perfomance tuning grates: cache dispatcher + better grate_inflight handling. (#992)

* init

* fix

* fix2

* fix3

* Add comments for new code. Removed unused commented out code.

* Fix typos

Author: stupendoussuperpowers

src/cage/src/cage.rs

@@ -322,6 +322,25 @@ pub fn get_cage_or_debug_panic(cageid: u64, caller: &str) -> Option<Arc<Cage>> {
     result
 }
 
+/// Borrows the `cageid` cage and applies a function `f` to it, return `Some(R)` or `None`
+/// depending on whether tha cage is out of range or does not exist.
+///
+/// Preferred over `get_cage` for synchronous operations where cloning the cage through an `Arc` is
+/// unnecessary.
+///
+/// SAFETY: Assumes that the cage cannot be removed while `f` is running (e.g. by ensuring
+/// `grate_inflight > 0`).
+pub fn with_cage<F, R>(cageid: u64, f: F) -> Option<R>
+where
+    F: FnOnce(&Cage) -> R,
+{
+    if cageid >= MAX_CAGEID as u64 {
+        return None;
+    }
+
+    unsafe { CAGE_MAP[cageid as usize].as_deref().map(f) }
+}
+
 #[allow(static_mut_refs)]
 // SAFETY: This code is single-threaded during teardown, and no other
 // mutable or immutable references to `CAGE_MAP` exist while this call executes.

src/lind-boot/src/lind_wasmtime/host.rs

@@ -1,18 +1,43 @@
 use crate::cli::CliOptions;
 use std::sync::Arc;
+use wasmtime::TypedFunc;
 use wasmtime_lind_common::LindEnviron;
 use wasmtime_lind_multi_process::{LindCtx, LindHost};
 use wasmtime_wasi_threads::WasiThreadsCtx;
 
+/// Function type for the `pass_fptr_to_wt` function used as an entry point for grate-run syscalls.
+pub type PassFptrTyped = TypedFunc<
+    (
+        u64,
+        u64,
+        u64,
+        u64,
+        u64,
+        u64,
+        u64,
+        u64,
+        u64,
+        u64,
+        u64,
+        u64,
+        u64,
+        u64,
+    ),
+    i32,
+>;
+
 /// The HostCtx host structure stores all relevant execution context objects:
 /// `lind_environ`: argv/environ data served by the 4 host functions in lind-common;
 /// `lind_fork_ctx`: the multi-process management structure, encapsulating fork/exec state;
 /// `wasi_threads`: which manages WASI thread-related capabilities.
+/// `pass_fptr_func`: the dispatcher function defined in the grate's WASM module. Cached after the
+/// first invocation.
 #[derive(Default, Clone)]
 pub struct HostCtx {
     pub lind_environ: Option<LindEnviron>,
     pub wasi_threads: Option<Arc<WasiThreadsCtx<HostCtx>>>,
     pub lind_fork_ctx: Option<LindCtx<HostCtx, CliOptions>>,
+    pub pass_fptr_func: Option<PassFptrTyped>,
 }
 
 impl HostCtx {
@@ -29,6 +54,7 @@ impl HostCtx {
             lind_environ: forked_lind_environ,
             lind_fork_ctx: forked_lind_fork_ctx,
             wasi_threads: self.wasi_threads.clone(),
+            pass_fptr_func: None,
         }
     }
 }

src/lind-boot/src/lind_wasmtime/trampoline.rs

@@ -1,3 +1,4 @@
+use crate::lind_wasmtime::host::PassFptrTyped;
 use crate::{cli::CliOptions, lind_wasmtime::host::HostCtx};
 use anyhow::anyhow;
 use threei::threei_const;
@@ -68,31 +69,27 @@ pub extern "C" fn grate_callback_trampoline(
                 caller: instance,
             } = caller;
 
-            // Resolve the unified entry function once per call
-            let entry_func = instance
-                .host_state()
-                .downcast_ref::<Instance>()
-                .ok_or_else(|| anyhow!("bad host_state Instance"))?
-                .get_export(&mut store, "pass_fptr_to_wt")
-                .and_then(|f| f.into_func())
-                .ok_or_else(|| anyhow!("missing export `pass_fptr_to_wt`"))?;
+            let typed_func: PassFptrTyped = match store.data().pass_fptr_func.clone() {
+                // Check if the dispatcher function is already cached.
+                Some(f) => f,
+                // Retrieve and store the dispatcher function. Done on the first syscall
+                // invocation.
+                None => {
+                    let entry_func = instance
+                        .host_state()
+                        .downcast_ref::<Instance>()
+                        .ok_or_else(|| anyhow!("bad host_state Instance"))?
+                        .get_export(&mut store, "pass_fptr_to_wt")
+                        .and_then(|f| f.into_func())
+                        .ok_or_else(|| anyhow!("missing export `pass_fptr_to_wt`"))?;
 
-            let typed_func = entry_func.typed::<(
-                u64,
-                u64,
-                u64,
-                u64,
-                u64,
-                u64,
-                u64,
-                u64,
-                u64,
-                u64,
-                u64,
-                u64,
-                u64,
-                u64,
-            ), i32>(&mut store)?;
+                    let typed = entry_func.typed(&mut store)?;
+
+                    store.data_mut().pass_fptr_func = Some(typed.clone());
+
+                    typed
+                }
+            };
 
             // Call the entry function with all arguments and in grate function pointer
             typed_func.call(

src/threei/src/threei.rs

@@ -1,5 +1,6 @@
 //! Threei (Three Interposition) module
 use cage::memory::{check_addr_read, check_addr_rw};
+use cage::with_cage;
 use core::panic;
 use dashmap::DashMap;
 use dashmap::DashSet;
@@ -495,21 +496,23 @@ pub fn make_syscall(
 
         // Track this in-flight grate dispatch so cage_finalize() waits
         // for it to complete before removing the cage.
-        if let Some(grate_cage) = cage::get_cage(grateid) {
-            grate_cage
+        //
+        // If the cage is already dead, return early.
+        let cage_dead = with_cage(grateid, |grate| {
+            grate
                 .grate_inflight
                 .fetch_add(1, std::sync::atomic::Ordering::AcqRel);
-            // After incrementing, recheck is_dead: if the cage started
-            // exiting between our earlier check and now, bail out.
-            if grate_cage
-                .is_dead
-                .load(std::sync::atomic::Ordering::Acquire)
-            {
-                grate_cage
+            if grate.is_dead.load(std::sync::atomic::Ordering::Acquire) {
+                grate
                     .grate_inflight
                     .fetch_sub(1, std::sync::atomic::Ordering::AcqRel);
                 return -(Errno::ESRCH as i32);
             }
+            0
+        });
+
+        if cage_dead != Some(0) {
+            return -(Errno::ESRCH as i32);
         }
 
         let grate_result = _call_grate_func(
@@ -530,11 +533,11 @@ pub fn make_syscall(
         );
 
         // Decrement the in-flight counter now that the dispatch returned.
-        if let Some(grate_cage) = cage::get_cage(grateid) {
-            grate_cage
+        with_cage(grateid, |grate| {
+            grate
                 .grate_inflight
                 .fetch_sub(1, std::sync::atomic::Ordering::AcqRel);
-        }
+        });
 
         if let Some(ret) = grate_result {
             return ret;