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        } else {
31            tracing::warn!(
32                "error.type" = "memory_limit_exceeded",
33                current,
34                desired,
35                maximum,
36                max_memory_size = self.max_memory_size,
37                "instance memory limit exceeded",
38            );
39        }
40        Ok(can_grow)
41    }
42
43    async fn table_growing(
44        &mut self,
45        _current: usize,
46        desired: usize,
47        _maximum: Option<usize>,
48    ) -> Result<bool> {
49        let can_grow = if let Some(limit) = self.max_table_elements {
50            desired <= limit
51        } else {
52            true
53        };
54        Ok(can_grow)
55    }
56}
57
58impl StoreLimitsAsync {
59    pub fn new(max_memory_size: Option<usize>, max_table_elements: Option<usize>) -> Self {
60        Self {
61            max_memory_size,
62            max_table_elements,
63            memory_consumed: 0,
64        }
65    }
66
67    /// How much memory has been consumed in bytes
68    pub fn memory_consumed(&self) -> u64 {
69        self.memory_consumed
70    }
71}
72
73#[cfg(test)]
74mod tests {
75    use super::*;
76
77    #[tokio::test]
78    async fn test_store_limits_memory() {
79        let mut limits = StoreLimitsAsync {
80            max_memory_size: Some(65536),
81            ..Default::default()
82        };
83        assert!(limits.memory_growing(0, 65536, None).await.unwrap());
84        assert_eq!(limits.memory_consumed, 65536);
85        assert!(!limits.memory_growing(65536, 131072, None).await.unwrap());
86        assert_eq!(limits.memory_consumed, 65536);
87    }
88
89    #[tokio::test]
90    async fn test_store_limits_table() {
91        let mut limits = StoreLimitsAsync {
92            max_table_elements: Some(10),
93            ..Default::default()
94        };
95        assert!(limits.table_growing(9, 10, None).await.unwrap());
96        assert!(!limits.table_growing(10, 11, None).await.unwrap());
97    }
98
99    #[tokio::test]
100    async fn test_memory_consumed() {
101        let engine = wasmtime::Engine::new(crate::Config::default().wasmtime_config()).unwrap();
102        let linker = wasmtime::component::Linker::new(&engine);
103        let component = wasmtime::component::Component::new(
104            &engine,
105            r#"
106            (component
107                (core module $m (memory 1))
108                (core instance $a (instantiate $m))
109            )
110            "#,
111        )
112        .unwrap();
113        let mut store = wasmtime::Store::new(&engine, StoreLimitsAsync::default());
114        store.limiter_async(|s| s);
115        let _ = linker
116            .instantiate_async(&mut store, &component)
117            .await
118            .unwrap();
119        assert_eq!(store.data().memory_consumed(), 1 << 16);
120    }
121}
spin_core/limits.rs

spin_core/
limits.rs
