|
| 1 | +//! Criterion benchmarks for WASM runtime performance. |
| 2 | +//! |
| 3 | +//! Run with: cargo bench --package mcp-wasm-runtime |
| 4 | +
|
| 5 | +use criterion::{Criterion, criterion_group, criterion_main}; |
| 6 | +use mcp_bridge::Bridge; |
| 7 | +use mcp_wasm_runtime::{ModuleCache, Runtime, SecurityConfig}; |
| 8 | +use std::hint::black_box; |
| 9 | +use std::sync::Arc; |
| 10 | + |
| 11 | +/// Benchmark WASM module compilation time. |
| 12 | +fn bench_module_compilation(c: &mut Criterion) { |
| 13 | + let wat = r#" |
| 14 | + (module |
| 15 | + (func $add (param $a i32) (param $b i32) (result i32) |
| 16 | + local.get $a |
| 17 | + local.get $b |
| 18 | + i32.add |
| 19 | + ) |
| 20 | + (func (export "main") (result i32) |
| 21 | + (i32.const 10) |
| 22 | + (i32.const 32) |
| 23 | + (call $add) |
| 24 | + ) |
| 25 | + ) |
| 26 | + "#; |
| 27 | + |
| 28 | + let wasm_bytes = wat::parse_str(wat).expect("Failed to parse WAT"); |
| 29 | + |
| 30 | + c.bench_function("module_compilation", |b| { |
| 31 | + b.iter(|| { |
| 32 | + let engine = wasmtime::Engine::default(); |
| 33 | + let module = wasmtime::Module::new(&engine, black_box(&wasm_bytes)) |
| 34 | + .expect("Failed to compile module"); |
| 35 | + black_box(module) |
| 36 | + }); |
| 37 | + }); |
| 38 | +} |
| 39 | + |
| 40 | +/// Benchmark cached module execution. |
| 41 | +fn bench_cached_execution(c: &mut Criterion) { |
| 42 | + let runtime = tokio::runtime::Runtime::new().unwrap(); |
| 43 | + |
| 44 | + let wat = r#" |
| 45 | + (module |
| 46 | + (func (export "main") (result i32) |
| 47 | + (i32.const 42) |
| 48 | + ) |
| 49 | + ) |
| 50 | + "#; |
| 51 | + |
| 52 | + let wasm_bytes = wat::parse_str(wat).expect("Failed to parse WAT"); |
| 53 | + |
| 54 | + let bridge = Arc::new(Bridge::new(1000)); |
| 55 | + let config = SecurityConfig::default(); |
| 56 | + let wasm_runtime = Runtime::new(bridge, config).expect("Failed to create runtime"); |
| 57 | + |
| 58 | + // Warm up cache |
| 59 | + runtime.block_on(async { |
| 60 | + let _ = wasm_runtime.execute(&wasm_bytes, "main", &Vec::new()).await; |
| 61 | + }); |
| 62 | + |
| 63 | + c.bench_function("cached_execution", |b| { |
| 64 | + b.iter(|| { |
| 65 | + runtime.block_on(async { |
| 66 | + let result = wasm_runtime |
| 67 | + .execute(black_box(&wasm_bytes), "main", &Vec::new()) |
| 68 | + .await; |
| 69 | + black_box(result) |
| 70 | + }) |
| 71 | + }); |
| 72 | + }); |
| 73 | +} |
| 74 | + |
| 75 | +/// Benchmark first-time execution (compilation + execution). |
| 76 | +fn bench_first_execution(c: &mut Criterion) { |
| 77 | + let runtime = tokio::runtime::Runtime::new().unwrap(); |
| 78 | + |
| 79 | + let wat = r#" |
| 80 | + (module |
| 81 | + (func (export "main") (result i32) |
| 82 | + (i32.const 42) |
| 83 | + ) |
| 84 | + ) |
| 85 | + "#; |
| 86 | + |
| 87 | + let wasm_bytes = wat::parse_str(wat).expect("Failed to parse WAT"); |
| 88 | + |
| 89 | + c.bench_function("first_execution", |b| { |
| 90 | + b.iter(|| { |
| 91 | + runtime.block_on(async { |
| 92 | + let bridge = Arc::new(Bridge::new(1000)); |
| 93 | + let config = SecurityConfig::default(); |
| 94 | + let wasm_runtime = Runtime::new(bridge, config).expect("Failed to create runtime"); |
| 95 | + |
| 96 | + let result = wasm_runtime |
| 97 | + .execute(black_box(&wasm_bytes), "main", &Vec::new()) |
| 98 | + .await; |
| 99 | + black_box(result) |
| 100 | + }) |
| 101 | + }); |
| 102 | + }); |
| 103 | +} |
| 104 | + |
| 105 | +/// Benchmark cache key generation. |
| 106 | +fn bench_cache_key_generation(c: &mut Criterion) { |
| 107 | + let wasm_bytes = vec![0u8; 1024]; // 1KB of data |
| 108 | + |
| 109 | + c.bench_function("cache_key_generation", |b| { |
| 110 | + b.iter(|| { |
| 111 | + let key = ModuleCache::cache_key_for_code(black_box(&wasm_bytes)); |
| 112 | + black_box(key) |
| 113 | + }); |
| 114 | + }); |
| 115 | +} |
| 116 | + |
| 117 | +/// Benchmark cache operations. |
| 118 | +fn bench_cache_operations(c: &mut Criterion) { |
| 119 | + let cache = ModuleCache::new(100); |
| 120 | + let engine = wasmtime::Engine::default(); |
| 121 | + let wat = "(module)"; |
| 122 | + let wasm = wat::parse_str(wat).expect("Failed to parse WAT"); |
| 123 | + let module = wasmtime::Module::new(&engine, &wasm).expect("Failed to create module"); |
| 124 | + |
| 125 | + c.bench_function("cache_insert", |b| { |
| 126 | + b.iter(|| { |
| 127 | + let key = ModuleCache::cache_key_for_code(b"test"); |
| 128 | + cache.insert(black_box(key), black_box(module.clone())); |
| 129 | + }); |
| 130 | + }); |
| 131 | + |
| 132 | + // Populate cache for get benchmark |
| 133 | + let test_key = ModuleCache::cache_key_for_code(b"test_get"); |
| 134 | + cache.insert(test_key.clone(), module.clone()); |
| 135 | + |
| 136 | + c.bench_function("cache_get", |b| { |
| 137 | + b.iter(|| { |
| 138 | + let result = cache.get(black_box(&test_key)); |
| 139 | + black_box(result) |
| 140 | + }); |
| 141 | + }); |
| 142 | +} |
| 143 | + |
| 144 | +/// Benchmark complex WASM execution. |
| 145 | +fn bench_complex_execution(c: &mut Criterion) { |
| 146 | + let runtime = tokio::runtime::Runtime::new().unwrap(); |
| 147 | + |
| 148 | + // More complex module with loops and function calls |
| 149 | + let wat = r#" |
| 150 | + (module |
| 151 | + (func $fibonacci (param $n i32) (result i32) |
| 152 | + (local $a i32) |
| 153 | + (local $b i32) |
| 154 | + (local $temp i32) |
| 155 | + (local $i i32) |
| 156 | +
|
| 157 | + (local.set $a (i32.const 0)) |
| 158 | + (local.set $b (i32.const 1)) |
| 159 | + (local.set $i (i32.const 0)) |
| 160 | +
|
| 161 | + (block $exit |
| 162 | + (loop $continue |
| 163 | + (br_if $exit (i32.ge_u (local.get $i) (local.get $n))) |
| 164 | +
|
| 165 | + (local.set $temp (local.get $b)) |
| 166 | + (local.set $b (i32.add (local.get $a) (local.get $b))) |
| 167 | + (local.set $a (local.get $temp)) |
| 168 | +
|
| 169 | + (local.set $i (i32.add (local.get $i) (i32.const 1))) |
| 170 | + (br $continue) |
| 171 | + ) |
| 172 | + ) |
| 173 | +
|
| 174 | + (local.get $a) |
| 175 | + ) |
| 176 | +
|
| 177 | + (func (export "main") (result i32) |
| 178 | + (i32.const 10) |
| 179 | + (call $fibonacci) |
| 180 | + ) |
| 181 | + ) |
| 182 | + "#; |
| 183 | + |
| 184 | + let wasm_bytes = wat::parse_str(wat).expect("Failed to parse WAT"); |
| 185 | + |
| 186 | + let bridge = Arc::new(Bridge::new(1000)); |
| 187 | + let config = SecurityConfig::default(); |
| 188 | + let wasm_runtime = Runtime::new(bridge, config).expect("Failed to create runtime"); |
| 189 | + |
| 190 | + // Warm up cache |
| 191 | + runtime.block_on(async { |
| 192 | + let _ = wasm_runtime.execute(&wasm_bytes, "main", &Vec::new()).await; |
| 193 | + }); |
| 194 | + |
| 195 | + c.bench_function("complex_execution", |b| { |
| 196 | + b.iter(|| { |
| 197 | + runtime.block_on(async { |
| 198 | + let result = wasm_runtime |
| 199 | + .execute(black_box(&wasm_bytes), "main", &Vec::new()) |
| 200 | + .await; |
| 201 | + black_box(result) |
| 202 | + }) |
| 203 | + }); |
| 204 | + }); |
| 205 | +} |
| 206 | + |
| 207 | +/// Benchmark host function calls (simplified - measures overhead). |
| 208 | +fn bench_host_function_overhead(c: &mut Criterion) { |
| 209 | + let runtime = tokio::runtime::Runtime::new().unwrap(); |
| 210 | + |
| 211 | + // Module that calls host function |
| 212 | + let wat = r#" |
| 213 | + (module |
| 214 | + (import "env" "host_add" (func $host_add (param i32 i32) (result i32))) |
| 215 | + (func (export "main") (result i32) |
| 216 | + (i32.const 10) |
| 217 | + (i32.const 32) |
| 218 | + (call $host_add) |
| 219 | + ) |
| 220 | + ) |
| 221 | + "#; |
| 222 | + |
| 223 | + let wasm_bytes = wat::parse_str(wat).expect("Failed to parse WAT"); |
| 224 | + |
| 225 | + let bridge = Arc::new(Bridge::new(1000)); |
| 226 | + let config = SecurityConfig::default(); |
| 227 | + let wasm_runtime = Runtime::new(bridge, config).expect("Failed to create runtime"); |
| 228 | + |
| 229 | + // Warm up cache |
| 230 | + runtime.block_on(async { |
| 231 | + let _ = wasm_runtime.execute(&wasm_bytes, "main", &Vec::new()).await; |
| 232 | + }); |
| 233 | + |
| 234 | + c.bench_function("host_function_overhead", |b| { |
| 235 | + b.iter(|| { |
| 236 | + runtime.block_on(async { |
| 237 | + let result = wasm_runtime |
| 238 | + .execute(black_box(&wasm_bytes), "main", &Vec::new()) |
| 239 | + .await; |
| 240 | + black_box(result) |
| 241 | + }) |
| 242 | + }); |
| 243 | + }); |
| 244 | +} |
| 245 | + |
| 246 | +criterion_group!( |
| 247 | + benches, |
| 248 | + bench_module_compilation, |
| 249 | + bench_cached_execution, |
| 250 | + bench_first_execution, |
| 251 | + bench_cache_key_generation, |
| 252 | + bench_cache_operations, |
| 253 | + bench_complex_execution, |
| 254 | + bench_host_function_overhead, |
| 255 | +); |
| 256 | + |
| 257 | +criterion_main!(benches); |
0 commit comments