introduce memory limiter

kitsune-soc · Nov 4, 2024 · abeb10f · abeb10f
1 parent f9d2b75
commit abeb10f
Show file tree

Hide file tree

Showing 3 changed files with 71 additions and 12 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/crates/kitsune-wasm-mrf/Cargo.toml b/crates/kitsune-wasm-mrf/Cargo.toml
@@ -16,6 +16,7 @@ fred = { workspace = true }
 futures-util = { version = "0.3.31", default-features = false, features = [
     "alloc",
 ] }
+human-size = "0.4.3"
 kitsune-config = { workspace = true }
 kitsune-derive = { workspace = true }
 kitsune-error = { workspace = true }

diff --git a/crates/kitsune-wasm-mrf/src/ctx.rs b/crates/kitsune-wasm-mrf/src/ctx.rs
@@ -1,4 +1,7 @@
-use crate::{kv_storage, mrf_wit::v1::fep::mrf::keyvalue};
+use crate::{
+    kv_storage,
+    mrf_wit::v1::fep::mrf::{http, keyvalue},
+};
 use slab::Slab;
 use triomphe::Arc;
 use wasmtime::{
@@ -7,6 +10,8 @@ use wasmtime::{
 };
 use wasmtime_wasi::{WasiCtx, WasiCtxBuilder, WasiView};
 
+const TABLE_ELEMENT_SIZE: usize = std::mem::size_of::<usize>();
+
 pub struct KvContext {
     pub module_name: Option<String>,
     pub storage: Arc<kv_storage::BackendDispatch>,
@@ -23,8 +28,49 @@ impl KvContext {
     }
 }
 
+pub struct HttpContext {
+    pub client: kitsune_http_client::Client,
+    pub bodies: Slab<todo!()>,
+}
+
+impl HttpContext {
+    #[inline]
+    pub fn get_body(&self, rep: &Resource<http::ResponseBody>) -> &todo!() {
+        &self.bodies[rep.rep() as usize]
+    }
+}
+
+pub struct ResourceLimiter {
+    pub max_memory_size: human_size::Size,
+    pub used_memory: usize,
+}
+
+impl wasmtime::ResourceLimiter for ResourceLimiter {
+    fn memory_growing(
+        &mut self,
+        current: usize,
+        desired: usize,
+        _maximum: Option<usize>,
+    ) -> wasmtime::Result<bool> {
+        self.used_memory += desired - current;
+        Ok(self.used_memory <= self.max_memory_size.to_bytes() as usize)
+    }
+
+    fn table_growing(
+        &mut self,
+        current: usize,
+        desired: usize,
+        _maximum: Option<usize>,
+    ) -> wasmtime::Result<bool> {
+        self.used_memory += TABLE_ELEMENT_SIZE * (desired - current);
+        Ok(self.used_memory <= self.max_memory_size.to_bytes() as usize)
+    }
+}
+
 pub struct Context {
+    pub http_ctx: HttpContext,
     pub kv_ctx: KvContext,
+    pub resource_limiter: ResourceLimiter,
     pub resource_table: ResourceTable,
     pub wasi_ctx: WasiCtx,
 }
@@ -50,16 +96,27 @@ pub fn construct_store(
         .allow_udp(false)
         .build();
 
-    Store::new(
-        engine,
-        Context {
-            kv_ctx: KvContext {
-                module_name: None,
-                storage,
-                buckets: Slab::new(),
-            },
-            resource_table: ResourceTable::new(),
-            wasi_ctx,
+    let data = Context {
+        http_ctx: HttpContext {
+            client: kitsune_http_client::Client::builder()
+                .content_length_limit(None)
+                .build(),
+            bodies: Slab::new(),
+        },
+        kv_ctx: KvContext {
+            module_name: None,
+            storage,
+            buckets: Slab::new(),
         },
-    )
+        resource_limiter: ResourceLimiter {
+            max_memory_size: human_size::Size::new(100, human_size::Any::Mebibyte).unwrap(),
+            used_memory: 0,
+        },
+        resource_table: ResourceTable::new(),
+        wasi_ctx,
+    };
+
+    let mut store = Store::new(engine, data);
+    store.limiter(|store| &mut store.resource_limiter);
+    store
 }