wip

Author: ndyakov

credentials_provider.go

@@ -2,6 +2,7 @@ package entraid
 
 import (
 	"fmt"
+	"sync"
 
 	"github.com/redis/go-redis/v9/auth"
 )
@@ -16,11 +17,18 @@ type entraidCredentialsProvider struct {
 	tokenManager       TokenManager
 	cancelTokenManager cancelFunc
 
+	// listeners is a slice of listeners that are notified when the token manager receives a new token.
 	listeners []auth.CredentialsListener
+
+	// rwLock is a mutex that is used to synchronize access to the listeners slice.
+	// It is used to ensure that only one goroutine can access the listeners slice at a time.
+	rwLock sync.RWMutex
 }
 
 // onTokenNext is a method that is called when the token manager receives a new token.
 func (e *entraidCredentialsProvider) onTokenNext(token *Token) {
+	e.rwLock.RLock()
+	defer e.rwLock.RUnlock()
 	// Notify all listeners with the new token.
 	for _, listener := range e.listeners {
 		listener.OnNext(&authCredentials{
@@ -34,6 +42,8 @@ func (e *entraidCredentialsProvider) onTokenNext(token *Token) {
 // onError is a method that is called when the token manager encounters an error.
 // It notifies all listeners with the error.
 func (e *entraidCredentialsProvider) onTokenError(err error) {
+	e.rwLock.RLock()
+	defer e.rwLock.RUnlock()
 	// Notify all listeners with the error.
 	for _, listener := range e.listeners {
 		listener.OnError(err)
@@ -46,6 +56,7 @@ func (e *entraidCredentialsProvider) onTokenError(err error) {
 // The listener will be notified with an error if there is an error obtaining the credentials.
 // The caller can cancel the subscription by calling the cancel function which is the second return value.
 func (e *entraidCredentialsProvider) Subscribe(listener auth.CredentialsListener) (auth.Credentials, auth.CancelProviderFunc, error) {
+	e.rwLock.Lock()
 	// Check if the listener is already in the list of listeners.
 	alreadySubscribed := false
 	for _, l := range e.listeners {
@@ -59,6 +70,7 @@ func (e *entraidCredentialsProvider) Subscribe(listener auth.CredentialsListener
 		// Get the token from the identity provider.
 		e.listeners = append(e.listeners, listener)
 	}
+	e.rwLock.Unlock()
 
 	token, err := e.tokenManager.GetToken()
 	if err != nil {
@@ -78,21 +90,32 @@ func (e *entraidCredentialsProvider) Subscribe(listener auth.CredentialsListener
 
 	cancel := func() error {
 		// Remove the listener from the list of listeners.
+		e.rwLock.Lock()
+		defer e.rwLock.Unlock()
 		for i, l := range e.listeners {
 			if l == listener {
 				e.listeners = append(e.listeners[:i], e.listeners[i+1:]...)
 				break
 			}
 		}
 		if len(e.listeners) == 0 {
-			e.cancelTokenManager()
+			e.close()
 		}
 		return nil
 	}
 
 	return credentials, cancel, nil
 }
 
+func (e *entraidCredentialsProvider) close() {
+	if e.cancelTokenManager != nil {
+		e.cancelTokenManager()
+	}
+	e.cancelTokenManager = nil
+	e.listeners = nil
+	e.rwLock = sync.RWMutex{}
+}
+
 type entraidTokenListener struct {
 	onTokenNext  func(token *Token)
 	onTokenError func(err error)

token_manager.go

@@ -48,6 +48,8 @@ type TokenManager interface {
 	GetToken() (*Token, error)
 	// Start starts the token manager and returns a channel that will receive updates.
 	Start(listener TokenListener) (cancelFunc, error)
+	// Close closes the token manager and releases any resources.
+	Close() error
 }
 
 // defaultTokenParser is a function that parses the raw token and returns Token object.
@@ -74,6 +76,7 @@ func NewTokenManager(idp IdentityProvider, options TokenManagerOptions) TokenMan
 	return &entraidTokenManager{
 		idp:          idp,
 		token:        nil,
+		closed:       make(chan struct{}),
 		tokenParser:  tokenParser,
 		retryOptions: retryOptions,
 	}
@@ -101,6 +104,8 @@ type entraidTokenManager struct {
 
 	// lock locks the listener to prevent concurrent access.
 	lock sync.Mutex
+
+	closed chan struct{}
 }
 
 func (e *entraidTokenManager) GetToken() (*Token, error) {
@@ -148,6 +153,7 @@ func (e *entraidTokenManager) Start(listener TokenListener) (cancelFunc, error)
 		return nil, fmt.Errorf("token manager already started")
 	}
 	e.listener = listener
+	e.closed = make(chan struct{})
 
 	token, err := e.GetToken()
 	if err != nil {
@@ -159,22 +165,61 @@ func (e *entraidTokenManager) Start(listener TokenListener) (cancelFunc, error)
 	go func(listener TokenListener) {
 		// Simulate token refresh
 		for {
-			time.Sleep(
-				time.Duration(e.retryOptions.InitialDelayMs) * time.Millisecond)
-			newToken, err := e.GetToken()
-			if err != nil {
-				listener.OnTokenError(err)
+			select {
+			case <-time.After(time.Duration(e.token.TTL) * time.Second):
+				// Token is about to expire, refresh it
+				for i := 0; i < e.retryOptions.MaxAttempts; i++ {
+					token, err := e.GetToken()
+					if err == nil {
+						listener.OnTokenNext(token)
+						break
+					}
+					// check if err is retryable
+					if err.Error() == "retryable error" {
+						// retry
+						continue
+					} else {
+						// not retryable
+						listener.OnTokenError(err)
+						return
+					}
+
+					// check if max attempts reached
+					if i == e.retryOptions.MaxAttempts-1 {
+						listener.OnTokenError(err)
+						return
+					}
+
+					// Exponential backoff
+					delay := time.Duration(e.retryOptions.InitialDelayMs) * time.Millisecond
+					if delay < time.Duration(e.retryOptions.MaxDelayMs)*time.Millisecond {
+						delay = time.Duration(float64(delay) * e.retryOptions.BackoffMultiplier)
+					}
+
+					time.Sleep(delay)
+
+					if delay > time.Duration(e.retryOptions.MaxDelayMs)*time.Millisecond {
+						delay = time.Duration(e.retryOptions.MaxDelayMs) * time.Millisecond
+					}
+				}
+			case <-e.closed:
+				// Token manager is closed, stop the loop
 				return
 			}
-			listener.OnTokenNext(newToken)
 		}
 	}(e.listener)
-	cancel := func() error {
-		// Stop the token manager.
-		return nil
-	}
 
-	return cancel, nil
+	return e.Close, nil
+}
+
+func (e *entraidTokenManager) Close() error {
+	e.lock.Lock()
+	defer e.lock.Unlock()
+	if e.listener != nil {
+		e.listener = nil
+	}
+	close(e.closed)
+	return nil
 }
 
 // defaultRetryOptionsOr returns the default retry options if the provided options are not set.