Skip redundant host function registration on reload (#422)

`get_loaded_sandbox` and `unload` each called `RegisterHostModules`,
although the base snapshot already holds that registration. A
load/unload cycle made four guest calls, but only `register_handler` and
`RunHandler` are needed. In hyperlight-dev/hyperlight-bench#46,
`kvm-hyperlight-js-reload` lost 15% throughput against 0.4.0.

This change registers host functions only when a restored snapshot
records different ones than the builder supplies. It then takes a new
base snapshot, so later cycles skip registration.

`full_lifecycle/add_load_handle_unload`:

| | main | this PR |
|---|---|---|
| MSHV | 415 µs | 315 µs |
| KVM | 248 µs | 203 µs |

KVM stays about 10% slower than 0.4.0 because each guest call costs more
in Hyperlight 0.18.

Breaking: the snapshot metadata schema is now 2. Snapshots from 0.5.0
can record host functions their base never registered, so they are
rejected and must be recreated.

---------

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

Author: ludfjig

docs/persistent-snapshots.md

@@ -155,6 +155,12 @@ The guest can call every supplied host function, including ones the snapshot
 does not record. Handlers loaded after `unload()` can import them too. Supply
 only the modules the guest should reach.
 
+If the builder supplies host functions the snapshot does not record,
+hyperlight-js registers them in the guest and copies the guest state. A
+`JSSandbox` does this in `build()` and in each `restore()` of such a snapshot.
+A `LoadedJSSandbox` does it in `unload()` after such a restore. Supply exactly
+the recorded functions to avoid this cost.
+
 A loaded snapshot records the host functions its evaluated handlers need. Saving
 it again after a restore keeps that record, even when the restoring builder
 supplied more functions. A `JSSandbox` snapshot records every function supplied
@@ -239,6 +245,8 @@ the same output routing is required.
   version for diagnostics. It also records the guest runtime ABI version.
   Unsupported metadata schema or guest ABI versions are rejected before starting
   guest execution.
+- Snapshots saved by hyperlight-js 0.5.0 use metadata schema 1 and are
+  rejected. Recreate them.
 - Digest verification detects accidental or malicious mutation, but does not
   authenticate who produced a snapshot. Authenticate snapshot artifacts through
   the distribution mechanism used by your application.

src/hyperlight-js/src/sandbox/host_fn.rs

@@ -68,6 +68,8 @@ pub(crate) fn register_host_function_manifest(
 }
 
 /// Return the registered host functions grouped by module.
+///
+/// Modules without functions are omitted because the guest has nothing to register for them.
 pub(crate) fn host_function_manifest(
     modules: &HashMap<String, HostModule>,
 ) -> HostFunctionManifest {
@@ -81,6 +83,7 @@ pub(crate) fn host_function_manifest(
             functions.sort();
             (module.clone(), functions)
         })
+        .filter(|(_, functions)| !functions.is_empty())
         .collect()
 }
 

src/hyperlight-js/src/sandbox/js_sandbox.rs

@@ -74,6 +74,23 @@ pub fn validate_namespace_not_reserved(namespace: &str) -> Result<()> {
     Ok(())
 }
 
+/// Return a base snapshot whose guest state registers `host_functions`.
+///
+/// `inner` must hold the guest state of `base`, which registers the host
+/// functions recorded in `metadata`.
+fn snapshot_base_with_host_functions(
+    inner: &mut Sandbox,
+    base: Arc<HyperlightSnapshot>,
+    metadata: &SnapshotMetadata,
+    host_functions: &HostFunctionManifest,
+) -> Result<Arc<HyperlightSnapshot>> {
+    if metadata.host_functions() == host_functions {
+        return Ok(base);
+    }
+    register_host_function_manifest(inner, host_functions)?;
+    inner.snapshot()
+}
+
 /// A Hyperlight Sandbox with a JavaScript run time loaded but no guest code.
 pub struct JSSandbox {
     pub(super) inner: Sandbox,
@@ -84,7 +101,8 @@ pub struct JSSandbox {
     modules: HashMap<String, Script>,
     // Host functions selected by SandboxBuilder.
     host_functions: HostFunctionManifest,
-    // Snapshot of state before any handlers are added.
+    // Snapshot of state before any handlers are added. Its guest state has
+    // `host_functions` registered and matches the current guest state.
     // This is used to restore state back to a neutral JSSandbox.
     snapshot: Arc<HyperlightSnapshot>,
     // metric drop guard to manage sandbox metric
@@ -109,10 +127,12 @@ impl JSSandbox {
     pub(crate) fn from_loaded(
         mut loaded: Sandbox,
         snapshot: Arc<HyperlightSnapshot>,
+        metadata: &SnapshotMetadata,
         host_functions: HostFunctionManifest,
     ) -> Result<Self> {
         loaded.restore(snapshot.clone())?;
-        register_host_function_manifest(&mut loaded, &host_functions)?;
+        let snapshot =
+            snapshot_base_with_host_functions(&mut loaded, snapshot, metadata, &host_functions)?;
         Ok(Self {
             inner: loaded,
             handlers: HashMap::new(),
@@ -124,16 +144,22 @@ impl JSSandbox {
     }
 
     pub(crate) fn from_snapshot(
-        inner: Sandbox,
+        mut inner: Sandbox,
         snapshot: &Snapshot<JSSandbox>,
         host_functions: HostFunctionManifest,
     ) -> Result<Self> {
+        let base = snapshot_base_with_host_functions(
+            &mut inner,
+            snapshot.inner(),
+            snapshot.metadata(),
+            &host_functions,
+        )?;
         Ok(Self {
             inner,
             handlers: snapshot.metadata().handlers().collect(),
             modules: snapshot.metadata().modules().collect(),
             host_functions,
-            snapshot: snapshot.inner(),
+            snapshot: base,
             _metric_guard: SandboxMetricsGuard::new(),
         })
     }
@@ -164,9 +190,23 @@ impl JSSandbox {
             .metadata()
             .validate_host_functions(&self.host_functions)?;
         self.inner.restore(snapshot.inner())?;
+        let base = snapshot_base_with_host_functions(
+            &mut self.inner,
+            snapshot.inner(),
+            snapshot.metadata(),
+            &self.host_functions,
+        );
+        self.snapshot = match base {
+            Ok(base) => base,
+            Err(error) => {
+                // Return the guest to the state the unchanged fields describe.
+                // Report the original error even if that fails.
+                let _ = self.inner.restore(self.snapshot.clone());
+                return Err(error);
+            }
+        };
         self.handlers = snapshot.metadata().handlers().collect();
         self.modules = snapshot.metadata().modules().collect();
-        self.snapshot = snapshot.inner();
         Ok(())
     }
 
@@ -377,10 +417,6 @@ impl JSSandbox {
             self.host_functions.clone(),
         ));
 
-        // Publish the builder-selected host functions before compiling modules
-        // and handlers that may import them.
-        register_host_function_manifest(&mut self.inner, &self.host_functions)?;
-
         // Register user modules first so that handlers can import them.
         // NOTE: HashMap iteration order is non-deterministic, but this is safe
         // because modules are lazily compiled by the UserModuleLoader when first

src/hyperlight-js/src/sandbox/loaded_js_sandbox.rs

@@ -38,6 +38,8 @@ use crate::sandbox::metrics::SandboxMetricsGuard;
 pub struct LoadedJSSandbox {
     inner: Sandbox,
     // Snapshot of state before the sandbox was loaded and before any handlers were added.
+    // Its guest state has the host functions in `snapshot_metadata` registered,
+    // which `unload` relies on.
     // This is used to restore state back to a JSSandbox.
     snapshot: Arc<HyperlightSnapshot>,
     // Metadata for the currently captured loaded guest state. Its host-function
@@ -205,8 +207,13 @@ impl LoadedJSSandbox {
     /// Unloads the Handlers from the sandbox and returns a `JSSandbox` with the JavaScript runtime loaded.
     #[instrument(err(Debug), skip_all, level=Level::DEBUG)]
     pub fn unload(self) -> Result<JSSandbox> {
-        JSSandbox::from_loaded(self.inner, self.snapshot, self.available_host_functions)
-            .inspect(|_| metrics::counter!(METRIC_SANDBOX_UNLOADS).increment(1))
+        JSSandbox::from_loaded(
+            self.inner,
+            self.snapshot,
+            &self.snapshot_metadata,
+            self.available_host_functions,
+        )
+        .inspect(|_| metrics::counter!(METRIC_SANDBOX_UNLOADS).increment(1))
     }
 
     /// Take a snapshot of the the current state of the sandbox.

src/hyperlight-js/src/sandbox/sandbox_builder.rs

@@ -424,6 +424,9 @@ impl<T: SnapshotTarget> SandboxBuilder<Snapshot<T>> {
     ///
     /// The snapshot's memory layout overrides heap, scratch, and transport
     /// settings. Supply every host function and module loader the snapshot uses.
+    /// Host functions the snapshot does not record are registered in the guest,
+    /// which costs one guest call and a copy of guest state. A `JSSandbox` pays
+    /// this in `build()`, and a `LoadedJSSandbox` in `unload()`.
     pub fn from_snapshot(snapshot: Snapshot<T>) -> Self {
         Self::with_source(snapshot)
     }

src/hyperlight-js/src/sandbox/snapshot.rs

@@ -30,7 +30,8 @@ use super::js_sandbox::{validate_module_identifier, validate_namespace_not_reser
 use crate::{new_error, JSSandbox, LoadedJSSandbox, Result, Script};
 
 const METADATA_NAMESPACE: &str = "dev.hyperlight.hyperlight-js.snapshot.v1";
-const METADATA_SCHEMA_VERSION: u32 = 1;
+// Version 2 guarantees that a snapshot's base registers its recorded host functions.
+const METADATA_SCHEMA_VERSION: u32 = 2;
 const BASE_TAG_SUFFIX: &str = ".hljs-base";
 const MAX_OCI_TAG_LENGTH: usize = 128;
 const MAX_LOADED_SNAPSHOT_TAG_LENGTH: usize = MAX_OCI_TAG_LENGTH - BASE_TAG_SUFFIX.len();
@@ -119,6 +120,10 @@ impl SnapshotMetadata {
         self.kind
     }
 
+    pub(crate) fn host_functions(&self) -> &HostFunctionManifest {
+        &self.host_functions
+    }
+
     pub(crate) fn handlers(&self) -> impl Iterator<Item = (String, Script)> + '_ {
         self.handlers
             .clone()
@@ -543,6 +548,16 @@ mod tests {
             .contains("produced by hyperlight-js 9.8.7"));
     }
 
+    #[test]
+    fn metadata_rejects_schema_version_1() {
+        let mut metadata = metadata(SnapshotKind::JsSandbox);
+        metadata.schema_version = 1;
+
+        let error = metadata.validate().unwrap_err();
+
+        assert!(error.to_string().contains("metadata version 1 "));
+    }
+
     #[test]
     fn metadata_rejects_other_guest_abi_version() {
         let mut metadata = metadata(SnapshotKind::JsSandbox);