Merge pull request #146 from Cysharp/fix-async

Fix Debounce, ThrottleFirstLast, ThrottleLast async overload does not handle correctly when func return immediately
Cysharp · Feb 29, 2024 · 81dfa96 · 81dfa96
2 parents ab76681 + c6cdd83
commit 81dfa96
Show file tree

Hide file tree

Showing 4 changed files with 21 additions and 15 deletions.
diff --git a/src/R3/Operators/Debounce.cs b/src/R3/Operators/Debounce.cs
@@ -112,7 +112,7 @@ sealed class _Debounce(Observer<T> observer, Func<T, CancellationToken, ValueTas
         readonly object gate = new object();
         T? latestValue;
         bool hasValue;
-        Task? runningTask;
+        bool isRunning;
         int taskId;
         CancellationTokenSource cancellationTokenSource = new();
 
@@ -123,15 +123,17 @@ protected override void OnNextCore(T value)
                 latestValue = value;
                 hasValue = true;
 
-                if (runningTask != null)
+                if (isRunning)
                 {
                     cancellationTokenSource.Cancel();
                     cancellationTokenSource = new CancellationTokenSource();
                 }
 
                 var newId = unchecked(taskId + 1);
                 Volatile.Write(ref taskId, newId);
-                runningTask = PublishOnNextAfterAsync(value, newId, cancellationTokenSource.Token);
+
+                isRunning = true;
+                PublishOnNextAfterAsync(value, newId, cancellationTokenSource.Token);
             }
         }
 
@@ -164,7 +166,7 @@ protected override void DisposeCore()
             cancellationTokenSource.Cancel();
         }
 
-        async Task PublishOnNextAfterAsync(T value, int id, CancellationToken cancellationToken)
+        async void PublishOnNextAfterAsync(T value, int id, CancellationToken cancellationToken)
         {
             try
             {
@@ -188,7 +190,7 @@ async Task PublishOnNextAfterAsync(T value, int id, CancellationToken cancellati
                     observer.OnNext(latestValue!);
                     hasValue = false;
                     latestValue = default;
-                    runningTask = null;
+                    isRunning = false;
 
                 END:
                     { }

diff --git a/src/R3/Operators/ThrottleFirstLast.cs b/src/R3/Operators/ThrottleFirstLast.cs
@@ -113,15 +113,16 @@ sealed class _ThrottleFirstLast(Observer<T> observer, Func<T, CancellationToken,
         readonly CancellationTokenSource cancellationTokenSource = new();
         T? lastValue;
         bool hasValue;
-        Task? taskRunner;
+        bool isRunning;
 
         protected override void OnNextCore(T value)
         {
             lock (gate)
             {
-                if (taskRunner == null)
+                if (!isRunning)
                 {
-                    taskRunner = RaiseOnNextAsync(value);
+                    isRunning = true;
+                    RaiseOnNextAsync(value);
                     observer.OnNext(value);
                 }
                 else
@@ -148,7 +149,7 @@ protected override void DisposeCore()
             cancellationTokenSource.Cancel();
         }
 
-        async Task RaiseOnNextAsync(T value)
+        async void RaiseOnNextAsync(T value)
         {
             try
             {
@@ -171,7 +172,7 @@ async Task RaiseOnNextAsync(T value)
                         observer.OnNext(lastValue!);
                         lastValue = default;
                         hasValue = false;
-                        taskRunner = null;
+                        isRunning = false;
                     }
                 }
             }

diff --git a/src/R3/Operators/ThrottleLast.cs b/src/R3/Operators/ThrottleLast.cs
@@ -105,16 +105,17 @@ sealed class _ThrottleLast(Observer<T> observer, Func<T, CancellationToken, Valu
         readonly object gate = new object();
         readonly CancellationTokenSource cancellationTokenSource = new();
         T? lastValue;
-        Task? taskRunner;
+        bool isRunning;
 
         protected override void OnNextCore(T value)
         {
             lock (gate)
             {
                 lastValue = value;
-                if (taskRunner == null)
+                if (!isRunning)
                 {
-                    taskRunner = RaiseOnNextAsync(value);
+                    isRunning = true;
+                    RaiseOnNextAsync(value);
                 }
             }
         }
@@ -135,7 +136,7 @@ protected override void DisposeCore()
             cancellationTokenSource.Cancel();
         }
 
-        async Task RaiseOnNextAsync(T value)
+        async void RaiseOnNextAsync(T value)
         {
             try
             {
@@ -155,7 +156,7 @@ async Task RaiseOnNextAsync(T value)
                 {
                     observer.OnNext(lastValue!);
                     lastValue = default;
-                    taskRunner = null;
+                    isRunning = false;
                 }
             }
         }

diff --git a/tests/R3.Tests/OperatorTests/ChunkTest.cs b/tests/R3.Tests/OperatorTests/ChunkTest.cs
@@ -257,6 +257,8 @@ public void ChunkFrameAndCount()
     [Fact]
     public void ChunkAsync()
     {
+        SynchronizationContext.SetSynchronizationContext(null);
+
         var publisher = new Subject<int>();
         var tp = new FakeTimeProvider();
         var list = publisher.Chunk(async (x, ct) =>