spin_core/
limits.rs

1use async_trait::async_trait;
2use wasmtime::ResourceLimiterAsync;
3
4/// Async implementation of wasmtime's `StoreLimits`: https://github.com/bytecodealliance/wasmtime/blob/main/crates/wasmtime/src/limits.rs
5/// Used to limit the memory use and table size of each Instance
6#[derive(Default)]
7pub struct StoreLimitsAsync {
8    max_memory_size: Option<usize>,
9    max_table_elements: Option<usize>,
10    memory_consumed: u64,
11}
12
13#[async_trait]
14impl ResourceLimiterAsync for StoreLimitsAsync {
15    async fn memory_growing(
16        &mut self,
17        current: usize,
18        desired: usize,
19        maximum: Option<usize>,
20    ) -> wasmtime::Result<bool> {
21        let can_grow = if let Some(limit) = self.max_memory_size {
22            desired <= limit
23        } else {
24            true
25        };
26        if can_grow {
27            self.memory_consumed =
28                (self.memory_consumed as i64 + (desired as i64 - current as i64)) as u64;
29        } else {
30            tracing::warn!(
31                "error.type" = "memory_limit_exceeded",
32                current,
33                desired,
34                maximum,
35                max_memory_size = self.max_memory_size,
36                "instance memory limit exceeded",
37            );
38        }
39        Ok(can_grow)
40    }
41
42    async fn table_growing(
43        &mut self,
44        _current: usize,
45        desired: usize,
46        _maximum: Option<usize>,
47    ) -> wasmtime::Result<bool> {
48        let can_grow = if let Some(limit) = self.max_table_elements {
49            desired <= limit
50        } else {
51            true
52        };
53        Ok(can_grow)
54    }
55}
56
57impl StoreLimitsAsync {
58    pub fn new(max_memory_size: Option<usize>, max_table_elements: Option<usize>) -> Self {
59        Self {
60            max_memory_size,
61            max_table_elements,
62            memory_consumed: 0,
63        }
64    }
65
66    /// How much memory has been consumed in bytes
67    pub fn memory_consumed(&self) -> u64 {
68        self.memory_consumed
69    }
70}
71
72#[cfg(test)]
73mod tests {
74    use super::*;
75
76    #[tokio::test]
77    async fn test_store_limits_memory() {
78        let mut limits = StoreLimitsAsync {
79            max_memory_size: Some(65536),
80            ..Default::default()
81        };
82        assert!(limits.memory_growing(0, 65536, None).await.unwrap());
83        assert_eq!(limits.memory_consumed, 65536);
84        assert!(!limits.memory_growing(65536, 131072, None).await.unwrap());
85        assert_eq!(limits.memory_consumed, 65536);
86    }
87
88    #[tokio::test]
89    async fn test_store_limits_table() {
90        let mut limits = StoreLimitsAsync {
91            max_table_elements: Some(10),
92            ..Default::default()
93        };
94        assert!(limits.table_growing(9, 10, None).await.unwrap());
95        assert!(!limits.table_growing(10, 11, None).await.unwrap());
96    }
97
98    #[tokio::test]
99    async fn test_memory_consumed() {
100        let engine = wasmtime::Engine::new(crate::Config::default().wasmtime_config()).unwrap();
101        let linker = wasmtime::component::Linker::new(&engine);
102        let component = wasmtime::component::Component::new(
103            &engine,
104            r#"
105            (component
106                (core module $m (memory 1))
107                (core instance $a (instantiate $m))
108            )
109            "#,
110        )
111        .unwrap();
112        let mut store = wasmtime::Store::new(&engine, StoreLimitsAsync::default());
113        store.limiter_async(|s| s);
114        let _ = linker
115            .instantiate_async(&mut store, &component)
116            .await
117            .unwrap();
118        assert_eq!(store.data().memory_consumed(), 1 << 16);
119    }
120}
spin_core/limits.rs

spin_core/
limits.rs
