Address PR comments, refactor

- Change type names to make more sense (e.g. Start / Stop )
- Add context.Context to the IDP RequestToken
- Add RequestTimeout to TokenManagerOptions which will be utilized by
  the context
- Change the LowerRefreshBoundMs from int64 to time.Duration and use
  better name (dropping the Ms suffix)

TODO:
 - Address changes in the documentation

Author: ndyakov

.github/workflows/build.yml

@@ -19,7 +19,7 @@ jobs:
     - name: Install dependencies
       run: go mod tidy
     - name: Run tests with coverage
-      run: go test ./... -coverprofile=./cover.out -covermode=atomic -race -count 2 -timeout 1m
+      run: go test ./... -coverprofile=./cover.out -covermode=atomic -race -count 2 -timeout 5m
     - name: Upload coverage
       uses: actions/upload-artifact@v4
       with:

README.md

@@ -231,7 +231,7 @@ type TokenManagerOptions struct {
 
     // Optional: Minimum time before expiration to refresh (ms)
     // Default: 10000 (10 seconds)
-    LowerRefreshBoundMs int64
+    LowerRefreshBounds int64
 
     // Optional: Configuration for retry behavior
     RetryOptions RetryOptions
@@ -350,7 +350,7 @@ options := entraid.CredentialsProviderOptions{
     ClientID: os.Getenv("AZURE_CLIENT_ID"),
     TokenManagerOptions: manager.TokenManagerOptions{
         ExpirationRefreshRatio: 0.7,
-        LowerRefreshBoundMs: 10000,
+        LowerRefreshBounds: 10000,
     },
 }
 ```
@@ -361,7 +361,7 @@ options := entraid.CredentialsProviderOptions{
     ClientID: os.Getenv("AZURE_CLIENT_ID"),
     TokenManagerOptions: manager.TokenManagerOptions{
         ExpirationRefreshRatio: 0.7,
-        LowerRefreshBoundMs: 10000,
+        LowerRefreshBounds: 10000,
         RetryOptions: manager.RetryOptions{
             MaxAttempts: 3,
             InitialDelayMs: 1000,
@@ -516,7 +516,7 @@ func main() {
     tokenManager, err := manager.NewTokenManager(customProvider, manager.TokenManagerOptions{
         // Configure token refresh behavior
         ExpirationRefreshRatio: 0.7,
-        LowerRefreshBoundMs:   10000,
+        LowerRefreshBounds:   10000,
     })
     if err != nil {
         log.Fatalf("Failed to create token manager: %v", err)

credentials_provider.go

@@ -19,8 +19,8 @@ var _ auth.StreamingCredentialsProvider = (*entraidCredentialsProvider)(nil)
 type entraidCredentialsProvider struct {
 	options CredentialsProviderOptions // Configuration options for the provider.
 
-	tokenManager      manager.TokenManager // Manages token retrieval.
-	closeTokenManager manager.CloseFunc    // Function to cancel the token manager.
+	tokenManager     manager.TokenManager // Manages token retrieval.
+	stopTokenManager manager.StopFunc     // Function to stop the token manager.
 
 	// listeners is a slice of listeners that are notified when the token manager receives a new token.
 	listeners []auth.CredentialsListener // Slice of listeners notified on token updates.
@@ -64,7 +64,7 @@ func (e *entraidCredentialsProvider) onTokenError(err error) {
 // - error: An error if the subscription fails, such as if the token cannot be retrieved.
 //
 // Note: If the listener is already subscribed, it will not receive duplicate notifications.
-func (e *entraidCredentialsProvider) Subscribe(listener auth.CredentialsListener) (auth.Credentials, auth.CancelProviderFunc, error) {
+func (e *entraidCredentialsProvider) Subscribe(listener auth.CredentialsListener) (auth.Credentials, auth.UnsubscribeFunc, error) {
 	// First try to get a token, only then subscribe the listener.
 	token, err := e.tokenManager.GetToken(false)
 	if err != nil {
@@ -87,7 +87,7 @@ func (e *entraidCredentialsProvider) Subscribe(listener auth.CredentialsListener
 	}
 	e.rwLock.Unlock()
 
-	cancel := func() error {
+	unsub := func() error {
 		// Remove the listener from the list of listeners.
 		e.rwLock.Lock()
 		defer e.rwLock.Unlock()
@@ -102,20 +102,20 @@ func (e *entraidCredentialsProvider) Subscribe(listener auth.CredentialsListener
 		// Clear the listeners slice if it's empty
 		if len(e.listeners) == 0 {
 			e.listeners = make([]auth.CredentialsListener, 0)
-			if e.closeTokenManager != nil {
-				err := e.closeTokenManager()
+			if e.stopTokenManager != nil {
+				err := e.stopTokenManager()
 				if err != nil {
 					return fmt.Errorf("couldn't cancel token manager: %w", err)
 				}
-				// Set the cancelTokenManager to nil to indicate that it has been canceled.
-				// This prevents multiple calls to cancelTokenManager.
-				e.closeTokenManager = nil
+				// Set the stopTokenManager to nil to indicate that it has been stopped.
+				// This prevents multiple calls to stopTokenManager.
+				e.stopTokenManager = nil
 			}
 		}
 		return nil
 	}
 
-	return token, cancel, nil
+	return token, unsub, nil
 }
 
 // NewCredentialsProvider creates a new credentials provider with the specified token manager and options.
@@ -134,10 +134,10 @@ func NewCredentialsProvider(tokenManager manager.TokenManager, options Credentia
 		options:      options,
 		listeners:    make([]auth.CredentialsListener, 0),
 	}
-	cancelTokenManager, err := cp.tokenManager.Start(tokenListenerFromCP(cp))
+	stopTM, err := cp.tokenManager.Start(tokenListenerFromCP(cp))
 	if err != nil {
 		return nil, fmt.Errorf("couldn't start token manager: %w", err)
 	}
-	cp.closeTokenManager = cancelTokenManager
+	cp.stopTokenManager = stopTM
 	return cp, nil
 }

credentials_provider_test.go

@@ -343,7 +343,7 @@ func TestCredentialsProviderSubscribe(t *testing.T) {
 		mtm.On("GetToken", false).Return(testToken, nil)
 		mtm.On("Start", mock.Anything).
 			Run(mockTokenManagerLoop(mtm, tokenExpiration, testToken, nil)).
-			Return(manager.CloseFunc(mtm.Close), nil)
+			Return(manager.StopFunc(mtm.Stop), nil)
 		provider, err := NewConfidentialCredentialsProvider(options)
 		require.NoError(t, err)
 		require.NotNil(t, provider)
@@ -396,13 +396,13 @@ func TestCredentialsProviderSubscribe(t *testing.T) {
 
 		mtm.On("Start", mock.Anything).
 			Run(mockTokenManagerLoop(mtm, tokenExpiration, nil, errTokenError)).
-			Return(manager.CloseFunc(mtm.Close), nil)
+			Return(manager.StopFunc(mtm.Stop), nil)
 		provider, err := NewConfidentialCredentialsProvider(options)
 		require.NoError(t, err)
 		require.NotNil(t, provider)
 		var wg sync.WaitGroup
 		listeners := make([]*mockCredentialsListener, numListeners)
-		cancels := make([]auth.CancelProviderFunc, numListeners)
+		cancels := make([]auth.UnsubscribeFunc, numListeners)
 
 		// Subscribe multiple listeners concurrently
 		for i := 0; i < numListeners; i++ {
@@ -467,7 +467,7 @@ func TestCredentialsProviderSubscribe(t *testing.T) {
 
 		mtm.On("Start", mock.Anything).
 			Run(mockTokenManagerLoop(mtm, tokenExpiration, nil, errTokenError)).
-			Return(manager.CloseFunc(mtm.Close), nil)
+			Return(manager.StopFunc(mtm.Stop), nil)
 		provider, err := NewConfidentialCredentialsProvider(options)
 		require.NoError(t, err)
 		require.NotNil(t, provider)
@@ -525,7 +525,7 @@ func TestCredentialsProviderSubscribe(t *testing.T) {
 		require.NotNil(t, provider)
 		var wg sync.WaitGroup
 		listeners := make([]*mockCredentialsListener, numListeners)
-		cancels := make([]auth.CancelProviderFunc, numListeners)
+		cancels := make([]auth.UnsubscribeFunc, numListeners)
 
 		// Subscribe multiple listeners concurrently
 		for i := 0; i < numListeners; i++ {

entraid_test.go

@@ -54,7 +54,7 @@ func (m *fakeTokenManager) GetToken(forceRefresh bool) (*token.Token, error) {
 	return m.token, m.err
 }
 
-func (m *fakeTokenManager) Start(listener manager.TokenListener) (manager.CloseFunc, error) {
+func (m *fakeTokenManager) Start(listener manager.TokenListener) (manager.StopFunc, error) {
 	if m.err != nil {
 		return nil, m.err
 	}
@@ -65,10 +65,10 @@ func (m *fakeTokenManager) Start(listener manager.TokenListener) (manager.CloseF
 			case <-time.After(tokenExpiration):
 				m.lock.Lock()
 				if m.err != nil {
-					listener.OnTokenError(m.err)
+					listener.OnError(m.err)
 					return
 				}
-				listener.OnTokenNext(m.token)
+				listener.OnNext(m.token)
 				m.lock.Unlock()
 			case <-done:
 				// Exit the loop if done channel is closed
@@ -84,7 +84,7 @@ func (m *fakeTokenManager) Start(listener manager.TokenListener) (manager.CloseF
 	}, nil
 }
 
-func (m *fakeTokenManager) Close() error {
+func (m *fakeTokenManager) Stop() error {
 	return nil
 }
 
@@ -147,7 +147,7 @@ func (m *mockTokenManager) GetToken(forceRefresh bool) (*token.Token, error) {
 	return args.Get(0).(*token.Token), args.Error(1)
 }
 
-func (m *mockTokenManager) Start(listener manager.TokenListener) (manager.CloseFunc, error) {
+func (m *mockTokenManager) Start(listener manager.TokenListener) (manager.StopFunc, error) {
 	args := m.Called(listener)
 	m.lock.Lock()
 	if m.done == nil {
@@ -161,13 +161,13 @@ func (m *mockTokenManager) Start(listener manager.TokenListener) (manager.CloseF
 		m.listener = listener
 	}
 	m.lock.Unlock()
-	return args.Get(0).(manager.CloseFunc), args.Error(1)
+	return args.Get(0).(manager.StopFunc), args.Error(1)
 }
-func (m *mockTokenManager) Close() error {
+func (m *mockTokenManager) Stop() error {
 	m.lock.Lock()
 	defer m.lock.Unlock()
 	if m.listener == nil {
-		return manager.ErrTokenManagerAlreadyClosed
+		return manager.ErrTokenManagerAlreadyStopped
 	}
 	if m.listener != nil {
 		m.listener = nil
@@ -200,9 +200,9 @@ func mockTokenManagerLoop(mtm *mockTokenManager, tokenExpiration time.Duration,
 				case <-time.After(tokenExpiration):
 					mtm.lock.Lock()
 					if err != nil {
-						mtm.listener.OnTokenError(err)
+						mtm.listener.OnError(err)
 					} else {
-						mtm.listener.OnTokenNext(testToken)
+						mtm.listener.OnNext(testToken)
 					}
 					mtm.lock.Unlock()
 				}

go.mod

@@ -4,12 +4,14 @@ go 1.18
 
 require (
 	github.com/Azure/azure-sdk-for-go/sdk/azidentity v1.8.0-beta.1
-	github.com/redis/go-redis/v9 v9.5.3-0.20250331212737-c248425ade4a
+	github.com/redis/go-redis/v9 v9.5.3-0.20250416091253-d0a8c76d8420
 	github.com/stretchr/testify v1.10.0
 )
 
 require (
+	github.com/cespare/xxhash/v2 v2.3.0 // indirect
 	github.com/davecgh/go-spew v1.1.1 // indirect
+	github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f // indirect
 	github.com/pmezard/go-difflib v1.0.0 // indirect
 	github.com/stretchr/objx v0.5.2 // indirect
 	gopkg.in/yaml.v3 v3.0.1 // indirect

go.sum

@@ -8,8 +8,12 @@ github.com/Azure/azure-sdk-for-go/sdk/internal v1.10.0/go.mod h1:iZDifYGJTIgIIkY
 github.com/AzureAD/microsoft-authentication-extensions-for-go/cache v0.1.1 h1:WJTmL004Abzc5wDB5VtZG2PJk5ndYDgVacGqfirKxjM=
 github.com/AzureAD/microsoft-authentication-library-for-go v1.4.1 h1:8BKxhZZLX/WosEeoCvWysmKUscfa9v8LIPEEU0JjE2o=
 github.com/AzureAD/microsoft-authentication-library-for-go v1.4.1/go.mod h1:wP83P5OoQ5p6ip3ScPr0BAq0BvuPAvacpEuSzyouqAI=
+github.com/cespare/xxhash/v2 v2.3.0 h1:UL815xU9SqsFlibzuggzjXhog7bL6oX9BbNZnL2UFvs=
+github.com/cespare/xxhash/v2 v2.3.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs=
 github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
+github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f h1:lO4WD4F/rVNCu3HqELle0jiPLLBs70cWOduZpkS1E78=
+github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f/go.mod h1:cuUVRXasLTGF7a8hSLbxyZXjz+1KgoB3wDUb6vlszIc=
 github.com/golang-jwt/jwt/v5 v5.2.1 h1:OuVbFODueb089Lh128TAcimifWaLhJwVflnrgM17wHk=
 github.com/golang-jwt/jwt/v5 v5.2.1/go.mod h1:pqrtFR0X4osieyHYxtmOUWsAWrfe1Q5UVIyoH402zdk=
 github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0=
@@ -25,6 +29,8 @@ github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZb
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
 github.com/redis/go-redis/v9 v9.5.3-0.20250331212737-c248425ade4a h1:R5xgk8m+CF7lVE0EGr+tLkT1eM3Zfd39BJfnANQqpKA=
 github.com/redis/go-redis/v9 v9.5.3-0.20250331212737-c248425ade4a/go.mod h1:huWgSWd8mW6+m0VPhJjSSQ+d6Nh1VICQ6Q5lHuCH/Iw=
+github.com/redis/go-redis/v9 v9.5.3-0.20250416091253-d0a8c76d8420 h1:/dxO9rhmlhKP5pyI7omDH3QQzC0AppWxHT1w5TBsdTU=
+github.com/redis/go-redis/v9 v9.5.3-0.20250416091253-d0a8c76d8420/go.mod h1:huWgSWd8mW6+m0VPhJjSSQ+d6Nh1VICQ6Q5lHuCH/Iw=
 github.com/rogpeppe/go-internal v1.12.0 h1:exVL4IDcn6na9z1rAb56Vxr+CgyK3nn3O+epU5NdKM8=
 github.com/stretchr/objx v0.5.2 h1:xuMeJ0Sdp5ZMRXx/aWO6RZxdr3beISkG5/G/aIRr3pY=
 github.com/stretchr/objx v0.5.2/go.mod h1:FRsXN1f5AsAjCGJKqEizvkpNtU+EGNCLh3NxZ/8L+MA=

identity/azure_default_identity_provider.go

@@ -58,7 +58,7 @@ func NewDefaultAzureIdentityProvider(opts DefaultAzureIdentityProviderOptions) (
 
 // RequestToken requests a token from the Azure Default Identity provider.
 // It returns the token, the expiration time, and an error if any.
-func (a *DefaultAzureIdentityProvider) RequestToken() (shared.IdentityProviderResponse, error) {
+func (a *DefaultAzureIdentityProvider) RequestToken(ctx context.Context) (shared.IdentityProviderResponse, error) {
 	credFactory := a.credFactory
 	if credFactory == nil {
 		credFactory = &defaultCredFactory{}
@@ -68,7 +68,7 @@ func (a *DefaultAzureIdentityProvider) RequestToken() (shared.IdentityProviderRe
 		return nil, fmt.Errorf("failed to create default azure credential: %w", err)
 	}
 
-	token, err := cred.GetToken(context.TODO(), policy.TokenRequestOptions{Scopes: a.scopes})
+	token, err := cred.GetToken(ctx, policy.TokenRequestOptions{Scopes: a.scopes})
 	if err != nil {
 		return nil, fmt.Errorf("failed to get token: %w", err)
 	}

identity/azure_default_identity_provider_test.go

@@ -1,6 +1,7 @@
 package identity
 
 import (
+	"context"
 	"testing"
 
 	"github.com/Azure/azure-sdk-for-go/sdk/azcore"
@@ -38,7 +39,7 @@ func TestAzureDefaultIdentityProvider_RequestToken(t *testing.T) {
 
 	// Request a token from the provider in incorrect environment
 	// should fail.
-	token, err := provider.RequestToken()
+	token, err := provider.RequestToken(context.Background())
 	assert.Nil(t, token, "token should be nil")
 	assert.Error(t, err, "failed to request token")
 
@@ -51,7 +52,7 @@ func TestAzureDefaultIdentityProvider_RequestToken(t *testing.T) {
 	mCredFactory := &mockCredFactory{}
 	mCredFactory.On("NewDefaultAzureCredential", mock.Anything).Return(mCreds, nil)
 	provider.credFactory = mCredFactory
-	token, err = provider.RequestToken()
+	token, err = provider.RequestToken(context.Background())
 	assert.NotNil(t, token, "token should not be nil")
 	assert.NoError(t, err, "failed to request token")
 	assert.Equal(t, shared.ResponseTypeAccessToken, token.Type(), "token type should be access token")
@@ -70,7 +71,7 @@ func TestAzureDefaultIdentityProvider_RequestTokenWithScopes(t *testing.T) {
 
 	t.Run("RequestToken with custom scopes", func(t *testing.T) {
 		// Request a token from the provider
-		token, err := provider.RequestToken()
+		token, err := provider.RequestToken(context.Background())
 		assert.Nil(t, token, "token should be nil")
 		assert.Error(t, err, "failed to request token")
 
@@ -83,7 +84,7 @@ func TestAzureDefaultIdentityProvider_RequestTokenWithScopes(t *testing.T) {
 		mCredFactory := &mockCredFactory{}
 		mCredFactory.On("NewDefaultAzureCredential", mock.Anything).Return(mCreds, nil)
 		provider.credFactory = mCredFactory
-		token, err = provider.RequestToken()
+		token, err = provider.RequestToken(context.Background())
 		assert.NotNil(t, token, "token should not be nil")
 		assert.NoError(t, err, "failed to request token")
 		assert.Equal(t, shared.ResponseTypeAccessToken, token.Type(), "token type should be access token")
@@ -94,7 +95,7 @@ func TestAzureDefaultIdentityProvider_RequestTokenWithScopes(t *testing.T) {
 		mCredFactory := &mockCredFactory{}
 		mCredFactory.On("NewDefaultAzureCredential", mock.Anything).Return(nil, assert.AnError)
 		provider.credFactory = mCredFactory
-		token, err := provider.RequestToken()
+		token, err := provider.RequestToken(context.Background())
 		assert.Nil(t, token, "token should be nil")
 		assert.Error(t, err, "failed to request token")
 	})

identity/confidential_identity_provider.go

@@ -155,12 +155,12 @@ func NewConfidentialIdentityProvider(opts ConfidentialIdentityProviderOptions) (
 
 // RequestToken requests a token from the identity provider.
 // It returns the identity provider response, including the auth result.
-func (c *ConfidentialIdentityProvider) RequestToken() (shared.IdentityProviderResponse, error) {
+func (c *ConfidentialIdentityProvider) RequestToken(ctx context.Context) (shared.IdentityProviderResponse, error) {
 	if c.client == nil {
 		return nil, fmt.Errorf("client is not initialized")
 	}
 
-	result, err := c.client.AcquireTokenByCredential(context.TODO(), c.scopes)
+	result, err := c.client.AcquireTokenByCredential(ctx, c.scopes)
 	if err != nil {
 		return nil, fmt.Errorf("failed to acquire token: %w", err)
 	}