Cleanup from cancellation for RabbitMQ

Author: phatboyg

src/Transports/MassTransit.RabbitMqTransport/RabbitMqTransport/Middleware/ConfigureRabbitMqTopologyFilter.cs

@@ -40,6 +40,13 @@ public async Task Send(ChannelContext context, IPipe<ChannelContext> next)
         }
     }
 
+    public void Probe(ProbeContext context)
+    {
+        var scope = context.CreateFilterScope("configureTopology");
+
+        _brokerTopology.Probe(scope);
+    }
+
     public async Task<OneTimeContext<ConfigureTopologyContext<TSettings>>> Configure(ChannelContext context, CancellationToken cancellationToken)
     {
         return await context.OneTimeSetup<ConfigureTopologyContext<TSettings>>(() =>
@@ -49,13 +56,6 @@ public async Task<OneTimeContext<ConfigureTopologyContext<TSettings>>> Configure
         }).ConfigureAwait(false);
     }
 
-    public void Probe(ProbeContext context)
-    {
-        var scope = context.CreateFilterScope("configureTopology");
-
-        _brokerTopology.Probe(scope);
-    }
-
     async Task ConfigureTopology(ChannelContext context, CancellationToken cancellationToken)
     {
         await Task.WhenAll(_brokerTopology.Queues.Select(queue => Declare(context, queue, cancellationToken))).ConfigureAwait(false);
@@ -71,7 +71,8 @@ static Task Declare(ChannelContext context, Exchange exchange, CancellationToken
     {
         RabbitMqLogMessages.DeclareExchange(exchange);
 
-        return context.ExchangeDeclare(exchange.ExchangeName, exchange.ExchangeType, exchange.Durable, exchange.AutoDelete, exchange.ExchangeArguments, cancellationToken);
+        return context.ExchangeDeclare(exchange.ExchangeName, exchange.ExchangeType, exchange.Durable, exchange.AutoDelete, exchange.ExchangeArguments,
+            cancellationToken);
     }
 
     static async Task Declare(ChannelContext context, Queue queue, CancellationToken cancellationToken)
@@ -103,6 +104,7 @@ static async Task Bind(ChannelContext context, ExchangeToQueueBinding binding, C
     {
         RabbitMqLogMessages.BindToQueue(binding);
 
-        await context.QueueBind(binding.Destination.QueueName, binding.Source.ExchangeName, binding.RoutingKey, binding.Arguments, cancellationToken).ConfigureAwait(false);
+        await context.QueueBind(binding.Destination.QueueName, binding.Source.ExchangeName, binding.RoutingKey, binding.Arguments, cancellationToken)
+            .ConfigureAwait(false);
     }
 }

src/Transports/MassTransit.RabbitMqTransport/RabbitMqTransport/RabbitMqChannelContext.cs

@@ -27,13 +27,11 @@ public RabbitMqChannelContext(ConnectionContext connectionContext, IChannel chan
         {
             ConnectionContext = connectionContext;
 
-            _cancellationToken = cancellationToken;
-            _tokenSource = CancellationTokenSource.CreateLinkedTokenSource(connectionContext.CancellationToken, cancellationToken);
-
             _channel = channel;
             _agent = agent;
 
-            _channel.ContinuationTimeout = ConnectionContext.ContinuationTimeout;
+            _cancellationToken = cancellationToken;
+            _tokenSource = CancellationTokenSource.CreateLinkedTokenSource(connectionContext.CancellationToken, cancellationToken);
         }
 
         public override CancellationToken CancellationToken => _tokenSource?.Token ?? _cancellationToken;
@@ -162,9 +160,8 @@ public async ValueTask DisposeAsync()
 
             await _channel.Cleanup(200, message, CancellationToken).ConfigureAwait(false);
 
-            var tokenSource = _tokenSource;
+            _tokenSource?.Dispose();
             _tokenSource = null;
-            tokenSource.Dispose();
         }
     }
 }

src/Transports/MassTransit.RabbitMqTransport/RabbitMqTransport/RabbitMqConnectionContext.cs

@@ -49,7 +49,11 @@ public async Task<IChannel> CreateChannel(ushort? concurrentMessageLimit, Cancel
         {
             var options = new CreateChannelOptions(PublisherConfirmation, PublisherConfirmation, consumerDispatchConcurrency: concurrentMessageLimit);
 
-            return await Connection.CreateChannelAsync(options, cancellationToken).ConfigureAwait(false);
+            var channel = await Connection.CreateChannelAsync(options, cancellationToken).ConfigureAwait(false);
+
+            channel.ContinuationTimeout = ContinuationTimeout;
+
+            return channel;
         }
 
         public async Task<ChannelContext> CreateChannelContext(IAgent agent, ushort? concurrentMessageLimit, CancellationToken cancellationToken)

src/Transports/MassTransit.RabbitMqTransport/RabbitMqTransport/ScopeChannelContext.cs

@@ -135,9 +135,8 @@ public void NotifyFaulted(Exception exception, Uri contextInputAddress)
 
         public void Dispose()
         {
-            var tokenSource = _tokenSource;
+            _tokenSource?.Dispose();
             _tokenSource = null;
-            tokenSource.Dispose();
         }
     }
 }

src/Transports/MassTransit.RabbitMqTransport/RabbitMqTransport/SharedChannelContext.cs

@@ -135,9 +135,8 @@ public void NotifyFaulted(Exception exception, Uri contextInputAddress)
 
         public void Dispose()
         {
-            var tokenSource = _tokenSource;
+            _tokenSource?.Dispose();
             _tokenSource = null;
-            tokenSource.Dispose();
         }
     }
 }

src/Transports/MassTransit.RabbitMqTransport/RabbitMqTransport/SharedConnectionContext.cs

@@ -22,6 +22,7 @@ public SharedConnectionContext(ConnectionContext context, CancellationToken canc
         {
             _context = context;
             _cancellationToken = cancellationToken;
+
             _tokenSource = CancellationTokenSource.CreateLinkedTokenSource(context.CancellationToken, cancellationToken);
         }
 
@@ -53,9 +54,8 @@ public async Task<ChannelContext> CreateChannelContext(IAgent agent, ushort? con
 
         public void Dispose()
         {
-            var tokenSource = _tokenSource;
+            _tokenSource?.Dispose();
             _tokenSource = null;
-            tokenSource.Dispose();
         }
     }
 }