Merge pull request #5635 from Particular/rhys/cancellation-endpointsettings

Add cancellation token to GetAllEndpointSettings

Author: rbev

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

@@ -2,7 +2,7 @@ namespace ServiceControl.Persistence.EFCore.Implementation;
 
 public class EndpointSettingsStore : IEndpointSettingsStore
 {
-    public IAsyncEnumerable<EndpointSettings> GetAllEndpointSettings() =>
+    public IAsyncEnumerable<EndpointSettings> GetAllEndpointSettings(CancellationToken cancellationToken) =>
         throw new NotImplementedException();
 
     public Task UpdateEndpointSettings(EndpointSettings settings, CancellationToken token) =>

src/ServiceControl.Persistence.RavenDB/EndpointSettingsStore.cs

@@ -1,6 +1,7 @@
 namespace ServiceControl.Persistence.RavenDB;
 
 using System.Collections.Generic;
+using System.Runtime.CompilerServices;
 using System.Threading;
 using System.Threading.Tasks;
 using Infrastructure;
@@ -12,14 +13,14 @@ class EndpointSettingsStore(IRavenSessionProvider sessionProvider) : IEndpointSe
     static string MakeDocumentId(string name) =>
         $"{EndpointSettings.CollectionName}/{DeterministicGuid.MakeId(name)}";
 
-    public async IAsyncEnumerable<EndpointSettings> GetAllEndpointSettings()
+    public async IAsyncEnumerable<EndpointSettings> GetAllEndpointSettings([EnumeratorCancellation] CancellationToken cancellationToken)
     {
-        using IAsyncDocumentSession session = await sessionProvider.OpenSession();
+        using IAsyncDocumentSession session = await sessionProvider.OpenSession(cancellationToken: cancellationToken);
         await using IAsyncEnumerator<StreamResult<EndpointSettings>> enumerator = await session
             .Advanced
-            .StreamAsync<EndpointSettings>($"{EndpointSettings.CollectionName}/");
+            .StreamAsync<EndpointSettings>($"{EndpointSettings.CollectionName}/", token: cancellationToken);
 
-        while (await enumerator.MoveNextAsync())
+        while (await enumerator.MoveNextAsync() && !cancellationToken.IsCancellationRequested)
         {
             yield return enumerator.Current.Document;
         }

src/ServiceControl.Persistence/IEndpointSettingsStore.cs

@@ -6,7 +6,7 @@
 
 public interface IEndpointSettingsStore
 {
-    IAsyncEnumerable<EndpointSettings> GetAllEndpointSettings();
+    IAsyncEnumerable<EndpointSettings> GetAllEndpointSettings(CancellationToken token);
 
     Task UpdateEndpointSettings(EndpointSettings settings, CancellationToken token);
     Task Delete(string name, CancellationToken cancellationToken);

src/ServiceControl.UnitTests/Monitoring/HeartbeatEndpointSettingsSyncHostedServiceTests.cs

@@ -242,7 +242,7 @@ public void RecordHeartbeat(EndpointInstanceId endpointInstanceId, DateTime time
 
     class MockEndpointSettingsStore(EndpointSettings[] settings) : IEndpointSettingsStore
     {
-        public IAsyncEnumerable<EndpointSettings> GetAllEndpointSettings() => settings.ToAsyncEnumerable();
+        public IAsyncEnumerable<EndpointSettings> GetAllEndpointSettings(CancellationToken token) => settings.ToAsyncEnumerable();
 
         public Task UpdateEndpointSettings(EndpointSettings settings, CancellationToken token)
         {

src/ServiceControl/Monitoring/HeartbeatEndpointSettingsSyncHostedService.cs

@@ -67,8 +67,7 @@ async Task PurgeMonitoringDataThatDoesNotNeedToBeTracked(CancellationToken cance
         ILookup<string, Guid> monitorEndpointsLookup = endpointsViews
             .Where(view => !view.IsSendingHeartbeats)
             .ToLookup(view => view.Name, view => view.Id);
-        await foreach (EndpointSettings endpointSetting in endpointSettingsStore.GetAllEndpointSettings()
-                           .WithCancellation(cancellationToken))
+        await foreach (EndpointSettings endpointSetting in endpointSettingsStore.GetAllEndpointSettings(cancellationToken))
         {
             if (!endpointSetting.TrackInstances)
             {
@@ -94,8 +93,7 @@ async Task InitialiseSettings(HashSet<string> monitorEndpoints, CancellationToke
         HashSet<string> settingsNames = [];
 
         // Delete any endpoints data that no longer exists
-        await foreach (EndpointSettings endpointSetting in endpointSettingsStore.GetAllEndpointSettings()
-                           .WithCancellation(cancellationToken))
+        await foreach (EndpointSettings endpointSetting in endpointSettingsStore.GetAllEndpointSettings(cancellationToken))
         {
             if (endpointSetting.Name == string.Empty)
             {

src/ServiceControl/Monitoring/Web/EndpointsSettingsController.cs

@@ -33,7 +33,7 @@ public class EndpointsSettingsController(
     public async IAsyncEnumerable<SettingsData> Endpoints([EnumeratorCancellation] CancellationToken token)
     {
         await using IAsyncEnumerator<EndpointSettings> enumerator =
-            dataStore.GetAllEndpointSettings().GetAsyncEnumerator(token);
+            dataStore.GetAllEndpointSettings(token).GetAsyncEnumerator(token);
         bool noResults = true;
         while (await enumerator.MoveNextAsync())
         {