Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion src/component/compute/compute_runtime.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down
4 changes: 4 additions & 0 deletions src/execute.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down
24 changes: 11 additions & 13 deletions src/linking.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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>) {
Expand All @@ -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);
Expand All @@ -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)
}
}
Expand All @@ -207,7 +207,6 @@ pub struct WasmCtx {
wasi_nn: WasiNnCtx,
session: Session,
guest_profiler: Option<Box<GuestProfiler>>,
limiter: Limiter,
}

impl WasmCtx {
Expand All @@ -228,7 +227,7 @@ impl WasmCtx {
}

pub fn limiter(&self) -> &Limiter {
&self.limiter
self.session.limiter()
}
}

Expand Down Expand Up @@ -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);
Expand All @@ -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)
}

Expand Down
28 changes: 28 additions & 0 deletions src/session.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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};

Expand Down Expand Up @@ -122,6 +123,8 @@ pub struct Session {
secrets_by_name: PrimaryMap<SecretHandle, SecretLookup>,
/// How many additional downstream requests have been receive by this Session.
next_req_accepted: usize,
/// Memory usage limiter to ensure the guest doesn't use over 128mb of heap.
limiter: Limiter,
}

impl Session {
Expand All @@ -144,6 +147,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,
Expand All @@ -165,6 +174,7 @@ impl Session {
secrets_by_name: PrimaryMap::new(),
downstream_pending_handle: None,
next_req_accepted: 0,
limiter,

ctx,
}
Expand Down Expand Up @@ -1325,6 +1335,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> {
Expand Down
4 changes: 4 additions & 0 deletions src/wiggle_abi/compute_runtime.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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())
}
}
12 changes: 12 additions & 0 deletions wasm_abi/compute-at-edge-abi/compute-at-edge.witx
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
1 change: 1 addition & 0 deletions wasm_abi/compute-at-edge-abi/typenames.witx
Original file line number Diff line number Diff line change
Expand Up @@ -491,6 +491,7 @@

(typename $body_length u64)
(typename $vcpu_ms u64)
(typename $memory_mib u32)

(typename $inspect_info_mask
(flags (@witx repr u32)
Expand Down
Loading