Skip to content

Commit 8b472d8

Browse files
kylejuliandevkylejuliandev
authored
refactor: Pass cancellation tokens to Provider Initialization functions (#640)
* Ensure we pass Cancellation Tokens to provider functions Signed-off-by: Kyle Julian <[email protected]> * Rename methods to include Async suffix Signed-off-by: Kyle Julian <[email protected]> * Add missing test case Signed-off-by: Kyle Julian <[email protected]> * Address gemini code review comments Signed-off-by: Kyle Julian <[email protected]> --------- Signed-off-by: Kyle Julian <[email protected]>
1 parent f65b099 commit 8b472d8

File tree

3 files changed

+127
-22
lines changed

3 files changed

+127
-22
lines changed

‎src/OpenFeature/Api.cs‎

Lines changed: 10 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -43,7 +43,8 @@ internal Api() { }
4343
public async Task SetProviderAsync(FeatureProvider featureProvider)
4444
{
4545
this._eventExecutor.RegisterDefaultFeatureProvider(featureProvider);
46-
await this._repository.SetProviderAsync(featureProvider, this.GetContext(), this.AfterInitialization, this.AfterError).ConfigureAwait(false);
46+
await this._repository.SetProviderAsync(featureProvider, this.GetContext(), this.AfterInitializationAsync, this.AfterErrorAsync)
47+
.ConfigureAwait(false);
4748

4849
}
4950

@@ -62,7 +63,8 @@ public async Task SetProviderAsync(string domain, FeatureProvider featureProvide
6263
throw new ArgumentNullException(nameof(domain));
6364
}
6465
this._eventExecutor.RegisterClientFeatureProvider(domain, featureProvider);
65-
await this._repository.SetProviderAsync(domain, featureProvider, this.GetContext(), this.AfterInitialization, this.AfterError).ConfigureAwait(false);
66+
await this._repository.SetProviderAsync(domain, featureProvider, this.GetContext(), this.AfterInitializationAsync, this.AfterErrorAsync)
67+
.ConfigureAwait(false);
6668
}
6769

6870
/// <summary>
@@ -324,7 +326,7 @@ internal void RemoveClientHandler(string client, ProviderEventTypes eventType, E
324326
/// <summary>
325327
/// Update the provider state to READY and emit a READY event after successful init.
326328
/// </summary>
327-
private async Task AfterInitialization(FeatureProvider provider)
329+
private async Task AfterInitializationAsync(FeatureProvider provider, CancellationToken cancellationToken = default)
328330
{
329331
provider.Status = ProviderStatus.Ready;
330332
var eventPayload = new ProviderEventPayload
@@ -334,13 +336,14 @@ private async Task AfterInitialization(FeatureProvider provider)
334336
ProviderName = provider.GetMetadata()?.Name,
335337
};
336338

337-
await this._eventExecutor.EventChannel.Writer.WriteAsync(new Event { Provider = provider, EventPayload = eventPayload }).ConfigureAwait(false);
339+
await this._eventExecutor.EventChannel.Writer.WriteAsync(new Event { Provider = provider, EventPayload = eventPayload }, cancellationToken)
340+
.ConfigureAwait(false);
338341
}
339342

340343
/// <summary>
341344
/// Update the provider state to ERROR and emit an ERROR after failed init.
342345
/// </summary>
343-
private async Task AfterError(FeatureProvider provider, Exception? ex)
346+
private async Task AfterErrorAsync(FeatureProvider provider, Exception? ex, CancellationToken cancellationToken = default)
344347
{
345348
provider.Status = typeof(ProviderFatalException) == ex?.GetType() ? ProviderStatus.Fatal : ProviderStatus.Error;
346349
var eventPayload = new ProviderEventPayload
@@ -350,7 +353,8 @@ private async Task AfterError(FeatureProvider provider, Exception? ex)
350353
ProviderName = provider.GetMetadata()?.Name,
351354
};
352355

353-
await this._eventExecutor.EventChannel.Writer.WriteAsync(new Event { Provider = provider, EventPayload = eventPayload }).ConfigureAwait(false);
356+
await this._eventExecutor.EventChannel.Writer.WriteAsync(new Event { Provider = provider, EventPayload = eventPayload }, cancellationToken)
357+
.ConfigureAwait(false);
354358
}
355359

356360
/// <summary>

‎src/OpenFeature/ProviderRepository.cs‎

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -55,8 +55,8 @@ public async ValueTask DisposeAsync()
5555
internal async Task SetProviderAsync(
5656
FeatureProvider? featureProvider,
5757
EvaluationContext context,
58-
Func<FeatureProvider, Task>? afterInitSuccess = null,
59-
Func<FeatureProvider, Exception, Task>? afterInitError = null,
58+
Func<FeatureProvider, CancellationToken, Task>? afterInitSuccess = null,
59+
Func<FeatureProvider, Exception, CancellationToken, Task>? afterInitError = null,
6060
CancellationToken cancellationToken = default)
6161
{
6262
// Cannot unset the feature provider.
@@ -93,8 +93,8 @@ await InitProviderAsync(this._defaultProvider, context, afterInitSuccess, afterI
9393
private static async Task InitProviderAsync(
9494
FeatureProvider? newProvider,
9595
EvaluationContext context,
96-
Func<FeatureProvider, Task>? afterInitialization,
97-
Func<FeatureProvider, Exception, Task>? afterError,
96+
Func<FeatureProvider, CancellationToken, Task>? afterInitialization,
97+
Func<FeatureProvider, Exception, CancellationToken, Task>? afterError,
9898
CancellationToken cancellationToken = default)
9999
{
100100
if (newProvider == null)
@@ -108,14 +108,14 @@ private static async Task InitProviderAsync(
108108
await newProvider.InitializeAsync(context, cancellationToken).ConfigureAwait(false);
109109
if (afterInitialization != null)
110110
{
111-
await afterInitialization.Invoke(newProvider).ConfigureAwait(false);
111+
await afterInitialization.Invoke(newProvider, cancellationToken).ConfigureAwait(false);
112112
}
113113
}
114114
catch (Exception ex)
115115
{
116116
if (afterError != null)
117117
{
118-
await afterError.Invoke(newProvider, ex).ConfigureAwait(false);
118+
await afterError.Invoke(newProvider, ex, cancellationToken).ConfigureAwait(false);
119119
}
120120
}
121121
}
@@ -138,8 +138,8 @@ private static async Task InitProviderAsync(
138138
internal async Task SetProviderAsync(string domain,
139139
FeatureProvider? featureProvider,
140140
EvaluationContext context,
141-
Func<FeatureProvider, Task>? afterInitSuccess = null,
142-
Func<FeatureProvider, Exception, Task>? afterInitError = null,
141+
Func<FeatureProvider, CancellationToken, Task>? afterInitSuccess = null,
142+
Func<FeatureProvider, Exception, CancellationToken, Task>? afterInitError = null,
143143
CancellationToken cancellationToken = default)
144144
{
145145
// Cannot set a provider for a null domain.

‎test/OpenFeature.Tests/ProviderRepositoryTests.cs‎

Lines changed: 109 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -39,7 +39,7 @@ public async Task AfterInitialization_Is_Invoked_For_Setting_Default_Provider()
3939
providerMock.Status.Returns(ProviderStatus.NotReady);
4040
var context = new EvaluationContextBuilder().Build();
4141
var callCount = 0;
42-
await repository.SetProviderAsync(providerMock, context, afterInitSuccess: (theProvider) =>
42+
await repository.SetProviderAsync(providerMock, context, afterInitSuccess: (theProvider, ct) =>
4343
{
4444
Assert.Equal(providerMock, theProvider);
4545
callCount++;
@@ -48,17 +48,42 @@ await repository.SetProviderAsync(providerMock, context, afterInitSuccess: (theP
4848
Assert.Equal(1, callCount);
4949
}
5050

51+
[Fact]
52+
public async Task AfterInitialization_Is_Invoked_With_CancellationToken()
53+
{
54+
var repository = new ProviderRepository();
55+
var providerMock = Substitute.For<FeatureProvider>();
56+
providerMock.Status.Returns(ProviderStatus.NotReady);
57+
58+
using var cancellationTokenSource = new CancellationTokenSource();
59+
var cancellationToken = cancellationTokenSource.Token;
60+
61+
var context = new EvaluationContextBuilder().Build();
62+
63+
var initCancellationToken = CancellationToken.None;
64+
await repository.SetProviderAsync(providerMock, context, afterInitSuccess: (theProvider, ct) =>
65+
{
66+
Assert.Equal(providerMock, theProvider);
67+
68+
initCancellationToken = ct;
69+
70+
return Task.CompletedTask;
71+
}, cancellationToken: cancellationToken);
72+
73+
Assert.Equal(cancellationToken, initCancellationToken);
74+
}
75+
5176
[Fact]
5277
public async Task AfterError_Is_Invoked_If_Initialization_Errors_Default_Provider()
5378
{
5479
var repository = new ProviderRepository();
5580
var providerMock = Substitute.For<FeatureProvider>();
5681
providerMock.Status.Returns(ProviderStatus.NotReady);
5782
var context = new EvaluationContextBuilder().Build();
58-
providerMock.When(x => x.InitializeAsync(context)).Throw(new Exception("BAD THINGS"));
83+
providerMock.When(x => x.InitializeAsync(context, Arg.Any<CancellationToken>())).Throw(new Exception("BAD THINGS"));
5984
var callCount = 0;
6085
Exception? receivedError = null;
61-
await repository.SetProviderAsync(providerMock, context, afterInitError: (theProvider, error) =>
86+
await repository.SetProviderAsync(providerMock, context, afterInitError: (theProvider, error, ct) =>
6287
{
6388
Assert.Equal(providerMock, theProvider);
6489
callCount++;
@@ -69,6 +94,32 @@ await repository.SetProviderAsync(providerMock, context, afterInitError: (thePro
6994
Assert.Equal(1, callCount);
7095
}
7196

97+
[Fact]
98+
public async Task AfterError_Is_Invoked_With_CancellationToken()
99+
{
100+
var repository = new ProviderRepository();
101+
var providerMock = Substitute.For<FeatureProvider>();
102+
providerMock.Status.Returns(ProviderStatus.NotReady);
103+
104+
using var cancellationTokenSource = new CancellationTokenSource();
105+
var cancellationToken = cancellationTokenSource.Token;
106+
107+
var context = new EvaluationContextBuilder().Build();
108+
providerMock.When(x => x.InitializeAsync(context, cancellationToken)).Throw(new Exception("BAD THINGS"));
109+
110+
var errorCancellationToken = CancellationToken.None;
111+
await repository.SetProviderAsync(providerMock, context, afterInitError: (theProvider, error, ct) =>
112+
{
113+
Assert.Equal(providerMock, theProvider);
114+
115+
errorCancellationToken = ct;
116+
117+
return Task.CompletedTask;
118+
}, cancellationToken: cancellationToken);
119+
120+
Assert.Equal(cancellationToken, errorCancellationToken);
121+
}
122+
72123
[Theory]
73124
[InlineData(ProviderStatus.Ready)]
74125
[InlineData(ProviderStatus.Stale)]
@@ -94,7 +145,7 @@ internal async Task AfterInitialize_Is_Not_Called_For_Ready_Provider(ProviderSta
94145
providerMock.Status.Returns(status);
95146
var context = new EvaluationContextBuilder().Build();
96147
var callCount = 0;
97-
await repository.SetProviderAsync(providerMock, context, afterInitSuccess: provider =>
148+
await repository.SetProviderAsync(providerMock, context, afterInitSuccess: (provider, ct) =>
98149
{
99150
callCount++;
100151
return Task.CompletedTask;
@@ -150,7 +201,7 @@ public async Task AfterInitialization_Is_Invoked_For_Setting_Named_Provider()
150201
providerMock.Status.Returns(ProviderStatus.NotReady);
151202
var context = new EvaluationContextBuilder().Build();
152203
var callCount = 0;
153-
await repository.SetProviderAsync("the-name", providerMock, context, afterInitSuccess: (theProvider) =>
204+
await repository.SetProviderAsync("the-name", providerMock, context, afterInitSuccess: (theProvider, ct) =>
154205
{
155206
Assert.Equal(providerMock, theProvider);
156207
callCount++;
@@ -159,17 +210,41 @@ await repository.SetProviderAsync("the-name", providerMock, context, afterInitSu
159210
Assert.Equal(1, callCount);
160211
}
161212

213+
[Fact]
214+
public async Task AfterInitialization_WithNamedProvider_Is_Invoked_With_CancellationToken()
215+
{
216+
var repository = new ProviderRepository();
217+
var providerMock = Substitute.For<FeatureProvider>();
218+
providerMock.Status.Returns(ProviderStatus.NotReady);
219+
220+
var context = new EvaluationContextBuilder().Build();
221+
using var cancellationTokenSource = new CancellationTokenSource();
222+
var cancellationToken = cancellationTokenSource.Token;
223+
224+
var initCancellationToken = CancellationToken.None;
225+
await repository.SetProviderAsync("the-name", providerMock, context, afterInitSuccess: (theProvider, ct) =>
226+
{
227+
Assert.Equal(providerMock, theProvider);
228+
229+
initCancellationToken = ct;
230+
231+
return Task.CompletedTask;
232+
}, cancellationToken: cancellationToken);
233+
234+
Assert.Equal(cancellationToken, initCancellationToken);
235+
}
236+
162237
[Fact]
163238
public async Task AfterError_Is_Invoked_If_Initialization_Errors_Named_Provider()
164239
{
165240
var repository = new ProviderRepository();
166241
var providerMock = Substitute.For<FeatureProvider>();
167242
providerMock.Status.Returns(ProviderStatus.NotReady);
168243
var context = new EvaluationContextBuilder().Build();
169-
providerMock.When(x => x.InitializeAsync(context)).Throw(new Exception("BAD THINGS"));
244+
providerMock.When(x => x.InitializeAsync(context, Arg.Any<CancellationToken>())).Throw(new Exception("BAD THINGS"));
170245
var callCount = 0;
171246
Exception? receivedError = null;
172-
await repository.SetProviderAsync("the-provider", providerMock, context, afterInitError: (theProvider, error) =>
247+
await repository.SetProviderAsync("the-provider", providerMock, context, afterInitError: (theProvider, error, ct) =>
173248
{
174249
Assert.Equal(providerMock, theProvider);
175250
callCount++;
@@ -180,6 +255,32 @@ await repository.SetProviderAsync("the-provider", providerMock, context, afterIn
180255
Assert.Equal(1, callCount);
181256
}
182257

258+
[Fact]
259+
public async Task AfterError_WithNamedProvider_Is_Invoked_With_CancellationToken()
260+
{
261+
var repository = new ProviderRepository();
262+
var providerMock = Substitute.For<FeatureProvider>();
263+
providerMock.Status.Returns(ProviderStatus.NotReady);
264+
265+
using var cancellationTokenSource = new CancellationTokenSource();
266+
var cancellationToken = cancellationTokenSource.Token;
267+
268+
var context = new EvaluationContextBuilder().Build();
269+
providerMock.When(x => x.InitializeAsync(context, cancellationToken)).Throw(new Exception("BAD THINGS"));
270+
271+
var errorCancellationToken = CancellationToken.None;
272+
await repository.SetProviderAsync("the-provider", providerMock, context, afterInitError: (theProvider, error, ct) =>
273+
{
274+
Assert.Equal(providerMock, theProvider);
275+
276+
errorCancellationToken = ct;
277+
278+
return Task.CompletedTask;
279+
}, cancellationToken: cancellationToken);
280+
281+
Assert.Equal(cancellationToken, errorCancellationToken);
282+
}
283+
183284
[Theory]
184285
[InlineData(ProviderStatus.Ready)]
185286
[InlineData(ProviderStatus.Stale)]
@@ -206,7 +307,7 @@ internal async Task AfterInitialize_Is_Not_Called_For_Ready_Named_Provider(Provi
206307
var context = new EvaluationContextBuilder().Build();
207308
var callCount = 0;
208309
await repository.SetProviderAsync("the-name", providerMock, context,
209-
afterInitSuccess: provider =>
310+
afterInitSuccess: (provider, ct) =>
210311
{
211312
callCount++;
212313
return Task.CompletedTask;

0 commit comments

Comments
 (0)