spin_core/
limits.rs

1use anyhow::Result;
2use async_trait::async_trait;
3use wasmtime::ResourceLimiterAsync;
4
5/// Async implementation of wasmtime's `StoreLimits`: https://github.com/bytecodealliance/wasmtime/blob/main/crates/wasmtime/src/limits.rs
6/// Used to limit the memory use and table size of each Instance
7#[derive(Default)]
8pub struct StoreLimitsAsync {
9    max_memory_size: Option<usize>,
10    max_table_elements: Option<usize>,
11    memory_consumed: u64,
12}
13
14#[async_trait]
15impl ResourceLimiterAsync for StoreLimitsAsync {
16    async fn memory_growing(
17        &mut self,
18        current: usize,
19        desired: usize,
20        _maximum: Option<usize>,
21    ) -> Result<bool> {
22        let can_grow = if let Some(limit) = self.max_memory_size {
23            desired <= limit
24        } else {
25            true
26        };
27        if can_grow {
28            self.memory_consumed =
29                (self.memory_consumed as i64 + (desired as i64 - current as i64)) as u64;
30        }
31        Ok(can_grow)
32    }
33
34    async fn table_growing(
35        &mut self,
36        _current: usize,
37        desired: usize,
38        _maximum: Option<usize>,
39    ) -> Result<bool> {
40        let can_grow = if let Some(limit) = self.max_table_elements {
41            desired <= limit
42        } else {
43            true
44        };
45        Ok(can_grow)
46    }
47}
48
49impl StoreLimitsAsync {
50    pub fn new(max_memory_size: Option<usize>, max_table_elements: Option<usize>) -> Self {
51        Self {
52            max_memory_size,
53            max_table_elements,
54            memory_consumed: 0,
55        }
56    }
57
58    /// How much memory has been consumed in bytes
59    pub fn memory_consumed(&self) -> u64 {
60        self.memory_consumed
61    }
62}
63
64#[cfg(test)]
65mod tests {
66    use super::*;
67
68    #[tokio::test]
69    async fn test_store_limits_memory() {
70        let mut limits = StoreLimitsAsync {
71            max_memory_size: Some(65536),
72            ..Default::default()
73        };
74        assert!(limits.memory_growing(0, 65536, None).await.unwrap());
75        assert_eq!(limits.memory_consumed, 65536);
76        assert!(!limits.memory_growing(65536, 131072, None).await.unwrap());
77        assert_eq!(limits.memory_consumed, 65536);
78    }
79
80    #[tokio::test]
81    async fn test_store_limits_table() {
82        let mut limits = StoreLimitsAsync {
83            max_table_elements: Some(10),
84            ..Default::default()
85        };
86        assert!(limits.table_growing(9, 10, None).await.unwrap());
87        assert!(!limits.table_growing(10, 11, None).await.unwrap());
88    }
89}