Bump hyperlight-host to new snapshot api

.github/workflows/Benchmarks.yml

@@ -71,7 +71,7 @@ jobs:
         # the component model benchmark depends on the wasm wit component
         just ensure-tools
         just compile-wit
-        just bench-ci dev release ${{ matrix.hypervisor == 'mshv3' && 'mshv3' || ''}}
+        just bench-ci dev release ${{ matrix.hypervisor == 'mshv' && 'mshv2' || ''}}
       working-directory: ./src/hyperlight_wasm
 
     - name: Upload Benchmarks

.github/workflows/dep_rust.yml

@@ -92,15 +92,15 @@ jobs:
         run: just clippy ${{ matrix.config }}
 
       - name: Build
-        run:  just build ${{ matrix.config }}  ${{ matrix.hypervisor == 'mshv3' && 'mshv3' || ''}}
+        run:  just build ${{ matrix.config }}  ${{ matrix.hypervisor == 'mshv' && 'mshv2' || ''}}
         working-directory: ./src/hyperlight_wasm
 
       - name: Build Rust Wasm examples
         run:  just build-rust-wasm-examples ${{ matrix.config }}
         working-directory: ./src/hyperlight_wasm
 
       - name: Test
-        run:  just test ${{ matrix.config }}  ${{ matrix.hypervisor == 'mshv3' && 'mshv3' || ''}}
+        run:  just test ${{ matrix.config }}  ${{ matrix.hypervisor == 'mshv' && 'mshv2' || ''}}
         working-directory: ./src/hyperlight_wasm
 
       - name: Install github-cli (Windows)
@@ -118,14 +118,14 @@ jobs:
         shell: pwsh
 
       - name: Test Examples
-        run: just examples-ci ${{ matrix.config }} ${{ matrix.hypervisor == 'mshv3' && 'mshv3' || ''}}
+        run: just examples-ci ${{ matrix.config }} ${{ matrix.hypervisor == 'mshv' && 'mshv2' || ''}}
         working-directory: ./src/hyperlight_wasm
         env:
           # required for gh cli when downloading
           GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
 
       - name: Test Component Model Examples
-        run: just examples-components ${{ matrix.config }} ${{ matrix.hypervisor == 'mshv3' && 'mshv3' || ''}}
+        run: just examples-components ${{ matrix.config }} ${{ matrix.hypervisor == 'mshv' && 'mshv2' || ''}}
         working-directory: ./src/hyperlight_wasm
 
       ### Benchmarks ###
@@ -141,6 +141,6 @@ jobs:
 
       - name: Run benchmarks
         run: |
-          just bench-ci dev ${{ matrix.config }} ${{ matrix.hypervisor == 'mshv3' && 'mshv3' || ''}}
+          just bench-ci dev ${{ matrix.config }} ${{ matrix.hypervisor == 'mshv' && 'mshv2' || ''}}
         working-directory: ./src/hyperlight_wasm
         if: ${{ matrix.config == 'release' }}

Cargo.toml

@@ -13,4 +13,4 @@ repository = "https://github.com/hyperlight-dev/hyperlight-wasm"
 readme = "README.md"
 
 [workspace.dependencies]
-hyperlight-host = { version = "0.7.0", git = "https://github.com/hyperlight-dev/hyperlight", rev = "ea6fa8f", default-features = false, features = ["executable_heap", "init-paging"] }
+hyperlight-host = { version = "0.7.0", git = "https://github.com/hyperlight-dev/hyperlight", rev = "172fcfa69b0f9064c7a0e48e512f8a86ae1fdbe1", default-features = false, features = ["executable_heap", "init-paging"] }

Justfile

@@ -89,16 +89,16 @@ test target=default-target features="": (test-seccomp target features)
     cargo test test_metrics {{ if features =="" {''} else if features=="no-default-features" {"--no-default-features" } else {"--no-default-features -F " + features } }}  --profile={{ if target == "debug" {"dev"} else { target } }} -- --ignored 
 
 test-seccomp target=default-target features="": 
-    cargo test {{ if features =="" {'--no-default-features -F "kvm,mshv2,seccomp"'} else {"--no-default-features -F seccomp," + features } }} --profile={{ if target == "debug" {"dev"} else { target } }} -- --test-threads=1
-    cargo test {{ if features =="" {'--no-default-features -F "kvm,mshv2,seccomp"'} else {"--no-default-features -F seccomp," + features } }} test_metrics --profile={{ if target == "debug" {"dev"} else { target } }} -- --ignored --test-threads=1 
-    cargo test {{ if features =="" {'--no-default-features -F "kvm,mshv2,seccomp"'} else {"--no-default-features -F seccomp," + features } }} test_gather_metrics --profile={{ if target == "debug" {"dev"} else { target } }} -- --ignored --test-threads=1 
+    cargo test {{ if features =="" {'--no-default-features -F "kvm,mshv3,seccomp"'} else {"--no-default-features -F seccomp," + features } }} --profile={{ if target == "debug" {"dev"} else { target } }} -- --test-threads=1
+    cargo test {{ if features =="" {'--no-default-features -F "kvm,mshv3,seccomp"'} else {"--no-default-features -F seccomp," + features } }} test_metrics --profile={{ if target == "debug" {"dev"} else { target } }} -- --ignored --test-threads=1 
+    cargo test {{ if features =="" {'--no-default-features -F "kvm,mshv3,seccomp"'} else {"--no-default-features -F seccomp," + features } }} test_gather_metrics --profile={{ if target == "debug" {"dev"} else { target } }} -- --ignored --test-threads=1 
 
 examples-ci target=default-target features="": (build-rust-wasm-examples target)
     cargo run {{ if features =="" {''} else {"--no-default-features -F " + features } }} --profile={{ if target == "debug" {"dev"} else { target } }} --example helloworld
     cargo run {{ if features =="" {''} else {"--no-default-features -F " + features } }} --profile={{ if target == "debug" {"dev"} else { target } }} --example hostfuncs
     cargo run {{ if features =="" {''} else {"--no-default-features -F " + features } }} --profile={{ if target == "debug" {"dev"} else { target } }} --example rust_wasm_examples
     cargo run {{ if features =="" {''} else {"--no-default-features -F function_call_metrics," + features } }} --profile={{ if target == "debug" {"dev"} else { target } }} --example metrics
-    cargo run {{ if features =="" {"--no-default-features --features kvm,mshv2"} else {"--no-default-features -F function_call_metrics," + features } }} --profile={{ if target == "debug" {"dev"} else { target } }} --example metrics 
+    cargo run {{ if features =="" {"--no-default-features --features kvm,mshv3"} else {"--no-default-features -F function_call_metrics," + features } }} --profile={{ if target == "debug" {"dev"} else { target } }} --example metrics 
 
 examples-components target=default-target features="": (build-rust-component-examples target) 
     {{ wit-world }} cargo run {{ if features =="" {''} else {"--no-default-features -F " + features } }} --profile={{ if target == "debug" {"dev"} else { target } }} --example component_example

src/hyperlight_wasm/Cargo.toml

@@ -58,7 +58,7 @@ windows = { version = "0.61", features = ["Win32_System_Threading"] }
 page_size = "0.6.0"
 
 [dev-dependencies]
-hyperlight-component-macro = { version = "0.7.0", git = "https://github.com/hyperlight-dev/hyperlight", rev = "b61265e4aa9e2ecf8d648b994022caeea0205352" }
+hyperlight-component-macro = { version = "0.7.0", git = "https://github.com/hyperlight-dev/hyperlight", rev = "172fcfa69b0f9064c7a0e48e512f8a86ae1fdbe1" }
 examples_common = { path = "../examples_common" }
 criterion = { version = "0.6.0", features = ["html_reports"] }
 crossbeam-queue = "0.3"
@@ -76,7 +76,7 @@ goblin = "0.10.0"
 tar = "0.4.44"
 
 [features]
-default = ["function_call_metrics", "kvm", "mshv2"]
+default = ["function_call_metrics", "kvm", "mshv3"]
 function_call_metrics = ["hyperlight-host/function_call_metrics"]
 seccomp = ["hyperlight-host/seccomp"]
 print_debug = ["hyperlight-host/print_debug"]

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);
+        }
     }
 }