feat: add benchmarking, optimize bounds checking in the kernel (#505)

- Adds benchmarking, run with `cargo bench` or `cargo criterion`
- Adds `MemoryRoot::pointer_in_bounds_fast` to do less precise bounds
checking in loads/stores

---------

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
Co-authored-by: zshipko <zshipko@users.noreply.github.com>

.github/workflows/ci-rust.yml

@@ -91,4 +91,29 @@ jobs:
         run: cargo test --all-features --release -p ${{ env.RUNTIME_CRATE }}
       - name: Test no features
         run: cargo test --no-default-features --release -p ${{ env.RUNTIME_CRATE }}
-
+  bench:
+    name: Benchmarking
+    runs-on: ${{ matrix.os }}
+    strategy:
+      matrix:
+        os: [ubuntu-latest]
+        rust:
+          - stable
+    steps:
+      - name: Checkout sources
+        uses: actions/checkout@v3
+      - name: Install Rust
+        uses: actions-rs/toolchain@v1
+        with:
+          toolchain: stable
+          override: true
+      - name: Cache Rust environment
+        uses: Swatinem/rust-cache@v1
+      - name: Cache target
+        id: cache-target
+        uses: actions/cache@v3
+        with:
+          path: target/**
+          key: ${{ runner.os }}-target-${{ github.sha }}
+      - run: cargo install cargo-criterion
+      - run: cargo criterion

Makefile

@@ -21,6 +21,9 @@ endif
 build:
 	cargo build --release $(FEATURE_FLAGS) --manifest-path libextism/Cargo.toml
 
+bench:
+	@(cargo criterion || echo 'For nicer output use cargo-criterion: `cargo install cargo-criterion` - using `cargo bench`') && cargo bench
+
 .PHONY: kernel
 kernel:
 	cd kernel && bash build.sh

kernel/build.sh

@@ -1,6 +1,21 @@
 #!/usr/bin/env bash
 
-cargo build --release --target wasm32-unknown-unknown --package extism-runtime-kernel --bin extism-runtime
+export CARGO_FLAGS=""
+
+while getopts d flag
+do
+    case "${flag}" in
+        d)
+            echo "Disabled bounds-checking";
+            export CARGO_FLAGS="--no-default-features";;
+        *)
+            echo "usage $0 [-d]"
+            echo "\t-d: build with bounds checking disabled"
+            exit 1
+    esac
+done
+
+cargo build --package extism-runtime-kernel --bin extism-runtime --release --target wasm32-unknown-unknown $CARGO_FLAGS
 cp target/wasm32-unknown-unknown/release/extism-runtime.wasm .
 wasm-strip extism-runtime.wasm
 mv extism-runtime.wasm ../runtime/src/extism-runtime.wasm

kernel/src/lib.rs

@@ -184,11 +184,21 @@ impl MemoryRoot {
     }
 
     #[inline(always)]
+    #[allow(unused)]
     fn pointer_in_bounds(&self, p: Pointer) -> bool {
         let start_ptr = self.blocks.as_ptr() as Pointer;
         p >= start_ptr && p < start_ptr + self.length.load(Ordering::Acquire) as Pointer
     }
 
+    #[inline(always)]
+    #[allow(unused)]
+    fn pointer_in_bounds_fast(p: Pointer) -> bool {
+        // Similar to `pointer_in_bounds` but less accurate on the upper bound. This uses the total memory size,
+        // instead of checking `MemoryRoot::length`
+        let end = core::arch::wasm32::memory_size(0) << 16;
+        p >= core::mem::size_of::<Self>() as Pointer && p <= end as Pointer
+    }
+
     // Find a block that is free to use, this can be a new block or an existing freed block. The `self_position` argument
     // is used to avoid loading the allocators position more than once when performing an allocation.
     unsafe fn find_free_block(
@@ -290,10 +300,7 @@ impl MemoryRoot {
 
     /// Finds the block at an offset in memory
     pub unsafe fn find_block(&mut self, offs: Pointer) -> Option<&mut MemoryBlock> {
-        let blocks_start = self.blocks.as_ptr() as Pointer;
-        if offs < blocks_start
-            || offs >= blocks_start + self.length.load(Ordering::Acquire) as Pointer
-        {
+        if !Self::pointer_in_bounds_fast(offs) {
             return None;
         }
         let ptr = offs - core::mem::size_of::<MemoryBlock>() as u64;
@@ -372,7 +379,7 @@ pub unsafe fn extism_length(p: Pointer) -> Length {
 #[no_mangle]
 pub unsafe fn extism_load_u8(p: Pointer) -> u8 {
     #[cfg(feature = "bounds-checking")]
-    if !MemoryRoot::new().pointer_in_bounds(p) {
+    if !MemoryRoot::pointer_in_bounds_fast(p) {
         return 0;
     }
     *(p as *mut u8)
@@ -382,7 +389,7 @@ pub unsafe fn extism_load_u8(p: Pointer) -> u8 {
 #[no_mangle]
 pub unsafe fn extism_load_u64(p: Pointer) -> u64 {
     #[cfg(feature = "bounds-checking")]
-    if !MemoryRoot::new().pointer_in_bounds(p + core::mem::size_of::<u64>() as Pointer - 1) {
+    if !MemoryRoot::pointer_in_bounds_fast(p + core::mem::size_of::<u64>() as u64 - 1) {
         return 0;
     }
     *(p as *mut u64)
@@ -412,7 +419,7 @@ pub unsafe fn extism_input_load_u64(p: Pointer) -> u64 {
 #[no_mangle]
 pub unsafe fn extism_store_u8(p: Pointer, x: u8) {
     #[cfg(feature = "bounds-checking")]
-    if !MemoryRoot::new().pointer_in_bounds(p) {
+    if !MemoryRoot::pointer_in_bounds_fast(p) {
         return;
     }
     *(p as *mut u8) = x;
@@ -422,7 +429,7 @@ pub unsafe fn extism_store_u8(p: Pointer, x: u8) {
 #[no_mangle]
 pub unsafe fn extism_store_u64(p: Pointer, x: u64) {
     #[cfg(feature = "bounds-checking")]
-    if !MemoryRoot::new().pointer_in_bounds(p) {
+    if !MemoryRoot::pointer_in_bounds_fast(p + core::mem::size_of::<u64>() as u64 - 1) {
         return;
     }
     *(p as *mut u64) = x;

runtime/Cargo.toml

@@ -34,3 +34,10 @@ http = ["ureq"]          # enables extism_http_request
 
 [build-dependencies]
 cbindgen = "0.26"
+
+[dev-dependencies]
+criterion = "0.5.1"
+
+[[bench]]
+name = "bench"
+harness = false

runtime/benches/bench.rs

@@ -0,0 +1,136 @@
+use criterion::{criterion_group, criterion_main, Criterion};
+use extism::*;
+
+const COUNT_VOWELS: &[u8] = include_bytes!("../../wasm/code.wasm");
+const REFLECT: &[u8] = include_bytes!("../../wasm/reflect.wasm");
+
+host_fn!(hello_world (a: String) -> String { a });
+
+pub fn basic(c: &mut Criterion) {
+    let mut g = c.benchmark_group("basic");
+    g.sample_size(300);
+    g.measurement_time(std::time::Duration::from_secs(6));
+    g.bench_function("basic", |b| {
+        let data = "a".repeat(4096);
+        b.iter(|| {
+            let mut plugin = Plugin::new(COUNT_VOWELS, [], true).unwrap();
+            let _: serde_json::Value = plugin.call("count_vowels", &data).unwrap();
+        })
+    });
+}
+
+pub fn create_plugin(c: &mut Criterion) {
+    let mut g = c.benchmark_group("create");
+    g.noise_threshold(1.0);
+    g.significance_level(0.2);
+    g.sample_size(300);
+    g.bench_function("create_plugin", |b| {
+        b.iter(|| {
+            let _plugin = PluginBuilder::new_with_module(COUNT_VOWELS)
+                .with_wasi(true)
+                .build()
+                .unwrap();
+        })
+    });
+}
+
+pub fn count_vowels(c: &mut Criterion) {
+    let mut g = c.benchmark_group("count_vowels");
+    g.sample_size(500);
+    let mut plugin = PluginBuilder::new_with_module(COUNT_VOWELS)
+        .with_wasi(true)
+        .build()
+        .unwrap();
+    let data = "a".repeat(4096);
+    g.bench_function("count_vowels(4096)", |b| {
+        b.iter(|| {
+            assert_eq!(
+                "{\"count\": 4096}",
+                plugin.call::<_, &str>("count_vowels", &data).unwrap()
+            );
+        })
+    });
+}
+
+pub fn reflect_1(c: &mut Criterion) {
+    let mut g = c.benchmark_group("reflect_1");
+    g.sample_size(500);
+    g.noise_threshold(1.0);
+    g.significance_level(0.2);
+    let mut plugin = PluginBuilder::new_with_module(REFLECT)
+        .with_wasi(true)
+        .with_function(
+            "host_reflect",
+            [ValType::I64],
+            [ValType::I64],
+            None,
+            hello_world,
+        )
+        .build()
+        .unwrap();
+    let data = "a".repeat(65536);
+    g.bench_function("reflect_1", |b| {
+        b.iter(|| {
+            assert_eq!(data, plugin.call::<_, &str>("reflect", &data).unwrap());
+        })
+    });
+}
+
+pub fn reflect_10(c: &mut Criterion) {
+    let mut g = c.benchmark_group("reflect_10");
+    g.sample_size(200);
+    g.noise_threshold(1.0);
+    g.significance_level(0.2);
+    let mut plugin = PluginBuilder::new_with_module(REFLECT)
+        .with_wasi(true)
+        .with_function(
+            "host_reflect",
+            [ValType::I64],
+            [ValType::I64],
+            None,
+            hello_world,
+        )
+        .build()
+        .unwrap();
+    let data = "a".repeat(65536 * 10);
+    g.bench_function("reflect_10", |b| {
+        b.iter(|| {
+            assert_eq!(data, plugin.call::<_, &str>("reflect", &data).unwrap());
+        })
+    });
+}
+
+pub fn reflect_100(c: &mut Criterion) {
+    let mut g = c.benchmark_group("reflect_100");
+    g.sample_size(50);
+    g.noise_threshold(1.0);
+    g.significance_level(0.2);
+    let mut plugin = PluginBuilder::new_with_module(REFLECT)
+        .with_wasi(true)
+        .with_function(
+            "host_reflect",
+            [ValType::I64],
+            [ValType::I64],
+            None,
+            hello_world,
+        )
+        .build()
+        .unwrap();
+    let data = "a".repeat(65536 * 100);
+    g.bench_function("reflect_100", |b| {
+        b.iter(|| {
+            assert_eq!(data, plugin.call::<_, &str>("reflect", &data).unwrap());
+        })
+    });
+}
+
+criterion_group!(
+    benches,
+    basic,
+    create_plugin,
+    count_vowels,
+    reflect_1,
+    reflect_10,
+    reflect_100
+);
+criterion_main!(benches);

runtime/src/tests/runtime.rs

@@ -300,7 +300,7 @@ fn test_fuzz_reflect_plugin() {
 
     for i in 1..65540 {
         let input = "a".repeat(i);
-        let output = plugin.call("reflect", input.clone());
+        let output = plugin.call("reflect", &input);
         let output = std::str::from_utf8(output.unwrap()).unwrap();
         assert_eq!(output, input);
     }