perf: reset scratch from what was written on restore

CHANGELOG.md

@@ -71,6 +71,21 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/).
   will be extended in the future.
 * MSHV: vCPU creation reads the partition's XSAVE format once. This removes
   one hypercall from every restore.
+* On KVM, restore zeroes the scratch pages that hold data and keeps them
+  mapped, so the next run takes no faults on them. Pages holding only zeros
+  are kept as well, and every 16th restore gives them back (64 runs of them
+  at most), so what stays resident follows what recent runs wrote. On Linux
+  6.7 or newer, a restore costs time in how much of scratch is in use rather
+  than in its size.
+* On WHP, and on MSHV on x86_64, restore zeroes only the scratch pages written
+  since the last restore, from the hypervisor's dirty-page log and a log of host
+  writes. WHP no longer replaces the scratch mapping on each restore, so the
+  pages a guest writes stay resident between restores. MSHV tracks scratch of
+  2 MiB or more. After a run writes more than a tenth of it (a quarter from 32
+  MiB), where zeroing all of it costs less, it stops tracking for 16 restores.
+* When Windows cannot tell which scratch pages were written, restore zeroes
+  scratch of up to 16 MiB in place, which is faster than replacing it below
+  that size, and replaces larger scratch.
 
 ### Removed
 * `RunPool` and the run-specific `AllocError::InvalidAlign` variant.

bench_report.toml

@@ -4,6 +4,21 @@
 # nothing fails the report.
 allowlist = [
   "^function_call_codec/decode_vec_bytes_copy$",
+  "^guest_calls/call_with_restore/default$",
+  "^guest_calls/call_with_restore/scratch_64mib$",
+  "^guest_calls/call_with_restore/scratch_64mib_write_1mib$",
+  "^guest_calls/call_with_restore/scratch_256mib$",
+  "^guest_calls/call_with_restore/scratch_256mib_write_1mib$",
+  "^guest_calls/call_with_restore/scratch_256mib_write_16mib$",
+  "^guest_calls/call_with_restore/32_after_32mib_scratch_64mib$",
+  "^guest_calls/call_after_restore/default$",
+  "^guest_calls/call_after_restore/scratch_256mib$",
+  "^guest_calls/call_after_restore/scratch_256mib_write_1mib$",
+  "^guest_calls/call_after_restore/scratch_256mib_write_16mib$",
+  "^snapshots/restore/default$",
+  "^snapshots/restore/scratch_256mib$",
+  "^snapshots/restore/scratch_256mib_write_1mib$",
+  "^snapshots/restore/scratch_256mib_write_16mib$",
   "^function_call_codec/encode_control/byte_chunks$",
   "^function_call_codec/encode_control/vec_bytes$",
   "^payload_allocation/slot_pool_segmented/262144$",

src/hyperlight_host/benches/benchmarks.rs

@@ -131,6 +131,103 @@ fn bench_guest_call_with_restore(b: &mut criterion::Bencher, size: SandboxSize)
     });
 }
 
+/// Sandboxes with large scratch, where resetting scratch dominates a
+/// restore: the name, MiB of scratch, and bytes of heap the guest writes a
+/// call (0 for only a few pages, an Echo).
+const LARGE_SCRATCH: [(&str, usize, u64); 5] = [
+    ("scratch_64mib", 64, 0),
+    ("scratch_64mib_write_1mib", 64, 1 << 20),
+    ("scratch_256mib", 256, 0),
+    ("scratch_256mib_write_1mib", 256, 1 << 20),
+    ("scratch_256mib_write_16mib", 256, 16 << 20),
+];
+
+/// A sandbox with `mib` MiB of scratch whose heap can allocate `write`
+/// bytes.
+fn large_scratch_sandbox(mib: usize, write: u64) -> Sandbox {
+    SandboxBuilder::from_file(simple_guest_as_pathbuf())
+        .heap_size((4 << 20).max(4 * write))
+        .scratch_size(mib << 20)
+        .build()
+        .unwrap()
+}
+
+/// A guest call that writes `write` bytes of heap, or an Echo at 0.
+fn call_writing(sbox: &mut Sandbox, write: u64) {
+    if write == 0 {
+        sbox.call::<String>("Echo", "hello\n".to_string()).unwrap();
+    } else {
+        sbox.call::<()>("AllocAndWritePattern", write).unwrap();
+    }
+}
+
+fn bench_large_scratch_call_with_restore(b: &mut criterion::Bencher, mib: usize, write: u64) {
+    let mut sbox = large_scratch_sandbox(mib, write);
+    let snapshot = sbox.snapshot().unwrap();
+
+    b.iter(|| {
+        call_writing(&mut sbox, write);
+        sbox.restore(snapshot.clone()).unwrap();
+    });
+}
+
+/// After one call that wrote `heavy` bytes, the next 32 calls that write
+/// little and their restores, as a pool of sandboxes sees after one
+/// large request. An iteration times those 32.
+fn bench_call_with_restore_after_heavy(b: &mut criterion::Bencher, mib: usize, heavy: u64) {
+    let mut sbox = large_scratch_sandbox(mib, heavy);
+    let snapshot = sbox.snapshot().unwrap();
+
+    b.iter_custom(|iters| {
+        let mut total = Duration::ZERO;
+        for _ in 0..iters {
+            call_writing(&mut sbox, heavy);
+            sbox.restore(snapshot.clone()).unwrap();
+            let start = Instant::now();
+            for _ in 0..32 {
+                call_writing(&mut sbox, 0);
+                sbox.restore(snapshot.clone()).unwrap();
+            }
+            total += start.elapsed();
+        }
+        total
+    });
+}
+
+/// Only the call after a restore, which pays for any page the restore
+/// left the guest to fault in again.
+fn bench_call_after_restore(b: &mut criterion::Bencher, mut sbox: Sandbox, write: u64) {
+    let snapshot = sbox.snapshot().unwrap();
+
+    b.iter_custom(|iters| {
+        let mut total = Duration::ZERO;
+        for _ in 0..iters {
+            sbox.restore(snapshot.clone()).unwrap();
+            let start = Instant::now();
+            call_writing(&mut sbox, write);
+            total += start.elapsed();
+        }
+        total
+    });
+}
+
+/// Only the restore after a call.
+fn bench_restore_after_call(b: &mut criterion::Bencher, mib: usize, write: u64) {
+    let mut sbox = large_scratch_sandbox(mib, write);
+    let snapshot = sbox.snapshot().unwrap();
+
+    b.iter_custom(|iters| {
+        let mut total = Duration::ZERO;
+        for _ in 0..iters {
+            call_writing(&mut sbox, write);
+            let start = Instant::now();
+            sbox.restore(snapshot.clone()).unwrap();
+            total += start.elapsed();
+        }
+        total
+    });
+}
+
 fn bench_guest_call_with_host_function(b: &mut criterion::Bencher, size: SandboxSize) {
     let mut multiuse_sandbox = size
         .builder()
@@ -238,6 +335,29 @@ fn guest_calls_benchmark(c: &mut Criterion) {
         });
     }
 
+    for (name, mib, write) in LARGE_SCRATCH {
+        group.bench_function(format!("call_with_restore/{name}"), |b| {
+            bench_large_scratch_call_with_restore(b, mib, write)
+        });
+    }
+
+    group.bench_function("call_with_restore/32_after_32mib_scratch_64mib", |b| {
+        bench_call_with_restore_after_heavy(b, 64, 32 << 20)
+    });
+
+    group.bench_function("call_after_restore/default", |b| {
+        bench_call_after_restore(
+            b,
+            create_multiuse_sandbox_with_size(SandboxSize::Default),
+            0,
+        )
+    });
+    for (name, mib, write) in LARGE_SCRATCH.into_iter().skip(2) {
+        group.bench_function(format!("call_after_restore/{name}"), |b| {
+            bench_call_after_restore(b, large_scratch_sandbox(mib, write), write)
+        });
+    }
+
     for size in SandboxSize::all() {
         group.bench_function(format!("call_with_host_function/{}", size.name()), |b| {
             bench_guest_call_with_host_function(b, size)
@@ -338,6 +458,12 @@ fn snapshots_benchmark(c: &mut Criterion) {
         });
     }
 
+    for (name, mib, write) in LARGE_SCRATCH.into_iter().skip(2) {
+        group.bench_function(format!("restore/{name}"), |b| {
+            bench_restore_after_call(b, mib, write)
+        });
+    }
+
     group.finish();
 }
 

src/hyperlight_host/src/hypervisor/hyperlight_vm/aarch64.rs

@@ -7,7 +7,7 @@ use std::sync::Arc;
 
 use super::{
     AccessPageTableError, CreateHyperlightVmError, DispatchGuestCallError, HyperlightVm,
-    InitializeError,
+    InitializeError, ScratchDirtyLog,
 };
 #[cfg(hvf)]
 use crate::hypervisor::HvfInterruptHandle;
@@ -100,6 +100,7 @@ impl HyperlightVm {
             snapshot_memory: None,
             scratch_slot,
             scratch_memory: None,
+            scratch_dirty: ScratchDirtyLog::default(),
 
             mmap_regions: Vec::new(),