fix(manager): wip, still have tests to resolve

Author: ndyakov

credentials_provider.go

@@ -67,26 +67,23 @@ func (e *entraidCredentialsProvider) onTokenError(err error) {
 //
 // Note: If the listener is already subscribed, it will not receive duplicate notifications.
 func (e *entraidCredentialsProvider) Subscribe(listener auth.CredentialsListener) (auth.Credentials, auth.UnsubscribeFunc, error) {
-	var token *token.Token
 	// check if the manager is working
 	// If the stopTokenManager is nil, the token manager is not started.
 	e.tmLock.Lock()
 	if e.stopTokenManager == nil {
-		t, stopTM, err := e.tokenManager.Start(tokenListenerFromCP(e))
+		stopTM, err := e.tokenManager.Start(tokenListenerFromCP(e))
 		if err != nil {
 			return nil, nil, fmt.Errorf("couldn't start token manager: %w", err)
 		}
 		e.stopTokenManager = stopTM
-		token = t
-	} else {
-		t, err := e.tokenManager.GetToken(false)
-		if err != nil {
-			return nil, nil, fmt.Errorf("couldn't get token: %w", err)
-		}
-		token = t
 	}
 	e.tmLock.Unlock()
 
+	token, err := e.tokenManager.GetToken(false)
+	if err != nil {
+		return nil, nil, fmt.Errorf("couldn't get token: %w", err)
+	}
+
 	e.rwLock.Lock()
 	// Check if the listener is already in the list of listeners.
 	alreadySubscribed := false
@@ -152,5 +149,11 @@ func NewCredentialsProvider(tokenManager manager.TokenManager, options Credentia
 		options:      options,
 		listeners:    make([]auth.CredentialsListener, 0),
 	}
+	// Start the token manager.
+	stop, err := tokenManager.Start(tokenListenerFromCP(cp))
+	if err != nil {
+		return nil, fmt.Errorf("couldn't start token manager: %w", err)
+	}
+	cp.stopTokenManager = stop
 	return cp, nil
 }

entraid_test.go

@@ -48,7 +48,7 @@ func (m *fakeTokenManager) GetToken(forceRefresh bool) (*token.Token, error) {
 			rawTokenString,
 			time.Now().Add(tokenExpiration),
 			time.Now(),
-			int64(100*time.Millisecond),
+			int64(tokenExpiration.Seconds()),
 		)
 	}
 	return m.token, m.err

manager/defaults.go

@@ -3,6 +3,7 @@ package manager
 import (
 	"errors"
 	"fmt"
+	"math"
 	"net"
 	"os"
 	"time"
@@ -105,7 +106,7 @@ func (*defaultIdentityProviderResponseParser) ParseResponse(response shared.Iden
 
 	var username, password, rawToken string
 	var expiresOn time.Time
-	now := time.Now().UTC().Truncate(time.Second)
+	now := time.Now().UTC().Truncate(time.Second).Add(time.Second)
 
 	switch response.Type() {
 	case shared.ResponseTypeAuthResult:
@@ -176,6 +177,6 @@ func (*defaultIdentityProviderResponseParser) ParseResponse(response shared.Iden
 		rawToken,
 		expiresOn,
 		now,
-		int64(time.Until(expiresOn).Seconds()),
+		int64(math.Ceil(time.Until(expiresOn).Seconds())),
 	), nil
 }

manager/entraid_manager.go

@@ -0,0 +1,267 @@
+package manager
+
+import (
+	"context"
+	"fmt"
+	"sync"
+	"time"
+
+	"github.com/redis-developer/go-redis-entraid/internal"
+	"github.com/redis-developer/go-redis-entraid/shared"
+	"github.com/redis-developer/go-redis-entraid/token"
+)
+
+// entraidTokenManager is a struct that implements the TokenManager interface.
+type entraidTokenManager struct {
+	// idp is the identity provider used to obtain the token.
+	idp shared.IdentityProvider
+
+	// token is the authentication token for the user which should be kept in memory if valid.
+	token *token.Token
+
+	// tokenRWLock is a read-write lock used to protect the token from concurrent access.
+	tokenRWLock sync.RWMutex
+
+	// identityProviderResponseParser is the parser used to parse the response from the identity provider.
+	// It`s ParseResponse method will be called to parse the response and return the token.
+	identityProviderResponseParser shared.IdentityProviderResponseParser
+
+	// retryOptions is a struct that contains the options for retrying the token request.
+	// It contains the maximum number of attempts, initial delay, maximum delay, and backoff multiplier.
+	// The default values are 3 attempts, 1000 ms initial delay, 10000 ms maximum delay, and 2.0 backoff multiplier.
+	// The values can be overridden by the user.
+	retryOptions RetryOptions
+
+	// listener is the single listener for the token manager.
+	// It is used to receive updates from the token manager.
+	// The token manager will call the listener's OnNext method with the updated token.
+	// If an error occurs, the token manager will call the listener's OnError method with the error.
+	// if listener is set, Start will fail
+	listener TokenListener
+
+	// lock locks the listener to prevent concurrent access.
+	lock sync.Mutex
+
+	// expirationRefreshRatio is the ratio of the token expiration time to refresh the token.
+	// It is used to determine when to refresh the token.
+	// The value should be between 0 and 1.
+	// For example, if the expiration time is 1 hour and the ratio is 0.75,
+	// the token will be refreshed after 45 minutes. (the token is refreshed when 75% of its lifetime has passed)
+	expirationRefreshRatio float64
+
+	// lowerBoundDuration is the lower bound for the refresh time in time.Duration.
+	lowerBoundDuration time.Duration
+
+	// closedChan is a channel that is closedChan when the token manager is closedChan.
+	// It is used to signal the token manager to stop requesting tokens.
+	closedChan chan struct{}
+
+	// context is the context used to request the token from the identity provider.
+	ctx context.Context
+
+	// ctxCancel is the cancel function for the context.
+	ctxCancel context.CancelFunc
+
+	// requestTimeout is the timeout for the request to the identity provider.
+	requestTimeout time.Duration
+}
+
+func (e *entraidTokenManager) GetToken(forceRefresh bool) (*token.Token, error) {
+	e.tokenRWLock.RLock()
+	// check if the token is nil and if it is not expired
+	t := e.token
+	duration := e.durationToRenewal(t)
+	if !forceRefresh && t != nil && duration > 0 {
+		e.tokenRWLock.RUnlock()
+		return t, nil
+	}
+	e.tokenRWLock.RUnlock()
+
+	// start the context early,
+	// since at heavy concurrent load
+	// locks may take some time to acquire
+	ctx, ctxCancel := context.WithTimeout(e.ctx, e.requestTimeout)
+	defer ctxCancel()
+
+	// Upgrade to write lock for token update
+	e.tokenRWLock.Lock()
+	defer e.tokenRWLock.Unlock()
+
+	// Double-check pattern to avoid unnecessary token refresh
+	t = e.token
+	duration = e.durationToRenewal(t)
+	if !forceRefresh && t != nil && duration > 0 {
+		return t, nil
+	}
+
+	// Request a new token from the identity provider
+	idpResult, err := e.idp.RequestToken(ctx)
+	if err != nil {
+		return nil, fmt.Errorf("failed to request token from idp: %w", err)
+	}
+
+	t, err = e.identityProviderResponseParser.ParseResponse(idpResult)
+	if err != nil {
+		return nil, fmt.Errorf("failed to parse token: %w", err)
+	}
+
+	if t == nil {
+		return nil, fmt.Errorf("failed to get token: token is nil")
+	}
+
+	// Store the token
+	e.token = t
+	// Return the token - no need to copy since it's immutable
+	return t, nil
+}
+
+// Start starts the token manager and returns cancelFunc to stop the token manager.
+// It takes a TokenListener as an argument, which is used to receive updates.
+// The token manager will call the listener's OnNext method with the updated token.
+// If an error occurs, the token manager will call the listener's OnError method with the error.
+func (e *entraidTokenManager) Start(listener TokenListener) (StopFunc, error) {
+	e.lock.Lock()
+	defer e.lock.Unlock()
+	if e.listener != nil {
+		return nil, ErrTokenManagerAlreadyStarted
+	}
+
+	if e.closedChan != nil && !internal.IsClosed(e.closedChan) {
+		// there is a hanging goroutine that is waiting for the closedChan to be closed
+		// if the closedChan is not nil and not closed, close it
+		close(e.closedChan)
+	}
+
+	ctx, ctxCancel := context.WithCancel(context.Background())
+	e.ctx = ctx
+	e.ctxCancel = ctxCancel
+
+	// make sure there is token in memory before starting the loop
+	_, err := e.GetToken(false)
+	if err != nil {
+		return nil, fmt.Errorf("failed to get token: %w", err)
+	}
+
+	e.closedChan = make(chan struct{})
+	e.listener = listener
+
+	go func(listener TokenListener, closed <-chan struct{}) {
+		maxDelay := e.retryOptions.MaxDelay
+		initialDelay := e.retryOptions.InitialDelay
+
+		for {
+			e.tokenRWLock.RLock()
+			timeToRenewal := e.durationToRenewal(e.token)
+			e.tokenRWLock.RUnlock()
+			select {
+			case <-closed:
+				return
+			case <-time.After(timeToRenewal):
+				if timeToRenewal == 0 {
+					// Token was requested immediately, guard against infinite loop
+					select {
+					case <-closed:
+						return
+					case <-time.After(initialDelay):
+						// continue to attempt
+					}
+				}
+
+				// Token is about to expire, refresh it
+				delay := initialDelay
+				for i := 0; i < e.retryOptions.MaxAttempts; i++ {
+					t, err := e.GetToken(true)
+					if err == nil {
+						listener.OnNext(t)
+						break
+					}
+
+					// check if err is retriable
+					if e.retryOptions.IsRetryable(err) {
+						if i == e.retryOptions.MaxAttempts-1 {
+							// last attempt, call OnError
+							listener.OnError(fmt.Errorf("max attempts reached: %w", err))
+							return
+						}
+
+						// Exponential backoff
+						if delay < maxDelay {
+							delay = time.Duration(float64(delay) * e.retryOptions.BackoffMultiplier)
+						}
+						if delay > maxDelay {
+							delay = maxDelay
+						}
+
+						select {
+						case <-closed:
+							return
+						case <-time.After(delay):
+							// continue to next attempt
+						}
+					} else {
+						// not retriable
+						listener.OnError(err)
+						return
+					}
+				}
+			}
+		}
+	}(listener, e.closedChan)
+
+	return e.stop, nil
+}
+
+// stop closes the token manager and releases any resources.
+func (e *entraidTokenManager) stop() (err error) {
+	e.lock.Lock()
+	defer e.lock.Unlock()
+	defer func() {
+		// recover from panic and return the error
+		if r := recover(); r != nil {
+			err = fmt.Errorf("failed to stop token manager: %s", r)
+		}
+	}()
+
+	if e.closedChan == nil || e.listener == nil {
+		return ErrTokenManagerAlreadyStopped
+	}
+
+	e.ctxCancel()
+	e.listener = nil
+	close(e.closedChan)
+
+	return nil
+}
+
+// durationToRenewal calculates the duration to the next token renewal.
+// It returns the duration to the next token renewal based on the expiration refresh ratio and the lower bound duration.
+// If the token is nil, it returns 0.
+// If the time till expiration is less than the lower bound duration, it returns 0 to renew the token now.
+func (e *entraidTokenManager) durationToRenewal(t *token.Token) time.Duration {
+	if t == nil {
+		return 0
+	}
+	expirationRefreshTime := t.ReceivedAt().Add(time.Duration(float64(t.TTL()) * float64(time.Second) * e.expirationRefreshRatio))
+	timeTillExpiration := time.Until(t.ExpirationOn())
+	now := time.Now().UTC()
+
+	if expirationRefreshTime.Before(now) {
+		return 0
+	}
+
+	// if the timeTillExpiration is less than the lower bound (or 0), return 0 to renew the token NOW
+	if timeTillExpiration <= e.lowerBoundDuration || timeTillExpiration <= 0 {
+		return 0
+	}
+
+	// Calculate the time to renew the token based on the expiration refresh ratio
+	duration := time.Until(expirationRefreshTime)
+
+	// if the duration will take us past the lower bound, return the duration to lower bound
+	if timeTillExpiration-e.lowerBoundDuration < duration {
+		return timeTillExpiration - e.lowerBoundDuration
+	}
+
+	// return the calculated duration
+	return duration
+}

manager/manager_test.go

@@ -61,7 +61,7 @@ var testTokenValid = token.New(
 	"test",
 	time.Now().Add(time.Hour),
 	time.Now(),
-	int64(time.Hour),
+	int64(time.Hour.Seconds()),
 )
 
 func newTestJWTToken(expiresOn time.Time) string {