From 1c2e51043103532743b9156e2c272699f792bf64 Mon Sep 17 00:00:00 2001 From: Tyler Rockwood Date: Mon, 16 Oct 2023 12:39:24 -0500 Subject: [PATCH] Support reset_fuel in store APIs (#7240) * Support set_fuel in store APIs Fixes: https://github.com/bytecodealliance/wasmtime/issues/5109 Signed-off-by: Tyler Rockwood * rename set_fuel to reset_fuel To make it more clear that consumed fuel is being reset. Signed-off-by: Tyler Rockwood * update out of date documentation for fuel in C API Signed-off-by: Tyler Rockwood --------- Signed-off-by: Tyler Rockwood --- crates/c-api/include/wasmtime/store.h | 25 ++++++++++++++-- crates/c-api/src/store.rs | 8 +++++ crates/wasmtime/src/func.rs | 7 +++++ crates/wasmtime/src/store.rs | 42 ++++++++++++++++++++++++++- tests/all/fuel.rs | 12 +++++++- 5 files changed, 90 insertions(+), 4 deletions(-) diff --git a/crates/c-api/include/wasmtime/store.h b/crates/c-api/include/wasmtime/store.h index c86ac57234e6..ac67d1a1dd82 100644 --- a/crates/c-api/include/wasmtime/store.h +++ b/crates/c-api/include/wasmtime/store.h @@ -156,14 +156,35 @@ WASM_API_EXTERN void wasmtime_context_gc(wasmtime_context_t* context); * This function must be called for the store to have * some fuel to allow WebAssembly to execute. * - * Note that at this time when fuel is entirely consumed it will cause - * wasm to trap. More usages of fuel are planned for the future. + * Note by default when fuel is entirely consumed it will cause + * wasm to trap. If async_support is enabled, you can use + * #wasmtime_context_out_of_fuel_async_yield and the async APIs to yield + * execution instead. * * If fuel is not enabled within this store then an error is returned. If fuel * is successfully added then NULL is returned. */ WASM_API_EXTERN wasmtime_error_t *wasmtime_context_add_fuel(wasmtime_context_t *store, uint64_t fuel); +/** + * \brief Set fuel to this context's store for wasm to consume while executing. + * + * For this method to work fuel consumption must be enabled via + * #wasmtime_config_consume_fuel_set. By default a store starts with 0 fuel + * for wasm to execute with (meaning it will immediately trap). + * This function must be called for the store to have + * some fuel to allow WebAssembly to execute. + * + * Note by default when fuel is entirely consumed it will cause + * wasm to trap. If async_support is enabled, you can use + * #wasmtime_context_out_of_fuel_async_yield and the async APIs to yield + * execution instead. + * + * If fuel is not enabled within this store then an error is returned. If fuel + * is successfully added then NULL is returned. + */ +WASM_API_EXTERN wasmtime_error_t *wasmtime_context_reset_fuel(wasmtime_context_t *store, uint64_t fuel); + /** * \brief Returns the amount of fuel consumed by this context's store execution * so far. diff --git a/crates/c-api/src/store.rs b/crates/c-api/src/store.rs index d56977619888..d6c95e00cbd3 100644 --- a/crates/c-api/src/store.rs +++ b/crates/c-api/src/store.rs @@ -216,6 +216,14 @@ pub extern "C" fn wasmtime_context_add_fuel( crate::handle_result(store.add_fuel(fuel), |()| {}) } +#[no_mangle] +pub extern "C" fn wasmtime_context_reset_fuel( + mut store: CStoreContextMut<'_>, + fuel: u64, +) -> Option> { + crate::handle_result(store.reset_fuel(fuel), |()| {}) +} + #[no_mangle] pub extern "C" fn wasmtime_context_fuel_consumed(store: CStoreContext<'_>, fuel: &mut u64) -> bool { match store.fuel_consumed() { diff --git a/crates/wasmtime/src/func.rs b/crates/wasmtime/src/func.rs index 8e3a39d004bf..a551b64daef6 100644 --- a/crates/wasmtime/src/func.rs +++ b/crates/wasmtime/src/func.rs @@ -1902,6 +1902,13 @@ impl Caller<'_, T> { self.store.fuel_consumed() } + /// Returns the remaining fuel in the store. + /// + /// For more information see [`Store::fuel_remaining`](crate::Store::fuel_remaining) + pub fn fuel_remaining(&self) -> Option { + self.store.fuel_remaining() + } + /// Inject more fuel into this store to be consumed when executing wasm code. /// /// For more information see [`Store::add_fuel`](crate::Store::add_fuel) diff --git a/crates/wasmtime/src/store.rs b/crates/wasmtime/src/store.rs index 48436c04f639..e4745d3adefa 100644 --- a/crates/wasmtime/src/store.rs +++ b/crates/wasmtime/src/store.rs @@ -792,7 +792,11 @@ impl Store { /// If fuel consumption is not enabled via /// [`Config::consume_fuel`](crate::Config::consume_fuel) then this /// function will return `None`. Also note that fuel, if enabled, must be - /// originally configured via [`Store::add_fuel`]. + /// originally configured via [`Store::add_fuel`] or [`Store::reset_fuel`]. + /// + /// Note that this function returns the amount of fuel consumed since the + /// last time [`reset_fuel`][Store::reset_fuel] was called, or since the creation + /// of the store if it's never been called. pub fn fuel_consumed(&self) -> Option { self.inner.fuel_consumed() } @@ -830,6 +834,20 @@ impl Store { pub fn add_fuel(&mut self, fuel: u64) -> Result<()> { self.inner.add_fuel(fuel) } + /// + /// Set the remaining fuel to this [`Store`] for wasm to consume while executing. + /// + /// See [`Store::add_fuel`] for more information about fuel. + /// + /// This method will also reset the amount of fuel that has been consumed. + /// + /// # Errors + /// + /// This function will return an error if fuel consumption is not enabled via + /// [`Config::consume_fuel`](crate::Config::consume_fuel). + pub fn reset_fuel(&mut self, fuel: u64) -> Result<()> { + self.inner.reset_fuel(fuel) + } /// Synthetically consumes fuel from this [`Store`]. /// @@ -1118,6 +1136,13 @@ impl<'a, T> StoreContextMut<'a, T> { self.0.consume_fuel(fuel) } + /// Set the fuel for this store to be consumed when executing wasm code. + /// + /// For more information see [`Store::reset_fuel`] + pub fn reset_fuel(&mut self, fuel: u64) -> Result<()> { + self.0.reset_fuel(fuel) + } + /// Configures this `Store` to trap whenever fuel runs out. /// /// For more information see [`Store::out_of_fuel_trap`] @@ -1551,6 +1576,10 @@ impl StoreOpaque { } fn consume_fuel(&mut self, fuel: u64) -> Result { + anyhow::ensure!( + self.engine().config().tunables.consume_fuel, + "fuel is not configured in this store" + ); let consumed_ptr = unsafe { &mut *self.runtime_limits.fuel_consumed.get() }; match i64::try_from(fuel) .ok() @@ -1564,6 +1593,17 @@ impl StoreOpaque { } } + fn reset_fuel(&mut self, fuel: u64) -> Result<()> { + anyhow::ensure!( + self.engine().config().tunables.consume_fuel, + "fuel is not configured in this store" + ); + let fuel = i64::try_from(fuel).unwrap_or(i64::max_value()); + self.fuel_adj = fuel; + unsafe { *self.runtime_limits.fuel_consumed.get() = -fuel }; + Ok(()) + } + #[inline] pub fn signal_handler(&self) -> Option<*const SignalHandler<'static>> { let handler = self.signal_handler.as_ref()?; diff --git a/tests/all/fuel.rs b/tests/all/fuel.rs index c3f300d09a33..df046763c95e 100644 --- a/tests/all/fuel.rs +++ b/tests/all/fuel.rs @@ -143,6 +143,14 @@ fn manual_fuel() { assert_eq!(store.consume_fuel(1).unwrap(), 0); assert_eq!(store.consume_fuel(0).unwrap(), 0); assert_eq!(store.fuel_remaining(), Some(0)); + store.add_fuel(5_000).unwrap(); + assert_eq!(store.fuel_consumed(), Some(10_000)); + assert_eq!(store.consume_fuel(2_500).unwrap(), 2_500); + assert_eq!(store.fuel_consumed(), Some(12_500)); + store.reset_fuel(5_000).unwrap(); + assert_eq!(store.fuel_consumed(), Some(0)); + assert_eq!(store.consume_fuel(2_500).unwrap(), 2_500); + assert_eq!(store.fuel_consumed(), Some(2_500)); } #[test] @@ -168,7 +176,9 @@ fn host_function_consumes_all() { store.add_fuel(FUEL).unwrap(); let func = Func::wrap(&mut store, |mut caller: Caller<'_, ()>| { let consumed = caller.fuel_consumed().unwrap(); - assert_eq!(caller.consume_fuel((FUEL - consumed) - 1).unwrap(), 1); + let remaining = caller.fuel_remaining().unwrap(); + assert_eq!(remaining, FUEL - consumed); + assert_eq!(caller.consume_fuel(remaining - 1).unwrap(), 1); }); let instance = Instance::new(&mut store, &module, &[func.into()]).unwrap();