[host] migrate to snapshot-based sandbox API

Replace transition-based sandbox evolution with direct snapshot/restore API:

- Remove EvolvableSandbox/DevolvableSandbox trait usage
- Replace transition callbacks with direct method calls
- Add snapshot/restore methods to LoadedWasmSandbox
- Store initial snapshots in WasmSandbox for efficient unloading
- Export Snapshot type in public API

This provides more direct control over sandbox state management
and enables features like checkpoint/restore.

Signed-off-by: Ludvig Liljenberg <4257730+ludfjig@users.noreply.github.com>

Author: ludfjig

src/hyperlight_wasm/src/lib.rs

@@ -39,6 +39,8 @@ pub type Result<T> = hyperlight_host::Result<T>;
 pub use hyperlight_host::is_hypervisor_present;
 /// Create a generic HyperlightError
 pub use hyperlight_host::new_error;
+// A snapshot of the memory of a sandbox at a given point in time.
+pub use hyperlight_host::sandbox::snapshot::Snapshot;
 
 /// Get the build information for this version of hyperlight-wasm
 pub fn get_build_info() -> BuildInfo {

src/hyperlight_wasm/src/sandbox/loaded_wasm_sandbox.rs

@@ -14,14 +14,14 @@ See the License for the specific language governing permissions and
 limitations under the License.
 */
 
+use std::fmt::Debug;
 use std::sync::Arc;
 
 use hyperlight_host::func::{ParameterTuple, SupportedReturnType};
 // re-export the InterruptHandle trait as it's part of the public API
 pub use hyperlight_host::hypervisor::InterruptHandle;
 use hyperlight_host::sandbox::Callable;
-use hyperlight_host::sandbox_state::sandbox::{DevolvableSandbox, Sandbox};
-use hyperlight_host::sandbox_state::transition::Noop;
+use hyperlight_host::sandbox::snapshot::Snapshot;
 use hyperlight_host::{MultiUseSandbox, Result, log_then_return, new_error};
 
 use super::metrics::METRIC_TOTAL_LOADED_WASM_SANDBOXES;
@@ -37,16 +37,15 @@ use crate::sandbox::metrics::{METRIC_ACTIVE_LOADED_WASM_SANDBOXES, METRIC_SANDBO
 /// memory context. If you want to "reset" the memory context, create
 /// a new `LoadedWasmSandbox` -- either from another `WasmSandbox` or by
 /// calling `my_loaded_wasm_sandbox.devolve()?.evolve()?`
-#[derive(Debug)]
 pub struct LoadedWasmSandbox {
     // inner is an Option<MultiUseSandbox> as we need to take ownership of it
     // We implement drop on the LoadedWasmSandbox to decrement the count of Sandboxes when it is dropped
     // because of this we cannot implement drop without making inner an Option (alternatively we could make MultiUseSandbox Copy but that would introduce other issues)
     inner: Option<MultiUseSandbox>,
+    // The state the sandbox was in before loading a wasm module. Used for transitioning back to a `WasmSandbox` (unloading the wasm module).
+    runtime_snapshot: Option<Snapshot>,
 }
 
-impl Sandbox for LoadedWasmSandbox {}
-
 impl LoadedWasmSandbox {
     /// Call the function in the guest with the name `fn_name`, passing
     /// parameters `params`.
@@ -64,17 +63,50 @@ impl LoadedWasmSandbox {
             None => log_then_return!("No inner MultiUseSandbox to call_guest_function"),
         }
     }
+
+    /// Take a snapshot of the current state of the sandbox.
+    pub fn snapshot(&mut self) -> Result<Snapshot> {
+        match &mut self.inner {
+            Some(inner) => inner.snapshot(),
+            None => log_then_return!("No inner MultiUseSandbox to snapshot"),
+        }
+    }
+
+    /// Restore the state of the sandbox to the state captured in the given snapshot.
+    pub fn restore(&mut self, snapshot: &Snapshot) -> Result<()> {
+        match &mut self.inner {
+            Some(inner) => inner.restore(snapshot),
+            None => log_then_return!("No inner MultiUseSandbox to restore"),
+        }
+    }
+
     /// unload the wasm module and return a `WasmSandbox` that can be used to load another module
-    pub fn unload_module(self) -> Result<WasmSandbox> {
-        self.devolve(Noop::default()).inspect(|_| {
+    pub fn unload_module(mut self) -> Result<WasmSandbox> {
+        let sandbox = self
+            .inner
+            .take()
+            .ok_or_else(|| new_error!("No inner MultiUseSandbox to unload"))?;
+
+        let snapshot = self
+            .runtime_snapshot
+            .take()
+            .ok_or_else(|| new_error!("No snapshot of the WasmSandbox to unload"))?;
+
+        WasmSandbox::new_from_loaded(sandbox, snapshot).inspect(|_| {
             metrics::counter!(METRIC_SANDBOX_UNLOADS).increment(1);
         })
     }
 
-    pub(super) fn new(inner: MultiUseSandbox) -> Result<LoadedWasmSandbox> {
+    pub(super) fn new(
+        inner: MultiUseSandbox,
+        runtime_snapshot: Snapshot,
+    ) -> Result<LoadedWasmSandbox> {
         metrics::gauge!(METRIC_ACTIVE_LOADED_WASM_SANDBOXES).increment(1);
         metrics::counter!(METRIC_TOTAL_LOADED_WASM_SANDBOXES).increment(1);
-        Ok(LoadedWasmSandbox { inner: Some(inner) })
+        Ok(LoadedWasmSandbox {
+            inner: Some(inner),
+            runtime_snapshot: Some(runtime_snapshot),
+        })
     }
 
     /// Get a handle to the interrupt handler for this sandbox,
@@ -100,26 +132,20 @@ impl Callable for LoadedWasmSandbox {
     }
 }
 
-/// Capability to transform a `LoadedWasmSandbox` back down to a
-/// `WasmSandbox`
-impl DevolvableSandbox<LoadedWasmSandbox, WasmSandbox, Noop<LoadedWasmSandbox, WasmSandbox>>
-    for LoadedWasmSandbox
-{
-    fn devolve(mut self, _: Noop<LoadedWasmSandbox, WasmSandbox>) -> Result<WasmSandbox> {
-        let new_inner: MultiUseSandbox = match self.inner.take() {
-            Some(inner) => inner.devolve(Noop::default())?,
-            None => log_then_return!("No inner MultiUseSandbox to devolve"),
-        };
-        Ok(WasmSandbox::new(new_inner))
-    }
-}
-
 impl Drop for LoadedWasmSandbox {
     fn drop(&mut self) {
         metrics::gauge!(METRIC_ACTIVE_LOADED_WASM_SANDBOXES).decrement(1);
     }
 }
 
+impl Debug for LoadedWasmSandbox {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("LoadedWasmSandbox")
+            .field("inner", &self.inner)
+            .finish()
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use std::sync::Arc;
@@ -161,7 +187,7 @@ mod tests {
         }
         .unwrap();
 
-        call_funcs(loaded_wasm_sandbox, 1000);
+        call_funcs(loaded_wasm_sandbox, 500);
     }
 
     #[test]

src/hyperlight_wasm/src/sandbox/metrics.rs

@@ -35,8 +35,16 @@ pub(crate) static METRIC_SANDBOX_UNLOADS: &str = "sandbox_unloads_total";
 #[cfg(test)]
 mod tests {
     use examples_common::get_wasm_module_path;
+    use hyperlight_host::HyperlightError;
 
-    use crate::{LoadedWasmSandbox, ProtoWasmSandbox};
+    use crate::{LoadedWasmSandbox, ProtoWasmSandbox, Result};
+
+    fn get_time_since_boot_microsecond() -> Result<i64> {
+        let res = std::time::SystemTime::now()
+            .duration_since(std::time::SystemTime::UNIX_EPOCH)?
+            .as_micros();
+        i64::try_from(res).map_err(HyperlightError::IntConversionFailure)
+    }
 
     #[test]
     #[ignore = "Needs to run separately to not get influenced by other tests"]
@@ -46,7 +54,13 @@ mod tests {
         recorder.install().unwrap();
 
         let snapshot = {
-            let sandbox = ProtoWasmSandbox::default();
+            let mut sandbox = ProtoWasmSandbox::default();
+            sandbox
+                .register(
+                    "GetTimeSinceBootMicrosecond",
+                    get_time_since_boot_microsecond,
+                )
+                .unwrap();
 
             let wasm_sandbox = sandbox.load_runtime().unwrap();
             let loaded_wasm_sandbox: LoadedWasmSandbox = {
@@ -57,6 +71,10 @@ mod tests {
             snapshotter.snapshot()
         };
         let snapshot = snapshot.into_vec();
-        assert_eq!(snapshot.len(), 8);
+        if cfg!(feature = "function_call_metrics") {
+            assert_eq!(snapshot.len(), 10);
+        } else {
+            assert_eq!(snapshot.len(), 8);
+        }
     }
 }

src/hyperlight_wasm/src/sandbox/proto_wasm_sandbox.rs

@@ -14,15 +14,11 @@ See the License for the specific language governing permissions and
 limitations under the License.
 */
 
-use hyperlight_host::func::call_ctx::MultiUseGuestCallContext;
 use hyperlight_host::func::{HostFunction, ParameterTuple, Registerable, SupportedReturnType};
-use hyperlight_host::sandbox::Callable;
 #[cfg(all(feature = "seccomp", target_os = "linux"))]
 use hyperlight_host::sandbox::ExtraAllowedSyscall;
 use hyperlight_host::sandbox::config::SandboxConfiguration;
-use hyperlight_host::sandbox_state::sandbox::{EvolvableSandbox, Sandbox};
-use hyperlight_host::sandbox_state::transition::{MultiUseContextCallback, Noop};
-use hyperlight_host::{GuestBinary, MultiUseSandbox, Result, UninitializedSandbox, new_error};
+use hyperlight_host::{GuestBinary, Result, UninitializedSandbox, new_error};
 
 use super::metrics::{METRIC_ACTIVE_PROTO_WASM_SANDBOXES, METRIC_TOTAL_PROTO_WASM_SANDBOXES};
 use super::sandbox_builder::SandboxBuilder;
@@ -39,8 +35,6 @@ pub struct ProtoWasmSandbox {
     pub(super) inner: Option<UninitializedSandbox>,
 }
 
-impl Sandbox for ProtoWasmSandbox {}
-
 impl Registerable for ProtoWasmSandbox {
     fn register_host_function<Args: ParameterTuple, Output: SupportedReturnType>(
         &mut self,
@@ -95,27 +89,20 @@ impl ProtoWasmSandbox {
     /// The returned `WasmSandbox` can be then be cached and used to load a different Wasm module.
     ///
     pub fn load_runtime(mut self) -> Result<WasmSandbox> {
-        let multi_use_sandbox: MultiUseSandbox = match self.inner.take() {
-            Some(s) => s.evolve(Noop::default())?,
+        let mut sandbox = match self.inner.take() {
+            Some(s) => s.evolve()?,
             None => return Err(new_error!("No inner sandbox found.")),
         };
 
-        let func = Box::new(move |call_ctx: &mut MultiUseGuestCallContext| {
-            let res: i32 = call_ctx.call("InitWasmRuntime", ())?;
-            if res != 0 {
-                return Err(new_error!(
-                    "InitWasmRuntime Failed  with error code {:?}",
-                    res
-                ));
-            }
-            Ok(())
-        });
-
-        let transition_func = MultiUseContextCallback::from(func);
-
-        let new_sbox: MultiUseSandbox = multi_use_sandbox.evolve(transition_func)?;
+        let res: i32 = sandbox.call_guest_function_by_name("InitWasmRuntime", ())?;
+        if res != 0 {
+            return Err(new_error!(
+                "InitWasmRuntime Failed  with error code {:?}",
+                res
+            ));
+        }
 
-        Ok(WasmSandbox::new(new_sbox))
+        WasmSandbox::new(sandbox)
     }
 
     /// Register the given host function `host_func` with `self` under