diff --git a/token_manager.go b/token_manager.go index 1361f636f..c6c4d7ddb 100644 --- a/token_manager.go +++ b/token_manager.go @@ -1,27 +1,37 @@ package redis import ( + "log" "sync" "time" ) type TokenManager struct { - token string - expiresAt time.Time - mutex sync.Mutex + token string + expiresAt time.Time + mutex sync.Mutex + refreshFunc func() (string, time.Duration, error) + stopChan chan struct{} } -func NewTokenManager() *TokenManager { - return &TokenManager{} +// NewTokenManager initializes a new TokenManager. +func NewTokenManager(refreshFunc func() (string, time.Duration, error)) *TokenManager { + return &TokenManager{ + refreshFunc: refreshFunc, + stopChan: make(chan struct{}), + } } +// SetToken updates the token and its expiration. func (tm *TokenManager) SetToken(token string, ttl time.Duration) { tm.mutex.Lock() defer tm.mutex.Unlock() tm.token = token tm.expiresAt = time.Now().Add(ttl) + log.Printf("Token updated with TTL: %s", ttl) } +// GetToken returns the current token if it's still valid. func (tm *TokenManager) GetToken() (string, bool) { tm.mutex.Lock() defer tm.mutex.Unlock() @@ -31,11 +41,73 @@ func (tm *TokenManager) GetToken() (string, bool) { return tm.token, true } -func (tm *TokenManager) RefreshToken(fetchToken func() (string, time.Duration, error)) error { - token, ttl, err := fetchToken() +// RefreshToken fetches a new token using the provided refresh function. +func (tm *TokenManager) RefreshToken() error { + if tm.refreshFunc == nil { + return nil + } + token, ttl, err := tm.refreshFunc() if err != nil { return err } tm.SetToken(token, ttl) return nil } + +// StartAutoRefresh starts a goroutine to proactively refresh the token. +func (tm *TokenManager) StartAutoRefresh() { + go func() { + ticker := time.NewTicker(1 * time.Minute) // Check periodically + defer ticker.Stop() + + for { + select { + case <-ticker.C: + if tm.shouldRefresh() { + log.Println("Proactively refreshing token...") + if err := tm.RefreshToken(); err != nil { + log.Printf("Failed to refresh token: %v", err) + } + } + case <-tm.stopChan: + log.Println("Stopping auto-refresh...") + return + } + } + }() +} + +// StopAutoRefresh stops the auto-refresh goroutine. +func (tm *TokenManager) StopAutoRefresh() { + close(tm.stopChan) +} + +// shouldRefresh checks if the token is nearing expiration. +func (tm *TokenManager) shouldRefresh() bool { + tm.mutex.Lock() + defer tm.mutex.Unlock() + remaining := time.Until(tm.expiresAt) + return remaining < 5*time.Minute // Refresh if less than 5 minutes remain +} + +// MonitorTelemetry adds monitoring for token usage and expiration. +func (tm *TokenManager) MonitorTelemetry() { + go func() { + ticker := time.NewTicker(30 * time.Second) // Adjust as needed + defer ticker.Stop() + + for { + select { + case <-ticker.C: + token, valid := tm.GetToken() + if !valid { + log.Println("Token has expired.") + } else { + log.Printf("Token is valid: %s, expires in: %s", token, time.Until(tm.expiresAt)) + } + case <-tm.stopChan: + return + } + } + }() +}