Add new benchmarks for varying sizes of sandboxes and guest function parameters

Signed-off-by: Ludvig Liljenberg <[email protected]>

Author: ludfjig

src/hyperlight_host/benches/benchmarks.rs

@@ -75,37 +75,65 @@ fn guest_call_benchmark(c: &mut Criterion) {
     group.finish();
 }
 
-fn guest_call_benchmark_large_param(c: &mut Criterion) {
+fn guest_call_benchmark_large_params(c: &mut Criterion) {
     let mut group = c.benchmark_group("guest_functions_with_large_parameters");
     #[cfg(target_os = "windows")]
     group.sample_size(10); // This benchmark is very slow on Windows, so we reduce the sample size to avoid long test runs.
 
-    // This benchmark includes time to first clone a vector and string, so it is not a "pure' benchmark of the guest call, but it's still useful
-    group.bench_function("guest_call_with_large_parameters", |b| {
-        const SIZE: usize = 50 * 1024 * 1024; // 50 MB
-        let large_vec = vec![0u8; SIZE];
-        let large_string = unsafe { String::from_utf8_unchecked(large_vec.clone()) }; // Safety: indeed above vec is valid utf8
+    // Helper function to create a benchmark for a specific size
+    let create_benchmark = |group: &mut criterion::BenchmarkGroup<_>, size_mb: usize| {
+        let benchmark_name = format!("guest_call_with_2_large_parameters_{}mb each", size_mb);
+        group.bench_function(&benchmark_name, |b| {
+            let size = size_mb * 1024 * 1024; // Convert MB to bytes
+            let large_vec = vec![0u8; size];
+            let large_string = unsafe { String::from_utf8_unchecked(large_vec.clone()) }; // Safety: indeed above vec is valid utf8
 
-        let mut config = SandboxConfiguration::default();
-        config.set_input_data_size(2 * SIZE + (1024 * 1024)); // 2 * SIZE + 1 MB, to allow 1MB for the rest of the serialized function call
-        config.set_heap_size(SIZE as u64 * 15);
+            let mut config = SandboxConfiguration::default();
+            config.set_input_data_size(2 * size + (1024 * 1024));
 
-        let sandbox = UninitializedSandbox::new(
-            GuestBinary::FilePath(simple_guest_as_string().unwrap()),
-            Some(config),
-        )
-        .unwrap();
-        let mut sandbox = sandbox.evolve().unwrap();
+            if size < 50 * 1024 * 1024 {
+                config.set_heap_size(size as u64 * 16);
+            } else {
+                config.set_heap_size(size as u64 * 11); // Set to 1GB for larger sizes
+            }
 
-        b.iter(|| {
-            sandbox
-                .call_guest_function_by_name::<()>(
-                    "LargeParameters",
-                    (large_vec.clone(), large_string.clone()),
-                )
-                .unwrap()
+            let sandbox = UninitializedSandbox::new(
+                GuestBinary::FilePath(simple_guest_as_string().unwrap()),
+                Some(config),
+            )
+            .unwrap();
+            let mut sandbox = sandbox.evolve().unwrap();
+
+            b.iter_custom(|iters| {
+                let mut total_duration = std::time::Duration::new(0, 0);
+
+                for _ in 0..iters {
+                    // Clone the data (not measured)
+                    let vec_clone = large_vec.clone();
+                    let string_clone = large_string.clone();
+
+                    // Measure only the guest function call
+                    let start = std::time::Instant::now();
+                    sandbox
+                        .call_guest_function_by_name::<()>(
+                            "LargeParameters",
+                            (vec_clone, string_clone),
+                        )
+                        .unwrap();
+                    total_duration += start.elapsed();
+                }
+
+                total_duration
+            });
         });
-    });
+    };
+
+    // Create benchmarks for different sizes
+    create_benchmark(&mut group, 5); // 5MB
+    create_benchmark(&mut group, 10); // 10MB
+    create_benchmark(&mut group, 20); // 20MB
+    create_benchmark(&mut group, 40); // 40MB
+    create_benchmark(&mut group, 60); // 60MB
 
     group.finish();
 }
@@ -138,9 +166,158 @@ fn sandbox_benchmark(c: &mut Criterion) {
     group.finish();
 }
 
+fn sandbox_heap_size_benchmark(c: &mut Criterion) {
+    let mut group = c.benchmark_group("sandbox_heap_sizes");
+
+    // Helper function to create sandbox with specific heap size
+    let create_sandbox_with_heap_size = |heap_size_mb: Option<u64>| {
+        let path = simple_guest_as_string().unwrap();
+        let config = if let Some(size_mb) = heap_size_mb {
+            let mut config = SandboxConfiguration::default();
+            config.set_heap_size(size_mb * 1024 * 1024); // Convert MB to bytes
+            Some(config)
+        } else {
+            None
+        };
+
+        let uninit_sandbox =
+            UninitializedSandbox::new(GuestBinary::FilePath(path), config).unwrap();
+        uninit_sandbox.evolve().unwrap()
+    };
+
+    // Benchmark sandbox creation with default heap size
+    group.bench_function("create_sandbox_default_heap", |b| {
+        b.iter_with_large_drop(|| create_sandbox_with_heap_size(None));
+    });
+
+    // Benchmark sandbox creation with 50MB heap
+    group.bench_function("create_sandbox_50mb_heap", |b| {
+        b.iter_with_large_drop(|| create_sandbox_with_heap_size(Some(50)));
+    });
+
+    // Benchmark sandbox creation with 100MB heap
+    group.bench_function("create_sandbox_100mb_heap", |b| {
+        b.iter_with_large_drop(|| create_sandbox_with_heap_size(Some(100)));
+    });
+
+    // Benchmark sandbox creation with 250MB heap
+    group.bench_function("create_sandbox_250mb_heap", |b| {
+        b.iter_with_large_drop(|| create_sandbox_with_heap_size(Some(250)));
+    });
+
+    // Benchmark sandbox creation with 500MB heap
+    group.bench_function("create_sandbox_500mb_heap", |b| {
+        b.iter_with_large_drop(|| create_sandbox_with_heap_size(Some(500)));
+    });
+
+    // Benchmark sandbox creation with 995MB heap (close to the limit of 1GB for a Sandbox )
+    group.bench_function("create_sandbox_995mb_heap", |b| {
+        b.iter_with_large_drop(|| create_sandbox_with_heap_size(Some(995)));
+    });
+
+    group.finish();
+}
+
+fn guest_call_heap_size_benchmark(c: &mut Criterion) {
+    let mut group = c.benchmark_group("guest_call_restore_heap_sizes");
+
+    // Helper function to create sandbox with specific heap size
+    let create_sandbox_with_heap_size = |heap_size_mb: Option<u64>| {
+        let path = simple_guest_as_string().unwrap();
+        let config = if let Some(size_mb) = heap_size_mb {
+            let mut config = SandboxConfiguration::default();
+            config.set_heap_size(size_mb * 1024 * 1024); // Convert MB to bytes
+            Some(config)
+        } else {
+            None
+        };
+
+        let uninit_sandbox =
+            UninitializedSandbox::new(GuestBinary::FilePath(path), config).unwrap();
+        uninit_sandbox.evolve().unwrap()
+    };
+
+    // Benchmark guest function call with default heap size
+    group.bench_function("guest_call_restore_default_heap", |b| {
+        let mut sandbox = create_sandbox_with_heap_size(None);
+        let snapshot = sandbox.snapshot().unwrap();
+
+        b.iter(|| {
+            sandbox
+                .call_guest_function_by_name::<String>("Echo", "hello\n".to_string())
+                .unwrap();
+            sandbox.restore(&snapshot).unwrap();
+        });
+    });
+
+    // Benchmark guest function call with 50MB heap
+    group.bench_function("guest_call_restore_50mb_heap", |b| {
+        let mut sandbox = create_sandbox_with_heap_size(Some(50));
+        let snapshot = sandbox.snapshot().unwrap();
+
+        b.iter(|| {
+            sandbox
+                .call_guest_function_by_name::<String>("Echo", "hello\n".to_string())
+                .unwrap();
+            sandbox.restore(&snapshot).unwrap();
+        });
+    });
+
+    // Benchmark guest function call with 100MB heap
+    group.bench_function("guest_call_restore_100mb_heap", |b| {
+        let mut sandbox = create_sandbox_with_heap_size(Some(100));
+        let snapshot = sandbox.snapshot().unwrap();
+
+        b.iter(|| {
+            sandbox
+                .call_guest_function_by_name::<String>("Echo", "hello\n".to_string())
+                .unwrap();
+            sandbox.restore(&snapshot).unwrap();
+        });
+    });
+
+    // Benchmark guest function call with 250MB heap
+    group.bench_function("guest_call_restore_250mb_heap", |b| {
+        let mut sandbox = create_sandbox_with_heap_size(Some(250));
+        let snapshot = sandbox.snapshot().unwrap();
+        b.iter(|| {
+            sandbox
+                .call_guest_function_by_name::<String>("Echo", "hello\n".to_string())
+                .unwrap();
+            sandbox.restore(&snapshot).unwrap();
+        });
+    });
+
+    // Benchmark guest function call with 500MB heap
+    group.bench_function("guest_call_restore_500mb_heap", |b| {
+        let mut sandbox = create_sandbox_with_heap_size(Some(500));
+        let snapshot = sandbox.snapshot().unwrap();
+        b.iter(|| {
+            sandbox
+                .call_guest_function_by_name::<String>("Echo", "hello\n".to_string())
+                .unwrap();
+            sandbox.restore(&snapshot).unwrap();
+        });
+    });
+
+    // Benchmark guest function call with 995MB heap
+    group.bench_function("guest_call_restore_995mb_heap", |b| {
+        let mut sandbox = create_sandbox_with_heap_size(Some(995));
+        let snapshot = sandbox.snapshot().unwrap();
+        b.iter(|| {
+            sandbox
+                .call_guest_function_by_name::<String>("Echo", "hello\n".to_string())
+                .unwrap();
+            sandbox.restore(&snapshot).unwrap();
+        });
+    });
+
+    group.finish();
+}
+
 criterion_group! {
     name = benches;
     config = Criterion::default();
-    targets = guest_call_benchmark, sandbox_benchmark, guest_call_benchmark_large_param
+    targets = guest_call_benchmark, sandbox_benchmark, sandbox_heap_size_benchmark, guest_call_benchmark_large_params, guest_call_heap_size_benchmark
 }
 criterion_main!(benches);