cleanup: add cost to each host function

The default for `timeout_ms` is currently 30s which was an arbitrary
decision but it forces a user to set `timeout_ms` to null to avoid
having plugins cancelled which is a little strange.

.github/workflows/release.yml

@@ -28,48 +28,56 @@ jobs:
             target: 'x86_64-apple-darwin'
             artifact: 'libextism.dylib'
             static-artifact: 'libextism.a'
+            static-dll-artifact: ''
             pc-in: 'extism.pc.in'
             static-pc-in: 'extism-static.pc.in'
           - os: 'macos'
             target: 'aarch64-apple-darwin'
             artifact: 'libextism.dylib'
             static-artifact: 'libextism.a'
+            static-dll-artifact: ''
             pc-in: 'extism.pc.in'
             static-pc-in: 'extism-static.pc.in'
           - os: 'ubuntu'
             target: 'aarch64-unknown-linux-gnu'
             artifact: 'libextism.so'
             static-artifact: 'libextism.a'
+            static-dll-artifact: ''
             pc-in: 'extism.pc.in'
             static-pc-in: 'extism-static.pc.in'
           - os: 'ubuntu'
             target: 'aarch64-unknown-linux-musl'
             artifact: 'libextism.so'
             static-artifact: 'libextism.a'
+            static-dll-artifact: ''
             pc-in: 'extism.pc.in'
             static-pc-in: 'extism-static.pc.in'
           - os: 'ubuntu'
             target: 'x86_64-unknown-linux-gnu'
             artifact: 'libextism.so'
             static-artifact: 'libextism.a'
+            static-dll-artifact: ''
             pc-in: 'extism.pc.in'
             static-pc-in: 'extism-static.pc.in'
           - os: 'ubuntu'
             target: 'x86_64-unknown-linux-musl'
             artifact: ''
             static-artifact: 'libextism.a'
+            static-dll-artifact: ''
             pc-in: ''
             static-pc-in: 'extism-static.pc.in'
           - os: 'windows'
             target: 'x86_64-pc-windows-gnu'
             artifact: 'extism.dll'
             static-artifact: 'libextism.a'
+            static-dll-artifact: 'libextism.dll.a'
             pc-in: 'extism.pc.in'
             static-pc-in: 'extism-static.pc.in'
           - os: 'windows'
             target: 'x86_64-pc-windows-msvc'
             artifact: 'extism.dll'
             static-artifact: 'extism.lib'
+            static-dll-artifact: 'extism.dll.lib'
             pc-in: ''
             static-pc-in: ''
 
@@ -158,7 +166,8 @@ jobs:
           cp LICENSE ${SRC_DIR}
           tar -C ${SRC_DIR} -czvf ${ARCHIVE} extism.h \
             ${{ matrix.artifact }} ${{ matrix.static-artifact }}  \
-            ${{ matrix.pc-in }} ${{ matrix.static-pc-in }}
+            ${{ matrix.pc-in }} ${{ matrix.static-pc-in }} \
+            ${{ matrix.static-dll-artifact }}
           ls -ll ${ARCHIVE}
 
           if &>/dev/null which shasum; then

.gitignore

@@ -1,6 +1,7 @@
 target
 runtime/ast.json
 runtime/target
+Cargo.lock
 .DS_Store
 .vscode
 **/libextism.dylib

convert/Cargo.toml

@@ -12,6 +12,7 @@ description = "Traits to make Rust types usable with Extism"
 [dependencies]
 anyhow = "1.0.75"
 base64 = "~0.21"
+bytemuck = {version = "1.14.0", optional = true }
 prost = { version = "0.12.0", optional = true }
 rmp-serde = { version = "1.1.2", optional = true }
 serde = "1.0.186"
@@ -21,6 +22,7 @@ serde_json = "1.0.105"
 serde = { version = "1.0.186", features = ["derive"] }
 
 [features]
-default = ["msgpack", "protobuf"]
+default = ["msgpack", "protobuf", "raw"]
 msgpack = ["rmp-serde"]
 protobuf = ["prost"]
+raw = ["bytemuck"]

convert/src/encoding.rs

@@ -138,3 +138,53 @@ impl<T: Default + prost::Message> FromBytesOwned for Protobuf<T> {
         Ok(Protobuf(T::decode(data)?))
     }
 }
+
+/// Raw does no conversion, it just copies the memory directly.
+/// Note: This will only work for types that implement [bytemuck::Pod](https://docs.rs/bytemuck/latest/bytemuck/trait.Pod.html)
+#[cfg(all(feature = "raw", target_endian = "little"))]
+pub struct Raw<'a, T: bytemuck::Pod>(pub &'a T);
+
+#[cfg(all(feature = "raw", target_endian = "little"))]
+impl<'a, T: bytemuck::Pod> ToBytes<'a> for Raw<'a, T> {
+    type Bytes = &'a [u8];
+
+    fn to_bytes(&self) -> Result<Self::Bytes, Error> {
+        Ok(bytemuck::bytes_of(self.0))
+    }
+}
+
+#[cfg(all(feature = "raw", target_endian = "little"))]
+impl<'a, T: bytemuck::Pod> FromBytes<'a> for Raw<'a, T> {
+    fn from_bytes(data: &'a [u8]) -> Result<Self, Error> {
+        let x = bytemuck::try_from_bytes(data).map_err(|x| Error::msg(x.to_string()))?;
+        Ok(Raw(x))
+    }
+}
+
+#[cfg(all(test, feature = "raw", target_endian = "little"))]
+mod tests {
+    use crate::*;
+
+    #[test]
+    fn test_raw() {
+        #[derive(Debug, Clone, Copy, PartialEq)]
+        struct TestRaw {
+            a: i32,
+            b: f64,
+            c: bool,
+        }
+        unsafe impl bytemuck::Pod for TestRaw {}
+        unsafe impl bytemuck::Zeroable for TestRaw {}
+        let x = TestRaw {
+            a: 123,
+            b: 45678.91011,
+            c: true,
+        };
+        let raw = Raw(&x).to_bytes().unwrap();
+        let y = Raw::from_bytes(&raw).unwrap();
+        assert_eq!(&x, y.0);
+
+        let y: Result<Raw<[u8; std::mem::size_of::<TestRaw>()]>, Error> = Raw::from_bytes(&raw);
+        assert!(y.is_ok());
+    }
+}

convert/src/lib.rs

@@ -21,6 +21,9 @@ pub use encoding::Msgpack;
 #[cfg(feature = "protobuf")]
 pub use encoding::Protobuf;
 
+#[cfg(all(feature = "raw", target_endian = "little"))]
+pub use encoding::Raw;
+
 pub use from_bytes::{FromBytes, FromBytesOwned};
 pub use memory_handle::MemoryHandle;
 pub use to_bytes::ToBytes;

flake.lock

@@ -1,94 +0,0 @@
-{
-  "nodes": {
-    "flake-utils": {
-      "inputs": {
-        "systems": "systems"
-      },
-      "locked": {
-        "lastModified": 1701680307,
-        "narHash": "sha256-kAuep2h5ajznlPMD9rnQyffWG8EM/C73lejGofXvdM8=",
-        "owner": "numtide",
-        "repo": "flake-utils",
-        "rev": "4022d587cbbfd70fe950c1e2083a02621806a725",
-        "type": "github"
-      },
-      "original": {
-        "owner": "numtide",
-        "repo": "flake-utils",
-        "type": "github"
-      }
-    },
-    "naersk": {
-      "inputs": {
-        "nixpkgs": "nixpkgs"
-      },
-      "locked": {
-        "lastModified": 1698420672,
-        "narHash": "sha256-/TdeHMPRjjdJub7p7+w55vyABrsJlt5QkznPYy55vKA=",
-        "owner": "nix-community",
-        "repo": "naersk",
-        "rev": "aeb58d5e8faead8980a807c840232697982d47b9",
-        "type": "github"
-      },
-      "original": {
-        "owner": "nix-community",
-        "repo": "naersk",
-        "type": "github"
-      }
-    },
-    "nixpkgs": {
-      "locked": {
-        "lastModified": 1701693815,
-        "narHash": "sha256-7BkrXykVWfkn6+c1EhFA3ko4MLi3gVG0p9G96PNnKTM=",
-        "owner": "NixOS",
-        "repo": "nixpkgs",
-        "rev": "09ec6a0881e1a36c29d67497693a67a16f4da573",
-        "type": "github"
-      },
-      "original": {
-        "id": "nixpkgs",
-        "type": "indirect"
-      }
-    },
-    "nixpkgs_2": {
-      "locked": {
-        "lastModified": 1701693815,
-        "narHash": "sha256-7BkrXykVWfkn6+c1EhFA3ko4MLi3gVG0p9G96PNnKTM=",
-        "owner": "NixOS",
-        "repo": "nixpkgs",
-        "rev": "09ec6a0881e1a36c29d67497693a67a16f4da573",
-        "type": "github"
-      },
-      "original": {
-        "owner": "NixOS",
-        "ref": "nixpkgs-unstable",
-        "repo": "nixpkgs",
-        "type": "github"
-      }
-    },
-    "root": {
-      "inputs": {
-        "flake-utils": "flake-utils",
-        "naersk": "naersk",
-        "nixpkgs": "nixpkgs_2"
-      }
-    },
-    "systems": {
-      "locked": {
-        "lastModified": 1681028828,
-        "narHash": "sha256-Vy1rq5AaRuLzOxct8nz4T6wlgyUR7zLU309k9mBC768=",
-        "owner": "nix-systems",
-        "repo": "default",
-        "rev": "da67096a3b9bf56a91d16901293e51ba5b49a27e",
-        "type": "github"
-      },
-      "original": {
-        "owner": "nix-systems",
-        "repo": "default",
-        "type": "github"
-      }
-    }
-  },
-  "root": "root",
-  "version": 7
-}

flake.nix

@@ -1,32 +0,0 @@
-{  
-  inputs = {
-    flake-utils.url = "github:numtide/flake-utils";
-    naersk.url = "github:nix-community/naersk";
-    nixpkgs.url = "github:NixOS/nixpkgs/nixpkgs-unstable";
-  };
-
-  outputs = { self, flake-utils, naersk, nixpkgs }:
-    flake-utils.lib.eachDefaultSystem (system:
-      let
-        pkgs = (import nixpkgs) {
-          inherit system;
-        };
-
-        naersk' = pkgs.callPackage naersk {};
-
-      in {
-        # For `nix build` & `nix run`:
-        defaultPackage = naersk'.buildPackage {
-          name = "libextism";
-          src = ./.;
-          copyLibs = true;
-          postInstall = "mkdir -p $out/include; cp runtime/extism.h $out/include";
-        };
-
-        # For `nix develop`:
-        devShell = pkgs.mkShell {
-          nativeBuildInputs = with pkgs; [ rustc cargo clippy gnumake ];
-        };
-      }
-    );
-}

kernel/Cargo.lock

@@ -1,7 +0,0 @@
-# This file is automatically @generated by Cargo.
-# It is not intended for manual editing.
-version = 3
-
-[[package]]
-name = "extism-runtime-kernel"
-version = "0.1.0"

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,
-            self.length.load(Ordering::Acquire) as usize,
+            MemoryStatus::Unused as u8,
+            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;
@@ -301,9 +306,26 @@ impl MemoryRoot {
         if !Self::pointer_in_bounds_fast(offs) {
             return None;
         }
-        let ptr = offs - core::mem::size_of::<MemoryBlock>() as u64;
-        let ptr = ptr as *mut MemoryBlock;
-        Some(&mut *ptr)
+
+        // Get the first block
+        let mut block = self.blocks.as_mut_ptr();
+
+        // Only loop while the block pointer is less then the current position
+        while (block as u64) < self.blocks.as_ptr() as u64 + offs {
+            let b = &mut *block;
+
+            // Get the block status, this lets us know if we are able to re-use it
+            let status = b.status.load(Ordering::Acquire);
+
+            if status == MemoryStatus::Active as u8 && b.data.as_ptr() as Pointer == offs {
+                return Some(b);
+            }
+
+            // Get the next block
+            block = b.next_ptr();
+        }
+
+        None
     }
 }
 

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.

manifest/src/lib.rs

@@ -231,20 +231,15 @@ pub struct Manifest {
     pub allowed_paths: Option<BTreeMap<PathBuf, PathBuf>>,
 
     /// The plugin timeout, by default this is set to 30s
-    #[serde(default = "default_timeout")]
+    #[serde(default)]
     pub timeout_ms: Option<u64>,
 }
 
-fn default_timeout() -> Option<u64> {
-    Some(30000)
-}
-
 impl Manifest {
     /// Create a new manifest
     pub fn new(wasm: impl IntoIterator<Item = impl Into<Wasm>>) -> Manifest {
         Manifest {
             wasm: wasm.into_iter().map(|x| x.into()).collect(),
-            timeout_ms: default_timeout(),
             ..Default::default()
         }
     }

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

@@ -97,6 +97,11 @@ typedef void (*ExtismFunctionType)(ExtismCurrentPlugin *plugin,
                                    ExtismSize n_outputs,
                                    void *data);
 
+/**
+ * Log drain callback
+ */
+typedef void (*ExtismLogDrainFunctionType)(const char *data, ExtismSize size);
+
 
 
 #ifdef __cplusplus
@@ -168,6 +173,13 @@ void extism_function_free(ExtismFunction *f);
  */
 void extism_function_set_namespace(ExtismFunction *ptr, const char *namespace_);
 
+/**
+ * Set the cost of an `ExtismFunction`, when set to 0 this has no effect, when set to `1` this will add
+ * the runtime of the function back to the plugin timer.
+ */
+void extism_function_set_cost(ExtismFunction *ptr,
+                              double cost);
+
 /**
  * Create a new plugin with host functions, the functions passed to this function no longer need to be manually freed using
  *
@@ -264,7 +276,12 @@ bool extism_log_custom(const char *log_level);
  * Calls the provided callback function for each buffered log line.
  * This is only needed when `extism_log_custom` is used.
  */
-void extism_log_drain(void (*handler)(const char*, uintptr_t));
+void extism_log_drain(ExtismLogDrainFunctionType handler);
+
+/**
+ * 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> {
-        let bytes = self.memory_bytes(handle)?;
+    pub fn memory_str_mut(&mut self, handle: MemoryHandle) -> Result<&mut str, Error> {
+        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/function.rs

@@ -155,7 +155,7 @@ unsafe impl<T> Sync for UserData<T> {}
 unsafe impl Send for CPtr {}
 unsafe impl Sync for CPtr {}
 
-type FunctionInner = dyn Fn(wasmtime::Caller<CurrentPlugin>, &[wasmtime::Val], &mut [wasmtime::Val]) -> Result<(), Error>
+pub(crate) type FunctionInner = dyn Fn(wasmtime::Caller<CurrentPlugin>, &[wasmtime::Val], &mut [wasmtime::Val]) -> Result<(), Error>
     + Sync
     + Send;
 
@@ -174,6 +174,9 @@ pub struct Function {
     /// Function handle
     pub(crate) f: Arc<FunctionInner>,
 
+    /// Function cost (in terms of time)
+    pub(crate) cost: f64,
+
     /// UserData
     pub(crate) _user_data: UserDataHandle,
 }
@@ -204,10 +207,12 @@ impl Function {
             ty,
             f: Arc::new(
                 move |mut caller: Caller<_>, inp: &[Val], outp: &mut [Val]| {
+                    let plugin = caller.data_mut();
                     let x = data.clone();
-                    f(caller.data_mut(), inp, outp, x)
+                    f(plugin, inp, outp, x)
                 },
             ),
+            cost: 0.0,
             namespace: None,
             _user_data: match &user_data {
                 UserData::C(ptr) => UserDataHandle::C(ptr.clone()),
@@ -239,6 +244,23 @@ impl Function {
         self
     }
 
+    /// Function cost
+    pub fn cost(&self) -> f64 {
+        self.cost
+    }
+
+    /// Set host function cost
+    pub fn set_cost(&mut self, cost: f64) {
+        trace!("Setting cost for {} to {cost}", self.name);
+        self.cost = cost;
+    }
+
+    /// Update host function cost
+    pub fn with_cost(mut self, cost: f64) -> Self {
+        self.set_cost(cost);
+        self
+    }
+
     /// Get function type
     pub fn ty(&self) -> &wasmtime::FuncType {
         &self.ty

runtime/src/lib.rs

@@ -14,6 +14,7 @@ pub(crate) mod manifest;
 pub(crate) mod pdk;
 mod plugin;
 mod plugin_builder;
+mod timeout_manager;
 mod timer;
 
 /// Extism C API
@@ -27,6 +28,7 @@ pub use plugin::{CancelHandle, Plugin, WasmInput, EXTISM_ENV_MODULE, EXTISM_USER
 pub use plugin_builder::{DebugOptions, PluginBuilder};
 
 pub(crate) use internal::{Internal, Wasi};
+pub(crate) use timeout_manager::TimeoutManager;
 pub(crate) use timer::{Timer, TimerAction};
 pub(crate) use tracing::{debug, error, trace, warn};
 

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,
 }
@@ -238,6 +238,20 @@ impl Plugin {
             CurrentPlugin::new(manifest, with_wasi, available_pages, id)?,
         );
         store.set_epoch_deadline(1);
+        store.call_hook(|data, hook| {
+            if hook.entering_host() {
+                let tx = Timer::tx();
+                tx.send(TimerAction::EnterHost {
+                    id: data.id.clone(),
+                })?;
+            } else if hook.exiting_host() {
+                let tx = Timer::tx();
+                tx.send(TimerAction::ExitHost {
+                    id: data.id.clone(),
+                })?;
+            }
+            Ok(())
+        });
 
         let mut linker = Linker::new(&engine);
         linker.allow_shadowing(true);
@@ -284,9 +298,13 @@ impl Plugin {
         for f in &mut imports {
             let name = f.name().to_string();
             let ns = f.namespace().unwrap_or(EXTISM_USER_MODULE);
-            unsafe {
-                linker.func_new(ns, &name, f.ty().clone(), &*(f.f.as_ref() as *const _))?;
-            }
+            let cost = f.cost();
+            let inner: &function::FunctionInner = unsafe { &*(f.f.as_ref() as *const _) };
+            linker.func_new(ns, &name, f.ty().clone(), move |caller, params, results| {
+                let _timeout =
+                    TimeoutManager::new(caller.data()).map(|x| x.with_cost(cost.clone()));
+                inner(caller, params, results)
+            })?;
         }
 
         let instance_pre = linker.instantiate_pre(main)?;
@@ -303,7 +321,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 +342,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 +373,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 +446,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") {
-            f.into_func().unwrap().call(&mut self.store, &[], &mut [])?;
-        } else {
-            error!(plugin = &id, "call to extism:host/env::reset failed");
-        }
-
+        self.reset()?;
         let handle = self.current_plugin_mut().memory_new(bytes)?;
 
         if let Some(f) = self
@@ -450,6 +463,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 +641,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) {
-                error!(
-                    plugin = self.id.to_string(),
-                    "call to Plugin::reset_store failed: {e:?}"
-                );
-            }
-            self.needs_reset = false;
+        if let Err(e) = self.reset_store(lock) {
+            error!(
+                plugin = self.id.to_string(),
+                "call to Plugin::reset_store failed: {e:?}"
+            );
         }
 
         self.instantiate(lock).map_err(|e| (e, -1))?;
@@ -667,7 +690,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

@@ -38,6 +38,9 @@ pub type ExtismFunctionType = extern "C" fn(
     data: *mut std::ffi::c_void,
 );
 
+/// Log drain callback
+pub type ExtismLogDrainFunctionType = extern "C" fn(data: *const std::ffi::c_char, size: Size);
+
 impl From<&wasmtime::Val> for ExtismVal {
     fn from(value: &wasmtime::Val) -> Self {
         match value.ty() {
@@ -255,6 +258,18 @@ pub unsafe extern "C" fn extism_function_set_namespace(
     }
 }
 
+/// Set the cost of an `ExtismFunction`, when set to 0 this has no effect, when set to `1` this will add
+/// the runtime of the function back to the plugin timer.
+#[no_mangle]
+pub unsafe extern "C" fn extism_function_set_cost(ptr: *mut ExtismFunction, cost: f64) {
+    let f = &mut *ptr;
+    if let Some(x) = f.0.get_mut() {
+        x.set_cost(cost);
+    } else {
+        debug!("Trying to set the cost of already registered function")
+    }
+}
+
 /// Create a new plugin with host functions, the functions passed to this function no longer need to be manually freed using
 ///
 /// `wasm`: is a WASM module (wat or wasm) or a JSON encoded manifest
@@ -667,11 +682,11 @@ unsafe fn set_log_buffer(filter: &str) -> Result<(), Error> {
 #[no_mangle]
 /// Calls the provided callback function for each buffered log line.
 /// This is only needed when `extism_log_custom` is used.
-pub unsafe extern "C" fn extism_log_drain(handler: extern "C" fn(*const std::ffi::c_char, usize)) {
+pub unsafe extern "C" fn extism_log_drain(handler: ExtismLogDrainFunctionType) {
     if let Some(buf) = &mut LOG_BUFFER {
         if let Ok(mut buf) = buf.buffer.lock() {
             for (line, len) in buf.drain(..) {
-                handler(line.as_ptr(), len);
+                handler(line.as_ptr(), len as u64);
             }
         }
     }
@@ -702,6 +717,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/kernel.rs

@@ -198,6 +198,10 @@ fn test_kernel_allocations() {
     // 512 bytes, test block re-use + splitting
     let p = extism_alloc(&mut store, instance, 512);
     assert_eq!(extism_length(&mut store, instance, p), 512);
+    assert_eq!(extism_length(&mut store, instance, p + 1), 0);
+    assert_eq!(extism_length(&mut store, instance, p + 2), 0);
+    assert_eq!(extism_length(&mut store, instance, p + 3), 0);
+    assert_eq!(extism_length(&mut store, instance, p + 4), 0);
     extism_free(&mut store, instance, p);
 
     // 128 bytes, should be split off the 512 byte block
@@ -210,7 +214,7 @@ fn test_kernel_allocations() {
     let r = extism_alloc(&mut store, instance, 128);
     assert!(p <= r && r < p + 512);
     assert!(r > p);
-    assert_eq!(extism_length(&mut store, instance, q), 128);
+    assert_eq!(extism_length(&mut store, instance, r), 128);
     extism_free(&mut store, instance, q);
 
     // 100 pages

runtime/src/tests/runtime.rs

@@ -235,6 +235,43 @@ fn test_timeout() {
     assert!(err == "timeout");
 }
 
+fn hello_world_timeout_manager(
+    plugin: &mut CurrentPlugin,
+    inputs: &[Val],
+    outputs: &mut [Val],
+    _user_data: UserData<()>,
+) -> Result<(), Error> {
+    let mgr = plugin.timeout_manager().map(|x| x.add_on_drop());
+    std::thread::sleep(std::time::Duration::from_secs(3));
+    outputs[0] = inputs[0].clone();
+    drop(mgr);
+    Ok(())
+}
+
+#[test]
+fn test_timeout_manager() {
+    let f = Function::new(
+        "hello_world",
+        [PTR],
+        [PTR],
+        UserData::default(),
+        hello_world_timeout_manager,
+    );
+
+    let manifest = Manifest::new([extism_manifest::Wasm::data(WASM)])
+        .with_timeout(std::time::Duration::from_secs(1));
+    let mut plugin = Plugin::new(manifest, [f], true).unwrap();
+
+    let start = std::time::Instant::now();
+    let output: Result<&[u8], Error> = plugin.call("count_vowels", "testing");
+    println!("Result {:?}", output);
+    assert!(output.is_ok());
+    let end = std::time::Instant::now();
+    let time = end - start;
+    println!("Plugin ran for {:?}", time);
+    assert!(time.as_secs() >= 3);
+}
+
 typed_plugin!(pub TestTypedPluginGenerics {
     count_vowels<T: FromBytes<'a>>(&str) -> T
 });
@@ -283,7 +320,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);
     }

runtime/src/timeout_manager.rs

@@ -0,0 +1,59 @@
+use crate::*;
+
+/// `TimeoutManager` is used to control `Plugin` timeouts from within host functions
+///
+/// It can be used to add or subtract time from a plug-in's timeout. If a plugin is not
+/// configured to have a timeout then this will have no effect.
+pub(crate) struct TimeoutManager {
+    start_time: std::time::Instant,
+    id: uuid::Uuid,
+    tx: std::sync::mpsc::Sender<TimerAction>,
+    cost: f64,
+}
+
+impl TimeoutManager {
+    /// Create a new `TimeoutManager` from the `CurrentPlugin`, this will return `None` if no timeout
+    /// is configured
+    pub fn new(plugin: &CurrentPlugin) -> Option<TimeoutManager> {
+        plugin.manifest.timeout_ms.map(|_| TimeoutManager {
+            start_time: std::time::Instant::now(),
+            id: plugin.id.clone(),
+            tx: Timer::tx(),
+            cost: 0.0,
+        })
+    }
+
+    /// Add the amount of time this value has existed to the configured timeout
+    pub fn add_elapsed(&mut self) -> Result<(), Error> {
+        let cost = self.cost.abs();
+        let d = self.start_time.elapsed().mul_f64(cost);
+        let mut d: timer::ExtendTimeout = d.into();
+        if self.cost.is_sign_negative() {
+            d = -d;
+        }
+        self.tx.send(TimerAction::Extend {
+            id: self.id.clone(),
+            duration: d,
+        })?;
+        self.start_time = std::time::Instant::now();
+        Ok(())
+    }
+
+    pub fn with_cost(mut self, cost: f64) -> Self {
+        self.cost = cost;
+        self
+    }
+}
+
+impl Drop for TimeoutManager {
+    fn drop(&mut self) {
+        let x = self.add_elapsed();
+        if let Err(e) = x {
+            error!(
+                plugin = self.id.to_string(),
+                "unable to extend timeout: {}",
+                e.to_string()
+            );
+        }
+    }
+}

runtime/src/timer.rs

@@ -1,17 +1,51 @@
 use crate::*;
 
+#[derive(Debug)]
+pub(crate) struct ExtendTimeout {
+    duration: std::time::Duration,
+    negative: bool,
+}
+
+impl From<std::time::Duration> for ExtendTimeout {
+    fn from(value: std::time::Duration) -> Self {
+        ExtendTimeout {
+            duration: value,
+            negative: false,
+        }
+    }
+}
+
+impl std::ops::Neg for ExtendTimeout {
+    type Output = ExtendTimeout;
+
+    fn neg(mut self) -> Self::Output {
+        self.negative = !self.negative;
+        self
+    }
+}
+
 pub(crate) enum TimerAction {
     Start {
         id: uuid::Uuid,
         engine: Engine,
         duration: Option<std::time::Duration>,
     },
+    Extend {
+        id: uuid::Uuid,
+        duration: ExtendTimeout,
+    },
     Stop {
         id: uuid::Uuid,
     },
     Cancel {
         id: uuid::Uuid,
     },
+    EnterHost {
+        id: uuid::Uuid,
+    },
+    ExitHost {
+        id: uuid::Uuid,
+    },
     Shutdown,
 }
 
@@ -31,6 +65,8 @@ extern "C" fn cleanup_timer() {
 
 static mut TIMER: std::sync::Mutex<Option<Timer>> = std::sync::Mutex::new(None);
 
+type TimerMap = std::collections::BTreeMap<uuid::Uuid, (Engine, Option<std::time::Instant>)>;
+
 impl Timer {
     pub(crate) fn tx() -> std::sync::mpsc::Sender<TimerAction> {
         let mut timer = match unsafe { TIMER.lock() } {
@@ -49,7 +85,8 @@ impl Timer {
     pub fn init(timer: &mut Option<Timer>) -> std::sync::mpsc::Sender<TimerAction> {
         let (tx, rx) = std::sync::mpsc::channel();
         let thread = std::thread::spawn(move || {
-            let mut plugins = std::collections::BTreeMap::new();
+            let mut plugins = TimerMap::new();
+            let mut in_host = TimerMap::new();
 
             macro_rules! handle {
                 ($x:expr) => {
@@ -70,12 +107,16 @@ impl Timer {
                         TimerAction::Stop { id } => {
                             trace!(plugin = id.to_string(), "handling stop event");
                             plugins.remove(&id);
+                            in_host.remove(&id);
                         }
                         TimerAction::Cancel { id } => {
                             trace!(plugin = id.to_string(), "handling cancel event");
                             if let Some((engine, _)) = plugins.remove(&id) {
                                 engine.increment_epoch();
                             }
+                            if let Some((engine, _)) = in_host.remove(&id) {
+                                engine.increment_epoch();
+                            }
                         }
                         TimerAction::Shutdown => {
                             trace!("Shutting down timer");
@@ -83,8 +124,42 @@ impl Timer {
                                 trace!(plugin = id.to_string(), "handling shutdown event");
                                 engine.increment_epoch();
                             }
+
+                            for (id, (engine, _)) in in_host.iter() {
+                                trace!(plugin = id.to_string(), "handling shutdown event");
+                                engine.increment_epoch();
+                            }
                             return;
                         }
+                        TimerAction::Extend { id, duration } => {
+                            if let Some((_engine, Some(timeout))) = plugins.get_mut(&id) {
+                                let x = if duration.negative {
+                                    timeout.checked_sub(duration.duration)
+                                } else {
+                                    timeout.checked_add(duration.duration)
+                                };
+                                if let Some(t) = x {
+                                    *timeout = t;
+                                } else {
+                                    error!(
+                                        plugin = id.to_string(),
+                                        "unable to extend timeout by {:?}", duration.duration
+                                    );
+                                }
+                            }
+                        }
+                        TimerAction::EnterHost { id } => {
+                            trace!(plugin = id.to_string(), "enter host function");
+                            if let Some(x) = plugins.remove(&id) {
+                                in_host.insert(id, x);
+                            }
+                        }
+                        TimerAction::ExitHost { id } => {
+                            trace!(plugin = id.to_string(), "exit host function");
+                            if let Some(x) = in_host.remove(&id) {
+                                plugins.insert(id, x);
+                            }
+                        }
                     }
                 };
             }
@@ -96,6 +171,10 @@ impl Timer {
                     }
                 }
 
+                for x in rx.try_iter() {
+                    handle!(x)
+                }
+
                 plugins = plugins
                     .into_iter()
                     .filter(|(_k, (engine, end))| {
@@ -109,10 +188,6 @@ impl Timer {
                         true
                     })
                     .collect();
-
-                for x in rx.try_iter() {
-                    handle!(x)
-                }
             }
         });
         *timer = Some(Timer {