Skip to content

Commit 38e89a9

Browse files
niemyjskiniemyjski
committed
Fixes dispose cancellation token usage
Ensures that cancellation tokens are properly used and disposed of in the HybridCacheClient, InMemoryCacheClient and MessageBusBase classes to prevent issues with orphaned tasks and resource leaks when the cache or message bus is disposed. Adds dispose check to prevent double dispose.
1 parent 564a857 commit 38e89a9

6 files changed

Lines changed: 66 additions & 35 deletions

File tree

‎src/Foundatio/Caching/HybridCacheClient.cs‎

Lines changed: 10 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -26,8 +26,10 @@ public class HybridCacheClient : IHybridCacheClient, IHaveTimeProvider, IHaveLog
2626
private readonly ILoggerFactory _loggerFactory;
2727
private readonly TimeProvider _timeProvider;
2828
private readonly IResiliencePolicyProvider _resiliencePolicyProvider;
29+
private readonly CancellationTokenSource _disposedCancellationTokenSource = new();
2930
private long _localCacheHits;
3031
private long _invalidateCacheCalls;
32+
private bool _isDisposed;
3133

3234
public HybridCacheClient(ICacheClient distributedCacheClient, IMessageBus messageBus, InMemoryCacheClientOptions localCacheOptions = null, ILoggerFactory loggerFactory = null)
3335
{
@@ -37,7 +39,7 @@ public HybridCacheClient(ICacheClient distributedCacheClient, IMessageBus messag
3739
_resiliencePolicyProvider = distributedCacheClient.GetResiliencePolicyProvider() ?? localCacheOptions?.ResiliencePolicyProvider;
3840
_distributedCache = distributedCacheClient;
3941
_messageBus = messageBus;
40-
_messageBus.SubscribeAsync<InvalidateCache>(OnRemoteCacheItemExpiredAsync).AnyContext().GetAwaiter().GetResult();
42+
_messageBus.SubscribeAsync<InvalidateCache>(OnRemoteCacheItemExpiredAsync, _disposedCancellationTokenSource.Token).AnyContext().GetAwaiter().GetResult();
4143
localCacheOptions ??= new InMemoryCacheClientOptions
4244
{
4345
TimeProvider = _timeProvider,
@@ -280,7 +282,7 @@ public async Task<bool> SetAsync<T>(string key, T value, TimeSpan? expiresIn = n
280282
public async Task<int> SetAllAsync<T>(IDictionary<string, T> values, TimeSpan? expiresIn = null)
281283
{
282284
ArgumentNullException.ThrowIfNull(values);
283-
285+
284286
if (values.Count is 0)
285287
return 0;
286288

@@ -738,9 +740,13 @@ public async Task<CacheValue<ICollection<T>>> GetListAsync<T>(string key, int? p
738740

739741
public virtual void Dispose()
740742
{
741-
_localCache.Dispose();
743+
if (_isDisposed)
744+
return;
742745

743-
// TODO: unsubscribe handler from messagebus.
746+
_isDisposed = true;
747+
_disposedCancellationTokenSource.Cancel();
748+
_disposedCancellationTokenSource.Dispose();
749+
_localCache.Dispose();
744750
}
745751

746752
public class InvalidateCache

‎src/Foundatio/Caching/InMemoryCacheClient.cs‎

Lines changed: 20 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,8 @@ public class InMemoryCacheClient : IMemoryCacheClient, IHaveTimeProvider, IHaveL
3636
private readonly ILogger _logger;
3737
private readonly ILoggerFactory _loggerFactory;
3838
private readonly AsyncLock _lock = new();
39+
private readonly CancellationTokenSource _disposedCancellationTokenSource = new();
40+
private bool _isDisposed;
3941

4042
public InMemoryCacheClient() : this(o => o)
4143
{
@@ -1058,7 +1060,7 @@ private async Task<bool> SetInternalAsync(string key, CacheEntry entry, bool add
10581060
public async Task<int> SetAllAsync<T>(IDictionary<string, T> values, TimeSpan? expiresIn = null)
10591061
{
10601062
ArgumentNullException.ThrowIfNull(values);
1061-
1063+
10621064
if (values.Count is 0)
10631065
return 0;
10641066

@@ -1421,18 +1423,21 @@ private async Task StartMaintenanceAsync(bool compactImmediately = false)
14211423
{
14221424
_logger.LogTrace("StartMaintenanceAsync called with compactImmediately={CompactImmediately}", compactImmediately);
14231425

1426+
if (_disposedCancellationTokenSource.IsCancellationRequested)
1427+
return;
1428+
14241429
var utcNow = _timeProvider.GetUtcNow().UtcDateTime;
14251430
if (compactImmediately)
14261431
await CompactAsync().AnyContext();
14271432

14281433
if (TimeSpan.FromMilliseconds(250) < utcNow - _lastMaintenance)
14291434
{
14301435
_lastMaintenance = utcNow;
1431-
_ = Task.Run(DoMaintenanceAsync);
1436+
_ = Task.Run(DoMaintenanceAsync, _disposedCancellationTokenSource.Token);
14321437
}
14331438
}
14341439

1435-
private bool ShouldCompact => (_maxItems.HasValue && _memory.Count > _maxItems) || (_shouldTrackMemory && _currentMemorySize > _maxMemorySize);
1440+
private bool ShouldCompact => !_disposedCancellationTokenSource.IsCancellationRequested && ((_maxItems.HasValue && _memory.Count > _maxItems) || (_shouldTrackMemory && _currentMemorySize > _maxMemorySize));
14361441

14371442
private async Task CompactAsync()
14381443
{
@@ -1445,7 +1450,7 @@ private async Task CompactAsync()
14451450
_logger.LogTrace("CompactAsync: Compacting cache");
14461451

14471452
var expiredKeys = new List<string>();
1448-
using (await _lock.LockAsync().AnyContext())
1453+
using (await _lock.LockAsync(_disposedCancellationTokenSource.Token).AnyContext())
14491454
{
14501455
int removalCount = 0;
14511456

@@ -1455,11 +1460,11 @@ private async Task CompactAsync()
14551460
const int absoluteMaxRemovals = 1000;
14561461

14571462
int itemOverLimitFactor = 1;
1458-
if (_maxItems.HasValue && _maxItems.Value > 0 && _memory.Count > _maxItems.Value)
1463+
if (_maxItems is > 0 && _memory.Count > _maxItems.Value)
14591464
itemOverLimitFactor = (int)Math.Ceiling((double)_memory.Count / _maxItems.Value);
14601465

14611466
int memoryOverLimitFactor = 1;
1462-
if (_shouldTrackMemory && _maxMemorySize.HasValue && _maxMemorySize.Value > 0 && _currentMemorySize > _maxMemorySize.Value)
1467+
if (_shouldTrackMemory && _maxMemorySize is > 0 && _currentMemorySize > _maxMemorySize.Value)
14631468
memoryOverLimitFactor = (int)Math.Ceiling((double)_currentMemorySize / _maxMemorySize.Value);
14641469

14651470
int overLimitFactor = Math.Max(itemOverLimitFactor, memoryOverLimitFactor);
@@ -1515,6 +1520,9 @@ private async Task CompactAsync()
15151520
}
15161521
}
15171522

1523+
if (_disposedCancellationTokenSource.IsCancellationRequested)
1524+
return;
1525+
15181526
// Notify about expired items
15191527
foreach (string expiredKey in expiredKeys)
15201528
OnItemExpired(expiredKey);
@@ -1661,7 +1669,13 @@ private async Task DoMaintenanceAsync()
16611669

16621670
public virtual void Dispose()
16631671
{
1672+
if (_isDisposed)
1673+
return;
1674+
1675+
_isDisposed = true;
16641676
_memory.Clear();
1677+
_disposedCancellationTokenSource.Cancel();
1678+
_disposedCancellationTokenSource.Dispose();
16651679
ItemExpired?.Dispose();
16661680
_sizeCalculator = null; // Allow GC to collect any captured closures
16671681
}

‎src/Foundatio/Messaging/MessageBusBase.cs‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
using System;
1+
using System;
22
using System.Collections.Concurrent;
33
using System.Collections.Generic;
44
using System.Diagnostics;
@@ -358,7 +358,7 @@ protected void SendDelayedMessage(Type messageType, object message, TimeSpan del
358358

359359
_logger.LogTrace("Sending delayed message scheduled for {SendTime:O} for type {MessageType}", sendTime, messageType);
360360
await PublishAsync(messageType, message).AnyContext();
361-
});
361+
}, _messageBusDisposedCancellationTokenSource.Token);
362362
}
363363

364364
public string MessageBusId { get; protected set; }

‎src/Foundatio/Queues/InMemoryQueue.cs‎

Lines changed: 13 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -107,7 +107,7 @@ protected override async Task<string> EnqueueImplAsync(T data, QueueEntryOptions
107107

108108
await OnEnqueuedAsync(entry).AnyContext();
109109
_logger.LogTrace("Enqueue done");
110-
}, _timeProvider, _queueDisposedCancellationTokenSource.Token);
110+
}, _timeProvider, DisposedCancellationToken);
111111
return id;
112112
}
113113

@@ -131,31 +131,32 @@ protected override void StartWorkingImpl(Func<IQueueEntry<T>, CancellationToken,
131131

132132
_logger.LogTrace("Queue {QueueName} start working", _options.Name);
133133

134+
var linkedCancellationTokenSource = GetLinkedDisposableCancellationTokenSource(cancellationToken);
134135
_workers.Add(Task.Run(async () =>
135136
{
136-
using var linkedCancellationToken = GetLinkedDisposableCancellationTokenSource(cancellationToken);
137+
using var _ = new DisposableAction(linkedCancellationTokenSource.Dispose);
137138
_logger.LogTrace("WorkerLoop Start {QueueName}", _options.Name);
138139

139-
while (!linkedCancellationToken.IsCancellationRequested)
140+
while (!linkedCancellationTokenSource.IsCancellationRequested)
140141
{
141142
_logger.LogTrace("WorkerLoop Signaled {QueueName}", _options.Name);
142143

143144
IQueueEntry<T> queueEntry = null;
144145
try
145146
{
146-
queueEntry = await DequeueImplAsync(linkedCancellationToken.Token).AnyContext();
147+
queueEntry = await DequeueImplAsync(linkedCancellationTokenSource.Token).AnyContext();
147148
}
148149
catch (Exception ex)
149150
{
150151
_logger.LogError(ex, "Error on Dequeue: {Message}", ex.Message);
151152
}
152153

153-
if (linkedCancellationToken.IsCancellationRequested || queueEntry == null)
154+
if (linkedCancellationTokenSource.IsCancellationRequested || queueEntry == null)
154155
return;
155156

156157
try
157158
{
158-
await handler(queueEntry, linkedCancellationToken.Token).AnyContext();
159+
await handler(queueEntry, linkedCancellationTokenSource.Token).AnyContext();
159160
}
160161
catch (Exception ex)
161162
{
@@ -165,7 +166,7 @@ protected override void StartWorkingImpl(Func<IQueueEntry<T>, CancellationToken,
165166
{
166167
try
167168
{
168-
await _resiliencePolicy.ExecuteAsync(async _ => await queueEntry.AbandonAsync(), linkedCancellationToken.Token).AnyContext();
169+
await _resiliencePolicy.ExecuteAsync(async _ => await queueEntry.AbandonAsync(), linkedCancellationTokenSource.Token).AnyContext();
169170
}
170171
catch (Exception abandonEx)
171172
{
@@ -180,7 +181,7 @@ protected override void StartWorkingImpl(Func<IQueueEntry<T>, CancellationToken,
180181
{
181182
try
182183
{
183-
await _resiliencePolicy.ExecuteAsync(async _ => await queueEntry.CompleteAsync(), linkedCancellationToken.Token).AnyContext();
184+
await _resiliencePolicy.ExecuteAsync(async _ => await queueEntry.CompleteAsync(), linkedCancellationTokenSource.Token).AnyContext();
184185
}
185186
catch (Exception ex)
186187
{
@@ -189,8 +190,8 @@ protected override void StartWorkingImpl(Func<IQueueEntry<T>, CancellationToken,
189190
}
190191
}
191192

192-
_logger.LogTrace("Worker exiting: {QueueName} Cancel Requested: {IsCancellationRequested}", _options.Name, linkedCancellationToken.IsCancellationRequested);
193-
}, GetLinkedDisposableCancellationTokenSource(cancellationToken).Token));
193+
_logger.LogTrace("Worker exiting: {QueueName} Cancel Requested: {IsCancellationRequested}", _options.Name, linkedCancellationTokenSource.IsCancellationRequested);
194+
}, linkedCancellationTokenSource.Token));
194195
}
195196

196197
protected override async Task<IQueueEntry<T>> DequeueImplAsync(CancellationToken linkedCancellationToken)
@@ -312,12 +313,12 @@ public override async Task AbandonAsync(IQueueEntry<T> queueEntry)
312313
if (_options.RetryDelay > TimeSpan.Zero)
313314
{
314315
_logger.LogTrace("Adding item to wait list for future retry: {QueueEntryId}", queueEntry.Id);
315-
var unawaited = Run.DelayedAsync(GetRetryDelay(targetEntry.Attempts), () => RetryAsync(targetEntry), _timeProvider, _queueDisposedCancellationTokenSource.Token);
316+
var unawaited = Run.DelayedAsync(GetRetryDelay(targetEntry.Attempts), () => RetryAsync(targetEntry), _timeProvider, DisposedCancellationToken);
316317
}
317318
else
318319
{
319320
_logger.LogTrace("Adding item back to queue for retry: {QueueEntryId}", queueEntry.Id);
320-
_ = Task.Run(() => RetryAsync(targetEntry));
321+
_ = Task.Run(() => RetryAsync(targetEntry), DisposedCancellationToken);
321322
}
322323
}
323324
else

‎src/Foundatio/Queues/QueueBase.cs‎

Lines changed: 5 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -36,7 +36,6 @@ public abstract class QueueBase<T, TOptions> : MaintenanceBase, IQueue<T>, IHave
3636
private readonly TagList _emptyTags = default;
3737

3838
private readonly List<IQueueBehavior<T>> _behaviors = new();
39-
protected readonly CancellationTokenSource _queueDisposedCancellationTokenSource;
4039
private bool _isDisposed;
4140
private QueueStats _queueStats;
4241
private DateTimeOffset _nextQueueStatsUpdate = DateTimeOffset.MinValue;
@@ -53,8 +52,6 @@ protected QueueBase(TOptions options) : base(options?.TimeProvider, options?.Log
5352
_serializer = options.Serializer ?? DefaultSerializer.Instance;
5453
options.Behaviors.ForEach(AttachBehavior);
5554

56-
_queueDisposedCancellationTokenSource = new CancellationTokenSource();
57-
5855
var resiliencePolicyProvider = _options.GetResiliencePolicyProvider() ?? DefaultResiliencePolicyProvider.Instance;
5956
_resiliencePolicy = resiliencePolicyProvider.GetPolicy<QueueBase<T, TOptions>, IQueue<T>, IQueue>(_logger, _timeProvider);
6057

@@ -130,7 +127,7 @@ public void AttachBehavior(IQueueBehavior<T> behavior)
130127
protected abstract Task<string> EnqueueImplAsync(T data, QueueEntryOptions options);
131128
public async Task<string> EnqueueAsync(T data, QueueEntryOptions options = null)
132129
{
133-
await EnsureQueueCreatedAsync(_queueDisposedCancellationTokenSource.Token).AnyContext();
130+
await EnsureQueueCreatedAsync(DisposedCancellationToken).AnyContext();
134131

135132
LastEnqueueActivity = _timeProvider.GetUtcNow();
136133
options ??= new QueueEntryOptions();
@@ -141,7 +138,7 @@ public async Task<string> EnqueueAsync(T data, QueueEntryOptions options = null)
141138
protected abstract Task<IQueueEntry<T>> DequeueImplAsync(CancellationToken linkedCancellationToken);
142139
public async Task<IQueueEntry<T>> DequeueAsync(CancellationToken cancellationToken)
143140
{
144-
await EnsureQueueCreatedAsync(_queueDisposedCancellationTokenSource.Token).AnyContext();
141+
await EnsureQueueCreatedAsync(DisposedCancellationToken).AnyContext();
145142

146143
LastDequeueActivity = _timeProvider.GetUtcNow();
147144
using var linkedCancellationTokenSource = GetLinkedDisposableCancellationTokenSource(cancellationToken);
@@ -163,7 +160,7 @@ public virtual async Task<IQueueEntry<T>> DequeueAsync(TimeSpan? timeout = null)
163160
protected abstract Task<IEnumerable<T>> GetDeadletterItemsImplAsync(CancellationToken cancellationToken);
164161
public async Task<IEnumerable<T>> GetDeadletterItemsAsync(CancellationToken cancellationToken = default)
165162
{
166-
await EnsureQueueCreatedAsync(_queueDisposedCancellationTokenSource.Token).AnyContext();
163+
await EnsureQueueCreatedAsync(DisposedCancellationToken).AnyContext();
167164
return await GetDeadletterItemsImplAsync(cancellationToken).AnyContext();
168165
}
169166

@@ -186,7 +183,7 @@ protected virtual QueueStats GetMetricsQueueStats()
186183
protected abstract void StartWorkingImpl(Func<IQueueEntry<T>, CancellationToken, Task> handler, bool autoComplete, CancellationToken cancellationToken);
187184
public async Task StartWorkingAsync(Func<IQueueEntry<T>, CancellationToken, Task> handler, bool autoComplete = false, CancellationToken cancellationToken = default)
188185
{
189-
await EnsureQueueCreatedAsync(_queueDisposedCancellationTokenSource.Token).AnyContext();
186+
await EnsureQueueCreatedAsync(DisposedCancellationToken).AnyContext();
190187
StartWorkingImpl(handler, autoComplete, cancellationToken);
191188
}
192189

@@ -384,7 +381,7 @@ protected string GetFullMetricName(string customMetricName, string name)
384381

385382
protected CancellationTokenSource GetLinkedDisposableCancellationTokenSource(CancellationToken cancellationToken)
386383
{
387-
return CancellationTokenSource.CreateLinkedTokenSource(cancellationToken, _queueDisposedCancellationTokenSource.Token);
384+
return CancellationTokenSource.CreateLinkedTokenSource(cancellationToken, DisposedCancellationToken);
388385
}
389386

390387
public override void Dispose()
@@ -397,8 +394,6 @@ public override void Dispose()
397394

398395
_isDisposed = true;
399396
_logger.LogTrace("Queue {QueueName} ({QueueId}) dispose", _options.Name, QueueId);
400-
_queueDisposedCancellationTokenSource?.Cancel();
401-
_queueDisposedCancellationTokenSource?.Dispose();
402397
base.Dispose();
403398

404399
Abandoned?.Dispose();

‎src/Foundatio/Utility/MaintenanceBase.cs‎

Lines changed: 16 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
1-
using System;
1+
using System;
2+
using System.Threading;
23
using System.Threading.Tasks;
34
using Microsoft.Extensions.Logging;
45
using Microsoft.Extensions.Logging.Abstractions;
@@ -11,6 +12,8 @@ public class MaintenanceBase : IDisposable
1112
protected readonly ILoggerFactory _loggerFactory;
1213
protected readonly TimeProvider _timeProvider;
1314
protected readonly ILogger _logger;
15+
private readonly CancellationTokenSource _disposedCancellationTokenSource = new();
16+
private bool _isDisposed;
1417

1518
public MaintenanceBase(TimeProvider timeProvider, ILoggerFactory loggerFactory)
1619
{
@@ -19,6 +22,12 @@ public MaintenanceBase(TimeProvider timeProvider, ILoggerFactory loggerFactory)
1922
_logger = _loggerFactory.CreateLogger(GetType());
2023
}
2124

25+
/// <summary>
26+
/// Gets a cancellation token that is canceled when this instance is disposed.
27+
/// Use this token to cancel background operations during shutdown.
28+
/// </summary>
29+
protected CancellationToken DisposedCancellationToken => _disposedCancellationTokenSource.Token;
30+
2231
protected void InitializeMaintenance(TimeSpan? dueTime = null, TimeSpan? intervalTime = null)
2332
{
2433
_maintenanceTimer = new ScheduledTimer(DoMaintenanceAsync, dueTime, intervalTime, _timeProvider, _loggerFactory);
@@ -36,6 +45,12 @@ protected void ScheduleNextMaintenance(DateTime utcDate)
3645

3746
public virtual void Dispose()
3847
{
48+
if (_isDisposed)
49+
return;
50+
51+
_isDisposed = true;
52+
_disposedCancellationTokenSource.Cancel();
53+
_disposedCancellationTokenSource.Dispose();
3954
_maintenanceTimer?.Dispose();
4055
}
4156
}

0 commit comments

Comments
 (0)