fix bug in pause/resume. (#2233)

Author: vwxyzh

src/Microsoft.Azure.SignalR.Common/Utilities/PauseHandler.cs

@@ -1,11 +1,13 @@
-// Copyright (c) Microsoft. All rights reserved.
+// Copyright (c) Microsoft. All rights reserved.
 // Licensed under the MIT license. See LICENSE file in the project root for full license information.
 
 using System.Threading;
 using System.Threading.Tasks;
 
 namespace Microsoft.Azure.SignalR;
 
+#nullable enable
+
 internal class PauseHandler
 {
     private readonly SemaphoreSlim _pauseSemaphore = new(1, 1);
@@ -14,9 +16,9 @@ internal class PauseHandler
 
     private volatile int _pauseAcked = 1;
 
-    public bool ShouldReplyAck => Interlocked.CompareExchange(ref _pauseAcked, 1, 0) == 0;
+    public bool ShouldReplyAck() => Interlocked.CompareExchange(ref _pauseAcked, 1, 0) == 0;
 
-    public async Task<bool> WaitAsync(int ms, CancellationToken ctoken) => _pauseSemaphore.Wait(0, ctoken) || await _pauseSemaphore.WaitAsync(ms, ctoken);
+    public Task<bool> TryAcquire(int millisecondsTimeout) => _pauseSemaphore.WaitAsync(millisecondsTimeout);
 
     public void Release() => _pauseSemaphore.Release();
 

src/Microsoft.Azure.SignalR/ClientConnections/ClientConnectionContext.cs

@@ -1,4 +1,4 @@
-// Copyright (c) Microsoft. All rights reserved.
+// Copyright (c) Microsoft. All rights reserved.
 // Licensed under the MIT license. See LICENSE file in the project root for full license information.
 
 using System;
@@ -261,7 +261,7 @@ public Task PauseAsync()
 
     public Task PauseAckAsync()
     {
-        if (_pauseHandler.ShouldReplyAck)
+        if (_pauseHandler.ShouldReplyAck())
         {
             Log.OutgoingTaskPauseAck(Logger, ConnectionId);
             var message = new ConnectionFlowControlMessage(ConnectionId, ConnectionFlowControlOperation.PauseAck);
@@ -334,18 +334,26 @@ internal async Task ProcessOutgoingMessagesAsync(SignalRProtocol.IHubProtocol pr
                     if (HandshakeResponseTask.IsCompleted)
                     {
                         var next = buffer;
-                        while (!buffer.IsEmpty && protocol.TryParseMessage(ref next, FakeInvocationBinder.Instance, out var message))
+
+                        // we still want messages to successfully going out when application completes
+                        int waitTime = 0;
+                        while (!await _pauseHandler.TryAcquire(1000))
                         {
-                            // we still want messages to successfully going out when application completes
-                            if (!await _pauseHandler.WaitAsync(StaticRandom.Next(500, 1500), OutgoingAborted))
+                            Log.OutgoingTaskPaused(Logger, ConnectionId);
+                            if (OutgoingAborted.IsCancellationRequested)
                             {
-                                Log.OutgoingTaskPaused(Logger, ConnectionId);
-                                buffer = buffer.Slice(0);
-                                break;
+                                waitTime++;
+                                if (waitTime > 5)
+                                {
+                                    OutgoingAborted.ThrowIfCancellationRequested();
+                                }
                             }
-
-                            try
+                        }
+                        try
+                        {
+                            while (!buffer.IsEmpty && protocol.TryParseMessage(ref next, FakeInvocationBinder.Instance, out var message))
                             {
+
                                 var messageType = message switch
                                 {
                                     SignalRProtocol.HubInvocationMessage => DataMessageType.Invocation,
@@ -363,10 +371,10 @@ internal async Task ProcessOutgoingMessagesAsync(SignalRProtocol.IHubProtocol pr
                                     _ => next,
                                 };
                             }
-                            finally
-                            {
-                                _pauseHandler.Release();
-                            }
+                        }
+                        finally
+                        {
+                            _pauseHandler.Release();
                         }
                     }
                 }

test/Microsoft.Azure.SignalR.Common.Tests/PauseHandlerTests.cs

@@ -1,4 +1,4 @@
-// Copyright (c) Microsoft. All rights reserved.
+// Copyright (c) Microsoft. All rights reserved.
 // Licensed under the MIT license. See LICENSE file in the project root for full license information.
 
 using System.Threading.Tasks;
@@ -13,22 +13,23 @@ public class PauseHandlerTests
     public async Task TestPauseAndResume()
     {
         var handler = new PauseHandler();
-        Assert.False(handler.ShouldReplyAck);
+        Assert.False(handler.ShouldReplyAck());
 
         await handler.PauseAsync();
-        Assert.False(await handler.WaitAsync(100, default));
+        Assert.False(await handler.TryAcquire(100));
 
-        Assert.True(handler.ShouldReplyAck);
-        Assert.False(handler.ShouldReplyAck); // ack only once
+        Assert.True(handler.ShouldReplyAck());
+        Assert.False(handler.ShouldReplyAck()); // ack only once
 
         await handler.PauseAsync(); // pause can be called multiple times.
-        Assert.False(await handler.WaitAsync(100, default));
-        Assert.False(handler.ShouldReplyAck); // already acked previously
+        Assert.False(await handler.TryAcquire(100));
+        Assert.False(handler.ShouldReplyAck()); // already acked previously
 
         await handler.ResumeAsync();
-        Assert.True(await handler.WaitAsync(100, default));
-        Assert.False(await handler.WaitAsync(100, default)); // only 1 parallel
+
+        Assert.True(await handler.TryAcquire(100));
+        Assert.False(await handler.TryAcquire(100)); // only 1 parallel
         handler.Release();
-        Assert.True(await handler.WaitAsync(100, default));
+        Assert.True(await handler.TryAcquire(100));
     }
 }