Propagate cancellation tokens through failed message editing and notification persistence

src/ServiceControl.AcceptanceTests/Monitoring/CustomChecks/When_email_notifications_are_enabled.cs

@@ -62,18 +62,18 @@ await Define<MyContext>(c =>
 
         class SetupNotificationSettings(INotificationsDataStore notificationsDataStore) : IHostedService
         {
-            public async Task StartAsync(CancellationToken cancellationToken)
+            public async Task StartAsync(CancellationToken cancellationToken = default)
             {
-                await using var notificationsManager = await notificationsDataStore.CreateNotificationsManager();
+                await using var notificationsManager = await notificationsDataStore.CreateNotificationsManager(cancellationToken);
 
-                var settings = await notificationsManager.LoadSettings();
+                var settings = await notificationsManager.LoadSettings(cancellationToken);
                 settings.Email.Enabled = true;
                 settings.Email.From = "YouServiceControl@particular.net";
                 settings.Email.To = "WhoeverMightBeConcerned@particular.net";
-                await notificationsManager.SaveChanges();
+                await notificationsManager.SaveChanges(cancellationToken);
             }
 
-            public Task StopAsync(CancellationToken cancellationToken) => Task.CompletedTask;
+            public Task StopAsync(CancellationToken cancellationToken = default) => Task.CompletedTask;
         }
 
         public class MyContext : ScenarioContext

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

@@ -5,7 +5,7 @@ namespace ServiceControl.Persistence.EFCore.Implementation;
 
 public class EditFailedMessagesDataStore(IServiceScopeFactory scopeFactory, TimeProvider timeProvider) : IEditFailedMessagesDataStore
 {
-    public Task<IEditFailedMessagesManager> CreateEditFailedMessageManager()
+    public Task<IEditFailedMessagesManager> CreateEditFailedMessageManager(CancellationToken cancellationToken = default)
     {
         var scope = scopeFactory.CreateAsyncScope();
         return Task.FromResult<IEditFailedMessagesManager>(

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

@@ -10,7 +10,7 @@ public class EditFailedMessagesManager(IAsyncDisposable scope, ServiceControlDbC
     FailedMessage? failedMessage;            // cached after GetFailedMessage
     FailedMessageEditEntity? editEntity;     // tracked after GetCurrentEditingRequestId / SetCurrentEditingRequestId
 
-    public async Task<FailedMessage?> GetFailedMessage(string failedMessageId)
+    public async Task<FailedMessage?> GetFailedMessage(string failedMessageId, CancellationToken cancellationToken = default)
     {
         if (!Guid.TryParse(failedMessageId, out var uniqueMessageId))
         {
@@ -19,7 +19,7 @@ public class EditFailedMessagesManager(IAsyncDisposable scope, ServiceControlDbC
 
         var entity = await dbContext.FailedMessages
             .AsNoTracking()
-            .SingleOrDefaultAsync(message => message.UniqueMessageId == uniqueMessageId);
+            .SingleOrDefaultAsync(message => message.UniqueMessageId == uniqueMessageId, cancellationToken);
 
         if (entity == null)
         {
@@ -34,7 +34,7 @@ public class EditFailedMessagesManager(IAsyncDisposable scope, ServiceControlDbC
         return result;
     }
 
-    public async Task<string?> GetCurrentEditingRequestId(string failedMessageId)
+    public async Task<string?> GetCurrentEditingRequestId(string failedMessageId, CancellationToken cancellationToken = default)
     {
         if (!Guid.TryParse(failedMessageId, out var uniqueMessageId))
         {
@@ -43,12 +43,12 @@ public class EditFailedMessagesManager(IAsyncDisposable scope, ServiceControlDbC
 
         editEntity = await dbContext.FailedMessageEdits
             .AsNoTracking()
-            .SingleOrDefaultAsync(e => e.UniqueMessageId == uniqueMessageId);
+            .SingleOrDefaultAsync(e => e.UniqueMessageId == uniqueMessageId, cancellationToken);
 
         return editEntity?.EditId;
     }
 
-    public Task SetCurrentEditingRequestId(string editingMessageId)
+    public Task SetCurrentEditingRequestId(string editingMessageId, CancellationToken cancellationToken = default)
     {
         if (failedMessage == null)
         {
@@ -65,7 +65,7 @@ public Task SetCurrentEditingRequestId(string editingMessageId)
         return Task.CompletedTask;
     }
 
-    public async Task SetFailedMessageAsResolved()
+    public async Task SetFailedMessageAsResolved(CancellationToken cancellationToken = default)
     {
         if (failedMessage == null)
         {
@@ -81,7 +81,7 @@ public async Task SetFailedMessageAsResolved()
         // swept immediately.
         var now = timeProvider.GetUtcNow().UtcDateTime;
         var entity = await dbContext.FailedMessages
-            .SingleOrDefaultAsync(m => m.UniqueMessageId == Guid.Parse(message.UniqueMessageId))
+            .SingleOrDefaultAsync(m => m.UniqueMessageId == Guid.Parse(message.UniqueMessageId), cancellationToken)
             ?? throw new InvalidOperationException("Failed message entity not found");
 
         entity.Status = FailedMessageStatus.Resolved;
@@ -90,7 +90,7 @@ public async Task SetFailedMessageAsResolved()
         message.Status = FailedMessageStatus.Resolved;
     }
 
-    public Task SaveChanges() => dbContext.SaveChangesAsync();
+    public Task SaveChanges(CancellationToken cancellationToken = default) => dbContext.SaveChangesAsync(cancellationToken);
 
     public ValueTask DisposeAsync() => scope.DisposeAsync();
 }

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

@@ -5,7 +5,7 @@ namespace ServiceControl.Persistence.EFCore.Implementation;
 
 public class NotificationsDataStore(IServiceProvider serviceProvider) : INotificationsDataStore
 {
-    public Task<INotificationsManager> CreateNotificationsManager()
+    public Task<INotificationsManager> CreateNotificationsManager(CancellationToken cancellationToken = default)
     {
         var scope = serviceProvider.CreateAsyncScope();
         ServiceControlDbContext serviceControlDbContext = scope.ServiceProvider.GetRequiredService<ServiceControlDbContext>();

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

@@ -8,19 +8,19 @@ public class NotificationsManager(IAsyncDisposable scope, ServiceControlDbContex
 {
     NotificationsSettings? _settings;
 
-    public Task SaveChanges()
+    public Task SaveChanges(CancellationToken cancellationToken = default)
     {
         if (_settings == null)
         {
             return Task.CompletedTask;
         }
 
-        return dbContext.StoreSetting(SettingKeys.NotificationEmails, _settings.Email, CancellationToken.None);
+        return dbContext.StoreSetting(SettingKeys.NotificationEmails, _settings.Email, cancellationToken);
     }
 
-    public async Task<NotificationsSettings> LoadSettings()
+    public async Task<NotificationsSettings> LoadSettings(CancellationToken cancellationToken = default)
     {
-        var settings = await dbContext.GetSetting<EmailNotifications>(SettingKeys.NotificationEmails, CancellationToken.None);
+        var settings = await dbContext.GetSetting<EmailNotifications>(SettingKeys.NotificationEmails, cancellationToken);
         _settings = new NotificationsSettings() { Email = settings ?? new EmailNotifications() };
         return _settings;
     }

src/ServiceControl.Persistence.RavenDB/Editing/EditFailedMessageManager.cs

@@ -1,6 +1,7 @@
 namespace ServiceControl.Persistence.RavenDB
 {
     using System;
+    using System.Threading;
     using System.Threading.Tasks;
     using Raven.Client.Documents.Session;
     using ServiceControl.MessageFailures;
@@ -19,19 +20,19 @@ public EditFailedMessageManager(IAsyncDocumentSession session, ExpirationManager
             this.expirationManager = expirationManager;
         }
 
-        public async Task<FailedMessage> GetFailedMessage(string failedMessageId)
+        public async Task<FailedMessage> GetFailedMessage(string failedMessageId, CancellationToken cancellationToken = default)
         {
-            failedMessage = await session.LoadAsync<FailedMessage>(FailedMessageIdGenerator.MakeDocumentId(failedMessageId));
+            failedMessage = await session.LoadAsync<FailedMessage>(FailedMessageIdGenerator.MakeDocumentId(failedMessageId), cancellationToken);
             return failedMessage;
         }
 
-        public async Task<string> GetCurrentEditingRequestId(string failedMessageId)
+        public async Task<string> GetCurrentEditingRequestId(string failedMessageId, CancellationToken cancellationToken = default)
         {
-            var edit = await session.LoadAsync<FailedMessageEdit>(FailedMessageEdit.MakeDocumentId(failedMessageId));
+            var edit = await session.LoadAsync<FailedMessageEdit>(FailedMessageEdit.MakeDocumentId(failedMessageId), cancellationToken);
             return edit?.EditId;
         }
 
-        public Task SetCurrentEditingRequestId(string editingMessageId)
+        public Task SetCurrentEditingRequestId(string editingMessageId, CancellationToken cancellationToken = default)
         {
             if (failedMessage == null)
             {
@@ -42,10 +43,10 @@ public Task SetCurrentEditingRequestId(string editingMessageId)
                 Id = FailedMessageEdit.MakeDocumentId(failedMessage.UniqueMessageId),
                 FailedMessageId = failedMessage.Id,
                 EditId = editingMessageId
-            });
+            }, cancellationToken);
         }
 
-        public Task SetFailedMessageAsResolved()
+        public Task SetFailedMessageAsResolved(CancellationToken cancellationToken = default)
         {
             // Instance is tracked by the document session
             failedMessage.Status = FailedMessageStatus.Resolved;

src/ServiceControl.Persistence.RavenDB/Editing/EditFailedMessagesDataStore.cs

@@ -1,11 +1,12 @@
 namespace ServiceControl.Persistence.RavenDB.Editing
 {
+    using System.Threading;
     using System.Threading.Tasks;
 
     class EditFailedMessagesDataStore(IRavenSessionProvider sessionProvider, ExpirationManager expirationManager) : IEditFailedMessagesDataStore
     {
-        public async Task<IEditFailedMessagesManager> CreateEditFailedMessageManager() =>
+        public async Task<IEditFailedMessagesManager> CreateEditFailedMessageManager(CancellationToken cancellationToken = default) =>
             // the edit failed message manager manages the lifetime of the session
-            new EditFailedMessageManager(await sessionProvider.OpenSession(), expirationManager);
+            new EditFailedMessageManager(await sessionProvider.OpenSession(cancellationToken: cancellationToken), expirationManager);
     }
 }

src/ServiceControl.Persistence.RavenDB/Editing/NotificationsDataStore.cs

@@ -1,11 +1,12 @@
 namespace ServiceControl.Persistence.RavenDB.Editing
 {
+    using System.Threading;
     using System.Threading.Tasks;
 
     class NotificationsDataStore(IRavenSessionProvider sessionProvider) : INotificationsDataStore
     {
-        public async Task<INotificationsManager> CreateNotificationsManager() =>
+        public async Task<INotificationsManager> CreateNotificationsManager(CancellationToken cancellationToken = default) =>
             // the notifications manager manages the lifetime of the session
-            new NotificationsManager(await sessionProvider.OpenSession());
+            new NotificationsManager(await sessionProvider.OpenSession(cancellationToken: cancellationToken));
     }
 }

src/ServiceControl.Persistence.RavenDB/Editing/NotificationsManager.cs

@@ -1,5 +1,6 @@
 namespace ServiceControl.Persistence.RavenDB.Editing
 {
+    using System.Threading;
     using System.Threading.Tasks;
     using Notifications;
     using Raven.Client.Documents.Session;
@@ -8,18 +9,18 @@ class NotificationsManager(IAsyncDocumentSession session) : AbstractSessionManag
     {
         const string SingleDocumentId = "NotificationsSettings/All";
 
-        public async Task<NotificationsSettings> LoadSettings()
+        public async Task<NotificationsSettings> LoadSettings(CancellationToken cancellationToken = default)
         {
             // Deliberately not aggressively cached. These settings are read rarely and edited by hand,
             // and aggressive caching invalidates asynchronously via the Changes API, so a read straight
             // after a save can return the pre-save document.
             var settings = await Session
-                .LoadAsync<NotificationsSettingsDocument>(SingleDocumentId);
+                .LoadAsync<NotificationsSettingsDocument>(SingleDocumentId, cancellationToken);
 
             if (settings == null)
             {
                 settings = new NotificationsSettingsDocument { Id = SingleDocumentId };
-                await Session.StoreAsync(settings);
+                await Session.StoreAsync(settings, cancellationToken);
             }
 
             return new NotificationsSettings()

src/ServiceControl.Persistence.RavenDB/Transactions/RavenTransactionalDataStore.cs

@@ -1,13 +1,14 @@
 namespace ServiceControl.Persistence.RavenDB
 {
+    using System.Threading;
     using System.Threading.Tasks;
     using Raven.Client.Documents.Session;
 
     abstract class AbstractSessionManager(IAsyncDocumentSession session) : IDataSessionManager
     {
         protected IAsyncDocumentSession Session { get; } = session;
 
-        public Task SaveChanges() => Session.SaveChangesAsync();
+        public Task SaveChanges(CancellationToken cancellationToken = default) => Session.SaveChangesAsync(cancellationToken);
         public ValueTask DisposeAsync()
         {
             Session.Dispose();

src/ServiceControl.Persistence/IDataSessionManager.cs

@@ -1,10 +1,11 @@
 namespace ServiceControl.Persistence
 {
     using System;
+    using System.Threading;
     using System.Threading.Tasks;
 
     public interface IDataSessionManager : IAsyncDisposable
     {
-        Task SaveChanges();
+        Task SaveChanges(CancellationToken cancellationToken = default);
     }
 }

src/ServiceControl.Persistence/IEditFailedMessagesDataStore.cs

@@ -1,9 +1,10 @@
 namespace ServiceControl.Persistence
 {
+    using System.Threading;
     using System.Threading.Tasks;
 
     public interface IEditFailedMessagesDataStore
     {
-        Task<IEditFailedMessagesManager> CreateEditFailedMessageManager();
+        Task<IEditFailedMessagesManager> CreateEditFailedMessageManager(CancellationToken cancellationToken = default);
     }
 }

src/ServiceControl.Persistence/IEditFailedMessagesManager.cs

@@ -1,14 +1,15 @@
 #nullable enable
 namespace ServiceControl.Persistence
 {
+    using System.Threading;
     using System.Threading.Tasks;
     using ServiceControl.MessageFailures;
 
     public interface IEditFailedMessagesManager : IDataSessionManager
     {
-        Task<FailedMessage?> GetFailedMessage(string failedMessageId);
-        Task<string?> GetCurrentEditingRequestId(string failedMessageId);
-        Task SetCurrentEditingRequestId(string editingMessageId);
-        Task SetFailedMessageAsResolved();
+        Task<FailedMessage?> GetFailedMessage(string failedMessageId, CancellationToken cancellationToken = default);
+        Task<string?> GetCurrentEditingRequestId(string failedMessageId, CancellationToken cancellationToken = default);
+        Task SetCurrentEditingRequestId(string editingMessageId, CancellationToken cancellationToken = default);
+        Task SetFailedMessageAsResolved(CancellationToken cancellationToken = default);
     }
 }

src/ServiceControl.Persistence/INotificationsDataStore.cs

@@ -1,9 +1,10 @@
 namespace ServiceControl.Persistence
 {
+    using System.Threading;
     using System.Threading.Tasks;
 
     public interface INotificationsDataStore
     {
-        Task<INotificationsManager> CreateNotificationsManager();
+        Task<INotificationsManager> CreateNotificationsManager(CancellationToken cancellationToken = default);
     }
 }

src/ServiceControl.Persistence/INotificationsManager.cs

@@ -1,10 +1,11 @@
 namespace ServiceControl.Persistence
 {
+    using System.Threading;
     using System.Threading.Tasks;
     using Notifications;
 
     public interface INotificationsManager : IDataSessionManager
     {
-        Task<NotificationsSettings> LoadSettings();
+        Task<NotificationsSettings> LoadSettings(CancellationToken cancellationToken = default);
     }
 }

src/ServiceControl.UnitTests/MessageFailures/EditFailedMessagesControllerAuditTests.cs

@@ -49,11 +49,11 @@ sealed class FakeEditFailedMessagesManager : IEditFailedMessagesManager
     {
         public string? CurrentEditingRequestId { get; set; }
 
-        public Task SaveChanges() => Task.CompletedTask;
-        public Task<FailedMessage?> GetFailedMessage(string failedMessageId) => Task.FromResult<FailedMessage?>(null);
-        public Task<string?> GetCurrentEditingRequestId(string failedMessageId) => Task.FromResult(CurrentEditingRequestId);
-        public Task SetCurrentEditingRequestId(string editingMessageId) => Task.CompletedTask;
-        public Task SetFailedMessageAsResolved() => Task.CompletedTask;
+        public Task SaveChanges(CancellationToken cancellationToken = default) => Task.CompletedTask;
+        public Task<FailedMessage?> GetFailedMessage(string failedMessageId, CancellationToken cancellationToken = default) => Task.FromResult<FailedMessage?>(null);
+        public Task<string?> GetCurrentEditingRequestId(string failedMessageId, CancellationToken cancellationToken = default) => Task.FromResult(CurrentEditingRequestId);
+        public Task SetCurrentEditingRequestId(string editingMessageId, CancellationToken cancellationToken = default) => Task.CompletedTask;
+        public Task SetFailedMessageAsResolved(CancellationToken cancellationToken = default) => Task.CompletedTask;
 
         public ValueTask DisposeAsync() => ValueTask.CompletedTask;
     }
@@ -63,7 +63,7 @@ sealed class StubErrorMessageDataStore : IFailedMessageQueryDataStore, IEditFail
         public FailedMessage? ErrorByResult { get; set; }
         public FakeEditFailedMessagesManager EditManager { get; } = new();
 
-        public Task<IEditFailedMessagesManager> CreateEditFailedMessageManager() => Task.FromResult<IEditFailedMessagesManager>(EditManager);
+        public Task<IEditFailedMessagesManager> CreateEditFailedMessageManager(CancellationToken cancellationToken = default) => Task.FromResult<IEditFailedMessagesManager>(EditManager);
         public Task<FailedMessage?> GetFailedMessage(string failedMessageId, CancellationToken cancellationToken = default) => Task.FromResult(ErrorByResult);
 
         public Task<FailedMessage[]> GetFailedMessagesByIds(Guid[] ids, CancellationToken cancellationToken = default) => throw new NotImplementedException();

src/ServiceControl/MessageFailures/Api/EditFailedMessagesController.cs

@@ -45,8 +45,8 @@ public async Task<ActionResult<EditRetryResponse>> Edit(string failedMessageId,
             }
 
             //HINT: This validation is the first one because we want to minimize the chance of two users concurrently execute an edit-retry.
-            var editManager = await editStore.CreateEditFailedMessageManager();
-            var editId = await editManager.GetCurrentEditingRequestId(failedMessageId);
+            var editManager = await editStore.CreateEditFailedMessageManager(cancellationToken);
+            var editId = await editManager.GetCurrentEditingRequestId(failedMessageId, cancellationToken);
             if (editId != null)
             {
                 logger.LogWarning("Cannot edit message {FailedMessageId} because it has already been edited", failedMessageId);