cancelImmediately

Cysharp · Dec 14, 2023 · 6096017 · 6096017
1 parent 3b92a99
commit 6096017
Showing 1 changed file with 33 additions and 7 deletions.
diff --git 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<Unit> EveryUpdate()
     {
-        return new EveryUpdate(EventSystem.DefaultFrameProvider, CancellationToken.None);
+        return new EveryUpdate(EventSystem.DefaultFrameProvider, CancellationToken.None, cancelImmediately: false);
     }
 
     public static Event<Unit> EveryUpdate(CancellationToken cancellationToken)
     {
-        return new EveryUpdate(EventSystem.DefaultFrameProvider, cancellationToken);
+        return new EveryUpdate(EventSystem.DefaultFrameProvider, cancellationToken, cancelImmediately: false);
     }
 
     public static Event<Unit> EveryUpdate(FrameProvider frameProvider)
     {
-        return new EveryUpdate(frameProvider, CancellationToken.None);
+        return new EveryUpdate(frameProvider, CancellationToken.None, cancelImmediately: false);
     }
 
     public static Event<Unit> EveryUpdate(FrameProvider frameProvider, CancellationToken cancellationToken)
     {
-        return new EveryUpdate(frameProvider, cancellationToken);
+        return new EveryUpdate(frameProvider, cancellationToken, cancelImmediately: false);
+    }
+
+    public static Event<Unit> EveryUpdate(FrameProvider frameProvider, CancellationToken cancellationToken, bool cancelImmediately)
+    {
+        return new EveryUpdate(frameProvider, cancellationToken, cancelImmediately: cancelImmediately);
     }
 }
 
 
 
-internal sealed class EveryUpdate(FrameProvider frameProvider, CancellationToken cancellationToken) : Event<Unit>
+internal sealed class EveryUpdate(FrameProvider frameProvider, CancellationToken cancellationToken, bool cancelImmediately) : Event<Unit>
 {
     protected override IDisposable SubscribeCore(Subscriber<Unit> subscriber)
     {
-        var runner = new EveryUpdateRunnerWorkItem(subscriber, cancellationToken);
+        var runner = new EveryUpdateRunnerWorkItem(subscriber, cancellationToken, cancelImmediately);
         frameProvider.Register(runner);
         return runner;
     }
 
-    class EveryUpdateRunnerWorkItem(Subscriber<Unit> subscriber, CancellationToken cancellationToken) : IFrameRunnerWorkItem, IDisposable
+    class EveryUpdateRunnerWorkItem : IFrameRunnerWorkItem, IDisposable
     {
+        Subscriber<Unit> subscriber;
+        CancellationToken cancellationToken;
+        CancellationTokenRegistration cancellationTokenRegistration;
         bool isDisposed;
 
+        public EveryUpdateRunnerWorkItem(Subscriber<Unit> 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();
         }
     }
 }