Propagate cancellation tokens through custom checks, event log and me…

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

@@ -8,7 +8,7 @@ namespace ServiceControl.Persistence.EFCore.Implementation;
 
 public class CustomCheckDataStore(IServiceScopeFactory scopeFactory) : DataStoreBase(scopeFactory), ICustomChecksDataStore
 {
-    public Task<CheckStateChange> UpdateCustomCheckStatus(CustomCheckDetail detail) => ExecuteWithDbContext(async context =>
+    public Task<CheckStateChange> UpdateCustomCheckStatus(CustomCheckDetail detail, CancellationToken cancellationToken = default) => ExecuteWithDbContext(async context =>
     {
         var status = CheckStateChange.Unchanged;
 
@@ -43,11 +43,12 @@ await context.UpsertAsync([detail.GetDeterministicId()],
                 entity.Category = detail.Category;
                 entity.ReportedAt = detail.ReportedAt;
                 entity.FailureReason = detail.FailureReason;
-            });
+            },
+            cancellationToken);
         return status;
     });
 
-    public Task<QueryResult<IList<CustomCheck>>> GetStats(PagingInfo paging, string? status = null) => ExecuteWithDbContext(async context =>
+    public Task<QueryResult<IList<CustomCheck>>> GetStats(PagingInfo paging, string? status = null, CancellationToken cancellationToken = default) => ExecuteWithDbContext(async context =>
     {
         var query = context.CustomChecks.AsQueryable().AsNoTracking();
 
@@ -62,7 +63,7 @@ public Task<QueryResult<IList<CustomCheck>>> GetStats(PagingInfo paging, string?
             .OrderBy(c => c.ReportedAt)
             .Skip(paging.Offset)
             .Take(paging.PageSize)
-            .ToListAsync();
+            .ToListAsync(cancellationToken);
 
         return new QueryResult<IList<CustomCheck>>(page.Select(c => new CustomCheck
         {
@@ -75,7 +76,7 @@ public Task<QueryResult<IList<CustomCheck>>> GetStats(PagingInfo paging, string?
         }).ToList(), new QueryStatsInfo("", page.Count, false));
     });
 
-    public Task DeleteCustomCheck(Guid id) => ExecuteWithDbContext(async context => await context.CustomChecks.AsNoTracking().Where(cc => cc.Id == id).ExecuteDeleteAsync());
+    public Task DeleteCustomCheck(Guid id, CancellationToken cancellationToken = default) => ExecuteWithDbContext(async context => await context.CustomChecks.AsNoTracking().Where(cc => cc.Id == id).ExecuteDeleteAsync(cancellationToken));
 
-    public Task<int> GetNumberOfFailedChecks() => ExecuteWithDbContext(async context => await context.CustomChecks.AsNoTracking().CountAsync(p => p.Status == Status.Fail));
+    public Task<int> GetNumberOfFailedChecks(CancellationToken cancellationToken = default) => ExecuteWithDbContext(async context => await context.CustomChecks.AsNoTracking().CountAsync(p => p.Status == Status.Fail, cancellationToken));
 }

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

@@ -8,7 +8,7 @@ namespace ServiceControl.Persistence.EFCore.Implementation;
 
 public class EventLogDataStore(IServiceScopeFactory scopeFactory) : DataStoreBase(scopeFactory), IEventLogDataStore
 {
-    public Task Add(EventLogItem logItem) =>
+    public Task Add(EventLogItem logItem, CancellationToken cancellationToken = default) =>
         ExecuteWithDbContext(async dbContext =>
         {
             dbContext.EventLogItems.Add(new EventLogItemEntity
@@ -21,11 +21,11 @@ public Task Add(EventLogItem logItem) =>
                 EventType = logItem.EventType
             });
 
-            await dbContext.SaveChangesAsync();
+            await dbContext.SaveChangesAsync(cancellationToken);
         });
 
     public Task<QueryResult<IList<EventLogItemView>>> GetEventLogItems(
-        PagingInfo pagingInfo, string? knownVersion = null) =>
+        PagingInfo pagingInfo, string? knownVersion = null, CancellationToken cancellationToken = default) =>
         ExecuteWithDbContext(async dbContext =>
         {
             var query = dbContext.EventLogItems.AsNoTracking();
@@ -40,7 +40,7 @@ public Task<QueryResult<IList<EventLogItemView>>> GetEventLogItems(
                     Newest = g.Max(e => (DateTime?)e.RaisedAt),
                     HighestId = g.Max(e => (long?)e.Id)
                 })
-                .FirstOrDefaultAsync();
+                .FirstOrDefaultAsync(cancellationToken);
 
             var total = stats?.Total ?? 0;
             var version = Version(total, stats?.Newest, stats?.HighestId);
@@ -71,7 +71,7 @@ public Task<QueryResult<IList<EventLogItemView>>> GetEventLogItems(
                     Category = e.Category,
                     EventType = e.EventType
                 })
-                .ToListAsync();
+                .ToListAsync(cancellationToken);
 
             return new QueryResult<IList<EventLogItemView>>(items, queryStats);
         });

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

@@ -7,12 +7,12 @@ namespace ServiceControl.Persistence.EFCore.Implementation;
 
 public class MessageRedirectsDataStore(IServiceScopeFactory scopeFactory) : DataStoreBase(scopeFactory), IMessageRedirectsDataStore
 {
-    public Task<IReadOnlyList<MessageRedirect>> GetRedirects() =>
+    public Task<IReadOnlyList<MessageRedirect>> GetRedirects(CancellationToken cancellationToken = default) =>
         ExecuteWithDbContext<IReadOnlyList<MessageRedirect>>(async dbContext =>
         {
             var rows = await dbContext.MessageRedirects
                 .AsNoTracking()
-                .ToListAsync();
+                .ToListAsync(cancellationToken);
 
             return [.. rows.Select(row => new MessageRedirect
             {
@@ -22,7 +22,7 @@ public Task<IReadOnlyList<MessageRedirect>> GetRedirects() =>
             })];
         });
 
-    public Task AddRedirect(MessageRedirect redirect) =>
+    public Task AddRedirect(MessageRedirect redirect, CancellationToken cancellationToken = default) =>
         ExecuteWithDbContext(async dbContext =>
         {
             dbContext.MessageRedirects.Add(new MessageRedirectEntity
@@ -32,18 +32,18 @@ public Task AddRedirect(MessageRedirect redirect) =>
                 LastModified = redirect.LastModified
             });
 
-            await dbContext.SaveChangesAsync();
+            await dbContext.SaveChangesAsync(cancellationToken);
         });
 
-    public Task UpdateRedirect(MessageRedirect redirect) =>
+    public Task UpdateRedirect(MessageRedirect redirect, CancellationToken cancellationToken = default) =>
         ExecuteWithDbContext(dbContext => dbContext.MessageRedirects
             .Where(row => row.FromPhysicalAddress == redirect.FromPhysicalAddress)
             .ExecuteUpdateAsync(setters => setters
                 .SetProperty(row => row.ToPhysicalAddress, redirect.ToPhysicalAddress)
-                .SetProperty(row => row.LastModified, redirect.LastModified)));
+                .SetProperty(row => row.LastModified, redirect.LastModified), cancellationToken));
 
-    public Task RemoveRedirect(MessageRedirect redirect) =>
+    public Task RemoveRedirect(MessageRedirect redirect, CancellationToken cancellationToken = default) =>
         ExecuteWithDbContext(dbContext => dbContext.MessageRedirects
             .Where(row => row.FromPhysicalAddress == redirect.FromPhysicalAddress)
-            .ExecuteDeleteAsync());
+            .ExecuteDeleteAsync(cancellationToken));
 }

src/ServiceControl.Persistence.RavenDB/EventLogDataStore.cs

@@ -2,40 +2,42 @@
 {
     using System;
     using System.Collections.Generic;
+    using System.Threading;
     using System.Threading.Tasks;
     using EventLog;
     using Persistence.Infrastructure;
     using Raven.Client.Documents;
 
     class EventLogDataStore(IRavenSessionProvider sessionProvider, ExpirationManager expirationManager) : IEventLogDataStore
     {
-        public async Task Add(EventLogItem logItem)
+        public async Task Add(EventLogItem logItem, CancellationToken cancellationToken = default)
         {
-            using var session = await sessionProvider.OpenSession();
+            using var session = await sessionProvider.OpenSession(cancellationToken: cancellationToken);
 
             // Version 7 rather than a random GUID so the final segment is time-ordered, which keeps
             // documents written together adjacent in the id index.
             await session.StoreAsync(
                 logItem,
-                EventLogItemIdGenerator.MakeDocumentId(logItem.Category, logItem.EventType, Guid.CreateVersion7()));
+                EventLogItemIdGenerator.MakeDocumentId(logItem.Category, logItem.EventType, Guid.CreateVersion7()),
+                cancellationToken);
 
             // Retention on RavenDB is per-document expiry metadata stamped at write time, not a
             // sweep. It has to be set here, on the only write path, or items never expire.
             expirationManager.EnableExpiration(session, logItem);
 
-            await session.SaveChangesAsync();
+            await session.SaveChangesAsync(cancellationToken);
         }
 
         public async Task<QueryResult<IList<EventLogItemView>>> GetEventLogItems(
-            PagingInfo pagingInfo, string knownVersion = null)
+            PagingInfo pagingInfo, string knownVersion = null, CancellationToken cancellationToken = default)
         {
-            using var session = await sessionProvider.OpenSession();
+            using var session = await sessionProvider.OpenSession(cancellationToken: cancellationToken);
             var documents = await session
                 .Query<EventLogItem>()
                 .Statistics(out var stats)
                 .OrderByDescending(p => p.RaisedAt)
                 .Paging(pagingInfo)
-                .ToListAsync();
+                .ToListAsync(cancellationToken);
 
             var queryStats = stats.ToQueryStatsInfo();
 

src/ServiceControl.Persistence.RavenDB/MessageRedirects/MessageRedirectsDataStore.cs

@@ -2,39 +2,40 @@
 {
     using System;
     using System.Collections.Generic;
+    using System.Threading;
     using System.Threading.Tasks;
     using ServiceControl.Persistence.MessageRedirects;
 
     class MessageRedirectsDataStore(IRavenSessionProvider sessionProvider) : IMessageRedirectsDataStore
     {
         public const string CollectionId = "messageredirects";
 
-        public async Task<IReadOnlyList<MessageRedirect>> GetRedirects()
+        public async Task<IReadOnlyList<MessageRedirect>> GetRedirects(CancellationToken cancellationToken = default)
         {
-            using var session = await sessionProvider.OpenSession();
-            var document = await session.LoadAsync<MessageRedirectsCollection>(CollectionId);
+            using var session = await sessionProvider.OpenSession(cancellationToken: cancellationToken);
+            var document = await session.LoadAsync<MessageRedirectsCollection>(CollectionId, cancellationToken);
 
             return document == null ? [] : document.ToRedirects();
         }
 
-        public Task AddRedirect(MessageRedirect redirect) => Mutate(document => document.Add(redirect));
+        public Task AddRedirect(MessageRedirect redirect, CancellationToken cancellationToken = default) => Mutate(document => document.Add(redirect), cancellationToken);
 
-        public Task UpdateRedirect(MessageRedirect redirect) => Mutate(document => document.Update(redirect));
+        public Task UpdateRedirect(MessageRedirect redirect, CancellationToken cancellationToken = default) => Mutate(document => document.Update(redirect), cancellationToken);
 
-        public Task RemoveRedirect(MessageRedirect redirect) => Mutate(document => document.Remove(redirect));
+        public Task RemoveRedirect(MessageRedirect redirect, CancellationToken cancellationToken = default) => Mutate(document => document.Remove(redirect), cancellationToken);
 
-        async Task Mutate(Action<MessageRedirectsCollection> mutate)
+        async Task Mutate(Action<MessageRedirectsCollection> mutate, CancellationToken cancellationToken)
         {
-            using var session = await sessionProvider.OpenSession();
-            var document = await session.LoadAsync<MessageRedirectsCollection>(CollectionId);
+            using var session = await sessionProvider.OpenSession(cancellationToken: cancellationToken);
+            var document = await session.LoadAsync<MessageRedirectsCollection>(CollectionId, cancellationToken);
             var changeVector = document == null ? null : session.Advanced.GetChangeVectorFor(document);
 
             document ??= new MessageRedirectsCollection();
 
             mutate(document);
 
-            await session.StoreAsync(document, changeVector, CollectionId);
-            await session.SaveChangesAsync();
+            await session.StoreAsync(document, changeVector, CollectionId, cancellationToken);
+            await session.SaveChangesAsync(cancellationToken);
         }
     }
 }

src/ServiceControl.Persistence.RavenDB/RavenCustomCheckDataStore.cs

@@ -2,6 +2,7 @@
 {
     using System;
     using System.Collections.Generic;
+    using System.Threading;
     using System.Threading.Tasks;
     using Raven.Client.Documents;
     using Raven.Client.Documents.Commands;
@@ -13,13 +14,13 @@
 
     class RavenCustomCheckDataStore(IRavenSessionProvider sessionProvider) : ICustomChecksDataStore
     {
-        public async Task<CheckStateChange> UpdateCustomCheckStatus(CustomCheckDetail detail)
+        public async Task<CheckStateChange> UpdateCustomCheckStatus(CustomCheckDetail detail, CancellationToken cancellationToken = default)
         {
             var status = CheckStateChange.Unchanged;
             var id = MakeId(detail.GetDeterministicId());
 
-            using var session = await sessionProvider.OpenSession();
-            var customCheck = await session.LoadAsync<CustomCheck>(id);
+            using var session = await sessionProvider.OpenSession(cancellationToken: cancellationToken);
+            var customCheck = await session.LoadAsync<CustomCheck>(id, cancellationToken);
 
             if (customCheck == null ||
                 (customCheck.Status == Status.Fail && !detail.HasFailed) ||
@@ -36,40 +37,40 @@ public async Task<CheckStateChange> UpdateCustomCheckStatus(CustomCheckDetail de
             customCheck.ReportedAt = detail.ReportedAt;
             customCheck.FailureReason = detail.FailureReason;
             customCheck.OriginatingEndpoint = detail.OriginatingEndpoint;
-            await session.StoreAsync(customCheck);
-            await session.SaveChangesAsync();
+            await session.StoreAsync(customCheck, cancellationToken);
+            await session.SaveChangesAsync(cancellationToken);
 
             return status;
         }
 
         static string MakeId(Guid id) => $"CustomChecks/{id}";
 
-        public async Task<QueryResult<IList<CustomCheck>>> GetStats(PagingInfo paging, string status = null)
+        public async Task<QueryResult<IList<CustomCheck>>> GetStats(PagingInfo paging, string status = null, CancellationToken cancellationToken = default)
         {
-            using var session = await sessionProvider.OpenSession();
+            using var session = await sessionProvider.OpenSession(cancellationToken: cancellationToken);
             var query =
                 session.Query<CustomCheck, CustomChecksIndex>().Statistics(out var stats);
 
             query = AddStatusFilter(query, status);
 
             var results = await query
                 .Paging(paging)
-                .ToListAsync();
+                .ToListAsync(cancellationToken);
 
             return new QueryResult<IList<CustomCheck>>(results, new QueryStatsInfo($"{stats.ResultEtag}", stats.TotalResults, stats.IsStale));
         }
 
-        public async Task DeleteCustomCheck(Guid id)
+        public async Task DeleteCustomCheck(Guid id, CancellationToken cancellationToken = default)
         {
             var documentId = MakeId(id);
-            using var session = await sessionProvider.OpenSession(new SessionOptions { NoTracking = true, NoCaching = true });
-            await session.Advanced.RequestExecutor.ExecuteAsync(new DeleteDocumentCommand(documentId, null), session.Advanced.Context);
+            using var session = await sessionProvider.OpenSession(new SessionOptions { NoTracking = true, NoCaching = true }, cancellationToken);
+            await session.Advanced.RequestExecutor.ExecuteAsync(new DeleteDocumentCommand(documentId, null), session.Advanced.Context, token: cancellationToken);
         }
 
-        public async Task<int> GetNumberOfFailedChecks()
+        public async Task<int> GetNumberOfFailedChecks(CancellationToken cancellationToken = default)
         {
-            using var session = await sessionProvider.OpenSession();
-            var failedCustomCheckCount = await session.Query<CustomCheck, CustomChecksIndex>().CountAsync(p => p.Status == Status.Fail);
+            using var session = await sessionProvider.OpenSession(cancellationToken: cancellationToken);
+            var failedCustomCheckCount = await session.Query<CustomCheck, CustomChecksIndex>().CountAsync(p => p.Status == Status.Fail, cancellationToken);
 
             return failedCustomCheckCount;
         }

src/ServiceControl.Persistence/ICustomChecksDataStore.cs

@@ -2,16 +2,17 @@ namespace ServiceControl.Persistence
 {
     using System;
     using System.Collections.Generic;
+    using System.Threading;
     using System.Threading.Tasks;
     using Infrastructure;
     using ServiceControl.Contracts.CustomChecks;
 
     public interface ICustomChecksDataStore
     {
-        Task<CheckStateChange> UpdateCustomCheckStatus(CustomCheckDetail detail);
+        Task<CheckStateChange> UpdateCustomCheckStatus(CustomCheckDetail detail, CancellationToken cancellationToken = default);
 
-        Task<QueryResult<IList<CustomCheck>>> GetStats(PagingInfo paging, string status = null);
-        Task DeleteCustomCheck(Guid id);
-        Task<int> GetNumberOfFailedChecks();
+        Task<QueryResult<IList<CustomCheck>>> GetStats(PagingInfo paging, string status = null, CancellationToken cancellationToken = default);
+        Task DeleteCustomCheck(Guid id, CancellationToken cancellationToken = default);
+        Task<int> GetNumberOfFailedChecks(CancellationToken cancellationToken = default);
     }
 }

src/ServiceControl.Persistence/IEventLogDataStore.cs

@@ -1,6 +1,7 @@
 namespace ServiceControl.Persistence
 {
     using System.Collections.Generic;
+    using System.Threading;
     using System.Threading.Tasks;
     using EventLog;
     using Infrastructure;
@@ -20,7 +21,7 @@ public interface IEventLogDataStore
         /// on <see cref="EventLogItemView.Id"/> when the item is read back.
         /// </summary>
         /// <param name="logItem">The item to store.</param>
-        Task Add(EventLogItem logItem);
+        Task Add(EventLogItem logItem, CancellationToken cancellationToken = default);
 
         /// <summary>
         /// Returns one page of event log items, newest <c>RaisedAt</c> first.
@@ -40,6 +41,6 @@ public interface IEventLogDataStore
         /// one is added, since nothing else tells a client its cached page is now wrong.
         /// </returns>
         Task<QueryResult<IList<EventLogItemView>>> GetEventLogItems(
-            PagingInfo pagingInfo, string knownVersion = null);
+            PagingInfo pagingInfo, string knownVersion = null, CancellationToken cancellationToken = default);
     }
 }

src/ServiceControl.Persistence/MessageRedirects/IMessageRedirectsDataStore.cs

@@ -1,13 +1,14 @@
 namespace ServiceControl.Persistence.MessageRedirects
 {
     using System.Collections.Generic;
+    using System.Threading;
     using System.Threading.Tasks;
 
     public interface IMessageRedirectsDataStore
     {
-        Task<IReadOnlyList<MessageRedirect>> GetRedirects();
-        Task AddRedirect(MessageRedirect redirect);
-        Task UpdateRedirect(MessageRedirect redirect);
-        Task RemoveRedirect(MessageRedirect redirect);
+        Task<IReadOnlyList<MessageRedirect>> GetRedirects(CancellationToken cancellationToken = default);
+        Task AddRedirect(MessageRedirect redirect, CancellationToken cancellationToken = default);
+        Task UpdateRedirect(MessageRedirect redirect, CancellationToken cancellationToken = default);
+        Task RemoveRedirect(MessageRedirect redirect, CancellationToken cancellationToken = default);
     }
 }