Support set_fuel in store APIs

Fixes: bytecodealliance#5109

Signed-off-by: Tyler Rockwood <rockwood@redpanda.com>

crates/c-api/include/wasmtime/store.h

@@ -164,6 +164,23 @@ WASM_API_EXTERN void wasmtime_context_gc(wasmtime_context_t* context);
  */
 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 that at this time when fuel is entirely consumed it will cause
+ * wasm to trap. More usages of fuel are planned for the future.
+ *
+ * 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_set_fuel(wasmtime_context_t *store, uint64_t fuel);
+
 /**
  * \brief Returns the amount of fuel consumed by this context's store execution
  * so far.

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_set_fuel(
+    mut store: CStoreContextMut<'_>,
+    fuel: u64,
+) -> Option<Box<wasmtime_error_t>> {
+    crate::handle_result(store.set_fuel(fuel), |()| {})
+}
+
 #[no_mangle]
 pub extern "C" fn wasmtime_context_fuel_consumed(store: CStoreContext<'_>, fuel: &mut u64) -> bool {
     match store.fuel_consumed() {

crates/wasmtime/src/func.rs

@@ -1902,6 +1902,13 @@ impl<T> 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<u64> {
+        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)

crates/wasmtime/src/store.rs

@@ -792,7 +792,11 @@ impl<T> Store<T> {
     /// 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::set_fuel`].
+    ///
+    /// Note that this function returns the amount of fuel consumed since the
+    /// last time [`set_fuel`][Store::set_fuel] was called, or since the creation
+    /// of the store if it's never been called.
     pub fn fuel_consumed(&self) -> Option<u64> {
         self.inner.fuel_consumed()
     }
@@ -830,6 +834,20 @@ impl<T> Store<T> {
     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 set_fuel(&mut self, fuel: u64) -> Result<()> {
+        self.inner.set_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::set_fuel`]
+    pub fn set_fuel(&mut self, fuel: u64) -> Result<()> {
+        self.0.set_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<u64> {
+        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 set_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()?;

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.set_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();