Add Store::consume_fuel to manually consume fuel (#3352)

This can be useful for host functions that want to consume fuel to
reflect their relative cost. Additionally it's a relatively easy
addition to have and someone's asking for it!

Closes #3315

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

@@ -139,6 +139,20 @@ WASM_API_EXTERN wasmtime_error_t *wasmtime_context_add_fuel(wasmtime_context_t *
  */
 WASM_API_EXTERN bool wasmtime_context_fuel_consumed(const wasmtime_context_t *context, uint64_t *fuel);
 
+/**
+ * \brief Attempt to manually consume fuel from the store.
+ *
+ * If fuel consumption is not enabled via #wasmtime_config_consume_fuel_set then
+ * this function will return an error. Otherwise this will attempt to consume
+ * the specified amount of `fuel` from the store. If successful the remaining
+ * amount of fuel is stored into `remaining`. If `fuel` couldn't be consumed
+ * then an error is returned.
+ *
+ * Also note that fuel, if enabled, must be originally configured via
+ * #wasmtime_context_add_fuel.
+ */
+WASM_API_EXTERN wasmtime_error_t *wasmtime_context_consume_fuel(wasmtime_context_t *context, uint64_t fuel, uint64_t *remaining);
+
 /**
  * \brief Configres WASI state within the specified store.
  *

crates/c-api/src/store.rs

@@ -145,6 +145,17 @@ pub extern "C" fn wasmtime_context_fuel_consumed(store: CStoreContext<'_>, fuel:
     }
 }
 
+#[no_mangle]
+pub extern "C" fn wasmtime_context_consume_fuel(
+    mut store: CStoreContextMut<'_>,
+    fuel: u64,
+    remaining_fuel: &mut u64,
+) -> Option<Box<wasmtime_error_t>> {
+    crate::handle_result(store.consume_fuel(fuel), |remaining| {
+        *remaining_fuel = remaining;
+    })
+}
+
 #[repr(C)]
 pub struct wasmtime_interrupt_handle_t {
     handle: InterruptHandle,

crates/wasmtime/src/func.rs

@@ -1621,6 +1621,13 @@ impl<T> Caller<'_, T> {
         self.store.add_fuel(fuel)
     }
 
+    /// Synthetically consumes fuel from the store.
+    ///
+    /// For more information see [`Store::consume_fuel`](crate::Store::consume_fuel)
+    pub fn consume_fuel(&mut self, fuel: u64) -> Result<u64> {
+        self.store.consume_fuel(fuel)
+    }
+
     /// Configures this `Store` to trap whenever fuel runs out.
     ///
     /// For more information see

crates/wasmtime/src/store.rs

@@ -614,6 +614,29 @@ impl<T> Store<T> {
         self.inner.add_fuel(fuel)
     }
 
+    /// Synthetically consumes fuel from this [`Store`].
+    ///
+    /// For this method to work fuel consumption must be enabled via
+    /// [`Config::consume_fuel`](crate::Config::consume_fuel).
+    ///
+    /// WebAssembly execution will automatically consume fuel but if so desired
+    /// the embedder can also consume fuel manually to account for relative
+    /// costs of host functions, for example.
+    ///
+    /// This function will attempt to consume `fuel` units of fuel from within
+    /// this store. If the remaining amount of fuel allows this then `Ok(N)` is
+    /// returned where `N` is the amount of remaining fuel. Otherwise an error
+    /// is returned and no fuel is consumed.
+    ///
+    /// # Errors
+    ///
+    /// This function will return an either either if fuel consumption via
+    /// [`Config`](crate::Config) is disabled or if `fuel` exceeds the amount
+    /// of remaining fuel within this store.
+    pub fn consume_fuel(&mut self, fuel: u64) -> Result<u64> {
+        self.inner.consume_fuel(fuel)
+    }
+
     /// Configures a [`Store`] to generate a [`Trap`] whenever it runs out of
     /// fuel.
     ///
@@ -748,6 +771,13 @@ impl<'a, T> StoreContextMut<'a, T> {
         self.0.add_fuel(fuel)
     }
 
+    /// Synthetically consume fuel from this store.
+    ///
+    /// For more information see [`Store::consume_fuel`]
+    pub fn consume_fuel(&mut self, fuel: u64) -> Result<u64> {
+        self.0.consume_fuel(fuel)
+    }
+
     /// Configures this `Store` to trap whenever fuel runs out.
     ///
     /// For more information see [`Store::out_of_fuel_trap`]
@@ -1028,6 +1058,20 @@ impl StoreOpaque {
         Ok(())
     }
 
+    fn consume_fuel(&mut self, fuel: u64) -> Result<u64> {
+        let consumed_ptr = unsafe { &mut *self.interrupts.fuel_consumed.get() };
+        match i64::try_from(fuel)
+            .ok()
+            .and_then(|fuel| consumed_ptr.checked_add(fuel))
+        {
+            Some(consumed) if consumed < 0 => {
+                *consumed_ptr = consumed;
+                Ok(u64::try_from(-consumed).unwrap())
+            }
+            _ => bail!("not enough fuel remaining in store"),
+        }
+    }
+
     #[inline]
     pub fn signal_handler(&self) -> Option<*const SignalHandler<'static>> {
         let handler = self.signal_handler.as_ref()?;

tests/all/fuel.rs

@@ -122,3 +122,70 @@ fn iloop() {
         );
     }
 }
+
+#[test]
+fn manual_fuel() {
+    let mut config = Config::new();
+    config.consume_fuel(true);
+    let engine = Engine::new(&config).unwrap();
+    let mut store = Store::new(&engine, ());
+    store.add_fuel(10_000).unwrap();
+    assert_eq!(store.fuel_consumed(), Some(0));
+    assert_eq!(store.consume_fuel(1).unwrap(), 9_999);
+    assert_eq!(store.fuel_consumed(), Some(1));
+    assert!(store.consume_fuel(10_000).is_err());
+    assert_eq!(store.consume_fuel(999).unwrap(), 9_000);
+    assert!(store.consume_fuel(10_000).is_err());
+    assert_eq!(store.consume_fuel(8998).unwrap(), 2);
+    assert!(store.consume_fuel(2).is_err());
+    assert_eq!(store.consume_fuel(1).unwrap(), 1);
+    assert!(store.consume_fuel(1).is_err());
+    assert_eq!(store.consume_fuel(0).unwrap(), 1);
+}
+
+#[test]
+fn host_function_consumes_all() {
+    const FUEL: u64 = 10_000;
+    let mut config = Config::new();
+    config.consume_fuel(true);
+    let engine = Engine::new(&config).unwrap();
+    let module = Module::new(
+        &engine,
+        r#"
+            (module
+                (import "" "" (func))
+                (func (export "")
+                    call 0
+                    call $other)
+                (func $other))
+        "#,
+    )
+    .unwrap();
+    let mut store = Store::new(&engine, ());
+    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 instance = Instance::new(&mut store, &module, &[func.into()]).unwrap();
+    let export = instance
+        .get_typed_func::<(), (), _>(&mut store, "")
+        .unwrap();
+    let trap = export.call(&mut store, ()).err().unwrap().to_string();
+    assert!(trap.contains("all fuel consumed"), "bad error: {}", trap);
+}
+
+#[test]
+fn manual_edge_cases() {
+    let mut config = Config::new();
+    config.consume_fuel(true);
+    let engine = Engine::new(&config).unwrap();
+    let mut store = Store::new(&engine, ());
+    store.add_fuel(u64::MAX).unwrap();
+    assert_eq!(store.fuel_consumed(), Some(0));
+    assert!(store.consume_fuel(u64::MAX).is_err());
+    assert!(store.consume_fuel(i64::MAX as u64).is_err());
+    assert!(store.consume_fuel(i64::MAX as u64 + 1).is_err());
+    assert_eq!(store.consume_fuel(i64::MAX as u64 - 1).unwrap(), 1);
+}