diff --git a/src/R3/Factories/_EventFactory.cs b/src/R3/Factories/_EventFactory.cs index ee9420b8..e33b1604 100644 --- a/src/R3/Factories/_EventFactory.cs +++ b/src/R3/Factories/_EventFactory.cs @@ -1,5 +1,6 @@  using System.Diagnostics.CodeAnalysis; +using System.Text; using System.Threading; namespace R3; @@ -30,40 +31,64 @@ public static partial class Event public static Event EveryUpdate() { - return new EveryUpdate(EventSystem.DefaultFrameProvider, CancellationToken.None); + return new EveryUpdate(EventSystem.DefaultFrameProvider, CancellationToken.None, cancelImmediately: false); } public static Event EveryUpdate(CancellationToken cancellationToken) { - return new EveryUpdate(EventSystem.DefaultFrameProvider, cancellationToken); + return new EveryUpdate(EventSystem.DefaultFrameProvider, cancellationToken, cancelImmediately: false); } public static Event EveryUpdate(FrameProvider frameProvider) { - return new EveryUpdate(frameProvider, CancellationToken.None); + return new EveryUpdate(frameProvider, CancellationToken.None, cancelImmediately: false); } public static Event EveryUpdate(FrameProvider frameProvider, CancellationToken cancellationToken) { - return new EveryUpdate(frameProvider, cancellationToken); + return new EveryUpdate(frameProvider, cancellationToken, cancelImmediately: false); + } + + public static Event EveryUpdate(FrameProvider frameProvider, CancellationToken cancellationToken, bool cancelImmediately) + { + return new EveryUpdate(frameProvider, cancellationToken, cancelImmediately: cancelImmediately); } } -internal sealed class EveryUpdate(FrameProvider frameProvider, CancellationToken cancellationToken) : Event +internal sealed class EveryUpdate(FrameProvider frameProvider, CancellationToken cancellationToken, bool cancelImmediately) : Event { protected override IDisposable SubscribeCore(Subscriber subscriber) { - var runner = new EveryUpdateRunnerWorkItem(subscriber, cancellationToken); + var runner = new EveryUpdateRunnerWorkItem(subscriber, cancellationToken, cancelImmediately); frameProvider.Register(runner); return runner; } - class EveryUpdateRunnerWorkItem(Subscriber subscriber, CancellationToken cancellationToken) : IFrameRunnerWorkItem, IDisposable + class EveryUpdateRunnerWorkItem : IFrameRunnerWorkItem, IDisposable { + Subscriber subscriber; + CancellationToken cancellationToken; + CancellationTokenRegistration cancellationTokenRegistration; bool isDisposed; + public EveryUpdateRunnerWorkItem(Subscriber subscriber, CancellationToken cancellationToken, bool cancelImmediately) + { + this.subscriber = subscriber; + this.cancellationToken = cancellationToken; + + if (cancelImmediately && cancellationToken.CanBeCanceled) + { + cancellationTokenRegistration = cancellationToken.UnsafeRegister(static state => + { + var s = (EveryUpdateRunnerWorkItem)state!; + s.subscriber.OnCompleted(); + s.Dispose(); + }, this); + } + } + public bool MoveNext(long frameCount) { if (isDisposed) @@ -85,6 +110,7 @@ public bool MoveNext(long frameCount) public void Dispose() { isDisposed = true; + cancellationTokenRegistration.Dispose(); } } }