diff --git a/src/ServiceControl.Persistence.EFCore/Implementation/EndpointSettingsStore.cs b/src/ServiceControl.Persistence.EFCore/Implementation/EndpointSettingsStore.cs index 7e79541e91..6de13c0fa4 100644 --- a/src/ServiceControl.Persistence.EFCore/Implementation/EndpointSettingsStore.cs +++ b/src/ServiceControl.Persistence.EFCore/Implementation/EndpointSettingsStore.cs @@ -2,7 +2,7 @@ namespace ServiceControl.Persistence.EFCore.Implementation; public class EndpointSettingsStore : IEndpointSettingsStore { - public IAsyncEnumerable GetAllEndpointSettings() => + public IAsyncEnumerable GetAllEndpointSettings(CancellationToken cancellationToken) => throw new NotImplementedException(); public Task UpdateEndpointSettings(EndpointSettings settings, CancellationToken token) => diff --git a/src/ServiceControl.Persistence.RavenDB/EndpointSettingsStore.cs b/src/ServiceControl.Persistence.RavenDB/EndpointSettingsStore.cs index d0442bdf9e..463a31efd0 100644 --- a/src/ServiceControl.Persistence.RavenDB/EndpointSettingsStore.cs +++ b/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 GetAllEndpointSettings() + public async IAsyncEnumerable GetAllEndpointSettings([EnumeratorCancellation] CancellationToken cancellationToken) { - using IAsyncDocumentSession session = await sessionProvider.OpenSession(); + using IAsyncDocumentSession session = await sessionProvider.OpenSession(cancellationToken: cancellationToken); await using IAsyncEnumerator> enumerator = await session .Advanced - .StreamAsync($"{EndpointSettings.CollectionName}/"); + .StreamAsync($"{EndpointSettings.CollectionName}/", token: cancellationToken); - while (await enumerator.MoveNextAsync()) + while (await enumerator.MoveNextAsync() && !cancellationToken.IsCancellationRequested) { yield return enumerator.Current.Document; } diff --git a/src/ServiceControl.Persistence/IEndpointSettingsStore.cs b/src/ServiceControl.Persistence/IEndpointSettingsStore.cs index dab2097718..d41296dd3a 100644 --- a/src/ServiceControl.Persistence/IEndpointSettingsStore.cs +++ b/src/ServiceControl.Persistence/IEndpointSettingsStore.cs @@ -6,7 +6,7 @@ public interface IEndpointSettingsStore { - IAsyncEnumerable GetAllEndpointSettings(); + IAsyncEnumerable GetAllEndpointSettings(CancellationToken token); Task UpdateEndpointSettings(EndpointSettings settings, CancellationToken token); Task Delete(string name, CancellationToken cancellationToken); diff --git a/src/ServiceControl.UnitTests/Monitoring/HeartbeatEndpointSettingsSyncHostedServiceTests.cs b/src/ServiceControl.UnitTests/Monitoring/HeartbeatEndpointSettingsSyncHostedServiceTests.cs index 6189431088..238021bf98 100644 --- a/src/ServiceControl.UnitTests/Monitoring/HeartbeatEndpointSettingsSyncHostedServiceTests.cs +++ b/src/ServiceControl.UnitTests/Monitoring/HeartbeatEndpointSettingsSyncHostedServiceTests.cs @@ -242,7 +242,7 @@ public void RecordHeartbeat(EndpointInstanceId endpointInstanceId, DateTime time class MockEndpointSettingsStore(EndpointSettings[] settings) : IEndpointSettingsStore { - public IAsyncEnumerable GetAllEndpointSettings() => settings.ToAsyncEnumerable(); + public IAsyncEnumerable GetAllEndpointSettings(CancellationToken token) => settings.ToAsyncEnumerable(); public Task UpdateEndpointSettings(EndpointSettings settings, CancellationToken token) { diff --git a/src/ServiceControl/Monitoring/HeartbeatEndpointSettingsSyncHostedService.cs b/src/ServiceControl/Monitoring/HeartbeatEndpointSettingsSyncHostedService.cs index ec5158a4f9..2856c00c94 100644 --- a/src/ServiceControl/Monitoring/HeartbeatEndpointSettingsSyncHostedService.cs +++ b/src/ServiceControl/Monitoring/HeartbeatEndpointSettingsSyncHostedService.cs @@ -67,8 +67,7 @@ async Task PurgeMonitoringDataThatDoesNotNeedToBeTracked(CancellationToken cance ILookup 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 monitorEndpoints, CancellationToke HashSet 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) { diff --git a/src/ServiceControl/Monitoring/Web/EndpointsSettingsController.cs b/src/ServiceControl/Monitoring/Web/EndpointsSettingsController.cs index a1d55138db..44bdd6f05e 100644 --- a/src/ServiceControl/Monitoring/Web/EndpointsSettingsController.cs +++ b/src/ServiceControl/Monitoring/Web/EndpointsSettingsController.cs @@ -33,7 +33,7 @@ public class EndpointsSettingsController( public async IAsyncEnumerable Endpoints([EnumeratorCancellation] CancellationToken token) { await using IAsyncEnumerator enumerator = - dataStore.GetAllEndpointSettings().GetAsyncEnumerator(token); + dataStore.GetAllEndpointSettings(token).GetAsyncEnumerator(token); bool noResults = true; while (await enumerator.MoveNextAsync()) {