From fbd14db19f78b36102e580ac3e94ebc00ece410a Mon Sep 17 00:00:00 2001 From: John Simons Date: Wed, 12 Aug 2026 07:04:03 +1000 Subject: [PATCH 1/2] Propagate cancellation tokens through ingestion and body storage persistence This change continues the effort to consistently apply optional `CancellationToken` parameters to asynchronous methods within the `IBodyStorage`, `IFailedErrorImportDataStore`, and ingestion unit of work interfaces, along with their EFCore and RavenDB implementations. Cancellation tokens are now propagated to underlying asynchronous database calls, ensuring proper cancellation of long-running operations. In the error processor, specific handling for `OperationCanceledException` was added to ensure message contexts are released during shutdown so that the ingestion loop can terminate rather than awaiting indefinitely. --- .../Implementation/BodyStorage/BodyStorage.cs | 16 ++++++++-------- .../FileSystemBodyStoragePersistence.cs | 3 +++ .../FailedErrorImportDataStore.cs | 11 +++++++---- .../UnitOfWork/EFIngestionUnitOfWork.cs | 2 +- .../UnitOfWork/EFIngestionUnitOfWorkFactory.cs | 2 +- .../EFMonitoringIngestionUnitOfWork.cs | 2 +- .../EFRecoverabilityIngestionUnitOfWork.cs | 7 ++++--- .../UnitOfWork/FailedMessageBatchWriter.cs | 2 +- .../ErrorMessagesDataStore.cs | 2 +- .../FailedErrorImportDataStore.cs | 18 +++++++++--------- .../RavenAttachmentsBodyStorage.cs | 15 ++++++++------- .../UnitOfWork/RavenIngestionUnitOfWork.cs | 2 +- .../RavenIngestionUnitOfWorkFactory.cs | 3 ++- .../RavenMonitoringIngestionUnitOfWork.cs | 3 ++- .../RavenRecoverabilityIngestionUnitOfWork.cs | 6 ++++-- .../EFCore/FailedErrorImportTests.cs | 8 ++++---- src/ServiceControl.Persistence/IBodyStorage.cs | 3 ++- .../IFailedErrorImportDataStore.cs | 6 +++--- .../UnitOfWork/FallbackIngestionUnitOfWork.cs | 2 +- .../FallbackIngestionUnitOfWorkFactory.cs | 7 ++++--- .../UnitOfWork/IIngestionUnitOfWork.cs | 2 +- .../UnitOfWork/IIngestionUnitOfWorkFactory.cs | 3 ++- .../IMonitoringIngestionUnitOfWork.cs | 3 ++- .../IRecoverabilityIngestionUnitOfWork.cs | 5 +++-- .../UnitOfWork/IngestionUnitOfWorkBase.cs | 5 ++++- .../Messages/GetMessagesController.cs | 2 +- .../FailedErrorImportCustomCheck.cs | 2 +- .../Operations/ErrorIngestionFaultPolicy.cs | 2 +- src/ServiceControl/Operations/ErrorIngestor.cs | 2 +- .../Operations/ErrorProcessor.cs | 12 ++++++++++-- .../Operations/ImportFailedErrors.cs | 4 ++-- .../Operations/RetryConfirmationProcessor.cs | 2 +- 32 files changed, 96 insertions(+), 68 deletions(-) diff --git a/src/ServiceControl.Persistence.EFCore/Implementation/BodyStorage/BodyStorage.cs b/src/ServiceControl.Persistence.EFCore/Implementation/BodyStorage/BodyStorage.cs index e62ac98259..700a978a0a 100644 --- a/src/ServiceControl.Persistence.EFCore/Implementation/BodyStorage/BodyStorage.cs +++ b/src/ServiceControl.Persistence.EFCore/Implementation/BodyStorage/BodyStorage.cs @@ -19,9 +19,9 @@ namespace ServiceControl.Persistence.EFCore.Implementation.BodyStorage; /// public class BodyStorage(IServiceScopeFactory scopeFactory, IBodyStoragePersistence storagePersistence) : DataStoreBase(scopeFactory), IBodyStorage { - public async Task TryFetch(string bodyId) + public async Task TryFetch(string bodyId, CancellationToken cancellationToken = default) { - var row = await ExecuteWithDbContext(dbContext => ResolveBody(dbContext, bodyId)); + var row = await ExecuteWithDbContext(dbContext => ResolveBody(dbContext, bodyId, cancellationToken)); if (row == null) { @@ -33,7 +33,7 @@ public class BodyStorage(IServiceScopeFactory scopeFactory, IBodyStoragePersiste if (row.BodyStoredExternally) { - var external = await storagePersistence.ReadBody(uniqueMessageId); + var external = await storagePersistence.ReadBody(uniqueMessageId, cancellationToken); return external == null ? new MessageBodyStreamResult { HasResult = false } @@ -64,21 +64,21 @@ public class BodyStorage(IServiceScopeFactory scopeFactory, IBodyStoragePersiste return new MessageBodyStreamResult { HasResult = false }; // Message exists but carries no body. } - static async Task ResolveBody(ServiceControlDbContext dbContext, string bodyId) + static async Task ResolveBody(ServiceControlDbContext dbContext, string bodyId, CancellationToken cancellationToken) { if (Guid.TryParse(bodyId, out var uniqueMessageId)) { - var byUniqueId = await Query(dbContext, message => message.UniqueMessageId == uniqueMessageId); + var byUniqueId = await Query(dbContext, message => message.UniqueMessageId == uniqueMessageId, cancellationToken); if (byUniqueId != null) { return byUniqueId; } } - return await Query(dbContext, message => message.MessageId == bodyId); + return await Query(dbContext, message => message.MessageId == bodyId, cancellationToken); } - static Task Query(ServiceControlDbContext dbContext, Expression> predicate) => + static Task Query(ServiceControlDbContext dbContext, Expression> predicate, CancellationToken cancellationToken) => dbContext.FailedMessages .AsNoTracking() .Where(predicate) @@ -90,7 +90,7 @@ public class BodyStorage(IServiceScopeFactory scopeFactory, IBodyStoragePersiste BodyStoredExternally = message.BodyStoredExternally, BodyContentType = message.BodyContentType }) - .FirstOrDefaultAsync(); + .FirstOrDefaultAsync(cancellationToken); sealed class BodyRow { diff --git a/src/ServiceControl.Persistence.EFCore/Implementation/BodyStorage/FileSystemBodyStoragePersistence.cs b/src/ServiceControl.Persistence.EFCore/Implementation/BodyStorage/FileSystemBodyStoragePersistence.cs index 062534f742..6f5ba321c7 100644 --- a/src/ServiceControl.Persistence.EFCore/Implementation/BodyStorage/FileSystemBodyStoragePersistence.cs +++ b/src/ServiceControl.Persistence.EFCore/Implementation/BodyStorage/FileSystemBodyStoragePersistence.cs @@ -72,11 +72,14 @@ public async Task WriteBody(string bodyId, ReadOnlyMemory body, string con TryDelete(tempFilePath); } } +#pragma warning disable PS0019 // The catch rethrows everything, cancellation included, so splitting + // out OperationCanceledException would duplicate the cleanup for no change in behaviour. catch { TryDelete(tempFilePath); throw; } +#pragma warning restore PS0019 } public Task ReadBody(string bodyId, CancellationToken cancellationToken = default) diff --git a/src/ServiceControl.Persistence.EFCore/Implementation/FailedErrorImportDataStore.cs b/src/ServiceControl.Persistence.EFCore/Implementation/FailedErrorImportDataStore.cs index d5fd7d7a70..a11bf466e5 100644 --- a/src/ServiceControl.Persistence.EFCore/Implementation/FailedErrorImportDataStore.cs +++ b/src/ServiceControl.Persistence.EFCore/Implementation/FailedErrorImportDataStore.cs @@ -22,14 +22,14 @@ public class FailedErrorImportDataStore( { const int BatchSize = 100; - public Task QueryContainsFailedImports() => + public Task QueryContainsFailedImports(CancellationToken cancellationToken = default) => ExecuteWithDbContext(dbContext => dbContext.FailedErrorImports.AsNoTracking().AnyAsync()); // Update-first, then insert. The dedupe key is deterministic, so a repeat failure updates the // existing row and concurrent writers that both miss it race only on the insert. The loser of // that race confirms the row is now present (the winner stored the same logical failure) and // otherwise rethrows, so the caller never treats a message as stored when it is not. - public Task StoreFailedErrorImport(FailedErrorImport failure) => + public Task StoreFailedErrorImport(FailedErrorImport failure, CancellationToken cancellationToken = default) => ExecuteWithDbContext(async dbContext => { var uniqueMessageId = FailedErrorImport.DeriveKey(failure.Message.Headers, failure.Message.Id); @@ -69,7 +69,7 @@ public Task StoreFailedErrorImport(FailedErrorImport failure) => // Replays oldest-first. Successful imports delete their row; failures are left in place, so the // count of failures so far is exactly the offset to the next unseen row. This walks the whole // set once without retrying a failure within the same run. - public async Task ProcessFailedErrorImports(Func processMessage, CancellationToken cancellationToken) + public async Task ProcessFailedErrorImports(Func processMessage, CancellationToken cancellationToken = default) { var succeeded = 0; var failed = 0; @@ -94,7 +94,7 @@ public async Task ProcessFailedErrorImports(Func p { var transportMessage = await ToTransportMessage(import, cancellationToken); - await processMessage(transportMessage); + await processMessage(transportMessage, cancellationToken); await DeleteImport(import, cancellationToken); @@ -192,10 +192,13 @@ async Task DeleteExternalBody(Guid uniqueMessageId, CancellationToken cancellati { await bodyStorage.DeleteBody(FailedErrorImportEntity.ExternalBodyId(uniqueMessageId), cancellationToken); } +#pragma warning disable PS0019 // The filter already excludes OperationCanceledException, so cancellation + // propagates; PS0019 only recognises a cancellationToken.IsCancellationRequested guard. catch (Exception ex) when (ex is not OperationCanceledException) { // Re-import must not stall on a missing or unavailable body. logger.LogWarning(ex, "Could not delete the external body for re-imported failed error {UniqueMessageId}", uniqueMessageId); } +#pragma warning restore PS0019 } } diff --git a/src/ServiceControl.Persistence.EFCore/Implementation/UnitOfWork/EFIngestionUnitOfWork.cs b/src/ServiceControl.Persistence.EFCore/Implementation/UnitOfWork/EFIngestionUnitOfWork.cs index 0f67fcc1af..acac97d01f 100644 --- a/src/ServiceControl.Persistence.EFCore/Implementation/UnitOfWork/EFIngestionUnitOfWork.cs +++ b/src/ServiceControl.Persistence.EFCore/Implementation/UnitOfWork/EFIngestionUnitOfWork.cs @@ -41,7 +41,7 @@ public EFIngestionUnitOfWork(IAsyncDisposable scope, ServiceControlDbContext dbC internal void RecordConfirmedRetry(Guid uniqueMessageId) => confirmedRetries.Enqueue(uniqueMessageId); - public async Task Complete(CancellationToken cancellationToken) + public async Task Complete(CancellationToken cancellationToken = default) { // External bodies are written before the rows that point at them await Task.WhenAll(bodyWrites); diff --git a/src/ServiceControl.Persistence.EFCore/Implementation/UnitOfWork/EFIngestionUnitOfWorkFactory.cs b/src/ServiceControl.Persistence.EFCore/Implementation/UnitOfWork/EFIngestionUnitOfWorkFactory.cs index c62ec80634..8e2e84bfa6 100644 --- a/src/ServiceControl.Persistence.EFCore/Implementation/UnitOfWork/EFIngestionUnitOfWorkFactory.cs +++ b/src/ServiceControl.Persistence.EFCore/Implementation/UnitOfWork/EFIngestionUnitOfWorkFactory.cs @@ -13,7 +13,7 @@ public class EFIngestionUnitOfWorkFactory( IFailedMessageIngestionSqlDialect dialect, TimeProvider timeProvider) : IIngestionUnitOfWorkFactory { - public ValueTask StartNew() + public ValueTask StartNew(CancellationToken cancellationToken = default) { var scope = serviceProvider.CreateAsyncScope(); var dbContext = scope.ServiceProvider.GetRequiredService(); diff --git a/src/ServiceControl.Persistence.EFCore/Implementation/UnitOfWork/EFMonitoringIngestionUnitOfWork.cs b/src/ServiceControl.Persistence.EFCore/Implementation/UnitOfWork/EFMonitoringIngestionUnitOfWork.cs index 19cdf86f54..6163e22525 100644 --- a/src/ServiceControl.Persistence.EFCore/Implementation/UnitOfWork/EFMonitoringIngestionUnitOfWork.cs +++ b/src/ServiceControl.Persistence.EFCore/Implementation/UnitOfWork/EFMonitoringIngestionUnitOfWork.cs @@ -4,7 +4,7 @@ namespace ServiceControl.Persistence.EFCore.Implementation.UnitOfWork; public class EFMonitoringIngestionUnitOfWork(EFIngestionUnitOfWork parentUnitOfWork) : IMonitoringIngestionUnitOfWork { - public Task RecordKnownEndpoint(KnownEndpoint knownEndpoint) + public Task RecordKnownEndpoint(KnownEndpoint knownEndpoint, CancellationToken cancellationToken = default) { parentUnitOfWork.Record(knownEndpoint); diff --git a/src/ServiceControl.Persistence.EFCore/Implementation/UnitOfWork/EFRecoverabilityIngestionUnitOfWork.cs b/src/ServiceControl.Persistence.EFCore/Implementation/UnitOfWork/EFRecoverabilityIngestionUnitOfWork.cs index 4a2067264b..8566b94bcf 100644 --- a/src/ServiceControl.Persistence.EFCore/Implementation/UnitOfWork/EFRecoverabilityIngestionUnitOfWork.cs +++ b/src/ServiceControl.Persistence.EFCore/Implementation/UnitOfWork/EFRecoverabilityIngestionUnitOfWork.cs @@ -14,7 +14,8 @@ public class EFRecoverabilityIngestionUnitOfWork(EFIngestionUnitOfWork parentUni { public Task RecordFailedProcessingAttempt(MessageContext context, FailedMessage.ProcessingAttempt processingAttempt, - List groups) + List groups, + CancellationToken cancellationToken = default) { var uniqueMessageId = context.Headers.UniqueId(); var contentType = context.Headers.GetValueOrDefault(Headers.ContentType, "text/plain"); @@ -24,7 +25,7 @@ public Task RecordFailedProcessingAttempt(MessageContext context, if (storeExternally) { parentUnitOfWork.RecordBodyWrite( - storagePersistence.WriteBody(uniqueMessageId, context.Body, contentType)); + storagePersistence.WriteBody(uniqueMessageId, context.Body, contentType, cancellationToken)); } var sendingEndpoint = GetMetadata(processingAttempt, "SendingEndpoint"); @@ -60,7 +61,7 @@ public Task RecordFailedProcessingAttempt(MessageContext context, return Task.CompletedTask; } - public Task RecordSuccessfulRetry(string retriedMessageUniqueId) + public Task RecordSuccessfulRetry(string retriedMessageUniqueId, CancellationToken cancellationToken = default) { parentUnitOfWork.RecordConfirmedRetry(Guid.Parse(retriedMessageUniqueId)); diff --git a/src/ServiceControl.Persistence.EFCore/Implementation/UnitOfWork/FailedMessageBatchWriter.cs b/src/ServiceControl.Persistence.EFCore/Implementation/UnitOfWork/FailedMessageBatchWriter.cs index 84d347494f..c451b0dd6b 100644 --- a/src/ServiceControl.Persistence.EFCore/Implementation/UnitOfWork/FailedMessageBatchWriter.cs +++ b/src/ServiceControl.Persistence.EFCore/Implementation/UnitOfWork/FailedMessageBatchWriter.cs @@ -17,7 +17,7 @@ public async Task Write( IReadOnlyCollection knownEndpoints, IReadOnlyCollection confirmedRetries, DateTime now, - CancellationToken cancellationToken) + CancellationToken cancellationToken = default) { var (failedMessages, groups) = Fold(attempts, now); var endpoints = BuildEndpointRows(knownEndpoints); diff --git a/src/ServiceControl.Persistence.RavenDB/ErrorMessagesDataStore.cs b/src/ServiceControl.Persistence.RavenDB/ErrorMessagesDataStore.cs index c15f3ffe1f..2e5bbce988 100644 --- a/src/ServiceControl.Persistence.RavenDB/ErrorMessagesDataStore.cs +++ b/src/ServiceControl.Persistence.RavenDB/ErrorMessagesDataStore.cs @@ -511,7 +511,7 @@ record struct FailedMessageProjection(string UniqueMessageId); public async Task GetFailedMessageBody(string uniqueMessageId, CancellationToken cancellationToken = default) { byte[] body = null; - var result = await bodyStorage.TryFetch(uniqueMessageId) + var result = await bodyStorage.TryFetch(uniqueMessageId, cancellationToken) ?? throw new InvalidOperationException("IBodyStorage.TryFetch result cannot be null"); if (result.HasResult) diff --git a/src/ServiceControl.Persistence.RavenDB/FailedErrorImportDataStore.cs b/src/ServiceControl.Persistence.RavenDB/FailedErrorImportDataStore.cs index 7b04debfe7..e6747de6f8 100644 --- a/src/ServiceControl.Persistence.RavenDB/FailedErrorImportDataStore.cs +++ b/src/ServiceControl.Persistence.RavenDB/FailedErrorImportDataStore.cs @@ -9,21 +9,21 @@ class FailedErrorImportDataStore(IRavenSessionProvider sessionProvider, ILogger logger) : IFailedErrorImportDataStore { - public async Task StoreFailedErrorImport(FailedErrorImport failure) + public async Task StoreFailedErrorImport(FailedErrorImport failure, CancellationToken cancellationToken = default) { - using var session = await sessionProvider.OpenSession(); + using var session = await sessionProvider.OpenSession(cancellationToken: cancellationToken); // This object's ID is generated externally, but is not in the RavenDB format // Check that's true to make sure that if it already is that it doesn't get double-formatted if (!failure.Id.StartsWith(CollectionName)) { failure.Id = MakeDocumentId(failure.Id); } - await session.StoreAsync(failure); + await session.StoreAsync(failure, cancellationToken); - await session.SaveChangesAsync(); + await session.SaveChangesAsync(cancellationToken); } - public async Task ProcessFailedErrorImports(Func processMessage, CancellationToken cancellationToken) + public async Task ProcessFailedErrorImports(Func processMessage, CancellationToken cancellationToken = default) { var succeeded = 0; var failed = 0; @@ -36,7 +36,7 @@ public async Task ProcessFailedErrorImports(Func p var transportMessage = stream.Current.Document.Message; try { - await processMessage(transportMessage); + await processMessage(transportMessage, cancellationToken); await session.Advanced.RequestExecutor.ExecuteAsync(new DeleteDocumentCommand(stream.Current.Id, null), session.Advanced.Context, token: cancellationToken); @@ -64,11 +64,11 @@ public async Task ProcessFailedErrorImports(Func p } } - public async Task QueryContainsFailedImports() + public async Task QueryContainsFailedImports(CancellationToken cancellationToken = default) { - using var session = await sessionProvider.OpenSession(); + using var session = await sessionProvider.OpenSession(cancellationToken: cancellationToken); var query = session.Query(); - await using var ie = await session.Advanced.StreamAsync(query); + await using var ie = await session.Advanced.StreamAsync(query, cancellationToken); return await ie.MoveNextAsync(); } diff --git a/src/ServiceControl.Persistence.RavenDB/RavenAttachmentsBodyStorage.cs b/src/ServiceControl.Persistence.RavenDB/RavenAttachmentsBodyStorage.cs index c2aa61ee80..64cc215ebc 100644 --- a/src/ServiceControl.Persistence.RavenDB/RavenAttachmentsBodyStorage.cs +++ b/src/ServiceControl.Persistence.RavenDB/RavenAttachmentsBodyStorage.cs @@ -2,6 +2,7 @@ { using System; using System.Linq; + using System.Threading; using System.Threading.Tasks; using Persistence.RavenDB; using Raven.Client.Documents; @@ -13,15 +14,15 @@ class RavenAttachmentsBodyStorage(IRavenSessionProvider sessionProvider) : IBody { public const string AttachmentName = "body"; - public async Task TryFetch(string bodyId) + public async Task TryFetch(string bodyId, CancellationToken cancellationToken = default) { - using var session = await sessionProvider.OpenSession(); + using var session = await sessionProvider.OpenSession(cancellationToken: cancellationToken); // BodyId could be a MessageID or a UniqueID, but if a UniqueID then it will be a DeterministicGuid of MessageID and endpoint name and be Guid-parseable // This is preferred, then we know we're getting the correct message body that is attached to the FailedMessage document if (Guid.TryParse(bodyId, out _)) { - var result = await ResultForUniqueId(session, bodyId); + var result = await ResultForUniqueId(session, bodyId, cancellationToken); if (result != null) { return result; @@ -34,21 +35,21 @@ public async Task TryFetch(string bodyId) .OfType() .Select(msg => msg.UniqueMessageId); - var uniqueId = await query.FirstOrDefaultAsync(); + var uniqueId = await query.FirstOrDefaultAsync(cancellationToken); if (uniqueId != null) { - return await ResultForUniqueId(session, uniqueId); + return await ResultForUniqueId(session, uniqueId, cancellationToken); } return null; } - async Task ResultForUniqueId(IAsyncDocumentSession session, string uniqueId) + async Task ResultForUniqueId(IAsyncDocumentSession session, string uniqueId, CancellationToken cancellationToken) { var documentId = FailedMessageIdGenerator.MakeDocumentId(uniqueId); - var result = await session.Advanced.Attachments.GetAsync(documentId, AttachmentName); + var result = await session.Advanced.Attachments.GetAsync(documentId, AttachmentName, cancellationToken); if (result == null) { diff --git a/src/ServiceControl.Persistence.RavenDB/UnitOfWork/RavenIngestionUnitOfWork.cs b/src/ServiceControl.Persistence.RavenDB/UnitOfWork/RavenIngestionUnitOfWork.cs index 64693498ff..7d39d56265 100644 --- a/src/ServiceControl.Persistence.RavenDB/UnitOfWork/RavenIngestionUnitOfWork.cs +++ b/src/ServiceControl.Persistence.RavenDB/UnitOfWork/RavenIngestionUnitOfWork.cs @@ -22,7 +22,7 @@ public RavenIngestionUnitOfWork(IRavenSessionProvider sessionProvider, Expiratio internal void AddCommand(ICommandData command) => commands.Enqueue(command); - public override async Task Complete(CancellationToken cancellationToken) + public override async Task Complete(CancellationToken cancellationToken = default) { using var session = await sessionProvider.OpenSession(cancellationToken: cancellationToken); // not really interested in the batch results since a batch is atomic diff --git a/src/ServiceControl.Persistence.RavenDB/UnitOfWork/RavenIngestionUnitOfWorkFactory.cs b/src/ServiceControl.Persistence.RavenDB/UnitOfWork/RavenIngestionUnitOfWorkFactory.cs index 2cd6237c1b..1ba74d0e11 100644 --- a/src/ServiceControl.Persistence.RavenDB/UnitOfWork/RavenIngestionUnitOfWorkFactory.cs +++ b/src/ServiceControl.Persistence.RavenDB/UnitOfWork/RavenIngestionUnitOfWorkFactory.cs @@ -1,5 +1,6 @@ namespace ServiceControl.Persistence.RavenDB { + using System.Threading; using System.Threading.Tasks; using ServiceControl.Persistence.UnitOfWork; @@ -10,7 +11,7 @@ class RavenIngestionUnitOfWorkFactory( RavenPersisterSettings settings) : IIngestionUnitOfWorkFactory { - public ValueTask StartNew() + public ValueTask StartNew(CancellationToken cancellationToken = default) => new(new RavenIngestionUnitOfWork(sessionProvider, expirationManager, settings)); public bool CanIngestMore() => customCheckState.CanIngestMore; diff --git a/src/ServiceControl.Persistence.RavenDB/UnitOfWork/RavenMonitoringIngestionUnitOfWork.cs b/src/ServiceControl.Persistence.RavenDB/UnitOfWork/RavenMonitoringIngestionUnitOfWork.cs index af690c8129..8d0b235c9a 100644 --- a/src/ServiceControl.Persistence.RavenDB/UnitOfWork/RavenMonitoringIngestionUnitOfWork.cs +++ b/src/ServiceControl.Persistence.RavenDB/UnitOfWork/RavenMonitoringIngestionUnitOfWork.cs @@ -1,5 +1,6 @@ namespace ServiceControl.Persistence.RavenDB { + using System.Threading; using System.Threading.Tasks; using Newtonsoft.Json.Linq; using Raven.Client.Documents.Commands.Batches; @@ -15,7 +16,7 @@ public RavenMonitoringIngestionUnitOfWork(RavenIngestionUnitOfWork parentUnitOfW this.parentUnitOfWork = parentUnitOfWork; } - public Task RecordKnownEndpoint(KnownEndpoint knownEndpoint) + public Task RecordKnownEndpoint(KnownEndpoint knownEndpoint, CancellationToken cancellationToken = default) { parentUnitOfWork.AddCommand(CreateKnownEndpointsPutCommand(knownEndpoint)); return Task.CompletedTask; diff --git a/src/ServiceControl.Persistence.RavenDB/UnitOfWork/RavenRecoverabilityIngestionUnitOfWork.cs b/src/ServiceControl.Persistence.RavenDB/UnitOfWork/RavenRecoverabilityIngestionUnitOfWork.cs index 882bd07d58..b53113beea 100644 --- a/src/ServiceControl.Persistence.RavenDB/UnitOfWork/RavenRecoverabilityIngestionUnitOfWork.cs +++ b/src/ServiceControl.Persistence.RavenDB/UnitOfWork/RavenRecoverabilityIngestionUnitOfWork.cs @@ -3,6 +3,7 @@ using System; using System.Collections.Generic; using System.Text; + using System.Threading; using System.Threading.Tasks; using NServiceBus; using NServiceBus.Transport; @@ -30,7 +31,8 @@ public RavenRecoverabilityIngestionUnitOfWork(RavenIngestionUnitOfWork parentUni public Task RecordFailedProcessingAttempt( MessageContext context, FailedMessage.ProcessingAttempt processingAttempt, - List groups) + List groups, + CancellationToken cancellationToken = default) { var uniqueMessageId = context.Headers.UniqueId(); var contentType = GetContentType(context.Headers, "text/plain"); @@ -66,7 +68,7 @@ public Task RecordFailedProcessingAttempt( return Task.CompletedTask; } - public Task RecordSuccessfulRetry(string retriedMessageUniqueId) + public Task RecordSuccessfulRetry(string retriedMessageUniqueId, CancellationToken cancellationToken = default) { var failedMessageDocumentId = FailedMessageIdGenerator.MakeDocumentId(retriedMessageUniqueId); var failedMessageRetryDocumentId = RetryDocumentDataStore.MakeFailedMessageRetriesDocumentId(retriedMessageUniqueId); diff --git a/src/ServiceControl.Persistence.Tests/EFCore/FailedErrorImportTests.cs b/src/ServiceControl.Persistence.Tests/EFCore/FailedErrorImportTests.cs index 44b26831c9..59d4143502 100644 --- a/src/ServiceControl.Persistence.Tests/EFCore/FailedErrorImportTests.cs +++ b/src/ServiceControl.Persistence.Tests/EFCore/FailedErrorImportTests.cs @@ -125,7 +125,7 @@ await Store( var replayed = new List(); await FailedImportStore.ProcessFailedErrorImports( - message => + (message, _) => { replayed.Add(message.Id); return message.Id == "native-2" ? throw new InvalidOperationException("boom") : Task.CompletedTask; @@ -140,7 +140,7 @@ await FailedImportStore.ProcessFailedErrorImports( var secondRun = new List(); await FailedImportStore.ProcessFailedErrorImports( - message => { secondRun.Add(message.Id); return Task.CompletedTask; }, + (message, _) => { secondRun.Add(message.Id); return Task.CompletedTask; }, TestContext.CurrentContext.CancellationToken); Assert.That(secondRun, Is.EqualTo(new[] { "native-2" })); @@ -176,7 +176,7 @@ await Store( var replayed = new List(); await FailedImportStore.ProcessFailedErrorImports( - message => + (message, _) => { replayed.Add(message.Id); cts.Cancel(); @@ -266,7 +266,7 @@ async Task> Replay() var replayed = new List(); await FailedImportStore.ProcessFailedErrorImports( - message => { replayed.Add(message); return Task.CompletedTask; }, + (message, _) => { replayed.Add(message); return Task.CompletedTask; }, TestContext.CurrentContext.CancellationToken); return replayed; diff --git a/src/ServiceControl.Persistence/IBodyStorage.cs b/src/ServiceControl.Persistence/IBodyStorage.cs index 5ade6139c6..4becd1ea2e 100644 --- a/src/ServiceControl.Persistence/IBodyStorage.cs +++ b/src/ServiceControl.Persistence/IBodyStorage.cs @@ -1,11 +1,12 @@ namespace ServiceControl.Operations.BodyStorage { using System.IO; + using System.Threading; using System.Threading.Tasks; public interface IBodyStorage { - Task TryFetch(string bodyId); + Task TryFetch(string bodyId, CancellationToken cancellationToken = default); } public class MessageBodyStreamResult diff --git a/src/ServiceControl.Persistence/IFailedErrorImportDataStore.cs b/src/ServiceControl.Persistence/IFailedErrorImportDataStore.cs index 4e579a560f..136fa1f675 100644 --- a/src/ServiceControl.Persistence/IFailedErrorImportDataStore.cs +++ b/src/ServiceControl.Persistence/IFailedErrorImportDataStore.cs @@ -7,8 +7,8 @@ public interface IFailedErrorImportDataStore { - Task StoreFailedErrorImport(FailedErrorImport failure); - Task ProcessFailedErrorImports(Func processMessage, CancellationToken cancellationToken); - Task QueryContainsFailedImports(); + Task StoreFailedErrorImport(FailedErrorImport failure, CancellationToken cancellationToken = default); + Task ProcessFailedErrorImports(Func processMessage, CancellationToken cancellationToken = default); + Task QueryContainsFailedImports(CancellationToken cancellationToken = default); } } \ No newline at end of file diff --git a/src/ServiceControl.Persistence/UnitOfWork/FallbackIngestionUnitOfWork.cs b/src/ServiceControl.Persistence/UnitOfWork/FallbackIngestionUnitOfWork.cs index 98921ae799..d257a15213 100644 --- a/src/ServiceControl.Persistence/UnitOfWork/FallbackIngestionUnitOfWork.cs +++ b/src/ServiceControl.Persistence/UnitOfWork/FallbackIngestionUnitOfWork.cs @@ -19,7 +19,7 @@ public FallbackIngestionUnitOfWork(IIngestionUnitOfWork primary, IIngestionUnitO Recoverability = primary.Recoverability ?? fallback.Recoverability; } - public override Task Complete(CancellationToken cancellationToken) + public override Task Complete(CancellationToken cancellationToken = default) => Task.WhenAll( primary.Complete(cancellationToken), fallback.Complete(cancellationToken) diff --git a/src/ServiceControl.Persistence/UnitOfWork/FallbackIngestionUnitOfWorkFactory.cs b/src/ServiceControl.Persistence/UnitOfWork/FallbackIngestionUnitOfWorkFactory.cs index a0c6c30bd4..c525774b52 100644 --- a/src/ServiceControl.Persistence/UnitOfWork/FallbackIngestionUnitOfWorkFactory.cs +++ b/src/ServiceControl.Persistence/UnitOfWork/FallbackIngestionUnitOfWorkFactory.cs @@ -1,5 +1,6 @@ namespace ServiceControl.Persistence.UnitOfWork { + using System.Threading; using System.Threading.Tasks; class FallbackIngestionUnitOfWorkFactory : IIngestionUnitOfWorkFactory @@ -13,10 +14,10 @@ public FallbackIngestionUnitOfWorkFactory(IIngestionUnitOfWorkFactory primary, I this.secondary = secondary; } - public async ValueTask StartNew() + public async ValueTask StartNew(CancellationToken cancellationToken = default) { - var primaryUnitOfWork = await primary.StartNew(); - var secondaryUnitOfWork = await secondary.StartNew(); + var primaryUnitOfWork = await primary.StartNew(cancellationToken); + var secondaryUnitOfWork = await secondary.StartNew(cancellationToken); return new FallbackIngestionUnitOfWork( primaryUnitOfWork, diff --git a/src/ServiceControl.Persistence/UnitOfWork/IIngestionUnitOfWork.cs b/src/ServiceControl.Persistence/UnitOfWork/IIngestionUnitOfWork.cs index ce5f245ac9..befc0c6280 100644 --- a/src/ServiceControl.Persistence/UnitOfWork/IIngestionUnitOfWork.cs +++ b/src/ServiceControl.Persistence/UnitOfWork/IIngestionUnitOfWork.cs @@ -8,6 +8,6 @@ public interface IIngestionUnitOfWork : IAsyncDisposable { IMonitoringIngestionUnitOfWork Monitoring { get; } IRecoverabilityIngestionUnitOfWork Recoverability { get; } - Task Complete(CancellationToken cancellationToken); + Task Complete(CancellationToken cancellationToken = default); } } \ No newline at end of file diff --git a/src/ServiceControl.Persistence/UnitOfWork/IIngestionUnitOfWorkFactory.cs b/src/ServiceControl.Persistence/UnitOfWork/IIngestionUnitOfWorkFactory.cs index 7fc0539338..5680bf430e 100644 --- a/src/ServiceControl.Persistence/UnitOfWork/IIngestionUnitOfWorkFactory.cs +++ b/src/ServiceControl.Persistence/UnitOfWork/IIngestionUnitOfWorkFactory.cs @@ -1,10 +1,11 @@ namespace ServiceControl.Persistence.UnitOfWork { + using System.Threading; using System.Threading.Tasks; public interface IIngestionUnitOfWorkFactory { - ValueTask StartNew(); + ValueTask StartNew(CancellationToken cancellationToken = default); bool CanIngestMore(); } } \ No newline at end of file diff --git a/src/ServiceControl.Persistence/UnitOfWork/IMonitoringIngestionUnitOfWork.cs b/src/ServiceControl.Persistence/UnitOfWork/IMonitoringIngestionUnitOfWork.cs index 3a270ea62c..55378abd1a 100644 --- a/src/ServiceControl.Persistence/UnitOfWork/IMonitoringIngestionUnitOfWork.cs +++ b/src/ServiceControl.Persistence/UnitOfWork/IMonitoringIngestionUnitOfWork.cs @@ -1,9 +1,10 @@ namespace ServiceControl.Persistence.UnitOfWork { + using System.Threading; using System.Threading.Tasks; public interface IMonitoringIngestionUnitOfWork { - Task RecordKnownEndpoint(KnownEndpoint knownEndpoint); + Task RecordKnownEndpoint(KnownEndpoint knownEndpoint, CancellationToken cancellationToken = default); } } \ No newline at end of file diff --git a/src/ServiceControl.Persistence/UnitOfWork/IRecoverabilityIngestionUnitOfWork.cs b/src/ServiceControl.Persistence/UnitOfWork/IRecoverabilityIngestionUnitOfWork.cs index 8d7b275421..6ea5d43621 100644 --- a/src/ServiceControl.Persistence/UnitOfWork/IRecoverabilityIngestionUnitOfWork.cs +++ b/src/ServiceControl.Persistence/UnitOfWork/IRecoverabilityIngestionUnitOfWork.cs @@ -1,6 +1,7 @@ namespace ServiceControl.Persistence.UnitOfWork { using System.Collections.Generic; + using System.Threading; using System.Threading.Tasks; using NServiceBus.Transport; using ServiceControl.MessageFailures; @@ -9,8 +10,8 @@ public interface IRecoverabilityIngestionUnitOfWork { Task RecordFailedProcessingAttempt(MessageContext context, FailedMessage.ProcessingAttempt processingAttempt, - List groups); + List groups, CancellationToken cancellationToken = default); - Task RecordSuccessfulRetry(string retriedMessageUniqueId); + Task RecordSuccessfulRetry(string retriedMessageUniqueId, CancellationToken cancellationToken = default); } } \ No newline at end of file diff --git a/src/ServiceControl.Persistence/UnitOfWork/IngestionUnitOfWorkBase.cs b/src/ServiceControl.Persistence/UnitOfWork/IngestionUnitOfWorkBase.cs index c348964646..5d55c742ac 100644 --- a/src/ServiceControl.Persistence/UnitOfWork/IngestionUnitOfWorkBase.cs +++ b/src/ServiceControl.Persistence/UnitOfWork/IngestionUnitOfWorkBase.cs @@ -6,7 +6,10 @@ namespace ServiceControl.Persistence.UnitOfWork public abstract class IngestionUnitOfWorkBase : IIngestionUnitOfWork { + // Mirrors IAsyncDisposable.DisposeAsync, which declares no token, so one cannot be added here. +#pragma warning disable PS0018 protected virtual ValueTask DisposeAsyncCore() => ValueTask.CompletedTask; +#pragma warning restore PS0018 public async ValueTask DisposeAsync() { @@ -16,6 +19,6 @@ public async ValueTask DisposeAsync() public IMonitoringIngestionUnitOfWork Monitoring { get; protected set; } public IRecoverabilityIngestionUnitOfWork Recoverability { get; protected set; } - public virtual Task Complete(CancellationToken cancellationToken) => Task.CompletedTask; + public virtual Task Complete(CancellationToken cancellationToken = default) => Task.CompletedTask; } } diff --git a/src/ServiceControl/CompositeViews/Messages/GetMessagesController.cs b/src/ServiceControl/CompositeViews/Messages/GetMessagesController.cs index bef095821d..a7d22eb027 100644 --- a/src/ServiceControl/CompositeViews/Messages/GetMessagesController.cs +++ b/src/ServiceControl/CompositeViews/Messages/GetMessagesController.cs @@ -88,7 +88,7 @@ public async Task Get(string id, [FromQuery(Name = "instance_id") { if (string.IsNullOrWhiteSpace(instanceId) || instanceId == settings.InstanceId) { - var result = await bodyStorage.TryFetch(id); + var result = await bodyStorage.TryFetch(id, cancellationToken); if (result == null) { diff --git a/src/ServiceControl/CustomChecks/FailedErrorImportCustomCheck.cs b/src/ServiceControl/CustomChecks/FailedErrorImportCustomCheck.cs index e8000d18f4..029bcc41c7 100644 --- a/src/ServiceControl/CustomChecks/FailedErrorImportCustomCheck.cs +++ b/src/ServiceControl/CustomChecks/FailedErrorImportCustomCheck.cs @@ -18,7 +18,7 @@ public FailedErrorImportCustomCheck(IFailedErrorImportDataStore store, ILogger PerformCheck(CancellationToken cancellationToken = default) { - var hasFailedImports = await store.QueryContainsFailedImports(); + var hasFailedImports = await store.QueryContainsFailedImports(cancellationToken); if (hasFailedImports) { diff --git a/src/ServiceControl/Operations/ErrorIngestionFaultPolicy.cs b/src/ServiceControl/Operations/ErrorIngestionFaultPolicy.cs index 6719738a29..93249bca49 100644 --- a/src/ServiceControl/Operations/ErrorIngestionFaultPolicy.cs +++ b/src/ServiceControl/Operations/ErrorIngestionFaultPolicy.cs @@ -75,7 +75,7 @@ async Task DoLogging(Exception exception, FailedErrorImport failure, Cancellatio logger.LogError(exception, "Failed importing error message"); // Write to data store - await store.StoreFailedErrorImport(failure); + await store.StoreFailedErrorImport(failure, cancellationToken); if (!AppEnvironment.RunningInContainer) { diff --git a/src/ServiceControl/Operations/ErrorIngestor.cs b/src/ServiceControl/Operations/ErrorIngestor.cs index 9b53d7a0b1..9e5040a601 100644 --- a/src/ServiceControl/Operations/ErrorIngestor.cs +++ b/src/ServiceControl/Operations/ErrorIngestor.cs @@ -120,7 +120,7 @@ async Task> PersistFailedMessages(List> Process(IReadOnlyList + await store.ProcessFailedErrorImports(async (transportMessage, token) => { var messageContext = new MessageContext( transportMessage.Id, @@ -32,7 +32,7 @@ await store.ProcessFailedErrorImports(async transportMessage => var taskCompletionSource = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); messageContext.SetTaskCompletionSource(taskCompletionSource); - await errorIngestor.Ingest([messageContext], cancellationToken); + await errorIngestor.Ingest([messageContext], token); await taskCompletionSource.Task; }, cancellationToken); } diff --git a/src/ServiceControl/Operations/RetryConfirmationProcessor.cs b/src/ServiceControl/Operations/RetryConfirmationProcessor.cs index 33a43180c1..8b948a7b91 100644 --- a/src/ServiceControl/Operations/RetryConfirmationProcessor.cs +++ b/src/ServiceControl/Operations/RetryConfirmationProcessor.cs @@ -23,7 +23,7 @@ public async Task Process(List contexts, IIngestionUnitOfWork un foreach (var context in contexts) { var retriedMessageUniqueId = context.Headers[RetryUniqueMessageIdHeader]; - await unitOfWork.Recoverability.RecordSuccessfulRetry(retriedMessageUniqueId); + await unitOfWork.Recoverability.RecordSuccessfulRetry(retriedMessageUniqueId, cancellationToken); } } From 39c0ea94d4c967a2b55a9ebf0ef72b710824e4d6 Mon Sep 17 00:00:00 2001 From: John Simons Date: Wed, 12 Aug 2026 13:37:37 +1000 Subject: [PATCH 2/2] Propagate cancellation tokens through the failed message retry data store Pass the cancellation token to the body storage fetch operation to ensure proper cancellation of database calls during failed message body retrieval. --- .../Implementation/FailedMessageRetryDataStore.cs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/ServiceControl.Persistence.EFCore/Implementation/FailedMessageRetryDataStore.cs b/src/ServiceControl.Persistence.EFCore/Implementation/FailedMessageRetryDataStore.cs index 44f928c403..f50e2b8be6 100644 --- a/src/ServiceControl.Persistence.EFCore/Implementation/FailedMessageRetryDataStore.cs +++ b/src/ServiceControl.Persistence.EFCore/Implementation/FailedMessageRetryDataStore.cs @@ -55,7 +55,7 @@ public Task ProcessPendingRetries(DateTime periodFrom, DateTime periodTo, string public async Task GetFailedMessageBody(string uniqueMessageId, CancellationToken cancellationToken) { - var result = await bodyStorage.TryFetch(uniqueMessageId) + var result = await bodyStorage.TryFetch(uniqueMessageId, cancellationToken) ?? throw new InvalidOperationException("IBodyStorage.TryFetch result cannot be null"); if (!result.HasResult)