1use anyhow::Result;
2use async_trait::async_trait;
3use wasmtime::ResourceLimiterAsync;
4
5#[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 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}