Propagate cancellation tokens through ingestion and body storage pers…

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/FailedMessageRetryDataStore.cs

@@ -55,7 +55,7 @@ public Task ProcessPendingRetries(DateTime periodFrom, DateTime periodTo, string
 
     public async Task<byte[]> 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)

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

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<MessageBodyStreamResult> TryFetch(string bodyId)
+        public async Task<MessageBodyStreamResult> 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<MessageBodyStreamResult> TryFetch(string bodyId)
                 .OfType<FailedMessage>()
                 .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<MessageBodyStreamResult> ResultForUniqueId(IAsyncDocumentSession session, string uniqueId)
+        async Task<MessageBodyStreamResult> 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)
             {

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

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<IIngestionUnitOfWork> StartNew()
+        public ValueTask<IIngestionUnitOfWork> StartNew(CancellationToken cancellationToken = default)
             => new(new RavenIngestionUnitOfWork(sessionProvider, expirationManager, settings));
 
         public bool CanIngestMore() => customCheckState.CanIngestMore;

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;

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