Allow GrainTimers to dispose themselves from their own callback (#9065)

dotnet · Jul 12, 2024 · 92e0bf3 · 92e0bf3
1 parent 77a187e
commit 92e0bf3
Show file tree

Hide file tree

Showing 4 changed files with 106 additions and 4 deletions.
diff --git a/src/Orleans.Runtime/Timers/GrainTimer.cs b/src/Orleans.Runtime/Timers/GrainTimer.cs
@@ -146,6 +146,12 @@ private void OnTickCompleted()
         // Schedule the next tick.
         try
         {
+            if (_cts.IsCancellationRequested)
+            {
+                // The instance has been disposed. No further ticks should be fired.
+                return;
+            }
+
             if (!_changed)
             {
                 // If the timer was not modified during the tick, schedule the next tick based on the period.
@@ -160,8 +166,10 @@ private void OnTickCompleted()
         catch (ObjectDisposedException)
         {
         }
-
-        _firing = false;
+        finally
+        {
+            _firing = false;
+        }
     }
 
     private Response OnCallbackException(Exception exc)

diff --git a/test/DefaultCluster.Tests/TimerOrleansTest.cs b/test/DefaultCluster.Tests/TimerOrleansTest.cs
@@ -197,6 +197,17 @@ public async Task GrainTimer_TestAllOverloads()
             await grain.TestCompletedTimerResults();
         }
 
+        [Fact, TestCategory("SlowBVT"), TestCategory("Timers")]
+        public async Task GrainTimer_DisposeFromCallback()
+        {
+            // Schedule a timer which disposes itself from its own callback.
+            var grain = GrainFactory.GetGrain<ITimerCallGrain>(GetRandomGrainId());
+            await grain.RunSelfDisposingTimer();
+
+            var pocoGrain = GrainFactory.GetGrain<IPocoTimerCallGrain>(GetRandomGrainId());
+            await pocoGrain.RunSelfDisposingTimer();
+        }
+
         [Fact, TestCategory("SlowBVT"), TestCategory("Timers")]
         public async Task NonReentrantGrainTimer_Test()
         {

diff --git a/test/Grains/TestGrainInterfaces/ITimerGrain.cs b/test/Grains/TestGrainInterfaces/ITimerGrain.cs
@@ -26,6 +26,7 @@ public interface ITimerCallGrain : IGrainWithIntegerKey
         Task RestartTimer(string name, TimeSpan dueTime);
         Task RestartTimer(string name, TimeSpan dueTime, TimeSpan period);
         Task StopTimer(string name);
+        Task RunSelfDisposingTimer();
     }
 
     public interface IPocoTimerCallGrain : ITimerCallGrain

diff --git a/test/Grains/TestInternalGrains/TimerGrain.cs b/test/Grains/TestInternalGrains/TimerGrain.cs
@@ -159,7 +159,7 @@ public Task StartTimer(string name, TimeSpan delay)
         {
             logger.LogInformation("StartTimer Name={Name} Delay={Delay}", name, delay);
             if (timer is not null) throw new InvalidOperationException("Expected timer to be null");
-            this.timer = this.RegisterGrainTimer(TimerTick, name, new(delay, Timeout.InfiniteTimeSpan)); // One shot timer
+            this.timer = this.RegisterGrainTimer(TimerTick, name, new(delay, Timeout.InfiniteTimeSpan) { Interleave = true }); // One shot timer
             this.timerName = name;
 
             return Task.CompletedTask;
@@ -170,7 +170,7 @@ public Task StartTimer(string name, TimeSpan delay, string operationType)
             logger.LogInformation("StartTimer Name={Name} Delay={Delay}", name, delay);
             if (timer is not null) throw new InvalidOperationException("Expected timer to be null");
             var state = Tuple.Create<string, object>(operationType, name);
-            this.timer = this.RegisterGrainTimer(TimerTickAdvanced, state, delay, Timeout.InfiniteTimeSpan); // One shot timer
+            this.timer = this.RegisterGrainTimer(TimerTickAdvanced, state, new(delay, Timeout.InfiniteTimeSpan) { Interleave = true }); // One shot timer
             this.timerName = name;
 
             return Task.CompletedTask;
@@ -208,6 +208,34 @@ public Task StopTimer(string name)
             return Task.CompletedTask;
         }
 
+        public async Task RunSelfDisposingTimer()
+        {
+            var tcs = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
+            var timer = new IGrainTimer[1];
+            timer[0] = this.RegisterGrainTimer(async (ct) =>
+            {
+                try
+                {
+                    Assert.False(ct.IsCancellationRequested);
+                    Assert.NotNull(timer[0]);
+                    timer[0].Dispose();
+                    Assert.True(ct.IsCancellationRequested);
+                    await Task.Delay(100);
+                    tcs.SetResult();
+                }
+                catch (Exception ex)
+                {
+                    tcs.TrySetException(ex);
+                }
+            },
+            new GrainTimerCreationOptions(TimeSpan.FromSeconds(1), TimeSpan.FromSeconds(1))
+            {
+                Interleave = true
+            });
+
+            await tcs.Task;
+        }
+
         private async Task TimerTick(object data)
         {
             try
@@ -837,6 +865,34 @@ public Task StopTimer(string name)
             return Task.CompletedTask;
         }
 
+        public async Task RunSelfDisposingTimer()
+        {
+            var tcs = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
+            var timer = new IGrainTimer[1];
+            timer[0] = this.RegisterGrainTimer(async (ct) =>
+            {
+                try
+                {
+                    Assert.False(ct.IsCancellationRequested);
+                    Assert.NotNull(timer[0]);
+                    timer[0].Dispose();
+                    Assert.True(ct.IsCancellationRequested);
+                    await Task.Delay(100);
+                    tcs.SetResult();
+                }
+                catch (Exception ex)
+                {
+                    tcs.TrySetException(ex);
+                }
+            },
+            new GrainTimerCreationOptions(TimeSpan.FromSeconds(1), TimeSpan.FromSeconds(1))
+            {
+                Interleave = true
+            });
+
+            await tcs.Task;
+        }
+
         private async Task TimerTick(object data)
         {
             try
@@ -1140,6 +1196,32 @@ public Task OnActivateAsync(CancellationToken cancellationToken)
             return Task.CompletedTask;
         }
 
+        public async Task RunSelfDisposingTimer()
+        {
+            var tcs = new TaskCompletionSource();
+            var timer = new IGrainTimer[1];
+            timer[0] = this.RegisterGrainTimer(async () =>
+            {
+                try
+                {
+                    Assert.NotNull(timer[0]);
+                    timer[0].Dispose();
+                    tcs.SetResult();
+                    await Task.Delay(100);
+                }
+                catch (Exception ex)
+                {
+                    tcs.TrySetException(ex);
+                }
+            },
+            new GrainTimerCreationOptions(TimeSpan.FromSeconds(1), TimeSpan.FromSeconds(1))
+            {
+                Interleave = true
+            });
+
+            await tcs.Task;
+        }
+
         public Task StartTimer(string name, TimeSpan delay, bool keepAlive)
         {
             _logger.LogInformation("StartTimer Name={Name} Delay={Delay}", name, delay);