* Set cancellation correctly for TaskCompletionSource in AsyncRpcContinuation

lukebakken · lukebakken · commit a3e2a697261f · 2024-12-27T11:32:33.000-08:00
diff --git a/projects/RabbitMQ.Client/ConsumerDispatching/ConsumerDispatcherChannelBase.cs b/projects/RabbitMQ.Client/ConsumerDispatching/ConsumerDispatcherChannelBase.cs
@@ -85,11 +85,23 @@ internal ConsumerDispatcherChannelBase(Impl.Channel channel, ushort concurrency)
 
         public ValueTask HandleBasicConsumeOkAsync(IAsyncBasicConsumer consumer, string consumerTag, CancellationToken cancellationToken)
         {
+            cancellationToken.ThrowIfCancellationRequested();
+
             if (false == _disposedValue && false == _quiesce)
             {
-                AddConsumer(consumer, consumerTag);
-                WorkStruct work = WorkStruct.CreateConsumeOk(consumer, consumerTag);
-                return _writer.WriteAsync(work, cancellationToken);
+                try
+                {
+                    AddConsumer(consumer, consumerTag);
+                    WorkStruct work = WorkStruct.CreateConsumeOk(consumer, consumerTag);
+
+                    cancellationToken.ThrowIfCancellationRequested();
+                    return _writer.WriteAsync(work, cancellationToken);
+                }
+                catch
+                {
+                    _ = GetAndRemoveConsumer(consumerTag);
+                    throw;
+                }
             }
             else
             {
@@ -101,10 +113,14 @@ public ValueTask HandleBasicDeliverAsync(string consumerTag, ulong deliveryTag,
             string exchange, string routingKey, IReadOnlyBasicProperties basicProperties, RentedMemory body,
             CancellationToken cancellationToken)
         {
+            cancellationToken.ThrowIfCancellationRequested();
+
             if (false == _disposedValue && false == _quiesce)
             {
                 IAsyncBasicConsumer consumer = GetConsumerOrDefault(consumerTag);
                 var work = WorkStruct.CreateDeliver(consumer, consumerTag, deliveryTag, redelivered, exchange, routingKey, basicProperties, body);
+
+                cancellationToken.ThrowIfCancellationRequested();
                 return _writer.WriteAsync(work, cancellationToken);
             }
             else
@@ -115,10 +131,14 @@ public ValueTask HandleBasicDeliverAsync(string consumerTag, ulong deliveryTag,
 
         public ValueTask HandleBasicCancelOkAsync(string consumerTag, CancellationToken cancellationToken)
         {
+            cancellationToken.ThrowIfCancellationRequested();
+
             if (false == _disposedValue && false == _quiesce)
             {
                 IAsyncBasicConsumer consumer = GetAndRemoveConsumer(consumerTag);
                 WorkStruct work = WorkStruct.CreateCancelOk(consumer, consumerTag);
+
+                cancellationToken.ThrowIfCancellationRequested();
                 return _writer.WriteAsync(work, cancellationToken);
             }
             else
@@ -129,10 +149,14 @@ public ValueTask HandleBasicCancelOkAsync(string consumerTag, CancellationToken
 
         public ValueTask HandleBasicCancelAsync(string consumerTag, CancellationToken cancellationToken)
         {
+            cancellationToken.ThrowIfCancellationRequested();
+
             if (false == _disposedValue && false == _quiesce)
             {
                 IAsyncBasicConsumer consumer = GetAndRemoveConsumer(consumerTag);
                 WorkStruct work = WorkStruct.CreateCancel(consumer, consumerTag);
+
+                cancellationToken.ThrowIfCancellationRequested();
                 return _writer.WriteAsync(work, cancellationToken);
             }
             else
diff --git a/projects/RabbitMQ.Client/Impl/AsyncRpcContinuations.cs b/projects/RabbitMQ.Client/Impl/AsyncRpcContinuations.cs
@@ -51,7 +51,7 @@ internal abstract class AsyncRpcContinuation<T> : IRpcContinuation
 
         private bool _disposedValue;
 
-        public AsyncRpcContinuation(TimeSpan continuationTimeout, CancellationToken cancellationToken)
+        public AsyncRpcContinuation(TimeSpan continuationTimeout, CancellationToken rpcCancellationToken)
         {
             /*
              * Note: we can't use an ObjectPool for these because the netstandard2.0
@@ -89,7 +89,7 @@ public AsyncRpcContinuation(TimeSpan continuationTimeout, CancellationToken canc
             _tcsConfiguredTaskAwaitable = _tcs.Task.ConfigureAwait(false);
 
             _linkedCancellationTokenSource = CancellationTokenSource.CreateLinkedTokenSource(
-                _continuationTimeoutCancellationTokenSource.Token, cancellationToken);
+                _continuationTimeoutCancellationTokenSource.Token, rpcCancellationToken);
         }
 
         public CancellationToken CancellationToken
@@ -105,7 +105,27 @@ public ConfiguredTaskAwaitable<T>.ConfiguredTaskAwaiter GetAwaiter()
             return _tcsConfiguredTaskAwaitable.GetAwaiter();
         }
 
-        public abstract Task HandleCommandAsync(IncomingCommand cmd);
+        public async Task HandleCommandAsync(IncomingCommand cmd)
+        {
+            try
+            {
+                await DoHandleCommandAsync(cmd)
+                    .ConfigureAwait(false);
+            }
+            catch (OperationCanceledException)
+            {
+                if (CancellationToken.IsCancellationRequested)
+                {
+                    _tcs.SetCanceled();
+                }
+                else
+                {
+                    throw;
+                }
+            }
+        }
+
+        protected abstract Task DoHandleCommandAsync(IncomingCommand cmd);
 
         public virtual void HandleChannelShutdown(ShutdownEventArgs reason)
         {
@@ -141,17 +161,17 @@ public ConnectionSecureOrTuneAsyncRpcContinuation(TimeSpan continuationTimeout,
         {
         }
 
-        public override Task HandleCommandAsync(IncomingCommand cmd)
+        protected override Task DoHandleCommandAsync(IncomingCommand cmd)
         {
             if (cmd.CommandId == ProtocolCommandId.ConnectionSecure)
             {
                 var secure = new ConnectionSecure(cmd.MethodSpan);
-                _tcs.TrySetResult(new ConnectionSecureOrTune(secure._challenge, default));
+                _tcs.SetResult(new ConnectionSecureOrTune(secure._challenge, default));
             }
             else if (cmd.CommandId == ProtocolCommandId.ConnectionTune)
             {
                 var tune = new ConnectionTune(cmd.MethodSpan);
-                _tcs.TrySetResult(new ConnectionSecureOrTune(default, new ConnectionTuneDetails
+                _tcs.SetResult(new ConnectionSecureOrTune(default, new ConnectionTuneDetails
                 {
                     m_channelMax = tune._channelMax,
                     m_frameMax = tune._frameMax,
@@ -178,11 +198,11 @@ public SimpleAsyncRpcContinuation(ProtocolCommandId expectedCommandId, TimeSpan
             _expectedCommandId = expectedCommandId;
         }
 
-        public override Task HandleCommandAsync(IncomingCommand cmd)
+        protected override Task DoHandleCommandAsync(IncomingCommand cmd)
         {
             if (cmd.CommandId == _expectedCommandId)
             {
-                _tcs.TrySetResult(true);
+                _tcs.SetResult(true);
             }
             else
             {
@@ -206,14 +226,14 @@ public BasicCancelAsyncRpcContinuation(string consumerTag, IConsumerDispatcher c
             _consumerDispatcher = consumerDispatcher;
         }
 
-        public override async Task HandleCommandAsync(IncomingCommand cmd)
+        protected override async Task DoHandleCommandAsync(IncomingCommand cmd)
         {
             if (cmd.CommandId == ProtocolCommandId.BasicCancelOk)
             {
-                _tcs.TrySetResult(true);
                 Debug.Assert(_consumerTag == new BasicCancelOk(cmd.MethodSpan)._consumerTag);
                 await _consumerDispatcher.HandleBasicCancelOkAsync(_consumerTag, CancellationToken)
                     .ConfigureAwait(false);
+                _tcs.SetResult(true);
             }
             else
             {
@@ -235,14 +255,16 @@ public BasicConsumeAsyncRpcContinuation(IAsyncBasicConsumer consumer, IConsumerD
             _consumerDispatcher = consumerDispatcher;
         }
 
-        public override async Task HandleCommandAsync(IncomingCommand cmd)
+        protected override async Task DoHandleCommandAsync(IncomingCommand cmd)
         {
             if (cmd.CommandId == ProtocolCommandId.BasicConsumeOk)
             {
                 var method = new BasicConsumeOk(cmd.MethodSpan);
-                _tcs.TrySetResult(method._consumerTag);
+
                 await _consumerDispatcher.HandleBasicConsumeOkAsync(_consumer, method._consumerTag, CancellationToken)
                     .ConfigureAwait(false);
+
+                _tcs.SetResult(method._consumerTag);
             }
             else
             {
@@ -264,7 +286,7 @@ public BasicGetAsyncRpcContinuation(Func<ulong, ulong> adjustDeliveryTag,
 
         internal DateTime StartTime { get; } = DateTime.UtcNow;
 
-        public override Task HandleCommandAsync(IncomingCommand cmd)
+        protected override Task DoHandleCommandAsync(IncomingCommand cmd)
         {
             if (cmd.CommandId == ProtocolCommandId.BasicGetOk)
             {
@@ -280,11 +302,11 @@ public override Task HandleCommandAsync(IncomingCommand cmd)
                     header,
                     cmd.Body.ToArray());
 
-                _tcs.TrySetResult(result);
+                _tcs.SetResult(result);
             }
             else if (cmd.CommandId == ProtocolCommandId.BasicGetEmpty)
             {
-                _tcs.TrySetResult(null);
+                _tcs.SetResult(null);
             }
             else
             {
@@ -325,7 +347,7 @@ public override void HandleChannelShutdown(ShutdownEventArgs reason)
 
         public Task OnConnectionShutdownAsync(object? sender, ShutdownEventArgs reason)
         {
-            _tcs.TrySetResult(true);
+            _tcs.SetResult(true);
             return Task.CompletedTask;
         }
     }
@@ -377,13 +399,13 @@ public QueueDeclareAsyncRpcContinuation(TimeSpan continuationTimeout, Cancellati
         {
         }
 
-        public override Task HandleCommandAsync(IncomingCommand cmd)
+        protected override Task DoHandleCommandAsync(IncomingCommand cmd)
         {
             if (cmd.CommandId == ProtocolCommandId.QueueDeclareOk)
             {
                 var method = new Client.Framing.QueueDeclareOk(cmd.MethodSpan);
                 var result = new QueueDeclareOk(method._queue, method._messageCount, method._consumerCount);
-                _tcs.TrySetResult(result);
+                _tcs.SetResult(result);
             }
             else
             {
@@ -417,12 +439,12 @@ public QueueDeleteAsyncRpcContinuation(TimeSpan continuationTimeout, Cancellatio
         {
         }
 
-        public override Task HandleCommandAsync(IncomingCommand cmd)
+        protected override Task DoHandleCommandAsync(IncomingCommand cmd)
         {
             if (cmd.CommandId == ProtocolCommandId.QueueDeleteOk)
             {
                 var method = new QueueDeleteOk(cmd.MethodSpan);
-                _tcs.TrySetResult(method._messageCount);
+                _tcs.SetResult(method._messageCount);
             }
             else
             {
@@ -440,12 +462,12 @@ public QueuePurgeAsyncRpcContinuation(TimeSpan continuationTimeout, Cancellation
         {
         }
 
-        public override Task HandleCommandAsync(IncomingCommand cmd)
+        protected override Task DoHandleCommandAsync(IncomingCommand cmd)
         {
             if (cmd.CommandId == ProtocolCommandId.QueuePurgeOk)
             {
                 var method = new QueuePurgeOk(cmd.MethodSpan);
-                _tcs.TrySetResult(method._messageCount);
+                _tcs.SetResult(method._messageCount);
             }
             else
             {
diff --git a/projects/Test/Integration/GH/TestGitHubIssues.cs b/projects/Test/Integration/GH/TestGitHubIssues.cs
@@ -0,0 +1,100 @@
+﻿// This source code is dual-licensed under the Apache License, version
+// 2.0, and the Mozilla Public License, version 2.0.
+//
+// The APL v2.0:
+//
+//---------------------------------------------------------------------------
+//   Copyright (c) 2007-2024 Broadcom. All Rights Reserved.
+//
+//   Licensed under the Apache License, Version 2.0 (the "License");
+//   you may not use this file except in compliance with the License.
+//   You may obtain a copy of the License at
+//
+//       https://www.apache.org/licenses/LICENSE-2.0
+//
+//   Unless required by applicable law or agreed to in writing, software
+//   distributed under the License is distributed on an "AS IS" BASIS,
+//   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+//   See the License for the specific language governing permissions and
+//   limitations under the License.
+//---------------------------------------------------------------------------
+//
+// The MPL v2.0:
+//
+//---------------------------------------------------------------------------
+// This Source Code Form is subject to the terms of the Mozilla Public
+// License, v. 2.0. If a copy of the MPL was not distributed with this
+// file, You can obtain one at https://mozilla.org/MPL/2.0/.
+//
+//  Copyright (c) 2007-2024 Broadcom. All Rights Reserved.
+//---------------------------------------------------------------------------
+
+using System.Threading;
+using System.Threading.Tasks;
+using RabbitMQ.Client;
+using RabbitMQ.Client.Events;
+using Xunit;
+using Xunit.Abstractions;
+
+#nullable enable
+
+namespace Test.Integration.GH
+{
+    public class TestGitHubIssues : IntegrationFixture
+    {
+        public TestGitHubIssues(ITestOutputHelper output) : base(output)
+        {
+        }
+
+        public override Task InitializeAsync()
+        {
+            // NB: nothing to do here since each test creates its own factory,
+            // connections and channels
+            Assert.Null(_connFactory);
+            Assert.Null(_conn);
+            Assert.Null(_channel);
+            return Task.CompletedTask;
+        }
+
+        [Fact]
+        public async Task TestBasicConsumeCancellation_GH1750()
+        {
+            /*
+             * Note:
+             * Testing that the task is actually canceled requires a hacked RabbitMQ server.
+             * Modify deps/rabbit/src/rabbit_channel.erl, handle_cast for basic.consume_ok
+             * Before send/2, add timer:sleep(1000), then `make run-broker`
+             *
+             * The _output line at the end of the test will print TaskCanceledException
+             */
+            Assert.Null(_connFactory);
+            Assert.Null(_conn);
+            Assert.Null(_channel);
+
+            _connFactory = CreateConnectionFactory();
+            _connFactory.AutomaticRecoveryEnabled = false;
+            _connFactory.TopologyRecoveryEnabled = false;
+
+            _conn = await _connFactory.CreateConnectionAsync();
+            _channel = await _conn.CreateChannelAsync();
+
+            QueueDeclareOk q = await _channel.QueueDeclareAsync();
+
+            var consumer = new AsyncEventingBasicConsumer(_channel);
+            consumer.ReceivedAsync += (o, a) =>
+            {
+                return Task.CompletedTask;
+            };
+
+            try
+            {
+                using var cts = new CancellationTokenSource(5);
+                await _channel.BasicConsumeAsync(q.QueueName, true, consumer, cts.Token);
+            }
+            catch (TaskCanceledException ex)
+            {
+                _output.WriteLine("ex: {0}", ex);
+            }
+        }
+    }
+}
diff --git a/projects/Test/Integration/TestToxiproxy.cs b/projects/Test/Integration/TestToxiproxy.cs

Original file line number	Diff line number	Diff line change
`@@ -85,11 +85,23 @@ internal ConsumerDispatcherChannelBase(Impl.Channel channel, ushort concurrency)`
`85`	`85`
`86`	`86`	`public ValueTask HandleBasicConsumeOkAsync(IAsyncBasicConsumer consumer, string consumerTag, CancellationToken cancellationToken)`
`87`	`87`	`{`
	`88`	`+ cancellationToken.ThrowIfCancellationRequested();`
	`89`	`+`
`88`	`90`	`if (false == _disposedValue && false == _quiesce)`
`89`	`91`	`{`
`90`		`- AddConsumer(consumer, consumerTag);`
`91`		`- WorkStruct work = WorkStruct.CreateConsumeOk(consumer, consumerTag);`
`92`		`- return _writer.WriteAsync(work, cancellationToken);`
	`92`	`+ try`
	`93`	`+ {`
	`94`	`+ AddConsumer(consumer, consumerTag);`
	`95`	`+ WorkStruct work = WorkStruct.CreateConsumeOk(consumer, consumerTag);`
	`96`	`+`
	`97`	`+ cancellationToken.ThrowIfCancellationRequested();`
	`98`	`+ return _writer.WriteAsync(work, cancellationToken);`
	`99`	`+ }`
	`100`	`+ catch`
	`101`	`+ {`
	`102`	`+ _ = GetAndRemoveConsumer(consumerTag);`
	`103`	`+ throw;`
	`104`	`+ }`
`93`	`105`	`}`
`94`	`106`	`else`
`95`	`107`	`{`
`@@ -101,10 +113,14 @@ public ValueTask HandleBasicDeliverAsync(string consumerTag, ulong deliveryTag,`
`101`	`113`	`string exchange, string routingKey, IReadOnlyBasicProperties basicProperties, RentedMemory body,`
`102`	`114`	`CancellationToken cancellationToken)`
`103`	`115`	`{`
	`116`	`+ cancellationToken.ThrowIfCancellationRequested();`
	`117`	`+`
`104`	`118`	`if (false == _disposedValue && false == _quiesce)`
`105`	`119`	`{`
`106`	`120`	`IAsyncBasicConsumer consumer = GetConsumerOrDefault(consumerTag);`
`107`	`121`	`var work = WorkStruct.CreateDeliver(consumer, consumerTag, deliveryTag, redelivered, exchange, routingKey, basicProperties, body);`
	`122`	`+`
	`123`	`+ cancellationToken.ThrowIfCancellationRequested();`
`108`	`124`	`return _writer.WriteAsync(work, cancellationToken);`
`109`	`125`	`}`
`110`	`126`	`else`
`@@ -115,10 +131,14 @@ public ValueTask HandleBasicDeliverAsync(string consumerTag, ulong deliveryTag,`
`115`	`131`
`116`	`132`	`public ValueTask HandleBasicCancelOkAsync(string consumerTag, CancellationToken cancellationToken)`
`117`	`133`	`{`
	`134`	`+ cancellationToken.ThrowIfCancellationRequested();`
	`135`	`+`
`118`	`136`	`if (false == _disposedValue && false == _quiesce)`
`119`	`137`	`{`
`120`	`138`	`IAsyncBasicConsumer consumer = GetAndRemoveConsumer(consumerTag);`
`121`	`139`	`WorkStruct work = WorkStruct.CreateCancelOk(consumer, consumerTag);`
	`140`	`+`
	`141`	`+ cancellationToken.ThrowIfCancellationRequested();`
`122`	`142`	`return _writer.WriteAsync(work, cancellationToken);`
`123`	`143`	`}`
`124`	`144`	`else`
`@@ -129,10 +149,14 @@ public ValueTask HandleBasicCancelOkAsync(string consumerTag, CancellationToken`
`129`	`149`
`130`	`150`	`public ValueTask HandleBasicCancelAsync(string consumerTag, CancellationToken cancellationToken)`
`131`	`151`	`{`
	`152`	`+ cancellationToken.ThrowIfCancellationRequested();`
	`153`	`+`
`132`	`154`	`if (false == _disposedValue && false == _quiesce)`
`133`	`155`	`{`
`134`	`156`	`IAsyncBasicConsumer consumer = GetAndRemoveConsumer(consumerTag);`
`135`	`157`	`WorkStruct work = WorkStruct.CreateCancel(consumer, consumerTag);`
	`158`	`+`
	`159`	`+ cancellationToken.ThrowIfCancellationRequested();`
`136`	`160`	`return _writer.WriteAsync(work, cancellationToken);`
`137`	`161`	`}`
`138`	`162`	`else`
Original file line number	Diff line number	Diff line change
`@@ -51,7 +51,7 @@ internal abstract class AsyncRpcContinuation<T> : IRpcContinuation`
`51`	`51`
`52`	`52`	`private bool _disposedValue;`
`53`	`53`
`54`		`- public AsyncRpcContinuation(TimeSpan continuationTimeout, CancellationToken cancellationToken)`
	`54`	`+ public AsyncRpcContinuation(TimeSpan continuationTimeout, CancellationToken rpcCancellationToken)`
`55`	`55`	`{`
`56`	`56`	`/*`
`57`	`57`	`* Note: we can't use an ObjectPool for these because the netstandard2.0`
`@@ -89,7 +89,7 @@ public AsyncRpcContinuation(TimeSpan continuationTimeout, CancellationToken canc`
`89`	`89`	`_tcsConfiguredTaskAwaitable = _tcs.Task.ConfigureAwait(false);`
`90`	`90`
`91`	`91`	`_linkedCancellationTokenSource = CancellationTokenSource.CreateLinkedTokenSource(`
`92`		`- _continuationTimeoutCancellationTokenSource.Token, cancellationToken);`
	`92`	`+ _continuationTimeoutCancellationTokenSource.Token, rpcCancellationToken);`
`93`	`93`	`}`
`94`	`94`
`95`	`95`	`public CancellationToken CancellationToken`
`@@ -105,7 +105,27 @@ public ConfiguredTaskAwaitable<T>.ConfiguredTaskAwaiter GetAwaiter()`
`105`	`105`	`return _tcsConfiguredTaskAwaitable.GetAwaiter();`
`106`	`106`	`}`
`107`	`107`
`108`		`- public abstract Task HandleCommandAsync(IncomingCommand cmd);`
	`108`	`+ public async Task HandleCommandAsync(IncomingCommand cmd)`
	`109`	`+ {`
	`110`	`+ try`
	`111`	`+ {`
	`112`	`+ await DoHandleCommandAsync(cmd)`
	`113`	`+ .ConfigureAwait(false);`
	`114`	`+ }`
	`115`	`+ catch (OperationCanceledException)`
	`116`	`+ {`
	`117`	`+ if (CancellationToken.IsCancellationRequested)`
	`118`	`+ {`
	`119`	`+ _tcs.SetCanceled();`
	`120`	`+ }`
	`121`	`+ else`
	`122`	`+ {`
	`123`	`+ throw;`
	`124`	`+ }`
	`125`	`+ }`
	`126`	`+ }`
	`127`	`+`
	`128`	`+ protected abstract Task DoHandleCommandAsync(IncomingCommand cmd);`
`109`	`129`
`110`	`130`	`public virtual void HandleChannelShutdown(ShutdownEventArgs reason)`
`111`	`131`	`{`
`@@ -141,17 +161,17 @@ public ConnectionSecureOrTuneAsyncRpcContinuation(TimeSpan continuationTimeout,`
`141`	`161`	`{`
`142`	`162`	`}`
`143`	`163`
`144`		`- public override Task HandleCommandAsync(IncomingCommand cmd)`
	`164`	`+ protected override Task DoHandleCommandAsync(IncomingCommand cmd)`
`145`	`165`	`{`
`146`	`166`	`if (cmd.CommandId == ProtocolCommandId.ConnectionSecure)`
`147`	`167`	`{`
`148`	`168`	`var secure = new ConnectionSecure(cmd.MethodSpan);`
`149`		`- _tcs.TrySetResult(new ConnectionSecureOrTune(secure._challenge, default));`
	`169`	`+ _tcs.SetResult(new ConnectionSecureOrTune(secure._challenge, default));`
`150`	`170`	`}`
`151`	`171`	`else if (cmd.CommandId == ProtocolCommandId.ConnectionTune)`
`152`	`172`	`{`
`153`	`173`	`var tune = new ConnectionTune(cmd.MethodSpan);`
`154`		`- _tcs.TrySetResult(new ConnectionSecureOrTune(default, new ConnectionTuneDetails`
	`174`	`+ _tcs.SetResult(new ConnectionSecureOrTune(default, new ConnectionTuneDetails`
`155`	`175`	`{`
`156`	`176`	`m_channelMax = tune._channelMax,`
`157`	`177`	`m_frameMax = tune._frameMax,`
`@@ -178,11 +198,11 @@ public SimpleAsyncRpcContinuation(ProtocolCommandId expectedCommandId, TimeSpan`
`178`	`198`	`_expectedCommandId = expectedCommandId;`
`179`	`199`	`}`
`180`	`200`
`181`		`- public override Task HandleCommandAsync(IncomingCommand cmd)`
	`201`	`+ protected override Task DoHandleCommandAsync(IncomingCommand cmd)`
`182`	`202`	`{`
`183`	`203`	`if (cmd.CommandId == _expectedCommandId)`
`184`	`204`	`{`
`185`		`- _tcs.TrySetResult(true);`
	`205`	`+ _tcs.SetResult(true);`
`186`	`206`	`}`
`187`	`207`	`else`
`188`	`208`	`{`
`@@ -206,14 +226,14 @@ public BasicCancelAsyncRpcContinuation(string consumerTag, IConsumerDispatcher c`
`206`	`226`	`_consumerDispatcher = consumerDispatcher;`
`207`	`227`	`}`
`208`	`228`
`209`		`- public override async Task HandleCommandAsync(IncomingCommand cmd)`
	`229`	`+ protected override async Task DoHandleCommandAsync(IncomingCommand cmd)`
`210`	`230`	`{`
`211`	`231`	`if (cmd.CommandId == ProtocolCommandId.BasicCancelOk)`
`212`	`232`	`{`
`213`		`- _tcs.TrySetResult(true);`
`214`	`233`	`Debug.Assert(_consumerTag == new BasicCancelOk(cmd.MethodSpan)._consumerTag);`
`215`	`234`	`await _consumerDispatcher.HandleBasicCancelOkAsync(_consumerTag, CancellationToken)`
`216`	`235`	`.ConfigureAwait(false);`
	`236`	`+ _tcs.SetResult(true);`
`217`	`237`	`}`
`218`	`238`	`else`
`219`	`239`	`{`
`@@ -235,14 +255,16 @@ public BasicConsumeAsyncRpcContinuation(IAsyncBasicConsumer consumer, IConsumerD`
`235`	`255`	`_consumerDispatcher = consumerDispatcher;`
`236`	`256`	`}`
`237`	`257`
`238`		`- public override async Task HandleCommandAsync(IncomingCommand cmd)`
	`258`	`+ protected override async Task DoHandleCommandAsync(IncomingCommand cmd)`
`239`	`259`	`{`
`240`	`260`	`if (cmd.CommandId == ProtocolCommandId.BasicConsumeOk)`
`241`	`261`	`{`
`242`	`262`	`var method = new BasicConsumeOk(cmd.MethodSpan);`
`243`		`- _tcs.TrySetResult(method._consumerTag);`
	`263`	`+`
`244`	`264`	`await _consumerDispatcher.HandleBasicConsumeOkAsync(_consumer, method._consumerTag, CancellationToken)`
`245`	`265`	`.ConfigureAwait(false);`
	`266`	`+`
	`267`	`+ _tcs.SetResult(method._consumerTag);`
`246`	`268`	`}`
`247`	`269`	`else`
`248`	`270`	`{`
`@@ -264,7 +286,7 @@ public BasicGetAsyncRpcContinuation(Func<ulong, ulong> adjustDeliveryTag,`
`264`	`286`
`265`	`287`	`internal DateTime StartTime { get; } = DateTime.UtcNow;`
`266`	`288`
`267`		`- public override Task HandleCommandAsync(IncomingCommand cmd)`
	`289`	`+ protected override Task DoHandleCommandAsync(IncomingCommand cmd)`
`268`	`290`	`{`
`269`	`291`	`if (cmd.CommandId == ProtocolCommandId.BasicGetOk)`
`270`	`292`	`{`
`@@ -280,11 +302,11 @@ public override Task HandleCommandAsync(IncomingCommand cmd)`
`280`	`302`	`header,`
`281`	`303`	`cmd.Body.ToArray());`
`282`	`304`
`283`		`- _tcs.TrySetResult(result);`
	`305`	`+ _tcs.SetResult(result);`
`284`	`306`	`}`
`285`	`307`	`else if (cmd.CommandId == ProtocolCommandId.BasicGetEmpty)`
`286`	`308`	`{`
`287`		`- _tcs.TrySetResult(null);`
	`309`	`+ _tcs.SetResult(null);`
`288`	`310`	`}`
`289`	`311`	`else`
`290`	`312`	`{`
`@@ -325,7 +347,7 @@ public override void HandleChannelShutdown(ShutdownEventArgs reason)`
`325`	`347`
`326`	`348`	`public Task OnConnectionShutdownAsync(object? sender, ShutdownEventArgs reason)`
`327`	`349`	`{`
`328`		`- _tcs.TrySetResult(true);`
	`350`	`+ _tcs.SetResult(true);`
`329`	`351`	`return Task.CompletedTask;`
`330`	`352`	`}`
`331`	`353`	`}`
`@@ -377,13 +399,13 @@ public QueueDeclareAsyncRpcContinuation(TimeSpan continuationTimeout, Cancellati`
`377`	`399`	`{`
`378`	`400`	`}`
`379`	`401`
`380`		`- public override Task HandleCommandAsync(IncomingCommand cmd)`
	`402`	`+ protected override Task DoHandleCommandAsync(IncomingCommand cmd)`
`381`	`403`	`{`
`382`	`404`	`if (cmd.CommandId == ProtocolCommandId.QueueDeclareOk)`
`383`	`405`	`{`
`384`	`406`	`var method = new Client.Framing.QueueDeclareOk(cmd.MethodSpan);`
`385`	`407`	`var result = new QueueDeclareOk(method._queue, method._messageCount, method._consumerCount);`
`386`		`- _tcs.TrySetResult(result);`
	`408`	`+ _tcs.SetResult(result);`
`387`	`409`	`}`
`388`	`410`	`else`
`389`	`411`	`{`
`@@ -417,12 +439,12 @@ public QueueDeleteAsyncRpcContinuation(TimeSpan continuationTimeout, Cancellatio`
`417`	`439`	`{`
`418`	`440`	`}`
`419`	`441`
`420`		`- public override Task HandleCommandAsync(IncomingCommand cmd)`
	`442`	`+ protected override Task DoHandleCommandAsync(IncomingCommand cmd)`
`421`	`443`	`{`
`422`	`444`	`if (cmd.CommandId == ProtocolCommandId.QueueDeleteOk)`
`423`	`445`	`{`
`424`	`446`	`var method = new QueueDeleteOk(cmd.MethodSpan);`
`425`		`- _tcs.TrySetResult(method._messageCount);`
	`447`	`+ _tcs.SetResult(method._messageCount);`
`426`	`448`	`}`
`427`	`449`	`else`
`428`	`450`	`{`
`@@ -440,12 +462,12 @@ public QueuePurgeAsyncRpcContinuation(TimeSpan continuationTimeout, Cancellation`
`440`	`462`	`{`
`441`	`463`	`}`
`442`	`464`
`443`		`- public override Task HandleCommandAsync(IncomingCommand cmd)`
	`465`	`+ protected override Task DoHandleCommandAsync(IncomingCommand cmd)`
`444`	`466`	`{`
`445`	`467`	`if (cmd.CommandId == ProtocolCommandId.QueuePurgeOk)`
`446`	`468`	`{`
`447`	`469`	`var method = new QueuePurgeOk(cmd.MethodSpan);`
`448`		`- _tcs.TrySetResult(method._messageCount);`
	`470`	`+ _tcs.SetResult(method._messageCount);`
`449`	`471`	`}`
`450`	`472`	`else`
`451`	`473`	`{`