diff --git a/src/Store.cs b/src/Store.cs index 405fe5f..326beb9 100644 --- a/src/Store.cs +++ b/src/Store.cs @@ -263,6 +263,91 @@ public void SetEpochDeadline(ulong ticksBeyondCurrent) /// public void SetData(object? data) => this.data = data; + /// + /// Represents a callback invoked when the epoch deadline of a store is reached. + /// + /// The store whose epoch deadline was reached. + /// The new deadline, in ticks beyond the current epoch, after which execution resumes. + /// + /// The callback runs on the thread executing the WebAssembly code. Throwing from it terminates + /// the execution; the exception becomes the InnerException of the resulting . + /// + public delegate ulong EpochDeadlineCallback(Store store); + + /// + /// Sets the callback invoked when WebAssembly code running in this store reaches its epoch + /// deadline, instead of trapping. Replaces any previously set callback. + /// + /// The callback to invoke. + /// + /// + /// For this to work epoch interruption must be enabled via . + /// + /// + /// The callback is kept alive until it is replaced or the store is disposed, so a callback that + /// captures this store keeps the store alive as well; prefer the store passed to the callback. + /// + /// + /// Thrown if callback is null + public void SetEpochDeadlineCallback(EpochDeadlineCallback callback) + { + if (callback is null) + { + throw new ArgumentNullException(nameof(callback)); + } + + unsafe + { + Native.WasmtimeEpochDeadlineCallback trampoline = + (context, data, epochDeadlineDelta, updateKind) => + InvokeEpochDeadlineCallback(callback, context, epochDeadlineDelta); + + // The GCHandle passed as callback data keeps the trampoline alive; Wasmtime runs the + // finalizer when the callback is replaced or the store is deleted. + Native.wasmtime_store_epoch_deadline_callback( + NativeHandle, + trampoline, + GCHandle.ToIntPtr(GCHandle.Alloc(trampoline)), + Finalizer + ); + } + } + + private static unsafe IntPtr InvokeEpochDeadlineCallback(EpochDeadlineCallback callback, IntPtr context, ulong* epochDeadlineDelta) + { + try + { + // The update kind is left at "continue"; yielding requires async support, which this binding does not enable. + *epochDeadlineDelta = callback(new StoreContext(context).Store); + return IntPtr.Zero; + } + catch (Exception ex) + { + return CreateEpochDeadlineError(ex); + } + } + + private static IntPtr CreateEpochDeadlineError(Exception ex) + { + try + { + // Store the exception as error cause, so that it becomes the WasmtimeException's + // InnerException when the error bubbles up. See Function.HandleCallbackException. + Function.CallbackErrorCause = ex is WasmtimeException wasmtimeException ? wasmtimeException.InnerException : ex; + + return Native.wasmtime_error_new(ex.Message); + } + catch (Exception separateException) + { + // We never must let .NET exceptions bubble through the native-to-managed transition; + // see Function.HandleCallbackException. + Environment.FailFast(separateException.Message, separateException); + + // Satisfy the control-flow analyzer; this line is never reached. + throw; + } + } + /// public void Dispose() { @@ -335,6 +420,14 @@ private static class Native [DllImport(Engine.LibraryName)] public static extern void wasmtime_store_limiter(Handle store, long memory_size, long table_elements, long instances, long tables, long memories); + + public unsafe delegate IntPtr WasmtimeEpochDeadlineCallback(IntPtr context, IntPtr data, ulong* epochDeadlineDelta, byte* updateKind); + + [DllImport(Engine.LibraryName)] + public static extern void wasmtime_store_epoch_deadline_callback(Handle store, WasmtimeEpochDeadlineCallback callback, IntPtr data, Finalizer? finalizer); + + [DllImport(Engine.LibraryName)] + public static extern IntPtr wasmtime_error_new([MarshalAs(Extensions.LPUTF8Str)] string message); } private readonly IntPtr contextHandle; diff --git a/tests/EpochInterruptionTests.cs b/tests/EpochInterruptionTests.cs index 03f0c37..0343cae 100644 --- a/tests/EpochInterruptionTests.cs +++ b/tests/EpochInterruptionTests.cs @@ -52,9 +52,116 @@ public void ItCanInterruptInfiniteLoop() .WithMessage("*wasm trap: interrupt*"); } + /// + /// Runs the given body while a background thread advances the engine epoch + /// + private void WhileEpochAdvances(Action body) + { + using var stop = new CancellationTokenSource(); + var ticker = new Thread(() => + { + while (!stop.Token.WaitHandle.WaitOne(TimeSpan.FromMilliseconds(5))) + { + Fixture.Engine.IncrementEpoch(); + } + }) + { + IsBackground = true + }; + + ticker.Start(); + + try + { + body(); + } + finally + { + stop.Cancel(); + ticker.Join(); + } + } + + [Fact] + public void ItRenewsTheDeadlineUntilTheCallbackThrows() + { + var invocations = 0; + Store observedStore = null; + var exceptionToThrow = new OperationCanceledException("the guest was cancelled"); + + Store.SetEpochDeadline(1); + Store.SetEpochDeadlineCallback(store => + { + observedStore = store; + + if (++invocations == 3) + { + throw exceptionToThrow; + } + + // Resume for one more tick, so the deadline is reached again. + return 1; + }); + + var instance = Linker.Instantiate(Store, Fixture.Module); + var run = instance.GetFunction("run"); + + WhileEpochAdvances(() => + { + var action = () => run.Invoke(); + + action.Should() + .Throw() + .Where(e => e.InnerException == exceptionToThrow) + .WithMessage("*the guest was cancelled*"); + }); + + invocations.Should().Be(3); + observedStore.Should().BeSameAs(Store); + } + + [Fact] + public void ItReplacesAPreviouslySetCallback() + { + var replacedInvocations = 0; + var invocations = 0; + + Store.SetEpochDeadline(1); + Store.SetEpochDeadlineCallback(_ => + { + replacedInvocations++; + return 1; + }); + + Store.SetEpochDeadlineCallback(_ => + { + invocations++; + throw new InvalidOperationException("stop"); + }); + + var instance = Linker.Instantiate(Store, Fixture.Module); + var run = instance.GetFunction("run"); + + WhileEpochAdvances(() => + { + var action = () => run.Invoke(); + + action.Should().Throw().WithMessage("*stop*"); + }); + + invocations.Should().Be(1); + replacedInvocations.Should().Be(0); + } + + [Fact] + public void ItThrowsForANullCallback() + { + Assert.Throws(() => Store.SetEpochDeadlineCallback(null!)); + } + public void Dispose() { Store.Dispose(); Linker.Dispose(); } -} \ No newline at end of file +}