From 0fefa2b2d504e7b4f8297f5d35396df0dd57c989 Mon Sep 17 00:00:00 2001
From: Trevor Elliott <telliott@fastly.com>
Date: Mon, 7 Aug 2023 16:35:28 -0700
Subject: [PATCH] Use a ResourceLimiter for tracking allocations (#300)

Co-authored-by: Adam C. Foltzer <acfoltzer@acfoltzer.net>
---
 lib/src/execute.rs |  7 ++-----
 lib/src/linking.rs | 37 +++++++++++++++++++++++++++++++++++++
 2 files changed, 39 insertions(+), 5 deletions(-)

diff --git a/lib/src/execute.rs b/lib/src/execute.rs
index 6cccb707..5d634014 100644
--- a/lib/src/execute.rs
+++ b/lib/src/execute.rs
@@ -350,16 +350,13 @@ impl ExecuteCtx {
         // sent during execution.
         store.data_mut().close_downstream_response_sender();
 
-        let heap_pages = instance
-            .get_memory(&mut store, "memory")
-            .expect("`memory` is exported")
-            .size(&store);
+        let heap_bytes = store.data().limiter().memory_allocated;
 
         let request_duration = Instant::now().duration_since(start_timestamp);
 
         info!(
             "request completed using {} of WebAssembly heap",
-            bytesize::ByteSize::kib(heap_pages * 64)
+            bytesize::ByteSize::b(heap_bytes as u64)
         );
 
         info!("request completed in {:.0?}", request_duration);
diff --git a/lib/src/linking.rs b/lib/src/linking.rs
index dcefd2f4..496efefb 100644
--- a/lib/src/linking.rs
+++ b/lib/src/linking.rs
@@ -13,11 +13,42 @@ use {
     wasmtime_wasi_nn::WasiNnCtx,
 };
 
+#[derive(Default)]
+pub struct Limiter {
+    /// Total memory allocated so far.
+    pub memory_allocated: usize,
+}
+
+impl wasmtime::ResourceLimiter for Limiter {
+    fn memory_growing(
+        &mut self,
+        current: usize,
+        desired: usize,
+        _maximum: Option<usize>,
+    ) -> anyhow::Result<bool> {
+        // Track the diff in memory allocated over time. As each instance will start with 0 and
+        // gradually resize, this will track the total allocations throughout the lifetime of the
+        // instance.
+        self.memory_allocated += desired - current;
+        Ok(true)
+    }
+
+    fn table_growing(
+        &mut self,
+        _current: u32,
+        _desired: u32,
+        _maximum: Option<u32>,
+    ) -> anyhow::Result<bool> {
+        Ok(true)
+    }
+}
+
 pub struct WasmCtx {
     wasi: WasiCtx,
     wasi_nn: WasiNnCtx,
     session: Session,
     guest_profiler: Option<Box<GuestProfiler>>,
+    limiter: Limiter,
 }
 
 impl WasmCtx {
@@ -36,6 +67,10 @@ impl WasmCtx {
     pub fn take_guest_profiler(&mut self) -> Option<Box<GuestProfiler>> {
         self.guest_profiler.take()
     }
+
+    pub fn limiter(&self) -> &Limiter {
+        &self.limiter
+    }
 }
 
 impl WasmCtx {
@@ -60,6 +95,7 @@ pub(crate) fn create_store(
         wasi_nn,
         session,
         guest_profiler: guest_profiler.map(Box::new),
+        limiter: Limiter::default(),
     };
     let mut store = Store::new(ctx.engine(), wasm_ctx);
     store.set_epoch_deadline(1);
@@ -70,6 +106,7 @@ pub(crate) fn create_store(
         }
         Ok(UpdateDeadline::Yield(1))
     });
+    store.limiter(|ctx| &mut ctx.limiter);
     Ok(store)
 }