Link cancellation tokens to disposal for prompt shutdown in MessageBus and Queue

- Add GetLinkedDisposableCancellationTokenSource to MessageBusBase so publish operations cancel when the bus is disposed mid-flight

- Use linked tokens in QueueBase.DequeueAsync, GetDeadletterItemsAsync, and StartWorkingAsync for EnsureQueueCreatedAsync calls

- Fix subscriber Action lambda to use its token parameter instead of a closure capture, preventing a subtle token mismatch bug

- Add DisposedCancellationToken check in SendMessageToSubscribersAsync to avoid dispatching to subscribers after the bus is disposed

- Pass cancellationToken to FolderFileStorage lock acquisitions so callers can cancel while waiting for the lock

- Remove unused DelayedMessage nested class from MessageBusBase

- Minor cleanups: TryGetValue, GetValueOrDefault, guard clause inversion

Co-authored-by: Cursor <cursoragent@cursor.com>

Author: niemyjski

src/Foundatio/Messaging/InMemoryMessageBus.cs

@@ -1,5 +1,6 @@
 using System;
 using System.Collections.Concurrent;
+using System.Collections.Generic;
 using System.Threading;
 using System.Threading.Tasks;
 using Foundatio.Utility;
@@ -23,12 +24,12 @@ public InMemoryMessageBus(Builder<InMemoryMessageBusOptionsBuilder, InMemoryMess
 
     public long GetMessagesSent(Type messageType)
     {
-        return _messageCounts.TryGetValue(GetMappedMessageType(messageType), out long count) ? count : 0;
+        return _messageCounts.GetValueOrDefault(GetMappedMessageType(messageType), 0);
     }
 
     public long GetMessagesSent<T>()
     {
-        return _messageCounts.TryGetValue(GetMappedMessageType(typeof(T)), out long count) ? count : 0;
+        return _messageCounts.GetValueOrDefault(GetMappedMessageType(typeof(T)), 0);
     }
 
     public void ResetMessagesSent()

src/Foundatio/Messaging/MessageBusBase.cs

@@ -49,6 +49,17 @@ public MessageBusBase(TOptions options)
     /// </summary>
     protected CancellationToken DisposedCancellationToken => _disposedCancellationTokenSource.Token;
 
+    /// <summary>
+    /// Creates a linked cancellation token source that combines the provided token with the disposal token.
+    /// This allows operations to be cancelled by either the caller or when this instance is disposed.
+    /// </summary>
+    /// <param name="cancellationToken">The caller's cancellation token to link.</param>
+    /// <returns>A new <see cref="CancellationTokenSource"/> that should be disposed by the caller.</returns>
+    protected CancellationTokenSource GetLinkedDisposableCancellationTokenSource(CancellationToken cancellationToken)
+    {
+        return CancellationTokenSource.CreateLinkedTokenSource(cancellationToken, DisposedCancellationToken);
+    }
+
     ILogger IHaveLogger.Logger => _logger;
     ILoggerFactory IHaveLoggerFactory.LoggerFactory => _loggerFactory;
     TimeProvider IHaveTimeProvider.TimeProvider => _timeProvider;
@@ -73,10 +84,11 @@ public async Task PublishAsync(Type messageType, object message, MessageOptions
                 options.Properties.Add("TraceState", Activity.Current.TraceStateString);
         }
 
+        using var linkedCancellationTokenSource = GetLinkedDisposableCancellationTokenSource(cancellationToken);
         try
         {
-            await EnsureTopicCreatedAsync(cancellationToken).AnyContext();
-            await PublishImplAsync(GetMappedMessageType(messageType), message, options, cancellationToken).AnyContext();
+            await EnsureTopicCreatedAsync(linkedCancellationTokenSource.Token).AnyContext();
+            await PublishImplAsync(GetMappedMessageType(messageType), message, options, linkedCancellationTokenSource.Token).AnyContext();
         }
         catch (Exception ex) when (ex is not OperationCanceledException and not MessageBusException)
         {
@@ -105,8 +117,8 @@ protected virtual Type GetMappedMessageType(string messageType)
 
         return _knownMessageTypesCache.GetOrAdd(messageType, type =>
         {
-            if (_options.MessageTypeMappings != null && _options.MessageTypeMappings.ContainsKey(type))
-                return _options.MessageTypeMappings[type];
+            if (_options.MessageTypeMappings != null && _options.MessageTypeMappings.TryGetValue(type, out Type typeMapping))
+                return typeMapping;
 
             try
             {
@@ -126,7 +138,6 @@ protected virtual Type GetMappedMessageType(string messageType)
                 catch (Exception ex)
                 {
                     _logger.LogError(ex, "Error getting message body type: {MessageType}", type);
-
                     return null;
                 }
             }
@@ -145,7 +156,7 @@ protected virtual Task SubscribeImplAsync<T>(Func<T, CancellationToken, Task> ha
             Action = (message, token) =>
             {
                 if (message is T typedMessage)
-                    return handler(typedMessage, cancellationToken);
+                    return handler(typedMessage, token);
 
                 _logger.LogTrace("Unable to call subscriber action: {MessageType} cannot be safely casted to {SubscriberType}", message.GetType(), typeof(T));
                 return Task.CompletedTask;
@@ -157,11 +168,11 @@ protected virtual Task SubscribeImplAsync<T>(Func<T, CancellationToken, Task> ha
             cancellationToken.Register(() =>
             {
                 _subscribers.TryRemove(subscriber.Id, out _);
-                if (_subscribers.Count == 0)
-                {
-                    _logger.LogDebug("Removing topic subscription for {MessageBusId}: No subscribers", MessageBusId);
-                    RemoveTopicSubscriptionAsync().GetAwaiter().GetResult();
-                }
+                if (_subscribers.Count > 0)
+                    return;
+
+                _logger.LogDebug("Removing topic subscription for {MessageBusId}: No subscribers", MessageBusId);
+                RemoveTopicSubscriptionAsync().GetAwaiter().GetResult();
             });
         }
 
@@ -265,7 +276,7 @@ protected async Task SendMessageToSubscribersAsync(IMessage message)
 
             return Task.Run(async () =>
             {
-                if (subscriber.CancellationToken.IsCancellationRequested)
+                if (DisposedCancellationToken.IsCancellationRequested || subscriber.CancellationToken.IsCancellationRequested)
                 {
                     _logger.LogTrace("The cancelled subscriber action will not be called: {SubscriberId}", subscriber.Id);
                     return;
@@ -422,14 +433,6 @@ public virtual void Dispose()
         _disposedCancellationTokenSource.Dispose();
     }
 
-    [DebuggerDisplay("MessageType: {MessageType} SendTime: {SendTime} Message: {Message}")]
-    protected class DelayedMessage
-    {
-        public DateTime SendTime { get; set; }
-        public Type MessageType { get; set; }
-        public object Message { get; set; }
-    }
-
     [DebuggerDisplay("Id: {Id} Type: {Type} CancellationToken: {CancellationToken}")]
     protected class Subscriber
     {

src/Foundatio/Queues/QueueBase.cs

@@ -139,10 +139,10 @@ public async Task<string> EnqueueAsync(T data, QueueEntryOptions options = null)
     protected abstract Task<IQueueEntry<T>> DequeueImplAsync(CancellationToken linkedCancellationToken);
     public async Task<IQueueEntry<T>> DequeueAsync(CancellationToken cancellationToken)
     {
-        await EnsureQueueCreatedAsync(DisposedCancellationToken).AnyContext();
+        using var linkedCancellationTokenSource = GetLinkedDisposableCancellationTokenSource(cancellationToken);
+        await EnsureQueueCreatedAsync(linkedCancellationTokenSource.Token).AnyContext();
 
         LastDequeueActivity = _timeProvider.GetUtcNow();
-        using var linkedCancellationTokenSource = GetLinkedDisposableCancellationTokenSource(cancellationToken);
         return await DequeueImplAsync(linkedCancellationTokenSource.Token).AnyContext();
     }
 
@@ -161,8 +161,9 @@ public virtual async Task<IQueueEntry<T>> DequeueAsync(TimeSpan? timeout = null)
     protected abstract Task<IEnumerable<T>> GetDeadletterItemsImplAsync(CancellationToken cancellationToken);
     public async Task<IEnumerable<T>> GetDeadletterItemsAsync(CancellationToken cancellationToken = default)
     {
-        await EnsureQueueCreatedAsync(DisposedCancellationToken).AnyContext();
-        return await GetDeadletterItemsImplAsync(cancellationToken).AnyContext();
+        using var linkedCancellationTokenSource = GetLinkedDisposableCancellationTokenSource(cancellationToken);
+        await EnsureQueueCreatedAsync(linkedCancellationTokenSource.Token).AnyContext();
+        return await GetDeadletterItemsImplAsync(linkedCancellationTokenSource.Token).AnyContext();
     }
 
     protected abstract Task<QueueStats> GetQueueStatsImplAsync();
@@ -191,7 +192,8 @@ public async Task DeleteQueueAsync()
     protected abstract void StartWorkingImpl(Func<IQueueEntry<T>, CancellationToken, Task> handler, bool autoComplete, CancellationToken cancellationToken);
     public async Task StartWorkingAsync(Func<IQueueEntry<T>, CancellationToken, Task> handler, bool autoComplete = false, CancellationToken cancellationToken = default)
     {
-        await EnsureQueueCreatedAsync(DisposedCancellationToken).AnyContext();
+        using var linkedCancellationTokenSource = GetLinkedDisposableCancellationTokenSource(cancellationToken);
+        await EnsureQueueCreatedAsync(linkedCancellationTokenSource.Token).AnyContext();
         StartWorkingImpl(handler, autoComplete, cancellationToken);
     }
 

src/Foundatio/Storage/FolderFileStorage.cs

@@ -151,7 +151,7 @@ public async Task<bool> RenameFileAsync(string path, string newPath, Cancellatio
 
         try
         {
-            using (await _lock.LockAsync().AnyContext())
+            using (await _lock.LockAsync(cancellationToken).AnyContext())
             {
                 string directory = Path.GetDirectoryName(normalizedNewPath);
                 if (directory != null)
@@ -198,7 +198,7 @@ public async Task<bool> CopyFileAsync(string path, string targetPath, Cancellati
 
         try
         {
-            using (await _lock.LockAsync().AnyContext())
+            using (await _lock.LockAsync(cancellationToken).AnyContext())
             {
                 string directory = Path.GetDirectoryName(normalizedTargetPath);
                 if (directory != null)

src/Foundatio/Utility/FoundatioDiagnostics.cs

@@ -69,4 +69,4 @@ public static Activity AddException(this Activity activity, Exception exception,
         return activity.AddEvent(new ActivityEvent(ExceptionEventName, timestamp, exceptionTags));
     }
 }
-#endif
+#endif