diff --git a/crates/wasmtime/src/store.rs b/crates/wasmtime/src/store.rs index 4c6eb299dc24..2669b340f012 100644 --- a/crates/wasmtime/src/store.rs +++ b/crates/wasmtime/src/store.rs @@ -97,7 +97,7 @@ pub(crate) struct StoreInner { enum OutOfGas { Trap, InjectFuel { - injection_count: u32, + injection_count: Option, fuel_to_inject: u64, }, } @@ -571,14 +571,19 @@ impl Store { /// will be consumed between yields of an async future. /// /// The `injection_count` parameter indicates how many times this fuel will - /// be injected. Multiplying the two parameters is the total amount of fuel - /// this store is allowed before wasm traps. + /// be injected. + /// If the parameter contains a value, then multiplying the two parameters + /// is the total amount of fuel this store is allowed before wasm traps. + /// If it is `None`, then additional fuel will be injected forever. + /// To still have the ability to stop a running store, you should can use + /// an [`InterruptHandle`]. See [`Store::interrupt_handle`]. + /// /// /// # Panics /// /// This method will panic if it is not called on a store associated with an [async /// config](crate::Config::async_support). - pub fn out_of_fuel_async_yield(&self, injection_count: u32, fuel_to_inject: u64) { + pub fn out_of_fuel_async_yield(&self, injection_count: Option, fuel_to_inject: u64) { assert!( self.async_support(), "cannot use `out_of_fuel_async_yield` without enabling async support in the config" @@ -883,13 +888,17 @@ unsafe impl TrapInfo for Store { injection_count, fuel_to_inject, } => { - if injection_count == 0 { - self.out_of_gas_trap(); + // If the injection count is limited, decrement the count or + // trap if it reaches 0. + if let Some(count) = injection_count { + if count == 0 { + self.out_of_gas_trap(); + } + self.inner.out_of_gas_behavior.set(OutOfGas::InjectFuel { + injection_count: Some(count - 1), + fuel_to_inject, + }); } - self.inner.out_of_gas_behavior.set(OutOfGas::InjectFuel { - injection_count: injection_count - 1, - fuel_to_inject, - }); self.out_of_gas_yield(fuel_to_inject); } #[cfg(not(feature = "async"))] diff --git a/examples/tokio/main.rs b/examples/tokio/main.rs index 5d645062892e..b453e2b7b4c9 100644 --- a/examples/tokio/main.rs +++ b/examples/tokio/main.rs @@ -137,7 +137,7 @@ async fn _run_wasm(inputs: Inputs) -> Result<(), Error> { // WebAssembly execution will be paused for an async yield every time it // consumes 10000 fuel. Fuel will be refilled u32::MAX times. - store.out_of_fuel_async_yield(u32::MAX, 10000); + store.out_of_fuel_async_yield(Some(u32::MAX), 10000); Wasi::set_context( &store, diff --git a/tests/all/async_functions.rs b/tests/all/async_functions.rs index d065fd4f3cfb..7790ea0dc072 100644 --- a/tests/all/async_functions.rs +++ b/tests/all/async_functions.rs @@ -375,7 +375,7 @@ fn dummy_waker() -> Waker { fn iloop_with_fuel() { let engine = Engine::new(Config::new().async_support(true).consume_fuel(true)).unwrap(); let store = Store::new(&engine); - store.out_of_fuel_async_yield(1_000, 10); + store.out_of_fuel_async_yield(Some(1_000), 10); let module = Module::new( &engine, " @@ -409,7 +409,7 @@ fn iloop_with_fuel() { fn fuel_eventually_finishes() { let engine = Engine::new(Config::new().async_support(true).consume_fuel(true)).unwrap(); let store = Store::new(&engine); - store.out_of_fuel_async_yield(u32::max_value(), 10); + store.out_of_fuel_async_yield(Some(u32::max_value()), 10); let module = Module::new( &engine, "