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 }
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 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}