From a91172d3ff1144ebe49bc266c7542b5a350d6ebb Mon Sep 17 00:00:00 2001 From: Tiziano Santoro Date: Thu, 18 Jan 2024 11:41:37 +0000 Subject: [PATCH] Add benchmark for Wasm invocations (#4664) Only create a single linker, and reuse that to create all subsequent instances. Add benchmarks to measure it. Before: ``` test wasm::tests::bench_invoke ... bench: 246,826 ns/iter (+/- 49,215) ``` After: ``` test wasm::tests::bench_invoke ... bench: 45,246 ns/iter (+/- 1,894) ``` Ref #3757 --- oak_functions_service/src/lib.rs | 2 + oak_functions_service/src/wasm/mod.rs | 37 +++++--------- oak_functions_service/src/wasm/tests.rs | 64 +++++++++++++++++++++++-- 3 files changed, 74 insertions(+), 29 deletions(-) diff --git a/oak_functions_service/src/lib.rs b/oak_functions_service/src/lib.rs index 1e782f286a4..613f9f10a5b 100644 --- a/oak_functions_service/src/lib.rs +++ b/oak_functions_service/src/lib.rs @@ -17,6 +17,8 @@ #![cfg_attr(not(feature = "std"), no_std)] #![feature(never_type)] #![feature(unwrap_infallible)] +// Required for enabling benchmark tests. +#![feature(test)] extern crate alloc; diff --git a/oak_functions_service/src/wasm/mod.rs b/oak_functions_service/src/wasm/mod.rs index d829aed921e..15928478462 100644 --- a/oak_functions_service/src/wasm/mod.rs +++ b/oak_functions_service/src/wasm/mod.rs @@ -32,7 +32,7 @@ use log::Level; use micro_rpc::StatusCode; use oak_functions_abi::{Request, Response}; use spinning_top::Spinlock; -use wasmi::{MemoryType, Store}; +use wasmi::Store; use crate::{ logger::{OakLogger, StandaloneLogger}, @@ -119,20 +119,9 @@ impl OakLinker where L: OakLogger, { - fn new(engine: &wasmi::Engine, store: &mut Store>) -> Self { + fn new(engine: &wasmi::Engine) -> Self { let mut linker: wasmi::Linker> = wasmi::Linker::new(engine); - // Add memory to linker. - // TODO(#3783): Find a sensible value for initial pages. - let initial_pages = 100; - let memory_type = - MemoryType::new(initial_pages, None).expect("failed to define Wasm memory type"); - let memory = - wasmi::Memory::new(store, memory_type).expect("failed to initialize Wasm memory"); - linker - .define(OAK_FUNCTIONS, MEMORY_NAME, wasmi::Extern::Memory(memory)) - .expect("failed to define Wasm memory in linker"); - linker .func_wrap( OAK_FUNCTIONS, @@ -218,13 +207,11 @@ where /// Instantiates the Oak Linker and checks whether the instance exports `main`, `alloc` and a /// memory is attached. - /// - /// Use the same store used when creating the linker. fn instantiate( - self, - mut store: Store>, + &self, + mut store: &mut Store>, module: Arc, - ) -> Result<(wasmi::Instance, Store>), micro_rpc::Status> { + ) -> Result { let instance = self .linker .instantiate(&mut store, &module) @@ -245,7 +232,7 @@ where // Check that the instance exports "main". let _ = &instance - .get_typed_func::<(), ()>(&store, MAIN_FUNCTION_NAME) + .get_typed_func::<(), ()>(&mut store, MAIN_FUNCTION_NAME) .map_err(|err| { micro_rpc::Status::new_with_message( micro_rpc::StatusCode::Internal, @@ -255,7 +242,7 @@ where // Check that the instance exports "alloc". let _ = &instance - .get_typed_func::(&store, ALLOC_FUNCTION_NAME) + .get_typed_func::(&mut store, ALLOC_FUNCTION_NAME) .map_err(|err| { micro_rpc::Status::new_with_message( micro_rpc::StatusCode::Internal, @@ -272,7 +259,7 @@ where ) })?; - Ok((instance, store)) + Ok(instance) } } @@ -423,9 +410,9 @@ where } // A request handler with a Wasm module for handling multiple requests. -#[derive(Clone)] pub struct WasmHandler { wasm_module: Arc, + linker: OakLinker, wasm_api_factory: Arc + Send + Sync>, logger: L, #[cfg_attr(not(feature = "std"), allow(dead_code))] @@ -466,8 +453,11 @@ where let module = wasmi::Module::new(&engine, wasm_module_bytes) .map_err(|err| anyhow::anyhow!("couldn't load module from buffer: {:?}", err))?; + let linker = OakLinker::new(module.engine()); + Ok(WasmHandler { wasm_module: Arc::new(module), + linker, wasm_api_factory, logger, observer, @@ -489,8 +479,7 @@ where let user_state = UserState::new(wasm_api.transport(), self.logger.clone()); // For isolated requests we need to create a new store for every request. let mut store = wasmi::Store::new(module.engine(), user_state); - let linker = OakLinker::new(module.engine(), &mut store); - let (instance, mut store) = linker.instantiate(store, module)?; + let instance = self.linker.instantiate(&mut store, module)?; instance.exports(&store).for_each(|export| { store diff --git a/oak_functions_service/src/wasm/tests.rs b/oak_functions_service/src/wasm/tests.rs index 6b744eae918..4d7ea5a83e0 100644 --- a/oak_functions_service/src/wasm/tests.rs +++ b/oak_functions_service/src/wasm/tests.rs @@ -14,11 +14,16 @@ // limitations under the License. // +extern crate test; + use alloc::{sync::Arc, vec::Vec}; +use std::time::Duration; use byteorder::{ByteOrder, LittleEndian}; use hashbrown::HashMap; +use oak_functions_abi::Request; use spinning_top::Spinlock; +use test::Bencher; use super::{ api::StdWasmApiFactory, OakLinker, UserState, WasmApiFactory, WasmHandler, ALLOC_FUNCTION_NAME, @@ -125,9 +130,57 @@ fn test_read_request() { assert_eq!(request_bytes, test_state.request.clone()); } +#[test] +fn test_invoke() { + let test_state = create_test_state(); + let data = b"Hello, world!"; + let response = test_state + .wasm_handler + .handle_invoke(Request { + body: data.to_vec(), + }) + .unwrap(); + assert_eq!(response.body, data.to_vec()); +} + +#[bench] +fn bench_invoke(bencher: &mut Bencher) { + let test_state = create_test_state(); + let data = b"Hello, world!"; + + let summary = bencher.bench(|bencher| { + bencher.iter(|| { + let response = test_state + .wasm_handler + .handle_invoke(Request { + body: data.to_vec(), + }) + .unwrap(); + assert_eq!(response.body, data.to_vec()); + }); + Ok(()) + }); + + // When running `cargo test` this benchmark test gets executed too, but `summary` will be `None` + // in that case. So, here we first check that `summary` is not empty. + if let Ok(Some(summary)) = summary { + // `summary.mean` is in nanoseconds, even though it is not explicitly documented in + // https://doc.rust-lang.org/test/stats/struct.Summary.html. + let elapsed = Duration::from_nanos(summary.mean as u64); + // We expect the `mean` time for loading the test Wasm module and running its main function + // to be less than a fixed threshold. + assert!( + elapsed < Duration::from_micros(100), + "elapsed time: {:.0?}", + elapsed + ); + } +} + struct TestState { instance: wasmi::Instance, store: wasmi::Store>, + wasm_handler: WasmHandler, request: Vec, } @@ -155,16 +208,17 @@ fn create_test_state() -> TestState { let user_state = UserState::new(wasm_api.transport(), logger.clone()); - let module = wasm_handler.wasm_module; + let module = wasm_handler.wasm_module.clone(); let mut store = wasmi::Store::new(module.engine(), user_state); - let linker = OakLinker::new(module.engine(), &mut store); - let (instance, store) = linker - .instantiate(store, module) + let linker = OakLinker::new(module.engine()); + let instance = linker + .instantiate(&mut store, module) .expect("couldn't instantiate Wasm module"); TestState { - store, instance, + store, + wasm_handler, request: request.clone(), } }