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}