Add fastly_compute_runtime::get_heap_mib hostcall

ulyssa · ulyssa · commit 025418ce2ae2 · 2026-01-16T13:54:28.000-05:00
diff --git a/src/component/compute/compute_runtime.rs b/src/component/compute/compute_runtime.rs
@@ -8,7 +8,7 @@ impl compute_runtime::Host for ComponentCtx {
     }
 
     fn get_heap_mib(&mut self) -> compute_runtime::MemoryMib {
-        0 // TODO
+        self.session().get_heap_usage_mib()
     }
 
     fn get_sandbox_id(&mut self) -> String {
diff --git a/src/execute.rs b/src/execute.rs
@@ -882,6 +882,10 @@ impl ExecuteCtx {
 
         Some(rx)
     }
+
+    pub fn is_component(&self) -> bool {
+        matches!(self.instance_pre.as_ref(), Instance::Component(_))
+    }
 }
 
 pub struct ExecuteCtxBuilder {
diff --git a/src/linking.rs b/src/linking.rs
@@ -19,13 +19,15 @@ pub struct Limiter {
     internal: StoreLimits,
 }
 
-impl Default for Limiter {
-    fn default() -> Self {
-        Limiter::new(1, 1, 1)
+impl Limiter {
+    pub fn for_wasip2() -> Self {
+        Self::new(100, 100, 100)
+    }
+
+    pub fn for_wasip1() -> Self {
+        Self::new(1, 1, 1)
     }
-}
 
-impl Limiter {
     fn new(max_instances: usize, max_memories: usize, max_tables: usize) -> Self {
         Limiter {
             memory_allocated: 0,
@@ -99,7 +101,6 @@ pub struct ComponentCtx {
     pub wasi_random: wasmtime_wasi::random::WasiRandomCtx,
     pub(crate) session: Session,
     guest_profiler: Option<Box<GuestProfiler>>,
-    limiter: Limiter,
 }
 
 /// An extension trait for users of `ComponentCtx` to access the session.
@@ -132,7 +133,7 @@ impl ComponentCtx {
     }
 
     pub fn limiter(&self) -> &Limiter {
-        &self.limiter
+        self.session.limiter()
     }
 
     pub fn close_downstream_response_sender(&mut self, resp: Response<Body>) {
@@ -159,7 +160,6 @@ impl ComponentCtx {
             wasi_random: wasmtime_wasi::random::WasiRandomCtx::default(),
             session,
             guest_profiler: guest_profiler.map(Box::new),
-            limiter: Limiter::new(100, 100, 100),
         };
         let mut store = Store::new(ctx.engine(), wasm_ctx);
         store.set_epoch_deadline(1);
@@ -182,7 +182,7 @@ impl ComponentCtx {
             Ok(UpdateDeadline::Yield(1))
         });
 
-        store.limiter(|ctx| &mut ctx.limiter);
+        store.limiter(|ctx| ctx.session.limiter_mut());
         Ok(store)
     }
 }
@@ -207,7 +207,6 @@ pub struct WasmCtx {
     wasi_nn: WasiNnCtx,
     session: Session,
     guest_profiler: Option<Box<GuestProfiler>>,
-    limiter: Limiter,
 }
 
 impl WasmCtx {
@@ -228,7 +227,7 @@ impl WasmCtx {
     }
 
     pub fn limiter(&self) -> &Limiter {
-        &self.limiter
+        self.session.limiter()
     }
 }
 
@@ -260,7 +259,6 @@ 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);
@@ -283,7 +281,7 @@ pub(crate) fn create_store(
         Ok(UpdateDeadline::Yield(1))
     });
 
-    store.limiter(|ctx| &mut ctx.limiter);
+    store.limiter(|ctx| ctx.session.limiter_mut());
     Ok(store)
 }
 
diff --git a/src/session.rs b/src/session.rs
@@ -18,6 +18,7 @@ use std::sync::{Arc, Mutex};
 use std::time::Duration;
 
 use crate::cache::{Cache, CacheEntry};
+use crate::linking::Limiter;
 use crate::object_store::KvStoreError;
 use crate::wiggle_abi::types::{CacheBusyHandle, CacheHandle, FramingHeadersMode};
 
@@ -122,6 +123,7 @@ pub struct Session {
     secrets_by_name: PrimaryMap<SecretHandle, SecretLookup>,
     /// How many additional downstream requests have been receive by this Session.
     next_req_accepted: usize,
+    limiter: Limiter,
 }
 
 impl Session {
@@ -144,6 +146,12 @@ impl Session {
         });
         let downstream_req_body_handle = async_items.push(Some(AsyncItem::Body(body))).into();
 
+        let limiter = if ctx.is_component() {
+            Limiter::for_wasip2()
+        } else {
+            Limiter::for_wasip1()
+        };
+
         Session {
             session_id,
             downstream_req_handle,
@@ -165,6 +173,7 @@ impl Session {
             secrets_by_name: PrimaryMap::new(),
             downstream_pending_handle: None,
             next_req_accepted: 0,
+            limiter,
 
             ctx,
         }
@@ -1325,6 +1334,24 @@ impl Session {
     pub fn ctx(&self) -> &Arc<ExecuteCtx> {
         &self.ctx
     }
+
+    /// Get the guest's heap usage in mebibytes.
+    ///
+    /// This rounds up to the nearest mebibyte, so that guests won't accidentally
+    /// rely on implementation details that may change over time.
+    pub fn get_heap_usage_mib(&self) -> u32 {
+        const MEBIBYTE: usize = 1024 * 1024;
+        let mb = self.limiter.memory_allocated.next_multiple_of(MEBIBYTE) / MEBIBYTE;
+        mb.try_into().unwrap_or(u32::MAX)
+    }
+
+    pub fn limiter(&self) -> &Limiter {
+        &self.limiter
+    }
+
+    pub fn limiter_mut(&mut self) -> &mut Limiter {
+        &mut self.limiter
+    }
 }
 
 pub struct SelectedTargets<'session> {
diff --git a/src/wiggle_abi/compute_runtime.rs b/src/wiggle_abi/compute_runtime.rs
@@ -11,4 +11,8 @@ impl FastlyComputeRuntime for Session {
         // try to minimize timing attacks.
         Ok(self.active_cpu_time_us.load(Ordering::SeqCst) / 1000)
     }
+
+    fn get_heap_mib(&mut self, _memory: &mut GuestMemory<'_>) -> Result<u32, Error> {
+        Ok(self.get_heap_usage_mib())
+    }
 }
diff --git a/wasm_abi/compute-at-edge-abi/compute-at-edge.witx b/wasm_abi/compute-at-edge-abi/compute-at-edge.witx
@@ -1352,6 +1352,18 @@
     (@interface func (export "get_vcpu_ms")
         (result $err (expected $vcpu_ms (error $fastly_status)))
     )
+
+    ;;; Get a snapshot of the current dynamic memory usage, rounded up to the nearest mebibyte (2^20).
+    ;;;
+    ;;; This includes usage from the Wasm linear memory (heap) and usage from host allocations
+    ;;; made on behalf of this sandbox, e.g. buffered bodies of HTTP responses.
+    ;;; The returned value is just a snapshot- it can change without any explicit action
+    ;;; by the sandbox (for instance, additional response data coming in from an HTTP response.)
+    ;;; It can also change over time / across runs, as the Compute platform's memory usage
+    ;;; changes. Consider the returned value with these uncertainties in mind.
+    (@interface func (export "get_heap_mib")
+        (result $err (expected $memory_mib (error $fastly_status)))
+    )
 )
 
 (module $fastly_acl
diff --git a/wasm_abi/compute-at-edge-abi/typenames.witx b/wasm_abi/compute-at-edge-abi/typenames.witx
@@ -491,6 +491,7 @@
 
 (typename $body_length u64)
 (typename $vcpu_ms u64)
+(typename $memory_mib u32)
 
 (typename $inspect_info_mask
     (flags (@witx repr u32)

Original file line number	Diff line number	Diff line change
`@@ -8,7 +8,7 @@ impl compute_runtime::Host for ComponentCtx {`
`8`	`8`	`}`
`9`	`9`
`10`	`10`	`fn get_heap_mib(&mut self) -> compute_runtime::MemoryMib {`
`11`		`- 0 // TODO`
	`11`	`+ self.session().get_heap_usage_mib()`
`12`	`12`	`}`
`13`	`13`
`14`	`14`	`fn get_sandbox_id(&mut self) -> String {`
Original file line number	Diff line number	Diff line change
`@@ -882,6 +882,10 @@ impl ExecuteCtx {`
`882`	`882`
`883`	`883`	`Some(rx)`
`884`	`884`	`}`
	`885`	`+`
	`886`	`+ pub fn is_component(&self) -> bool {`
	`887`	`+ matches!(self.instance_pre.as_ref(), Instance::Component(_))`
	`888`	`+ }`
`885`	`889`	`}`
`886`	`890`
`887`	`891`	`pub struct ExecuteCtxBuilder {`
Original file line number	Diff line number	Diff line change
`@@ -11,4 +11,8 @@ impl FastlyComputeRuntime for Session {`
`11`	`11`	`// try to minimize timing attacks.`
`12`	`12`	`Ok(self.active_cpu_time_us.load(Ordering::SeqCst) / 1000)`
`13`	`13`	`}`
	`14`	`+`
	`15`	`+ fn get_heap_mib(&mut self, _memory: &mut GuestMemory<'_>) -> Result<u32, Error> {`
	`16`	`+ Ok(self.get_heap_usage_mib())`
	`17`	`+ }`
`14`	`18`	`}`