Pass cancellation through the dispatcher

danielmarbach · danielmarbach · commit c6825300f44d · 2025-08-21T22:10:09.000+02:00
diff --git a/projects/RabbitMQ.Client/ConsumerDispatching/AsyncConsumerDispatcher.cs b/projects/RabbitMQ.Client/ConsumerDispatching/AsyncConsumerDispatcher.cs
@@ -34,20 +34,20 @@ protected override async Task ProcessChannelAsync()
                                         {
                                             await work.Consumer.HandleBasicDeliverAsync(
                                                 work.ConsumerTag!, work.DeliveryTag, work.Redelivered,
-                                                work.Exchange!, work.RoutingKey!, work.BasicProperties!, work.Body.Memory)
+                                                work.Exchange!, work.RoutingKey!, work.BasicProperties!, work.Body.Memory, work.CancellationToken)
                                                 .ConfigureAwait(false);
                                         }
                                         break;
                                     case WorkType.Cancel:
-                                        await work.Consumer.HandleBasicCancelAsync(work.ConsumerTag!)
+                                        await work.Consumer.HandleBasicCancelAsync(work.ConsumerTag!, work.CancellationToken)
                                             .ConfigureAwait(false);
                                         break;
                                     case WorkType.CancelOk:
-                                        await work.Consumer.HandleBasicCancelOkAsync(work.ConsumerTag!)
+                                        await work.Consumer.HandleBasicCancelOkAsync(work.ConsumerTag!, work.CancellationToken)
                                             .ConfigureAwait(false);
                                         break;
                                     case WorkType.ConsumeOk:
-                                        await work.Consumer.HandleBasicConsumeOkAsync(work.ConsumerTag!)
+                                        await work.Consumer.HandleBasicConsumeOkAsync(work.ConsumerTag!, work.CancellationToken)
                                             .ConfigureAwait(false);
                                         break;
                                     case WorkType.Shutdown:
diff --git a/projects/RabbitMQ.Client/ConsumerDispatching/ConsumerDispatcherChannelBase.cs b/projects/RabbitMQ.Client/ConsumerDispatching/ConsumerDispatcherChannelBase.cs
@@ -46,6 +46,7 @@ internal abstract class ConsumerDispatcherChannelBase : ConsumerDispatcherBase,
         private readonly ushort _concurrency;
         private long _isQuiescing;
         private bool _disposed;
+        private readonly CancellationTokenSource _shutdownCts = new CancellationTokenSource();
 
         internal ConsumerDispatcherChannelBase(Impl.Channel channel, ushort concurrency)
         {
@@ -92,7 +93,7 @@ public async ValueTask HandleBasicConsumeOkAsync(IAsyncBasicConsumer consumer, s
                 try
                 {
                     AddConsumer(consumer, consumerTag);
-                    WorkStruct work = WorkStruct.CreateConsumeOk(consumer, consumerTag);
+                    WorkStruct work = WorkStruct.CreateConsumeOk(consumer, consumerTag, _shutdownCts);
                     await _writer.WriteAsync(work, cancellationToken)
                         .ConfigureAwait(false);
                 }
@@ -113,7 +114,7 @@ public async ValueTask HandleBasicDeliverAsync(string consumerTag, ulong deliver
             if (false == _disposed && false == IsQuiescing)
             {
                 IAsyncBasicConsumer consumer = GetConsumerOrDefault(consumerTag);
-                var work = WorkStruct.CreateDeliver(consumer, consumerTag, deliveryTag, redelivered, exchange, routingKey, basicProperties, body);
+                var work = WorkStruct.CreateDeliver(consumer, consumerTag, deliveryTag, redelivered, exchange, routingKey, basicProperties, body, _shutdownCts);
                 await _writer.WriteAsync(work, cancellationToken)
                     .ConfigureAwait(false);
             }
@@ -126,7 +127,7 @@ public async ValueTask HandleBasicCancelOkAsync(string consumerTag, Cancellation
             if (false == _disposed && false == IsQuiescing)
             {
                 IAsyncBasicConsumer consumer = GetAndRemoveConsumer(consumerTag);
-                WorkStruct work = WorkStruct.CreateCancelOk(consumer, consumerTag);
+                WorkStruct work = WorkStruct.CreateCancelOk(consumer, consumerTag, _shutdownCts);
                 await _writer.WriteAsync(work, cancellationToken)
                     .ConfigureAwait(false);
             }
@@ -139,7 +140,7 @@ public async ValueTask HandleBasicCancelAsync(string consumerTag, CancellationTo
             if (false == _disposed && false == IsQuiescing)
             {
                 IAsyncBasicConsumer consumer = GetAndRemoveConsumer(consumerTag);
-                WorkStruct work = WorkStruct.CreateCancel(consumer, consumerTag);
+                WorkStruct work = WorkStruct.CreateCancel(consumer, consumerTag, _shutdownCts);
                 await _writer.WriteAsync(work, cancellationToken)
                     .ConfigureAwait(false);
             }
@@ -148,6 +149,14 @@ await _writer.WriteAsync(work, cancellationToken)
         public void Quiesce()
         {
             Interlocked.Exchange(ref _isQuiescing, 1);
+            try
+            {
+                _shutdownCts.Cancel();
+            }
+            catch
+            {
+                // ignore
+            }
         }
 
         public async Task WaitForShutdownAsync()
@@ -214,7 +223,7 @@ protected bool IsQuiescing
 
         protected sealed override void ShutdownConsumer(IAsyncBasicConsumer consumer, ShutdownEventArgs reason)
         {
-            _writer.TryWrite(WorkStruct.CreateShutdown(consumer, reason));
+            _writer.TryWrite(WorkStruct.CreateShutdown(consumer, reason, _shutdownCts));
         }
 
         protected override Task InternalShutdownAsync()
@@ -237,25 +246,32 @@ protected override Task InternalShutdownAsync()
             public readonly RentedMemory Body;
             public readonly ShutdownEventArgs? Reason;
             public readonly WorkType WorkType;
+            public readonly CancellationToken CancellationToken;
+            private readonly CancellationTokenSource? _cancellationTokenSource;
 
-            private WorkStruct(WorkType type, IAsyncBasicConsumer consumer, string consumerTag)
+            private WorkStruct(WorkType type, IAsyncBasicConsumer consumer, string consumerTag, CancellationToken cancellationToken)
                 : this()
             {
                 WorkType = type;
                 Consumer = consumer;
                 ConsumerTag = consumerTag;
+                CancellationToken = cancellationToken;
+                _cancellationTokenSource = null;
             }
 
-            private WorkStruct(IAsyncBasicConsumer consumer, ShutdownEventArgs reason)
+            private WorkStruct(IAsyncBasicConsumer consumer, ShutdownEventArgs reason, CancellationTokenSource? cancellationTokenSource)
                 : this()
             {
                 WorkType = WorkType.Shutdown;
                 Consumer = consumer;
                 Reason = reason;
+                CancellationToken = cancellationTokenSource?.Token ?? CancellationToken.None;
+                this._cancellationTokenSource = cancellationTokenSource;
             }
 
             private WorkStruct(IAsyncBasicConsumer consumer, string consumerTag, ulong deliveryTag, bool redelivered,
-                string exchange, string routingKey, IReadOnlyBasicProperties basicProperties, RentedMemory body)
+                string exchange, string routingKey, IReadOnlyBasicProperties basicProperties, RentedMemory body,
+                CancellationToken cancellationToken)
             {
                 WorkType = WorkType.Deliver;
                 Consumer = consumer;
@@ -266,37 +282,62 @@ private WorkStruct(IAsyncBasicConsumer consumer, string consumerTag, ulong deliv
                 RoutingKey = routingKey;
                 BasicProperties = basicProperties;
                 Body = body;
-                Reason = default;
+                Reason = null;
+                CancellationToken = cancellationToken;
+                _cancellationTokenSource = null;
             }
 
-            public static WorkStruct CreateCancel(IAsyncBasicConsumer consumer, string consumerTag)
+            public static WorkStruct CreateCancel(IAsyncBasicConsumer consumer, string consumerTag, CancellationTokenSource cancellationTokenSource)
             {
-                return new WorkStruct(WorkType.Cancel, consumer, consumerTag);
+                return new WorkStruct(WorkType.Cancel, consumer, consumerTag, cancellationTokenSource.Token);
             }
 
-            public static WorkStruct CreateCancelOk(IAsyncBasicConsumer consumer, string consumerTag)
+            public static WorkStruct CreateCancelOk(IAsyncBasicConsumer consumer, string consumerTag, CancellationTokenSource cancellationTokenSource)
             {
-                return new WorkStruct(WorkType.CancelOk, consumer, consumerTag);
+                return new WorkStruct(WorkType.CancelOk, consumer, consumerTag, cancellationTokenSource.Token);
             }
 
-            public static WorkStruct CreateConsumeOk(IAsyncBasicConsumer consumer, string consumerTag)
+            public static WorkStruct CreateConsumeOk(IAsyncBasicConsumer consumer, string consumerTag, CancellationTokenSource cancellationTokenSource)
             {
-                return new WorkStruct(WorkType.ConsumeOk, consumer, consumerTag);
+                return new WorkStruct(WorkType.ConsumeOk, consumer, consumerTag, cancellationTokenSource.Token);
             }
 
-            public static WorkStruct CreateShutdown(IAsyncBasicConsumer consumer, ShutdownEventArgs reason)
+            public static WorkStruct CreateShutdown(IAsyncBasicConsumer consumer, ShutdownEventArgs reason, CancellationTokenSource cancellationTokenSource)
             {
-                return new WorkStruct(consumer, reason);
+                // Create a linked CTS so the shutdown args token reflects both dispatcher cancellation and any upstream token.
+                CancellationTokenSource? linked = null;
+                try
+                {
+                    if (reason.CancellationToken.CanBeCanceled)
+                    {
+                        linked = CancellationTokenSource.CreateLinkedTokenSource(cancellationTokenSource.Token, reason.CancellationToken);
+                    }
+                }
+                catch
+                {
+                    linked = null;
+                }
+
+                CancellationToken token = linked?.Token ?? cancellationTokenSource.Token;
+                ShutdownEventArgs argsWithToken = reason.Exception != null ?
+                    new ShutdownEventArgs(reason.Initiator, reason.ReplyCode, reason.ReplyText, reason.Exception, token) :
+                    new ShutdownEventArgs(reason.Initiator, reason.ReplyCode, reason.ReplyText, reason.ClassId, reason.MethodId, reason.Cause, token);
+
+                return new WorkStruct(consumer, argsWithToken, linked);
             }
 
             public static WorkStruct CreateDeliver(IAsyncBasicConsumer consumer, string consumerTag, ulong deliveryTag, bool redelivered,
-                string exchange, string routingKey, IReadOnlyBasicProperties basicProperties, RentedMemory body)
+                string exchange, string routingKey, IReadOnlyBasicProperties basicProperties, RentedMemory body, CancellationTokenSource cancellationTokenSource)
             {
                 return new WorkStruct(consumer, consumerTag, deliveryTag, redelivered,
-                    exchange, routingKey, basicProperties, body);
+                    exchange, routingKey, basicProperties, body, cancellationTokenSource.Token);
             }
 
-            public void Dispose() => Body.Dispose();
+            public void Dispose()
+            {
+                Body.Dispose();
+                _cancellationTokenSource?.Dispose();
+            }
         }
 
         protected enum WorkType : byte
@@ -317,6 +358,7 @@ protected virtual void Dispose(bool disposing)
                     if (disposing)
                     {
                         Quiesce();
+                        _shutdownCts.Dispose();
                     }
                 }
                 catch

Original file line number	Diff line number	Diff line change
`@@ -46,6 +46,7 @@ internal abstract class ConsumerDispatcherChannelBase : ConsumerDispatcherBase,`
`46`	`46`	`private readonly ushort _concurrency;`
`47`	`47`	`private long _isQuiescing;`
`48`	`48`	`private bool _disposed;`
	`49`	`+ private readonly CancellationTokenSource _shutdownCts = new CancellationTokenSource();`
`49`	`50`
`50`	`51`	`internal ConsumerDispatcherChannelBase(Impl.Channel channel, ushort concurrency)`
`51`	`52`	`{`
`@@ -92,7 +93,7 @@ public async ValueTask HandleBasicConsumeOkAsync(IAsyncBasicConsumer consumer, s`
`92`	`93`	`try`
`93`	`94`	`{`
`94`	`95`	`AddConsumer(consumer, consumerTag);`
`95`		`- WorkStruct work = WorkStruct.CreateConsumeOk(consumer, consumerTag);`
	`96`	`+ WorkStruct work = WorkStruct.CreateConsumeOk(consumer, consumerTag, _shutdownCts);`
`96`	`97`	`await _writer.WriteAsync(work, cancellationToken)`
`97`	`98`	`.ConfigureAwait(false);`
`98`	`99`	`}`
`@@ -113,7 +114,7 @@ public async ValueTask HandleBasicDeliverAsync(string consumerTag, ulong deliver`
`113`	`114`	`if (false == _disposed && false == IsQuiescing)`
`114`	`115`	`{`
`115`	`116`	`IAsyncBasicConsumer consumer = GetConsumerOrDefault(consumerTag);`
`116`		`- var work = WorkStruct.CreateDeliver(consumer, consumerTag, deliveryTag, redelivered, exchange, routingKey, basicProperties, body);`
	`117`	`+ var work = WorkStruct.CreateDeliver(consumer, consumerTag, deliveryTag, redelivered, exchange, routingKey, basicProperties, body, _shutdownCts);`
`117`	`118`	`await _writer.WriteAsync(work, cancellationToken)`
`118`	`119`	`.ConfigureAwait(false);`
`119`	`120`	`}`
`@@ -126,7 +127,7 @@ public async ValueTask HandleBasicCancelOkAsync(string consumerTag, Cancellation`
`126`	`127`	`if (false == _disposed && false == IsQuiescing)`
`127`	`128`	`{`
`128`	`129`	`IAsyncBasicConsumer consumer = GetAndRemoveConsumer(consumerTag);`
`129`		`- WorkStruct work = WorkStruct.CreateCancelOk(consumer, consumerTag);`
	`130`	`+ WorkStruct work = WorkStruct.CreateCancelOk(consumer, consumerTag, _shutdownCts);`
`130`	`131`	`await _writer.WriteAsync(work, cancellationToken)`
`131`	`132`	`.ConfigureAwait(false);`
`132`	`133`	`}`
`@@ -139,7 +140,7 @@ public async ValueTask HandleBasicCancelAsync(string consumerTag, CancellationTo`
`139`	`140`	`if (false == _disposed && false == IsQuiescing)`
`140`	`141`	`{`
`141`	`142`	`IAsyncBasicConsumer consumer = GetAndRemoveConsumer(consumerTag);`
`142`		`- WorkStruct work = WorkStruct.CreateCancel(consumer, consumerTag);`
	`143`	`+ WorkStruct work = WorkStruct.CreateCancel(consumer, consumerTag, _shutdownCts);`
`143`	`144`	`await _writer.WriteAsync(work, cancellationToken)`
`144`	`145`	`.ConfigureAwait(false);`
`145`	`146`	`}`
`@@ -148,6 +149,14 @@ await _writer.WriteAsync(work, cancellationToken)`
`148`	`149`	`public void Quiesce()`
`149`	`150`	`{`
`150`	`151`	`Interlocked.Exchange(ref _isQuiescing, 1);`
	`152`	`+ try`
	`153`	`+ {`
	`154`	`+ _shutdownCts.Cancel();`
	`155`	`+ }`
	`156`	`+ catch`
	`157`	`+ {`
	`158`	`+ // ignore`
	`159`	`+ }`
`151`	`160`	`}`
`152`	`161`
`153`	`162`	`public async Task WaitForShutdownAsync()`
`@@ -214,7 +223,7 @@ protected bool IsQuiescing`
`214`	`223`
`215`	`224`	`protected sealed override void ShutdownConsumer(IAsyncBasicConsumer consumer, ShutdownEventArgs reason)`
`216`	`225`	`{`
`217`		`- _writer.TryWrite(WorkStruct.CreateShutdown(consumer, reason));`
	`226`	`+ _writer.TryWrite(WorkStruct.CreateShutdown(consumer, reason, _shutdownCts));`
`218`	`227`	`}`
`219`	`228`
`220`	`229`	`protected override Task InternalShutdownAsync()`
`@@ -237,25 +246,32 @@ protected override Task InternalShutdownAsync()`
`237`	`246`	`public readonly RentedMemory Body;`
`238`	`247`	`public readonly ShutdownEventArgs? Reason;`
`239`	`248`	`public readonly WorkType WorkType;`
	`249`	`+ public readonly CancellationToken CancellationToken;`
	`250`	`+ private readonly CancellationTokenSource? _cancellationTokenSource;`
`240`	`251`
`241`		`- private WorkStruct(WorkType type, IAsyncBasicConsumer consumer, string consumerTag)`
	`252`	`+ private WorkStruct(WorkType type, IAsyncBasicConsumer consumer, string consumerTag, CancellationToken cancellationToken)`
`242`	`253`	`: this()`
`243`	`254`	`{`
`244`	`255`	`WorkType = type;`
`245`	`256`	`Consumer = consumer;`
`246`	`257`	`ConsumerTag = consumerTag;`
	`258`	`+ CancellationToken = cancellationToken;`
	`259`	`+ _cancellationTokenSource = null;`
`247`	`260`	`}`
`248`	`261`
`249`		`- private WorkStruct(IAsyncBasicConsumer consumer, ShutdownEventArgs reason)`
	`262`	`+ private WorkStruct(IAsyncBasicConsumer consumer, ShutdownEventArgs reason, CancellationTokenSource? cancellationTokenSource)`
`250`	`263`	`: this()`
`251`	`264`	`{`
`252`	`265`	`WorkType = WorkType.Shutdown;`
`253`	`266`	`Consumer = consumer;`
`254`	`267`	`Reason = reason;`
	`268`	`+ CancellationToken = cancellationTokenSource?.Token ?? CancellationToken.None;`
	`269`	`+ this._cancellationTokenSource = cancellationTokenSource;`
`255`	`270`	`}`
`256`	`271`
`257`	`272`	`private WorkStruct(IAsyncBasicConsumer consumer, string consumerTag, ulong deliveryTag, bool redelivered,`
`258`		`- string exchange, string routingKey, IReadOnlyBasicProperties basicProperties, RentedMemory body)`
	`273`	`+ string exchange, string routingKey, IReadOnlyBasicProperties basicProperties, RentedMemory body,`
	`274`	`+ CancellationToken cancellationToken)`
`259`	`275`	`{`
`260`	`276`	`WorkType = WorkType.Deliver;`
`261`	`277`	`Consumer = consumer;`
`@@ -266,37 +282,62 @@ private WorkStruct(IAsyncBasicConsumer consumer, string consumerTag, ulong deliv`
`266`	`282`	`RoutingKey = routingKey;`
`267`	`283`	`BasicProperties = basicProperties;`
`268`	`284`	`Body = body;`
`269`		`- Reason = default;`
	`285`	`+ Reason = null;`
	`286`	`+ CancellationToken = cancellationToken;`
	`287`	`+ _cancellationTokenSource = null;`
`270`	`288`	`}`
`271`	`289`
`272`		`- public static WorkStruct CreateCancel(IAsyncBasicConsumer consumer, string consumerTag)`
	`290`	`+ public static WorkStruct CreateCancel(IAsyncBasicConsumer consumer, string consumerTag, CancellationTokenSource cancellationTokenSource)`
`273`	`291`	`{`
`274`		`- return new WorkStruct(WorkType.Cancel, consumer, consumerTag);`
	`292`	`+ return new WorkStruct(WorkType.Cancel, consumer, consumerTag, cancellationTokenSource.Token);`
`275`	`293`	`}`
`276`	`294`
`277`		`- public static WorkStruct CreateCancelOk(IAsyncBasicConsumer consumer, string consumerTag)`
	`295`	`+ public static WorkStruct CreateCancelOk(IAsyncBasicConsumer consumer, string consumerTag, CancellationTokenSource cancellationTokenSource)`
`278`	`296`	`{`
`279`		`- return new WorkStruct(WorkType.CancelOk, consumer, consumerTag);`
	`297`	`+ return new WorkStruct(WorkType.CancelOk, consumer, consumerTag, cancellationTokenSource.Token);`
`280`	`298`	`}`
`281`	`299`
`282`		`- public static WorkStruct CreateConsumeOk(IAsyncBasicConsumer consumer, string consumerTag)`
	`300`	`+ public static WorkStruct CreateConsumeOk(IAsyncBasicConsumer consumer, string consumerTag, CancellationTokenSource cancellationTokenSource)`
`283`	`301`	`{`
`284`		`- return new WorkStruct(WorkType.ConsumeOk, consumer, consumerTag);`
	`302`	`+ return new WorkStruct(WorkType.ConsumeOk, consumer, consumerTag, cancellationTokenSource.Token);`
`285`	`303`	`}`
`286`	`304`
`287`		`- public static WorkStruct CreateShutdown(IAsyncBasicConsumer consumer, ShutdownEventArgs reason)`
	`305`	`+ public static WorkStruct CreateShutdown(IAsyncBasicConsumer consumer, ShutdownEventArgs reason, CancellationTokenSource cancellationTokenSource)`
`288`	`306`	`{`
`289`		`- return new WorkStruct(consumer, reason);`
	`307`	`+ // Create a linked CTS so the shutdown args token reflects both dispatcher cancellation and any upstream token.`
	`308`	`+ CancellationTokenSource? linked = null;`
	`309`	`+ try`
	`310`	`+ {`
	`311`	`+ if (reason.CancellationToken.CanBeCanceled)`
	`312`	`+ {`
	`313`	`+ linked = CancellationTokenSource.CreateLinkedTokenSource(cancellationTokenSource.Token, reason.CancellationToken);`
	`314`	`+ }`
	`315`	`+ }`
	`316`	`+ catch`
	`317`	`+ {`
	`318`	`+ linked = null;`
	`319`	`+ }`
	`320`	`+`
	`321`	`+ CancellationToken token = linked?.Token ?? cancellationTokenSource.Token;`
	`322`	`+ ShutdownEventArgs argsWithToken = reason.Exception != null ?`
	`323`	`+ new ShutdownEventArgs(reason.Initiator, reason.ReplyCode, reason.ReplyText, reason.Exception, token) :`
	`324`	`+ new ShutdownEventArgs(reason.Initiator, reason.ReplyCode, reason.ReplyText, reason.ClassId, reason.MethodId, reason.Cause, token);`
	`325`	`+`
	`326`	`+ return new WorkStruct(consumer, argsWithToken, linked);`
`290`	`327`	`}`
`291`	`328`
`292`	`329`	`public static WorkStruct CreateDeliver(IAsyncBasicConsumer consumer, string consumerTag, ulong deliveryTag, bool redelivered,`
`293`		`- string exchange, string routingKey, IReadOnlyBasicProperties basicProperties, RentedMemory body)`
	`330`	`+ string exchange, string routingKey, IReadOnlyBasicProperties basicProperties, RentedMemory body, CancellationTokenSource cancellationTokenSource)`
`294`	`331`	`{`
`295`	`332`	`return new WorkStruct(consumer, consumerTag, deliveryTag, redelivered,`
`296`		`- exchange, routingKey, basicProperties, body);`
	`333`	`+ exchange, routingKey, basicProperties, body, cancellationTokenSource.Token);`
`297`	`334`	`}`
`298`	`335`
`299`		`- public void Dispose() => Body.Dispose();`
	`336`	`+ public void Dispose()`
	`337`	`+ {`
	`338`	`+ Body.Dispose();`
	`339`	`+ _cancellationTokenSource?.Dispose();`
	`340`	`+ }`
`300`	`341`	`}`
`301`	`342`
`302`	`343`	`protected enum WorkType : byte`
`@@ -317,6 +358,7 @@ protected virtual void Dispose(bool disposing)`
`317`	`358`	`if (disposing)`
`318`	`359`	`{`
`319`	`360`	`Quiesce();`
	`361`	`+ _shutdownCts.Dispose();`
`320`	`362`	`}`
`321`	`363`	`}`
`322`	`364`	`catch`