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, + ) -> anyhow::Result { + // 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, + ) -> anyhow::Result { + Ok(true) + } +} + pub struct WasmCtx { wasi: WasiCtx, wasi_nn: WasiNnCtx, session: Session, guest_profiler: Option>, + limiter: Limiter, } impl WasmCtx { @@ -36,6 +67,10 @@ impl WasmCtx { pub fn take_guest_profiler(&mut self) -> Option> { 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) }