Allow CancellationToken to be used within OnDisconnected

Cysharp · Sep 6, 2024 · f70c2d0 · f70c2d0
1 parent 52f1cb9
commit f70c2d0
Show file tree

Hide file tree

Showing 3 changed files with 18 additions and 2 deletions.
diff --git a/src/MagicOnion.Server/Hubs/StreamingHub.cs b/src/MagicOnion.Server/Hubs/StreamingHub.cs
@@ -134,10 +134,13 @@ internal async Task<DuplexStreamingResult<StreamingHubPayload, StreamingHubPaylo
         {
             Metrics.StreamingHubConnectionDecrement(Context.Metrics, Context.MethodHandler.ServiceName);
 
-            heartbeatHandle.Dispose();
             StreamingServiceContext.CompleteStreamingHub();
+            heartbeatHandle.Unregister(); // NOTE: To be able to use CancellationToken within OnDisconnected event, separate the calls to Dispose and Unregister.
+
             await OnDisconnected();
+
             await this.Group.DisposeAsync();
+            heartbeatHandle.Dispose();
             remoteClientResultPendingTasks.Dispose();
         }
 

diff --git a/src/MagicOnion.Server/Hubs/StreamingHubHeartbeatManager.cs b/src/MagicOnion.Server/Hubs/StreamingHubHeartbeatManager.cs
@@ -23,6 +23,7 @@ internal class StreamingHubHeartbeatHandle : IDisposable
     readonly CancellationTokenSource timeoutToken;
     readonly TimeSpan timeoutDuration;
     bool disposed;
+    bool unregistered;
     short waitingSequence = -1;
     bool timeoutTimerIsRunning;
     DateTimeOffset lastSentAt;
@@ -99,17 +100,24 @@ public void SetAckCallback(Action<TimeSpan>? callbackAction)
 
     public void Unregister()
     {
+        if (unregistered) return;
+
         manager.Unregister(ServiceContext);
         timeoutToken.CancelAfter(Timeout.InfiniteTimeSpan);
         timeoutTimerIsRunning = false;
+        unregistered = true;
     }
 
     public void Dispose()
     {
         if (disposed) return;
+
         disposed = true;
         onAckCallback = null;
-        manager.Unregister(ServiceContext);
+        if (!unregistered)
+        {
+            manager.Unregister(ServiceContext);
+        }
         timeoutToken.Dispose();
     }
 }

diff --git a/tests/MagicOnion.Server.Tests/StreamingHubHeartbeat/StreamingHubServerHeartbeatTest.cs b/tests/MagicOnion.Server.Tests/StreamingHubHeartbeat/StreamingHubServerHeartbeatTest.cs
@@ -134,6 +134,7 @@ public async Task Timeout()
 
         // Assert
         Assert.True((bool)Fixture.Items.GetValueOrDefault("Disconnected"));
+        Assert.True((bool)Fixture.Items.GetValueOrDefault("Heartbeat/TimeoutToken/IsCancellationRequested"));
         Assert.True(client.WaitForDisconnect().IsCompleted);
     }
 
@@ -269,7 +270,11 @@ public class StreamingHubServerHeartbeatTestHub_TimeoutBehavior([FromKeyedServic
 {
     protected override ValueTask OnDisconnected()
     {
+        var httpContext = Context.CallContext.GetHttpContext();
+        var heartbeatFeature = httpContext.Features.GetRequiredFeature<IMagicOnionHeartbeatFeature>();
+
         items["Disconnected"] = true;
+        items["Heartbeat/TimeoutToken/IsCancellationRequested"] = heartbeatFeature.TimeoutToken.IsCancellationRequested;
         return base.OnDisconnected();
     }
 }