New cancellation token allows further improvement to reduce shutdown duration to close within 30 seconds

ramonsmits · ramonsmits · commit 21ffe9cba8bb · 2025-02-04T15:50:25.000+01:00
diff --git a/src/ServiceControl.Audit.Persistence.InMemory/InMemoryAuditIngestionUnitOfWorkFactory.cs b/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));
diff --git a/src/ServiceControl.Audit.Persistence.RavenDB/UnitOfWork/RavenAuditUnitOfWorkFactory.cs b/src/ServiceControl.Audit.Persistence.RavenDB/UnitOfWork/RavenAuditUnitOfWorkFactory.cs
@@ -13,9 +13,10 @@ 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 timedCancellationSource = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
+            timedCancellationSource.CancelAfter(databaseConfiguration.BulkInsertCommitTimeout);
             var bulkInsert = (await documentStoreProvider.GetDocumentStore(timedCancellationSource.Token))
                 .BulkInsert(new BulkInsertOptions { SkipOverwriteIfUnchanged = true, }, timedCancellationSource.Token);
 
diff --git a/src/ServiceControl.Audit.Persistence/UnitOfWork/IAuditIngestionUnitOfWorkFactory.cs b/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 = default); //Throws if not enough space or some other problem preventing from writing data
         bool CanIngestMore();
     }
 }
diff --git a/src/ServiceControl.Audit/Auditing/AuditIngestion.cs b/src/ServiceControl.Audit/Auditing/AuditIngestion.cs
@@ -129,7 +129,7 @@ async Task EnsureStarted(CancellationToken cancellationToken = default)
 
                 queueIngestor = transportInfrastructure.Receivers[inputEndpoint];
 
-                await auditIngestor.VerifyCanReachForwardingAddress();
+                await auditIngestor.VerifyCanReachForwardingAddress(cancellationToken);
 
                 await queueIngestor.StartReceive(cancellationToken);
 
@@ -230,7 +230,7 @@ async Task Loop(CancellationToken cancellationToken)
                         auditBatchSize.Record(contexts.Count);
                         var sw = Stopwatch.StartNew();
 
-                        await auditIngestor.Ingest(contexts);
+                        await auditIngestor.Ingest(contexts, cancellationToken);
                         auditBatchDuration.Record(sw.ElapsedMilliseconds);
                     }
                     catch (OperationCanceledException e) when (e.CancellationToken == cancellationToken)
diff --git a/src/ServiceControl.Audit/Auditing/AuditIngestor.cs b/src/ServiceControl.Audit/Auditing/AuditIngestor.cs
@@ -3,6 +3,7 @@
     using System;
     using System.Collections.Generic;
     using System.Linq;
+    using System.Threading;
     using System.Threading.Tasks;
     using Infrastructure.Settings;
     using Monitoring;
@@ -41,14 +42,14 @@ ITransportCustomization transportCustomization
             );
         }
 
-        public async Task Ingest(List<MessageContext> contexts)
+        public async Task Ingest(List<MessageContext> contexts, CancellationToken cancellationToken)
         {
             if (Log.IsDebugEnabled)
             {
                 Log.Debug($"Ingesting {contexts.Count} message contexts");
             }
 
-            var stored = await auditPersister.Persist(contexts);
+            var stored = await auditPersister.Persist(contexts, cancellationToken);
 
             try
             {
@@ -59,7 +60,7 @@ public async Task Ingest(List<MessageContext> contexts)
                         Log.Debug($"Forwarding {stored.Count} messages");
                     }
 
-                    await Forward(stored, logQueueAddress);
+                    await Forward(stored, logQueueAddress, cancellationToken);
                     if (Log.IsDebugEnabled)
                     {
                         Log.Debug("Forwarded messages");
@@ -86,7 +87,7 @@ public async Task Ingest(List<MessageContext> contexts)
             }
         }
 
-        Task Forward(IReadOnlyCollection<MessageContext> messageContexts, string forwardingAddress)
+        Task Forward(IReadOnlyCollection<MessageContext> messageContexts, string forwardingAddress, CancellationToken cancellationToken)
         {
             var transportOperations = new TransportOperation[messageContexts.Count]; //We could allocate based on the actual number of ProcessedMessages but this should be OK
             var index = 0;
@@ -103,7 +104,8 @@ Task Forward(IReadOnlyCollection<MessageContext> messageContexts, string forward
                 var outgoingMessage = new OutgoingMessage(
                     messageContext.NativeMessageId,
                     messageContext.Headers,
-                    messageContext.Body);
+                    messageContext.Body
+                );
 
                 // Forwarded messages should last as long as possible
                 outgoingMessage.Headers.Remove(Headers.TimeToBeReceived);
@@ -115,12 +117,13 @@ Task Forward(IReadOnlyCollection<MessageContext> messageContexts, string forward
             return anyContext != null
                 ? messageDispatcher.Value.Dispatch(
                     new TransportOperations(transportOperations),
-                    anyContext.TransportTransaction
+                    anyContext.TransportTransaction,
+                    cancellationToken
                 )
                 : Task.CompletedTask;
         }
 
-        public async Task VerifyCanReachForwardingAddress()
+        public async Task VerifyCanReachForwardingAddress(CancellationToken cancellationToken)
         {
             if (!settings.ForwardAuditMessages)
             {
@@ -137,7 +140,7 @@ public async Task VerifyCanReachForwardingAddress()
                     )
                 );
 
-                await messageDispatcher.Value.Dispatch(transportOperations, new TransportTransaction());
+                await messageDispatcher.Value.Dispatch(transportOperations, new TransportTransaction(), cancellationToken);
             }
             catch (Exception e)
             {
diff --git a/src/ServiceControl.Audit/Auditing/AuditPersister.cs b/src/ServiceControl.Audit/Auditing/AuditPersister.cs
@@ -5,6 +5,7 @@
     using System.Diagnostics;
     using System.Diagnostics.Metrics;
     using System.Text.Json;
+    using System.Threading;
     using System.Threading.Tasks;
     using Infrastructure;
     using Monitoring;
@@ -23,7 +24,7 @@ class AuditPersister(IAuditIngestionUnitOfWorkFactory unitOfWorkFactory,
         IMessageSession messageSession,
         Lazy<IMessageDispatcher> messageDispatcher)
     {
-        public async Task<IReadOnlyList<MessageContext>> Persist(IReadOnlyList<MessageContext> contexts)
+        public async Task<IReadOnlyList<MessageContext>> Persist(IReadOnlyList<MessageContext> contexts, CancellationToken cancellationToken)
         {
             var stopwatch = Stopwatch.StartNew();
 
@@ -37,7 +38,7 @@ public async Task<IReadOnlyList<MessageContext>> Persist(IReadOnlyList<MessageCo
             try
             {
                 // deliberately not using the using statement because we dispose async explicitly
-                unitOfWork = await unitOfWorkFactory.StartNew(contexts.Count);
+                unitOfWork = await unitOfWorkFactory.StartNew(contexts.Count, cancellationToken);
                 var inserts = new List<Task>(contexts.Count);
                 foreach (var context in contexts)
                 {
diff --git a/src/ServiceControl.Audit/Auditing/ImportFailedAudits.cs b/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;
@@ -37,7 +37,7 @@ await failedAuditStore.ProcessFailedMessages(
                             var taskCompletionSource = new TaskCompletionSource<bool>(TaskCreationOptions.RunContinuationsAsynchronously);
                             messageContext.SetTaskCompletionSource(taskCompletionSource);
 
-                            await auditIngestor.Ingest([messageContext]);
+                            await auditIngestor.Ingest([messageContext], cancellationToken);
 
                             await taskCompletionSource.Task;
 

Original file line number	Diff line number	Diff line change
`@@ -1,5 +1,6 @@`
`1`	`1`	`namespace ServiceControl.Audit.Persistence.InMemory`
`2`	`2`	`{`
	`3`	`+ using System.Threading;`
`3`	`4`	`using System.Threading.Tasks;`
`4`	`5`	`using ServiceControl.Audit.Auditing.BodyStorage;`
`5`	`6`	`using ServiceControl.Audit.Persistence.UnitOfWork;`
`@@ -12,7 +13,7 @@ public InMemoryAuditIngestionUnitOfWorkFactory(InMemoryAuditDataStore dataStore,`
`12`	`13`	`bodyStorageEnricher = new BodyStorageEnricher(bodyStorage, settings);`
`13`	`14`	`}`
`14`	`15`
`15`		`- public ValueTask<IAuditIngestionUnitOfWork> StartNew(int batchSize)`
	`16`	`+ public ValueTask<IAuditIngestionUnitOfWork> StartNew(int batchSize, CancellationToken cancellationToken)`
`16`	`17`	`{`
`17`	`18`	`//The batchSize argument is ignored: the in-memory storage implementation doesn't support batching.`
`18`	`19`	`return new ValueTask<IAuditIngestionUnitOfWork>(new InMemoryAuditIngestionUnitOfWork(dataStore, bodyStorageEnricher));`
Original file line number	Diff line number	Diff line change
`@@ -13,9 +13,10 @@ class RavenAuditIngestionUnitOfWorkFactory(`
`13`	`13`	`MinimumRequiredStorageState customCheckState)`
`14`	`14`	`: IAuditIngestionUnitOfWorkFactory`
`15`	`15`	`{`
`16`		`- public async ValueTask<IAuditIngestionUnitOfWork> StartNew(int batchSize)`
	`16`	`+ public async ValueTask<IAuditIngestionUnitOfWork> StartNew(int batchSize, CancellationToken cancellationToken)`
`17`	`17`	`{`
`18`		`- var timedCancellationSource = new CancellationTokenSource(databaseConfiguration.BulkInsertCommitTimeout);`
	`18`	`+ var timedCancellationSource = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);`
	`19`	`+ timedCancellationSource.CancelAfter(databaseConfiguration.BulkInsertCommitTimeout);`
`19`	`20`	`var bulkInsert = (await documentStoreProvider.GetDocumentStore(timedCancellationSource.Token))`
`20`	`21`	`.BulkInsert(new BulkInsertOptions { SkipOverwriteIfUnchanged = true, }, timedCancellationSource.Token);`
`21`	`22`
Original file line number	Diff line number	Diff line change
`@@ -1,10 +1,11 @@`
`1`	`1`	`namespace ServiceControl.Audit.Persistence.UnitOfWork`
`2`	`2`	`{`
	`3`	`+ using System.Threading;`
`3`	`4`	`using System.Threading.Tasks;`
`4`	`5`
`5`	`6`	`public interface IAuditIngestionUnitOfWorkFactory`
`6`	`7`	`{`
`7`		`- ValueTask<IAuditIngestionUnitOfWork> StartNew(int batchSize); //Throws if not enough space or some other problem preventing from writing data`
	`8`	`+ ValueTask<IAuditIngestionUnitOfWork> StartNew(int batchSize, CancellationToken cancellationToken = default); //Throws if not enough space or some other problem preventing from writing data`
`8`	`9`	`bool CanIngestMore();`
`9`	`10`	`}`
`10`	`11`	`}`
Original file line number	Diff line number	Diff line change
`@@ -3,6 +3,7 @@`
`3`	`3`	`using System;`
`4`	`4`	`using System.Collections.Generic;`
`5`	`5`	`using System.Linq;`
	`6`	`+ using System.Threading;`
`6`	`7`	`using System.Threading.Tasks;`
`7`	`8`	`using Infrastructure.Settings;`
`8`	`9`	`using Monitoring;`
`@@ -41,14 +42,14 @@ ITransportCustomization transportCustomization`
`41`	`42`	`);`
`42`	`43`	`}`
`43`	`44`
`44`		`- public async Task Ingest(List<MessageContext> contexts)`
	`45`	`+ public async Task Ingest(List<MessageContext> contexts, CancellationToken cancellationToken)`
`45`	`46`	`{`
`46`	`47`	`if (Log.IsDebugEnabled)`
`47`	`48`	`{`
`48`	`49`	`Log.Debug($"Ingesting {contexts.Count} message contexts");`
`49`	`50`	`}`
`50`	`51`
`51`		`- var stored = await auditPersister.Persist(contexts);`
	`52`	`+ var stored = await auditPersister.Persist(contexts, cancellationToken);`
`52`	`53`
`53`	`54`	`try`
`54`	`55`	`{`
`@@ -59,7 +60,7 @@ public async Task Ingest(List<MessageContext> contexts)`
`59`	`60`	`Log.Debug($"Forwarding {stored.Count} messages");`
`60`	`61`	`}`
`61`	`62`
`62`		`- await Forward(stored, logQueueAddress);`
	`63`	`+ await Forward(stored, logQueueAddress, cancellationToken);`
`63`	`64`	`if (Log.IsDebugEnabled)`
`64`	`65`	`{`
`65`	`66`	`Log.Debug("Forwarded messages");`
`@@ -86,7 +87,7 @@ public async Task Ingest(List<MessageContext> contexts)`
`86`	`87`	`}`
`87`	`88`	`}`
`88`	`89`
`89`		`- Task Forward(IReadOnlyCollection<MessageContext> messageContexts, string forwardingAddress)`
	`90`	`+ Task Forward(IReadOnlyCollection<MessageContext> messageContexts, string forwardingAddress, CancellationToken cancellationToken)`
`90`	`91`	`{`
`91`	`92`	`var transportOperations = new TransportOperation[messageContexts.Count]; //We could allocate based on the actual number of ProcessedMessages but this should be OK`
`92`	`93`	`var index = 0;`
`@@ -103,7 +104,8 @@ Task Forward(IReadOnlyCollection<MessageContext> messageContexts, string forward`
`103`	`104`	`var outgoingMessage = new OutgoingMessage(`
`104`	`105`	`messageContext.NativeMessageId,`
`105`	`106`	`messageContext.Headers,`
`106`		`- messageContext.Body);`
	`107`	`+ messageContext.Body`
	`108`	`+ );`
`107`	`109`
`108`	`110`	`// Forwarded messages should last as long as possible`
`109`	`111`	`outgoingMessage.Headers.Remove(Headers.TimeToBeReceived);`
`@@ -115,12 +117,13 @@ Task Forward(IReadOnlyCollection<MessageContext> messageContexts, string forward`
`115`	`117`	`return anyContext != null`
`116`	`118`	`? messageDispatcher.Value.Dispatch(`
`117`	`119`	`new TransportOperations(transportOperations),`
`118`		`- anyContext.TransportTransaction`
	`120`	`+ anyContext.TransportTransaction,`
	`121`	`+ cancellationToken`
`119`	`122`	`)`
`120`	`123`	`: Task.CompletedTask;`
`121`	`124`	`}`
`122`	`125`
`123`		`- public async Task VerifyCanReachForwardingAddress()`
	`126`	`+ public async Task VerifyCanReachForwardingAddress(CancellationToken cancellationToken)`
`124`	`127`	`{`
`125`	`128`	`if (!settings.ForwardAuditMessages)`
`126`	`129`	`{`
`@@ -137,7 +140,7 @@ public async Task VerifyCanReachForwardingAddress()`
`137`	`140`	`)`
`138`	`141`	`);`
`139`	`142`
`140`		`- await messageDispatcher.Value.Dispatch(transportOperations, new TransportTransaction());`
	`143`	`+ await messageDispatcher.Value.Dispatch(transportOperations, new TransportTransaction(), cancellationToken);`
`141`	`144`	`}`
`142`	`145`	`catch (Exception e)`
`143`	`146`	`{`
Original file line number	Diff line number	Diff line change
`@@ -5,6 +5,7 @@`
`5`	`5`	`using System.Diagnostics;`
`6`	`6`	`using System.Diagnostics.Metrics;`
`7`	`7`	`using System.Text.Json;`
	`8`	`+ using System.Threading;`
`8`	`9`	`using System.Threading.Tasks;`
`9`	`10`	`using Infrastructure;`
`10`	`11`	`using Monitoring;`
`@@ -23,7 +24,7 @@ class AuditPersister(IAuditIngestionUnitOfWorkFactory unitOfWorkFactory,`
`23`	`24`	`IMessageSession messageSession,`
`24`	`25`	`Lazy<IMessageDispatcher> messageDispatcher)`
`25`	`26`	`{`
`26`		`- public async Task<IReadOnlyList<MessageContext>> Persist(IReadOnlyList<MessageContext> contexts)`
	`27`	`+ public async Task<IReadOnlyList<MessageContext>> Persist(IReadOnlyList<MessageContext> contexts, CancellationToken cancellationToken)`
`27`	`28`	`{`
`28`	`29`	`var stopwatch = Stopwatch.StartNew();`
`29`	`30`
`@@ -37,7 +38,7 @@ public async Task<IReadOnlyList<MessageContext>> Persist(IReadOnlyList<MessageCo`
`37`	`38`	`try`
`38`	`39`	`{`
`39`	`40`	`// deliberately not using the using statement because we dispose async explicitly`
`40`		`- unitOfWork = await unitOfWorkFactory.StartNew(contexts.Count);`
	`41`	`+ unitOfWork = await unitOfWorkFactory.StartNew(contexts.Count, cancellationToken);`
`41`	`42`	`var inserts = new List<Task>(contexts.Count);`
`42`	`43`	`foreach (var context in contexts)`
`43`	`44`	`{`