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.

Author: johnsimons

src/ServiceControl.Persistence.EFCore/Implementation/BodyStorage/BodyStorage.cs

@@ -19,9 +19,9 @@ namespace ServiceControl.Persistence.EFCore.Implementation.BodyStorage;
 /// </remarks>
 public class BodyStorage(IServiceScopeFactory scopeFactory, IBodyStoragePersistence storagePersistence) : DataStoreBase(scopeFactory), IBodyStorage
 {
-    public async Task<MessageBodyStreamResult?> TryFetch(string bodyId)
+    public async Task<MessageBodyStreamResult?> 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<BodyRow?> ResolveBody(ServiceControlDbContext dbContext, string bodyId)
+    static async Task<BodyRow?> 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<BodyRow?> Query(ServiceControlDbContext dbContext, Expression<Func<FailedMessageEntity, bool>> predicate) =>
+    static Task<BodyRow?> Query(ServiceControlDbContext dbContext, Expression<Func<FailedMessageEntity, bool>> 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
     {

src/ServiceControl.Persistence.EFCore/Implementation/BodyStorage/FileSystemBodyStoragePersistence.cs

@@ -72,11 +72,14 @@ public async Task WriteBody(string bodyId, ReadOnlyMemory<byte> 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<MessageBodyFileResult?> ReadBody(string bodyId, CancellationToken cancellationToken = default)

src/ServiceControl.Persistence.EFCore/Implementation/FailedErrorImportDataStore.cs

@@ -22,14 +22,14 @@ public class FailedErrorImportDataStore(
 {
     const int BatchSize = 100;
 
-    public Task<bool> QueryContainsFailedImports() =>
+    public Task<bool> 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<FailedTransportMessage, Task> processMessage, CancellationToken cancellationToken)
+    public async Task ProcessFailedErrorImports(Func<FailedTransportMessage, CancellationToken, Task> processMessage, CancellationToken cancellationToken = default)
     {
         var succeeded = 0;
         var failed = 0;
@@ -94,7 +94,7 @@ public async Task ProcessFailedErrorImports(Func<FailedTransportMessage, Task> 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
     }
 }

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);

src/ServiceControl.Persistence.EFCore/Implementation/UnitOfWork/EFIngestionUnitOfWorkFactory.cs

@@ -13,7 +13,7 @@ public class EFIngestionUnitOfWorkFactory(
     IFailedMessageIngestionSqlDialect dialect,
     TimeProvider timeProvider) : IIngestionUnitOfWorkFactory
 {
-    public ValueTask<IIngestionUnitOfWork> StartNew()
+    public ValueTask<IIngestionUnitOfWork> StartNew(CancellationToken cancellationToken = default)
     {
         var scope = serviceProvider.CreateAsyncScope();
         var dbContext = scope.ServiceProvider.GetRequiredService<ServiceControlDbContext>();

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);
 

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<FailedMessage.FailureGroup> groups)
+        List<FailedMessage.FailureGroup> 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<EndpointDetails>(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));
 

src/ServiceControl.Persistence.EFCore/Implementation/UnitOfWork/FailedMessageBatchWriter.cs

@@ -17,7 +17,7 @@ public async Task Write(
         IReadOnlyCollection<KnownEndpoint> knownEndpoints,
         IReadOnlyCollection<Guid> confirmedRetries,
         DateTime now,
-        CancellationToken cancellationToken)
+        CancellationToken cancellationToken = default)
     {
         var (failedMessages, groups) = Fold(attempts, now);
         var endpoints = BuildEndpointRows(knownEndpoints);

src/ServiceControl.Persistence.RavenDB/ErrorMessagesDataStore.cs

@@ -511,7 +511,7 @@ record struct FailedMessageProjection(string UniqueMessageId);
         public async Task<byte[]> 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)

src/ServiceControl.Persistence.RavenDB/FailedErrorImportDataStore.cs

@@ -9,21 +9,21 @@
 
     class FailedErrorImportDataStore(IRavenSessionProvider sessionProvider, ILogger<FailedErrorImportDataStore> 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<FailedTransportMessage, Task> processMessage, CancellationToken cancellationToken)
+        public async Task ProcessFailedErrorImports(Func<FailedTransportMessage, CancellationToken, Task> processMessage, CancellationToken cancellationToken = default)
         {
             var succeeded = 0;
             var failed = 0;
@@ -36,7 +36,7 @@ public async Task ProcessFailedErrorImports(Func<FailedTransportMessage, Task> 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<FailedTransportMessage, Task> p
             }
         }
 
-        public async Task<bool> QueryContainsFailedImports()
+        public async Task<bool> QueryContainsFailedImports(CancellationToken cancellationToken = default)
         {
-            using var session = await sessionProvider.OpenSession();
+            using var session = await sessionProvider.OpenSession(cancellationToken: cancellationToken);
             var query = session.Query<FailedErrorImport, FailedErrorImportIndex>();
-            await using var ie = await session.Advanced.StreamAsync(query);
+            await using var ie = await session.Advanced.StreamAsync(query, cancellationToken);
             return await ie.MoveNextAsync();
         }