Support full range of fuel by stashing overflow in reserves

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

crates/wasmtime/src/store.rs

@@ -1108,6 +1108,31 @@ impl<T> StoreInner<T> {
     }
 }
 
+fn get_fuel(consumed_fuel: i64, fuel_reserve: u64) -> u64 {
+    u64::try_from(-consumed_fuel)
+        .unwrap_or(0)
+        .saturating_add(fuel_reserve)
+}
+
+fn set_fuel(consumed_ptr: &mut i64, fuel_reserve: &mut u64, yield_interval: Option<NonZeroU64>, new_fuel_amount: u64) {
+        // Fuel is stored as an i64, and users can pass in more than that. So in order to support
+        // easier accounting, we put the overflow into `fuel_reserve`, as `fuel_reserve` has the
+        // full u64 range.
+        let overflow = new_fuel_amount.saturating_sub(i64::MAX as u64);
+        let fuel = i64::try_from(new_fuel_amount).unwrap_or(i64::MAX);
+        if yield_interval.is_none() {
+            *fuel_reserve = overflow;
+            *consumed_ptr = -fuel;
+            return;
+        }
+        let interval = i64::try_from(yield_interval.unwrap().get()).unwrap_or(i64::MAX);
+        let injected = std::cmp::min(interval, fuel);
+        *fuel_reserve = u64::try_from(fuel - injected)
+            .unwrap_or(0)
+            .saturating_add(overflow);
+        *consumed_ptr = -injected;
+}
+
 #[doc(hidden)]
 impl StoreOpaque {
     pub fn id(&self) -> StoreId {
@@ -1355,31 +1380,16 @@ impl StoreOpaque {
             "fuel is not configured in this store"
         );
         let consumed = unsafe { *self.runtime_limits.fuel_consumed.get() };
-        Ok(u64::try_from(-consumed)
-           .unwrap_or(0)
-           .saturating_add(self.fuel_reserve))
+        Ok(get_fuel(consumed, self.fuel_reserve))
     }
 
     pub fn set_fuel(&mut self, fuel: u64) -> Result<()> {
         anyhow::ensure!(
             self.engine().config().tunables.consume_fuel,
             "fuel is not configured in this store"
         );
-        // Fuel is stored as an i64, so we need to cast it. If the provided fuel
-        // value overflows that just assume that i64::MAX will suffice. Wasm
-        // execution isn't fast enough to burn through i64::MAX fuel in any
-        // reasonable amount of time anyway.
-        let fuel = i64::try_from(fuel).unwrap_or(i64::MAX);
         let consumed_ptr = unsafe { &mut *self.runtime_limits.fuel_consumed.get() };
-        if self.fuel_yield_interval.is_none() {
-            self.fuel_reserve = 0;
-            *consumed_ptr = -fuel;
-            return Ok(());
-        }
-        let interval = i64::try_from(self.fuel_yield_interval.unwrap().get()).unwrap_or(i64::MAX);
-        let injected = std::cmp::min(interval, fuel);
-        self.fuel_reserve = u64::try_from(fuel - injected).unwrap_or(0);
-        *consumed_ptr = -injected;
+        set_fuel(consumed_ptr, &mut self.fuel_reserve, self.fuel_yield_interval, fuel);
         Ok(())
     }
 
@@ -1388,6 +1398,10 @@ impl StoreOpaque {
             self.engine().config().tunables.consume_fuel,
             "fuel is not configured in this store"
         );
+        anyhow::ensure!(
+            self.engine().config().async_support,
+            "async support is not configured in this store"
+        );
         self.fuel_yield_interval = interval;
         // Reset the fuel active + reserve states by resetting the amount.
         self.set_fuel(self.get_fuel()?)
@@ -2052,7 +2066,10 @@ unsafe impl<T> wasmtime_runtime::Store for StoreInner<T> {
             return Err(Trap::OutOfFuel.into());
         }
         self.set_fuel(remaining)?;
-        self.async_yield_impl()?;
+        #[cfg(feature = "async")]
+        if self.fuel_yield_interval.is_some() {
+            self.async_yield_impl()?;
+        }
         Ok(())
     }
 
@@ -2218,3 +2235,71 @@ impl<T: Copy> Drop for Reset<T> {
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use std::num::NonZeroU64;
+    use super::{get_fuel, set_fuel};
+
+    struct FuelTank {
+        pub consumed_fuel: i64,
+        pub reserve_fuel: u64,
+        pub yield_interval: Option<NonZeroU64>,
+    }
+
+    impl FuelTank {
+        fn new() -> Self {
+            FuelTank { consumed_fuel: 0, reserve_fuel: 0, yield_interval: None }
+        }
+        fn get_fuel(&self) -> u64 {
+            get_fuel(self.consumed_fuel, self.reserve_fuel)
+        }
+
+        fn set_fuel(&mut self, fuel: u64) {
+            set_fuel(&mut self.consumed_fuel, &mut self.reserve_fuel, self.yield_interval, fuel);
+        }
+    }
+
+    #[test]
+    fn smoke() {
+        let mut tank = FuelTank::new();
+        tank.set_fuel(10);
+        assert_eq!(tank.consumed_fuel, -10);
+        assert_eq!(tank.reserve_fuel, 0);
+
+        tank.yield_interval = NonZeroU64::new(10);
+        tank.set_fuel(25);
+        assert_eq!(tank.consumed_fuel, -10);
+        assert_eq!(tank.reserve_fuel, 15);
+    }
+
+    #[test]
+    fn does_not_lose_precision() {
+        let mut tank = FuelTank::new();
+        tank.set_fuel(u64::MAX);
+        assert_eq!(tank.get_fuel(), u64::MAX);
+
+        tank.set_fuel(i64::MAX as u64);
+        assert_eq!(tank.get_fuel(), i64::MAX as u64);
+
+        tank.set_fuel(i64::MAX as u64 + 1);
+        assert_eq!(tank.get_fuel(), i64::MAX as u64 + 1);
+    }
+
+    #[test]
+    fn yielding_does_not_lose_precision() {
+        let mut tank = FuelTank::new();
+
+        tank.yield_interval = NonZeroU64::new(10);
+        tank.set_fuel(u64::MAX);
+        assert_eq!(tank.get_fuel(), u64::MAX);
+        assert_eq!(tank.consumed_fuel, -10);
+        assert_eq!(tank.reserve_fuel, u64::MAX - 10);
+
+        tank.yield_interval = NonZeroU64::new(u64::MAX);
+        tank.set_fuel(u64::MAX);
+        assert_eq!(tank.get_fuel(), u64::MAX);
+        assert_eq!(tank.consumed_fuel, -i64::MAX);
+        assert_eq!(tank.reserve_fuel, u64::MAX - (i64::MAX as u64));
+    }
+}

tests/all/fuel.rs

@@ -1,3 +1,5 @@
+use std::num::NonZeroU64;
+
 use anyhow::Result;
 use wasmtime::*;
 use wast::parser::{self, Parse, ParseBuffer, Parser};
@@ -174,10 +176,27 @@ fn manual_edge_cases() {
     let engine = Engine::new(&config).unwrap();
     let mut store = Store::new(&engine, ());
     store.set_fuel(u64::MAX).unwrap();
-    assert_eq!(store.get_fuel().unwrap(), i64::MAX as u64);
-    assert!(store.set_fuel(i64::MAX as u64).is_ok());
-    assert_eq!(store.get_fuel().unwrap(), i64::MAX as u64);
-    assert!(store.set_fuel(i64::MAX as u64).is_ok());
+    assert_eq!(store.get_fuel().unwrap(), u64::MAX);
+}
+
+#[test]
+fn manual_async_edge_cases() {
+    let mut config = Config::new();
+    config.consume_fuel(true).async_support(true);
+    let engine = Engine::new(&config).unwrap();
+    let mut store = Store::new(&engine, ());
+
+    store.fuel_async_yield_interval(NonZeroU64::new(u64::MAX)).unwrap();
+    store.set_fuel(u64::MAX).unwrap();
+    assert_eq!(store.get_fuel().unwrap(), u64::MAX);
+
+    store.fuel_async_yield_interval(NonZeroU64::new(1)).unwrap();
+    store.set_fuel(u64::MAX).unwrap();
+    assert_eq!(store.get_fuel().unwrap(), u64::MAX);
+
+    store.fuel_async_yield_interval(NonZeroU64::new(i64::MAX as u64)).unwrap();
+    store.set_fuel(u64::MAX).unwrap();
+    assert_eq!(store.get_fuel().unwrap(), u64::MAX);
 }
 
 #[test]