fix(kernel): improve performance after large allocations, add extism_plugin_reset to give users more control when dealing with large allocations (#627)

See https://github.com/extism/cpp-sdk/issues/15

- Limits a call to memset in the kernel to the size of the current
memory offset instead of the total size of memory.
- Adds `extism_plugin_reset` to the C API and `extism::Plugin::reset` to
Rust

---------

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

kernel/src/lib.rs

@@ -139,7 +139,9 @@ impl MemoryRoot {
         }
 
         // Ensure that at least one page is allocated to store the `MemoryRoot` data
-        if core::arch::wasm32::memory_size(0) == 0 && core::arch::wasm32::memory_grow(0, 1) == usize::MAX {
+        if core::arch::wasm32::memory_size(0) == 0
+            && core::arch::wasm32::memory_grow(0, 1) == usize::MAX
+        {
             core::arch::wasm32::unreachable()
         }
 
@@ -168,12 +170,15 @@ impl MemoryRoot {
 
     /// Resets the position of the allocator and zeroes out all allocations
     pub unsafe fn reset(&mut self) {
+        // Clear allocated data
+        let self_position = self.position.fetch_and(0, Ordering::SeqCst);
         core::ptr::write_bytes(
             self.blocks.as_mut_ptr() as *mut u8,
-            0,
+            MemoryStatus::Unused as u8,
-            self.length.load(Ordering::Acquire) as usize,
+            self_position as usize,
         );
-        self.position.store(0, Ordering::Release);
+
+        // Clear extism runtime metadata
         self.error.store(0, Ordering::Release);
         self.input_offset = 0;
         self.input_length = 0;

libextism/API.md

@@ -90,6 +90,14 @@ Get the plugin's output data.
 const uint8_t *extism_plugin_output_data(ExtismPlugin *plugin);
 ```
 
+### `extism_plugin_reset`
+
+Reset the Extism runtime, this will invalidate all allocated memory.
+
+```c
+bool extism_plugin_reset(ExtismPlugin *plugin);
+```
+
 ### `extism_log_file`
 
 Set log file and level.

runtime/benches/bench.rs

@@ -118,9 +118,7 @@ pub fn echo(c: &mut Criterion) {
 
 pub fn reflect(c: &mut Criterion) {
     let mut g = c.benchmark_group("reflect");
-    g.sample_size(500);
+
-    g.noise_threshold(1.0);
-    g.significance_level(0.2);
     let mut plugin = PluginBuilder::new(REFLECT)
         .with_wasi(true)
         .with_function(
@@ -132,22 +130,23 @@ pub fn reflect(c: &mut Criterion) {
         )
         .build()
         .unwrap();
-
     for (i, elements) in [
         b"a".repeat(65536),
         b"a".repeat(65536 * 10),
         b"a".repeat(65536 * 100),
+        b"a".repeat(65536),
     ]
     .iter()
     .enumerate()
     {
         g.throughput(criterion::Throughput::Bytes(elements.len() as u64));
         g.bench_with_input(
-            format!("reflect {} bytes", 10u32.pow(i as u32) * 65536),
+            format!("{i}: reflect {} bytes", elements.len()),
             elements,
             |b, elems| {
                 b.iter(|| {
                     assert_eq!(elems, plugin.call::<_, &[u8]>("reflect", &elems).unwrap());
+                    // plugin.reset().unwrap();
                 });
             },
         );

runtime/extism.h

@@ -266,6 +266,11 @@ bool extism_log_custom(const char *log_level);
  */
 void extism_log_drain(void (*handler)(const char*, uintptr_t));
 
+/**
+ * Reset the Extism runtime, this will invalidate all allocated memory
+ */
+bool extism_plugin_reset(ExtismPlugin *plugin);
+
 /**
  * Get the Extism version string
  */

runtime/src/current_plugin.rs

@@ -87,18 +87,27 @@ impl CurrentPlugin {
     }
 
     /// Access memory bytes as `str`
-    pub fn memory_str(&mut self, handle: MemoryHandle) -> Result<&mut str, Error> {
+    pub fn memory_str_mut(&mut self, handle: MemoryHandle) -> Result<&mut str, Error> {
-        let bytes = self.memory_bytes(handle)?;
+        let bytes = self.memory_bytes_mut(handle)?;
         let s = std::str::from_utf8_mut(bytes)?;
         Ok(s)
     }
 
+    pub fn memory_str(&mut self, handle: MemoryHandle) -> Result<&str, Error> {
+        let bytes = self.memory_bytes(handle)?;
+        let s = std::str::from_utf8(bytes)?;
+        Ok(s)
+    }
+
     /// Allocate a handle large enough for the encoded Rust type and copy it into Extism memory
     pub fn memory_new<'a, T: ToBytes<'a>>(&mut self, t: T) -> Result<MemoryHandle, Error> {
         let data = t.to_bytes()?;
         let data = data.as_ref();
+        if data.is_empty() {
+            return Ok(MemoryHandle::null());
+        }
         let handle = self.memory_alloc(data.len() as u64)?;
-        let bytes = self.memory_bytes(handle)?;
+        let bytes = self.memory_bytes_mut(handle)?;
         bytes.copy_from_slice(data.as_ref());
         Ok(handle)
     }
@@ -144,7 +153,7 @@ impl CurrentPlugin {
         Ok(())
     }
 
-    pub fn memory_bytes(&mut self, handle: MemoryHandle) -> Result<&mut [u8], Error> {
+    pub fn memory_bytes_mut(&mut self, handle: MemoryHandle) -> Result<&mut [u8], Error> {
         let (linker, mut store) = self.linker_and_store();
         if let Some(mem) = linker.get(&mut store, EXTISM_ENV_MODULE, "memory") {
             let mem = mem.into_memory().unwrap();
@@ -158,6 +167,20 @@ impl CurrentPlugin {
         anyhow::bail!("{} unable to locate extism memory", self.id)
     }
 
+    pub fn memory_bytes(&mut self, handle: MemoryHandle) -> Result<&[u8], Error> {
+        let (linker, mut store) = self.linker_and_store();
+        if let Some(mem) = linker.get(&mut store, EXTISM_ENV_MODULE, "memory") {
+            let mem = mem.into_memory().unwrap();
+            let ptr = unsafe { mem.data_ptr(&store).add(handle.offset() as usize) };
+            if ptr.is_null() {
+                return Ok(&[]);
+            }
+            return Ok(unsafe { std::slice::from_raw_parts(ptr, handle.len()) });
+        }
+
+        anyhow::bail!("{} unable to locate extism memory", self.id)
+    }
+
     pub fn memory_alloc(&mut self, n: u64) -> Result<MemoryHandle, Error> {
         if n == 0 {
             return Ok(MemoryHandle {

runtime/src/plugin.rs

@@ -74,7 +74,7 @@ pub struct Plugin {
 
     /// Set to `true` when de-initializarion may have occured (i.e.a call to `_start`),
     /// in this case we need to re-initialize the entire module.
-    pub(crate) needs_reset: bool,
+    pub(crate) store_needs_reset: bool,
 
     pub(crate) debug_options: DebugOptions,
 }
@@ -303,7 +303,7 @@ impl Plugin {
             cancel_handle: CancelHandle { id, timer_tx },
             instantiations: 0,
             output: Output::default(),
-            needs_reset: false,
+            store_needs_reset: false,
             debug_options,
             _functions: imports,
         };
@@ -324,7 +324,7 @@ impl Plugin {
         &mut self,
         instance_lock: &mut std::sync::MutexGuard<Option<Instance>>,
     ) -> Result<(), Error> {
-        if self.instantiations > 100 {
+        if self.store_needs_reset {
             let engine = self.store.engine().clone();
             let internal = self.current_plugin_mut();
             self.store = Store::new(
@@ -355,9 +355,9 @@ impl Plugin {
             }
             self.instantiations = 0;
             self.instance_pre = self.linker.instantiate_pre(main)?;
+            **instance_lock = None;
+            self.store_needs_reset = false;
         }
-
-        **instance_lock = None;
         Ok(())
     }
 
@@ -428,12 +428,7 @@ impl Plugin {
         let bytes = unsafe { std::slice::from_raw_parts(input, len) };
         debug!(plugin = &id, "input size: {}", bytes.len());
 
-        if let Some(f) = self.linker.get(&mut self.store, EXTISM_ENV_MODULE, "reset") {
+        self.reset()?;
-            f.into_func().unwrap().call(&mut self.store, &[], &mut [])?;
-        } else {
-            error!(plugin = &id, "call to extism:host/env::reset failed");
-        }
-
         let handle = self.current_plugin_mut().memory_new(bytes)?;
 
         if let Some(f) = self
@@ -450,6 +445,19 @@ impl Plugin {
         Ok(())
     }
 
+    /// Reset Extism runtime, this will invalidate all allocated memory
+    pub fn reset(&mut self) -> Result<(), Error> {
+        let id = self.id.to_string();
+
+        if let Some(f) = self.linker.get(&mut self.store, EXTISM_ENV_MODULE, "reset") {
+            f.into_func().unwrap().call(&mut self.store, &[], &mut [])?;
+        } else {
+            error!(plugin = &id, "call to extism:host/env::reset failed");
+        }
+
+        Ok(())
+    }
+
     /// Determine if wasi is enabled
     pub fn has_wasi(&self) -> bool {
         self.current_plugin().wasi.is_some()
@@ -615,14 +623,11 @@ impl Plugin {
         let name = name.as_ref();
         let input = input.as_ref();
 
-        if self.needs_reset {
+        if let Err(e) = self.reset_store(lock) {
-            if let Err(e) = self.reset_store(lock) {
+            error!(
-                error!(
+                plugin = self.id.to_string(),
-                    plugin = self.id.to_string(),
+                "call to Plugin::reset_store failed: {e:?}"
-                    "call to Plugin::reset_store failed: {e:?}"
+            );
-                );
-            }
-            self.needs_reset = false;
         }
 
         self.instantiate(lock).map_err(|e| (e, -1))?;
@@ -667,7 +672,7 @@ impl Plugin {
         self.store
             .epoch_deadline_callback(|_| Ok(UpdateDeadline::Continue(1)));
         let _ = self.timer_tx.send(TimerAction::Stop { id: self.id });
-        self.needs_reset = name == "_start";
+        self.store_needs_reset = name == "_start";
 
         // Get extism error
         self.get_output_after_call().map_err(|x| (x, -1))?;

runtime/src/sdk.rs

@@ -702,6 +702,30 @@ impl std::io::Write for LogBuffer {
     }
 }
 
+/// Reset the Extism runtime, this will invalidate all allocated memory
+#[no_mangle]
+pub unsafe extern "C" fn extism_plugin_reset(plugin: *mut Plugin) -> bool {
+    let plugin = &mut *plugin;
+
+    if let Err(e) = plugin.reset() {
+        error!(
+            plugin = plugin.id.to_string(),
+            "unable to reset plugin: {}",
+            e.to_string()
+        );
+        if let Err(e) = plugin.current_plugin_mut().set_error(e.to_string()) {
+            error!(
+                plugin = plugin.id.to_string(),
+                "unable to set error after failed plugin reset: {}",
+                e.to_string()
+            );
+        }
+        false
+    } else {
+        true
+    }
+}
+
 /// Get the Extism version string
 #[no_mangle]
 pub unsafe extern "C" fn extism_version() -> *const c_char {

runtime/src/tests/runtime.rs

@@ -283,7 +283,7 @@ fn test_multiple_instantiations() {
 #[test]
 fn test_globals() {
     let mut plugin = Plugin::new(WASM_GLOBALS, [], true).unwrap();
-    for i in 0..1000 {
+    for i in 0..100000 {
         let Json(count) = plugin.call::<_, Json<Count>>("globals", "").unwrap();
         assert_eq!(count.count, i);
     }