feat: Add optional CancellationToken parameter to SetProviderAsync (#638)

* Add CancellationToken optional parameter to SetProviderAsync

Signed-off-by: Kyle Julian <[email protected]>

* Address gemini review comments

Signed-off-by: Kyle Julian <[email protected]>

* Address API review comments

Signed-off-by: Kyle Julian <[email protected]>

---------

Signed-off-by: Kyle Julian <[email protected]>

Author: kylejuliandev

src/OpenFeature/Api.cs

@@ -40,12 +40,39 @@ internal Api() { }
     /// </summary>
     /// <remarks>The provider cannot be set to null. Attempting to set the provider to null has no effect. May throw an exception if <paramref name="featureProvider"/> cannot be initialized.</remarks>
     /// <param name="featureProvider">Implementation of <see cref="FeatureProvider"/></param>
-    public async Task SetProviderAsync(FeatureProvider featureProvider)
+    /// <returns>A <see cref="Task"/> that completes once Provider initialization is complete.</returns>
+    public Task SetProviderAsync(FeatureProvider featureProvider)
+    {
+        return this.SetProviderAsync(featureProvider, CancellationToken.None);
+    }
+
+    /// <summary>
+    /// Sets the default feature provider. In order to wait for the provider to be set, and initialization to complete,
+    /// await the returned task.
+    /// </summary>
+    /// <remarks>The provider cannot be set to null. Attempting to set the provider to null has no effect. May throw an exception if <paramref name="featureProvider"/> cannot be initialized.</remarks>
+    /// <param name="featureProvider">Implementation of <see cref="FeatureProvider"/></param>
+    /// <param name="cancellationToken">Propagates notification that the provider initialization should be canceled.</param>
+    /// <returns>A <see cref="Task"/> that completes once Provider initialization is complete.</returns>
+    public async Task SetProviderAsync(FeatureProvider featureProvider, CancellationToken cancellationToken)
     {
         this._eventExecutor.RegisterDefaultFeatureProvider(featureProvider);
-        await this._repository.SetProviderAsync(featureProvider, this.GetContext(), this.AfterInitializationAsync, this.AfterErrorAsync)
+        await this._repository.SetProviderAsync(featureProvider, this.GetContext(), this.AfterInitializationAsync, this.AfterErrorAsync, cancellationToken)
             .ConfigureAwait(false);
+    }
 
+    /// <summary>
+    /// Binds the feature provider to the given domain. In order to wait for the provider to be set, and
+    /// initialization to complete, await the returned task.
+    /// </summary>
+    /// <remarks>The provider cannot be set to null. Attempting to set the provider to null has no effect. May throw an exception if <paramref name="featureProvider"/> cannot be initialized.</remarks>
+    /// <param name="domain">An identifier which logically binds clients with providers</param>
+    /// <param name="featureProvider">Implementation of <see cref="FeatureProvider"/></param>
+    /// <exception cref="ArgumentNullException">domain cannot be null or empty</exception>
+    /// <returns>A <see cref="Task"/> that completes once Provider initialization is complete.</returns>
+    public Task SetProviderAsync(string domain, FeatureProvider featureProvider)
+    {
+        return this.SetProviderAsync(domain, featureProvider, CancellationToken.None);
     }
 
     /// <summary>
@@ -55,15 +82,17 @@ await this._repository.SetProviderAsync(featureProvider, this.GetContext(), this
     /// <remarks>The provider cannot be set to null. Attempting to set the provider to null has no effect. May throw an exception if <paramref name="featureProvider"/> cannot be initialized.</remarks>
     /// <param name="domain">An identifier which logically binds clients with providers</param>
     /// <param name="featureProvider">Implementation of <see cref="FeatureProvider"/></param>
+    /// <param name="cancellationToken">Propagates notification that the provider initialization should be canceled.</param>
     /// <exception cref="ArgumentNullException">domain cannot be null or empty</exception>
-    public async Task SetProviderAsync(string domain, FeatureProvider featureProvider)
+    /// <returns>A <see cref="Task"/> that completes once Provider initialization is complete.</returns>
+    public async Task SetProviderAsync(string domain, FeatureProvider featureProvider, CancellationToken cancellationToken)
     {
         if (string.IsNullOrWhiteSpace(domain))
         {
             throw new ArgumentNullException(nameof(domain));
         }
         this._eventExecutor.RegisterClientFeatureProvider(domain, featureProvider);
-        await this._repository.SetProviderAsync(domain, featureProvider, this.GetContext(), this.AfterInitializationAsync, this.AfterErrorAsync)
+        await this._repository.SetProviderAsync(domain, featureProvider, this.GetContext(), this.AfterInitializationAsync, this.AfterErrorAsync, cancellationToken)
             .ConfigureAwait(false);
     }
 

test/OpenFeature.Tests/OpenFeatureTests.cs

@@ -34,6 +34,43 @@ public async Task OpenFeature_Should_Initialize_Provider()
         await providerMockNamed.Received(1).InitializeAsync(Api.Instance.GetContext());
     }
 
+    [Fact]
+    public async Task OpenFeature_Should_Initialize_Provider_WithCancellationToken()
+    {
+        var providerMockDefault = Substitute.For<FeatureProvider>();
+        providerMockDefault.Status.Returns(ProviderStatus.NotReady);
+
+        using var cancellationTokenSource = new CancellationTokenSource();
+        var cancellationToken = cancellationTokenSource.Token;
+
+        await Api.Instance.SetProviderAsync(providerMockDefault, cancellationToken);
+        await providerMockDefault.Received(1).InitializeAsync(Api.Instance.GetContext(), cancellationToken);
+
+        var providerMockNamed = Substitute.For<FeatureProvider>();
+        providerMockNamed.Status.Returns(ProviderStatus.NotReady);
+
+        await Api.Instance.SetProviderAsync("the-name", providerMockNamed, cancellationToken);
+        await providerMockNamed.Received(1).InitializeAsync(Api.Instance.GetContext(), cancellationToken);
+    }
+
+    [Fact]
+    public async Task OpenFeature_Should_Handle_Cancellation_During_Initialization()
+    {
+        using var cancellationTokenSource = new CancellationTokenSource();
+        cancellationTokenSource.Cancel();
+        var cancellationToken = cancellationTokenSource.Token;
+
+        var providerMockDefault = Substitute.For<FeatureProvider>();
+        providerMockDefault.InitializeAsync(Arg.Any<EvaluationContext>(), cancellationToken)
+            .Returns(ci => Task.FromCanceled(cancellationToken));
+
+        await Assert.ThrowsAsync<TaskCanceledException>(() =>
+            Api.Instance.SetProviderAsync(providerMockDefault, cancellationToken));
+
+        await providerMockDefault.Received(1).InitializeAsync(Api.Instance.GetContext(), cancellationToken);
+        Assert.Equal(ProviderStatus.Error, providerMockDefault.Status);
+    }
+
     [Fact]
     [Specification("1.1.2.3",
         "The provider mutator function MUST invoke the shutdown function on the previously registered provider once it's no longer being used to resolve flag values.")]