diff --git a/src/Winton.Extensions.Configuration.Consul/ConsulConfigurationProvider.cs b/src/Winton.Extensions.Configuration.Consul/ConsulConfigurationProvider.cs index 00b484c..0cd3aff 100644 --- a/src/Winton.Extensions.Configuration.Consul/ConsulConfigurationProvider.cs +++ b/src/Winton.Extensions.Configuration.Consul/ConsulConfigurationProvider.cs @@ -45,7 +45,6 @@ public ConsulConfigurationProvider( public void Dispose() { _cancellationTokenSource.Cancel(); - _pollTask?.Wait(500); _cancellationTokenSource.Dispose(); } @@ -57,20 +56,22 @@ public override void Load() return; } - DoLoad().GetAwaiter().GetResult(); + var cancellationToken = _cancellationTokenSource.Token; + + DoLoad(cancellationToken).GetAwaiter().GetResult(); // Polling starts after the initial load to ensure no concurrent access to the key from this instance if (_source.ReloadOnChange) { - _pollTask = Task.Run(PollingLoop); + _pollTask = Task.Run(() => PollingLoop(cancellationToken), cancellationToken); } } - private async Task DoLoad() + private async Task DoLoad(CancellationToken cancellationToken) { try { - var result = await GetKvPairs(false).ConfigureAwait(false); + var result = await GetKvPairs(false, cancellationToken).ConfigureAwait(false); if (result.HasValue()) { @@ -94,7 +95,7 @@ private async Task DoLoad() } } - private async Task> GetKvPairs(bool waitForChange) + private async Task> GetKvPairs(bool waitForChange, CancellationToken cancellationToken) { using var consulClient = _consulClientFactory.Create(); var queryOptions = new QueryOptions @@ -106,27 +107,25 @@ private async Task> GetKvPairs(bool waitForChange) var result = await consulClient .KV - .List(_source.Key, queryOptions, _cancellationTokenSource.Token) + .List(_source.Key, queryOptions, cancellationToken) .ConfigureAwait(false); return result.StatusCode switch { HttpStatusCode.OK => result, HttpStatusCode.NotFound => result, - _ => - throw - new Exception($"Error loading configuration from consul. Status code: {result.StatusCode}.") - }; + _ => throw new Exception($"Error loading configuration from consul. Status code: {result.StatusCode}.") + }; } - private async Task PollingLoop() + private async Task PollingLoop(CancellationToken cancellationToken) { var consecutiveFailureCount = 0; - while (!_cancellationTokenSource.Token.IsCancellationRequested) + while (!cancellationToken.IsCancellationRequested) { try { - var result = await GetKvPairs(true).ConfigureAwait(false); + var result = await GetKvPairs(true, cancellationToken).ConfigureAwait(false); if (result.HasValue() && result.LastIndex > _lastIndex) { @@ -143,7 +142,7 @@ private async Task PollingLoop() _source.OnWatchException?.Invoke( new ConsulWatchExceptionContext(exception, ++consecutiveFailureCount, _source)) ?? TimeSpan.FromSeconds(5); - await Task.Delay(wait, _cancellationTokenSource.Token); + await Task.Delay(wait, cancellationToken); } } } diff --git a/test/Winton.Extensions.Configuration.Consul.Test/ConsulConfigurationProviderTests.cs b/test/Winton.Extensions.Configuration.Consul.Test/ConsulConfigurationProviderTests.cs index f8738a9..611f291 100644 --- a/test/Winton.Extensions.Configuration.Consul.Test/ConsulConfigurationProviderTests.cs +++ b/test/Winton.Extensions.Configuration.Consul.Test/ConsulConfigurationProviderTests.cs @@ -72,8 +72,8 @@ public sealed class Dispose : ConsulConfigurationProviderTests private async Task ShouldCancelPollingTaskWhenReloading() { var expectedKvCalls = 0; - var pollingCancelled = new TaskCompletionSource(); - CancellationToken cancellationToken = default; + var pollingCancelled = new TaskCompletionSource(); + _source.ReloadOnChange = true; _source.Optional = true; _kvEndpoint @@ -87,18 +87,21 @@ private async Task ShouldCancelPollingTaskWhenReloading() } expectedKvCalls++; - if (cancellationToken == default) + if (token.CanBeCanceled) { - cancellationToken = token; + token.Register(() => pollingCancelled.TrySetResult(true)); } }) - .Returns( - pollingCancelled.Task.IsCompleted - ? Task.FromCanceled>(pollingCancelled.Task.Result) - : Task.FromResult(new QueryResult { StatusCode = HttpStatusCode.OK })); + .Returns( + (_, __, token) => + token.IsCancellationRequested + ? Task.FromCanceled>(token) + : Task.Delay(5).ContinueWith(t => new QueryResult { StatusCode = HttpStatusCode.OK })); _provider.Load(); - cancellationToken.Register(() => pollingCancelled.SetResult(cancellationToken)); + + // allow polling loop to spin up + await Task.Delay(25); _provider.Dispose();