Properly "refuel" for async yields

prtest:full

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

crates/wasmtime/src/store.rs

@@ -1109,9 +1109,25 @@ 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)
+    fuel_reserve.saturating_add_signed(-consumed_fuel)
+}
+
+// Add remaining fuel from the reserve into the active fuel if there is any left.
+//
+// We use this "special" API for the async interval because consumed_ptr can be > 0 (due to some
+// instructions costing more than a single fuel unit).
+fn refuel(
+    consumed_ptr: &mut i64,
+    fuel_reserve: &mut u64,
+    yield_interval: Option<NonZeroU64>,
+) -> bool {
+    let fuel = get_fuel(*consumed_ptr, *fuel_reserve);
+    if fuel > 0 {
+        set_fuel(consumed_ptr, fuel_reserve, yield_interval, fuel);
+        true
+    } else {
+        false
+    }
 }
 
 fn set_fuel(
@@ -1388,6 +1404,15 @@ impl StoreOpaque {
         Ok(get_fuel(consumed, self.fuel_reserve))
     }
 
+    fn refuel(&mut self) -> bool {
+        let consumed_ptr = unsafe { &mut *self.runtime_limits.fuel_consumed.get() };
+        refuel(
+            consumed_ptr,
+            &mut self.fuel_reserve,
+            self.fuel_yield_interval,
+        )
+    }
+
     pub fn set_fuel(&mut self, fuel: u64) -> Result<()> {
         anyhow::ensure!(
             self.engine().config().tunables.consume_fuel,
@@ -2071,11 +2096,9 @@ unsafe impl<T> wasmtime_runtime::Store for StoreInner<T> {
     }
 
     fn out_of_gas(&mut self) -> Result<()> {
-        let remaining = self.fuel_reserve;
-        if remaining == 0 {
+        if !self.refuel() {
             return Err(Trap::OutOfFuel.into());
         }
-        self.set_fuel(remaining)?;
         #[cfg(feature = "async")]
         if self.fuel_yield_interval.is_some() {
             self.async_yield_impl()?;
@@ -2248,7 +2271,7 @@ impl<T: Copy> Drop for Reset<T> {
 
 #[cfg(test)]
 mod tests {
-    use super::{get_fuel, set_fuel};
+    use super::{get_fuel, set_fuel, refuel};
     use std::num::NonZeroU64;
 
     struct FuelTank {
@@ -2268,7 +2291,9 @@ mod tests {
         fn get_fuel(&self) -> u64 {
             get_fuel(self.consumed_fuel, self.reserve_fuel)
         }
-
+        fn refuel(&mut self) -> bool {
+            refuel(&mut self.consumed_fuel, &mut self.reserve_fuel, self.yield_interval)
+        }
         fn set_fuel(&mut self, fuel: u64) {
             set_fuel(
                 &mut self.consumed_fuel,
@@ -2321,4 +2346,37 @@ mod tests {
         assert_eq!(tank.consumed_fuel, -i64::MAX);
         assert_eq!(tank.reserve_fuel, u64::MAX - (i64::MAX as u64));
     }
+
+    #[test]
+    fn refueling() {
+        // It's possible to fuel to have consumed over the limit as some instructions can consume
+        // multiple units of fuel at once. Refueling should be strict in it's consumption and not
+        // add more fuel than there is.
+        let mut tank = FuelTank::new();
+
+        tank.yield_interval = NonZeroU64::new(10);
+        tank.reserve_fuel = 42;
+        tank.consumed_fuel = 4;
+        assert!(tank.refuel());
+        assert_eq!(tank.reserve_fuel, 28);
+        assert_eq!(tank.consumed_fuel, -10);
+
+        tank.yield_interval = NonZeroU64::new(1);
+        tank.reserve_fuel = 8;
+        tank.consumed_fuel = 4;
+        assert_eq!(tank.get_fuel(), 4);
+        assert!(tank.refuel());
+        assert_eq!(tank.reserve_fuel, 3);
+        assert_eq!(tank.consumed_fuel, -1);
+        assert_eq!(tank.get_fuel(), 4);
+
+        tank.yield_interval = NonZeroU64::new(10);
+        tank.reserve_fuel = 3;
+        tank.consumed_fuel = 4;
+        assert_eq!(tank.get_fuel(), 0);
+        assert!(!tank.refuel());
+        assert_eq!(tank.reserve_fuel, 3);
+        assert_eq!(tank.consumed_fuel, 4);
+        assert_eq!(tank.get_fuel(), 0);
+    }
 }