Support cancellation of audit ingestion

Author: ramonsmits

src/ServiceControl.Audit.Persistence.InMemory/InMemoryAuditIngestionUnitOfWorkFactory.cs

@@ -1,5 +1,6 @@
 namespace ServiceControl.Audit.Persistence.InMemory
 {
+    using System.Threading;
     using System.Threading.Tasks;
     using ServiceControl.Audit.Auditing.BodyStorage;
     using ServiceControl.Audit.Persistence.UnitOfWork;
@@ -12,7 +13,7 @@ public InMemoryAuditIngestionUnitOfWorkFactory(InMemoryAuditDataStore dataStore,
             bodyStorageEnricher = new BodyStorageEnricher(bodyStorage, settings);
         }
 
-        public ValueTask<IAuditIngestionUnitOfWork> StartNew(int batchSize)
+        public ValueTask<IAuditIngestionUnitOfWork> StartNew(int batchSize, CancellationToken cancellationToken)
         {
             //The batchSize argument is ignored: the in-memory storage implementation doesn't support batching.
             return new ValueTask<IAuditIngestionUnitOfWork>(new InMemoryAuditIngestionUnitOfWork(dataStore, bodyStorageEnricher));

src/ServiceControl.Audit.Persistence.RavenDB/UnitOfWork/RavenAuditUnitOfWorkFactory.cs

@@ -13,7 +13,7 @@ class RavenAuditIngestionUnitOfWorkFactory(
         MinimumRequiredStorageState customCheckState)
         : IAuditIngestionUnitOfWorkFactory
     {
-        public async ValueTask<IAuditIngestionUnitOfWork> StartNew(int batchSize)
+        public async ValueTask<IAuditIngestionUnitOfWork> StartNew(int batchSize, CancellationToken cancellationToken)
         {
             var timedCancellationSource = new CancellationTokenSource(databaseConfiguration.BulkInsertCommitTimeout);
             var bulkInsert = (await documentStoreProvider.GetDocumentStore(timedCancellationSource.Token))

src/ServiceControl.Audit.Persistence.Tests/PersistenceTestFixture.cs

@@ -64,7 +64,7 @@ protected string GetManifestPath()
             configuration.AuditIngestionUnitOfWorkFactory;
 
         protected ValueTask<IAuditIngestionUnitOfWork> StartAuditUnitOfWork(int batchSize) =>
-            AuditIngestionUnitOfWorkFactory.StartNew(batchSize);
+            AuditIngestionUnitOfWorkFactory.StartNew(batchSize, TestContext.CurrentContext.CancellationToken);
 
         protected IServiceProvider ServiceProvider => configuration.ServiceProvider;
 

src/ServiceControl.Audit.Persistence/UnitOfWork/IAuditIngestionUnitOfWorkFactory.cs

@@ -1,10 +1,11 @@
 namespace ServiceControl.Audit.Persistence.UnitOfWork
 {
+    using System.Threading;
     using System.Threading.Tasks;
 
     public interface IAuditIngestionUnitOfWorkFactory
     {
-        ValueTask<IAuditIngestionUnitOfWork> StartNew(int batchSize); //Throws if not enough space or some other problem preventing from writing data
+        ValueTask<IAuditIngestionUnitOfWork> StartNew(int batchSize, CancellationToken cancellationToken); //Throws if not enough space or some other problem preventing from writing data
         bool CanIngestMore();
     }
 }

src/ServiceControl.Audit/Auditing/ImportFailedAudits.cs

@@ -23,7 +23,7 @@ public ImportFailedAudits(
 
         public async Task Run(CancellationToken cancellationToken = default)
         {
-            await auditIngestor.VerifyCanReachForwardingAddress();
+            await auditIngestor.VerifyCanReachForwardingAddress(cancellationToken);
 
             var succeeded = 0;
             var failed = 0;
@@ -33,11 +33,18 @@ await failedAuditStore.ProcessFailedMessages(
                     {
                         try
                         {
-                            var messageContext = new MessageContext(transportMessage.Id, transportMessage.Headers, transportMessage.Body, EmptyTransaction, settings.AuditQueue, EmptyContextBag);
+                            var messageContext = new MessageContext(
+                                transportMessage.Id,
+                                transportMessage.Headers,
+                                transportMessage.Body,
+                                EmptyTransaction,
+                                settings.AuditQueue,
+                                EmptyContextBag
+                            );
                             var taskCompletionSource = new TaskCompletionSource<bool>(TaskCreationOptions.RunContinuationsAsynchronously);
                             messageContext.SetTaskCompletionSource(taskCompletionSource);
 
-                            await auditIngestor.Ingest([messageContext]);
+                            await auditIngestor.Ingest([messageContext], cancellationToken);
 
                             await taskCompletionSource.Task;