[DRAFT] Impl scheduler

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

@@ -1,4 +1,5 @@
 use crate::lind_wasmtime::host::DylinkMetadata;
+use crate::lind_wasmtime::host::{init_grate_pool, register_grate_handler_for_cage, cleanup_grate_handler};
 use crate::{cli::CliOptions, lind_wasmtime::host::HostCtx, lind_wasmtime::trampoline::*};
 use anyhow::{Context, Result, anyhow, bail};
 use cage::signal::{lind_signal_init, signal_may_trigger};
@@ -17,7 +18,8 @@ use wasmtime::{
     AsContextMut, Engine, Export, Func, InstantiateType, Linker, Module, Precompiled, SharedMemory,
     Store, Val, ValType, WasmBacktraceDetails,
 };
-use wasmtime_lind_3i::{VmCtxWrapper, init_vmctx_pool, rm_vmctx, set_vmctx, set_vmctx_thread};
+// use wasmtime_lind_3i::{VmCtxWrapper, init_vmctx_pool, rm_vmctx, set_vmctx, set_vmctx_thread};
+use wasmtime_lind_3i::*;
 use wasmtime_lind_common::LindEnviron;
 use wasmtime_lind_dylink::DynamicLoader;
 use wasmtime_lind_multi_process::{
@@ -61,13 +63,13 @@ pub fn execute_wasmtime(lindboot_cli: CliOptions) -> anyhow::Result<i32> {
     // new cage is created
     lind_manager.increment();
 
-    // Initialize vmctx pool
-    init_vmctx_pool();
+    let grate_cleanup_funcptr = cleanup_grate_handler as *const () as usize as u64;
     // Initialize trampoline entry function pointer for wasmtime runtime.
     // This is for grate calls to re-enter wasmtime runtime.
     threei::register_trampoline(
         threei_const::RUNTIME_TYPE_WASMTIME,
         grate_callback_trampoline,
+        grate_cleanup_funcptr,
     );
 
     // Register syscall handlers (clone/exec/exit) with 3i
@@ -632,6 +634,9 @@ fn load_main_module(
     // see comments at signal_may_trigger for more details
     signal_may_trigger(cageid);
 
+    init_grate_pool();
+    init_vmctx_pool();
+
     // The main challenge in enabling dynamic syscall interposition between grates and 3i lies in Rust’s
     // strict lifetime and ownership system, which makes retrieving the Wasmtime runtime context across
     // instance boundaries particularly difficult. To overcome this, the design employs low-level context
@@ -651,48 +656,41 @@ fn load_main_module(
     // This function will be called at either the first cage or exec-ed cages.
     set_vmctx_thread(cageid, THREAD_START_ID as u64, vmctx_wrapper);
 
+    let grate_template = GrateTemplate {
+        engine: module.engine().clone(),
+        module: module.clone(),
+        linker: linker_guard.clone(),
+    };
+
+    let host = store.data().clone();
+
+    // 4) register grate workers for this cage
+    register_grate_handler_for_cage(&grate_template, host, cageid)
+        .with_context(|| format!("failed to register grate workers for cage {}", cageid))?;
+
     // 4) Notify threei of the cage runtime type
     threei::set_cage_runtime(cageid, threei_const::RUNTIME_TYPE_WASMTIME);
 
     let mut linker = linker_guard.clone();
     linker.define_weak_imports_as_traps(&module);
 
-    // 5) Create backup instances to populate the vmctx pool
-    // See more comments in lind-3i/lib.rs
-    for _ in 0..INSTANCE_NUMBER {
-        let (instance, backup_cage_instanceid) = linker
-            .instantiate_with_lind_thread(&mut *store, &module, false)
-            .context(format!("failed to instantiate"))?;
-
-        // Extract vmctx pointer
-        let backup_cage_storeopaque = store.inner_mut();
-        let backup_cage_instancehandler = backup_cage_storeopaque.instance(backup_cage_instanceid);
-        let backup_vmctx_ptr: *mut c_void = backup_cage_instancehandler.vmctx().cast();
-
-        // Put vmctx in a Send+Sync wrapper
-        let backup_vmctx_wrapper = VmCtxWrapper {
-            vmctx: NonNull::new(backup_vmctx_ptr).ok_or_else(|| anyhow!("null vmctx"))?,
-        };
-
-        // Store the vmctx wrapper in the global table for later retrieval during grate calls
-        set_vmctx(cageid, backup_vmctx_wrapper);
-    }
-
     // must drop linker before jump into wasm
     drop(linker);
     drop(linker_guard);
 
+    println!("[lind-boot] Starting main module in cage {}", cageid);
+
     let ret = match func {
         Some(func) => invoke_func(store, func, &args),
         None => Ok(vec![]),
     };
 
-    if !rm_vmctx(cageid) {
-        panic!(
-            "[lind-boot] Failed to remove existing VMContext for cage_id {}",
-            cageid
-        );
-    }
+    // if !rm_vmctx(cageid) {
+    //     panic!(
+    //         "[lind-boot] Failed to remove existing VMContext for cage_id {}",
+    //         cageid
+    //     );
+    // }
 
     ret
 }
@@ -806,7 +804,7 @@ pub fn precompile_module(cli: &CliOptions) -> Result<()> {
         .with_context(|| format!("failed to read {}", wasm_path.display()))?;
     let cwasm_bytes = engine
         .precompile_module(&wasm_bytes)
-        .context("failed to precompile module")?;
+        .with_context(|| format!("failed to precompile module {}", wasm_path.display()))?;
     std::fs::write(&cwasm_path, cwasm_bytes)
         .with_context(|| format!("failed to write {}", cwasm_path.display()))?;
 
@@ -827,8 +825,8 @@ fn read_wasm_or_cwasm(engine: &Engine, path: &Path) -> Result<Module> {
     // We can therefore not call .context()? on this function since that would unwind and not run the Module::from_file()
     match engine.detect_precompiled_file(path) {
         Ok(_) => unsafe { Module::deserialize_file(engine, path) }
-            .context("failed to deserialize precompiled module"),
-        Err(_) => Module::from_file(engine, path).context("failed to compile module"),
+            .with_context(|| format!("failed to deserialize precompiled module {}", path.display())),
+        Err(_) => Module::from_file(engine, path).with_context(|| format!("failed to compile module {}", path.display())),
     }
 }
 

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

@@ -1,42 +1,19 @@
 use crate::cli::CliOptions;
-use std::sync::{Arc, Mutex};
+use std::sync::{Arc, Mutex, OnceLock};
 use wasmtime::{Table, TypedFunc};
 use wasmtime_lind_common::LindEnviron;
 use wasmtime_lind_multi_process::{LindCtx, LindHost};
 use wasmtime_lind_utils::LindGOT;
-
-/// 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,
->;
+use wasmtime_lind_3i::*;
+use sysdefs::constants::lind_platform_const;
 
 /// 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 lind_fork_ctx: Option<LindCtx<HostCtx, CliOptions>>,
-    pub pass_fptr_func: Option<PassFptrTyped>,
 }
 
 impl HostCtx {
@@ -52,7 +29,6 @@ impl HostCtx {
         Self {
             lind_environ: forked_lind_environ,
             lind_fork_ctx: forked_lind_fork_ctx,
-            pass_fptr_func: None,
         }
     }
 }
@@ -80,3 +56,114 @@ impl DylinkMetadata {
         }
     }
 }
+
+
+static GRATE_POOL: OnceLock<Vec<Mutex<Option<Arc<GrateHandler<HostCtx>>>>>> = OnceLock::new();
+
+/// Initialize the global `GrateHandler` pool.
+///
+/// This function must be called exactly once during lind-wasm startup, before any `GrateHandler` is
+/// pushed to or retrieved from the pool. It eagerly allocates one empty queue per possible `cage_id`.
+pub fn init_grate_pool() {
+    GRATE_POOL.get_or_init(|| {
+        (0..lind_platform_const::MAX_CAGEID)
+            .map(|_| Mutex::new(None))
+            .collect()
+    });
+}
+
+pub fn register_grate_handler_for_cage(
+    template: &GrateTemplate<HostCtx>,
+    host: HostCtx,
+    cageid: u64,
+) -> anyhow::Result<()>
+{
+    let handler = create_handler_for_cage(
+        template,
+        host,
+        cageid,
+        ConcurrencyMode::Serialized,
+    )?;
+
+    let pool = GRATE_POOL
+        .get()
+        .ok_or_else(|| anyhow::anyhow!("GRATE_POOL is not initialized"))?;
+
+    let slot = pool
+        .get(cageid as usize)
+        .ok_or_else(|| anyhow::anyhow!("invalid cageid {}", cageid))?;
+
+    let mut guard = slot.lock().unwrap();
+
+    if guard.is_some() {
+        anyhow::bail!("GrateHandler for cageid {} already exists", cageid);
+    }
+
+    *guard = Some(Arc::new(handler));
+    Ok(())
+}
+
+fn get_grate_handler(grate_id: u64) -> anyhow::Result<Arc<GrateHandler<HostCtx>>> {
+    println!("[lind-boot] Retrieving grate handler for grate_id {}", grate_id);
+    let pool = GRATE_POOL
+        .get()
+        .ok_or_else(|| anyhow::anyhow!("GRATE_POOL is not initialized"))?;
+
+    let slot = pool
+        .get(grate_id as usize)
+        .ok_or_else(|| anyhow::anyhow!("invalid grate_id {}", grate_id))?;
+
+    let guard = slot.lock().unwrap();
+
+    println!("[lind-boot] Grate handler for grate_id {} is {}", grate_id, if guard.is_some() { "present" } else { "absent" });
+
+    guard
+        .as_ref()
+        .cloned()
+        .ok_or_else(|| anyhow::anyhow!("grate handler {} not found", grate_id))
+}
+
+pub fn submit_grate_request(grate_id: u64, req: GrateRequest) -> anyhow::Result<i32> {
+    println!("[lind-boot] Submitting grate request to cage {}, handler_addr: {:#x}", req.cageid, req.handler_addr);
+    let handler = match get_grate_handler(grate_id) {
+        Ok(handler) => {
+            println!("[lind-boot] got handler");
+            handler
+        }
+        Err(e) => {
+            panic!("[lind-boot] get_grate_handler failed: {:?}", e);
+        }
+    };
+    handler.submit(req)
+}
+
+pub fn unregister_grate_handler(
+    grate_id: u64,
+) -> anyhow::Result<Arc<GrateHandler<HostCtx>>> {
+    let pool = GRATE_POOL
+        .get()
+        .ok_or_else(|| anyhow::anyhow!("GRATE_POOL is not initialized"))?;
+
+    let slot = pool
+        .get(grate_id as usize)
+        .ok_or_else(|| anyhow::anyhow!("invalid grate_id {}", grate_id))?;
+
+    let mut guard = slot.lock().unwrap();
+
+    guard
+        .take()
+        .ok_or_else(|| anyhow::anyhow!("grate handler {} not found", grate_id))
+}
+
+pub fn cleanup_grate_handler(grate_id: u64) -> anyhow::Result<()> {
+    let handler = unregister_grate_handler(grate_id)?;
+
+    // 1. Deactivate the handler to prevent new requests from being accepted
+    // and interrupt other running stores
+    handler.begin_shutdown();
+
+    // 2. Wait for the handler to finish processing any in-flight requests and become idle.
+    handler.wait_for_idle();
+
+    Ok(())
+}