fix: use singleflight and add backoff for token refresh

go.mod

@@ -6,6 +6,7 @@ require (
 	github.com/coder/websocket v1.8.15
 	github.com/stretchr/testify v1.11.1
 	golang.org/x/oauth2 v0.36.0
+	golang.org/x/sync v0.22.0
 	google.golang.org/grpc v1.81.1
 	google.golang.org/protobuf v1.36.11
 )

go.sum

@@ -34,6 +34,8 @@ golang.org/x/net v0.51.0 h1:94R/GTO7mt3/4wIKpcR5gkGmRLOuE/2hNGeWq/GBIFo=
 golang.org/x/net v0.51.0/go.mod h1:aamm+2QF5ogm02fjy5Bb7CQ0WMt1/WVM7FtyaTLlA9Y=
 golang.org/x/oauth2 v0.36.0 h1:peZ/1z27fi9hUOFCAZaHyrpWG5lwe0RJEEEeH0ThlIs=
 golang.org/x/oauth2 v0.36.0/go.mod h1:YDBUJMTkDnJS+A4BP4eZBjCqtokkg1hODuPjwiGPO7Q=
+golang.org/x/sync v0.22.0 h1:SZjpbeLmrCk4xhRSZFNZW5gFUeCeFgjekvI/+gfScek=
+golang.org/x/sync v0.22.0/go.mod h1:9xrNwdLfx4jkKbNva9FpL6vEN7evnE43NNNJQ2LF3+0=
 golang.org/x/sys v0.42.0 h1:omrd2nAlyT5ESRdCLYdm3+fMfNFE/+Rf4bDIQImRJeo=
 golang.org/x/sys v0.42.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw=
 golang.org/x/text v0.34.0 h1:oL/Qq0Kdaqxa1KbNeMKwQq0reLCCaFtqu2eNuSeNHbk=

openshell/v1/auth_refresh.go

@@ -10,11 +10,16 @@ import (
 	"time"
 
 	"golang.org/x/oauth2"
+	"golang.org/x/sync/singleflight"
 
 	"github.com/rhuss/openshell-sdk-go/openshell/v1/types"
 )
 
-const defaultLeeway = 10 * time.Second
+const (
+	defaultLeeway  = 10 * time.Second
+	initialBackoff = 1 * time.Second
+	maxBackoff     = 30 * time.Second
+)
 
 var errNilTokenSource = errors.New("openshell: TokenSource must not be nil")
 
@@ -52,11 +57,14 @@ func WithLogger(l types.Logger) RefreshOption {
 }
 
 type refreshableAuth struct {
-	source oauth2.TokenSource
-	mu     sync.RWMutex
-	tok    *oauth2.Token
-	leeway time.Duration
-	logger types.Logger
+	source     oauth2.TokenSource
+	mu         sync.Mutex
+	group      singleflight.Group
+	tok        *oauth2.Token
+	leeway     time.Duration
+	logger     types.Logger
+	nextRetry  time.Time
+	backoff    time.Duration
 }
 
 func (r *refreshableAuth) isTokenValid() bool {
@@ -70,25 +78,46 @@ func (r *refreshableAuth) isTokenValid() bool {
 }
 
 func (r *refreshableAuth) GetRequestMetadata(_ context.Context, _ ...string) (map[string]string, error) {
-	// Fast path: RLock, return cached token if valid.
-	r.mu.RLock()
+	r.mu.Lock()
 	if r.isTokenValid() {
 		tok := r.tok.AccessToken
-		r.mu.RUnlock()
+		r.mu.Unlock()
 		return map[string]string{"authorization": "Bearer " + tok}, nil
 	}
-	r.mu.RUnlock()
 
-	// Slow path: Lock, re-check, fetch if still stale.
+	if !r.nextRetry.IsZero() && time.Now().Before(r.nextRetry) {
+		if r.tok != nil {
+			tok := r.tok.AccessToken
+			r.mu.Unlock()
+			return map[string]string{"authorization": "Bearer " + tok}, nil
+		}
+		r.mu.Unlock()
+		return nil, errors.New("openshell: token refresh failed and backoff is active")
+	}
+	r.mu.Unlock()
+
+	val, err, _ := r.group.Do("refresh", func() (any, error) {
+		return r.source.Token()
+	})
+
 	r.mu.Lock()
 	defer r.mu.Unlock()
 
-	if r.isTokenValid() {
-		return map[string]string{"authorization": "Bearer " + r.tok.AccessToken}, nil
-	}
-
-	newTok, err := r.source.Token()
 	if err != nil {
+		if r.nextRetry.IsZero() || !time.Now().Before(r.nextRetry) {
+			bo := r.backoff
+			if bo == 0 {
+				bo = initialBackoff
+			} else {
+				bo *= 2
+				if bo > maxBackoff {
+					bo = maxBackoff
+				}
+			}
+			r.backoff = bo
+			r.nextRetry = time.Now().Add(bo)
+		}
+
 		if r.tok != nil {
 			if r.logger != nil {
 				r.logger.Error(err, "token refresh failed, using cached token")
@@ -98,7 +127,13 @@ func (r *refreshableAuth) GetRequestMetadata(_ context.Context, _ ...string) (ma
 		return nil, err
 	}
 
-	r.tok = newTok
+	tok, ok := val.(*oauth2.Token)
+	if !ok || tok == nil {
+		return nil, errors.New("openshell: token source returned nil token without error")
+	}
+	r.tok = tok
+	r.backoff = 0
+	r.nextRetry = time.Time{}
 	return map[string]string{"authorization": "Bearer " + r.tok.AccessToken}, nil
 }
 
@@ -108,7 +143,8 @@ func (r *refreshableAuth) RequireTransportSecurity() bool {
 
 // RefreshableToken returns an AuthProvider that caches tokens from src
 // and refreshes them before expiry. Concurrent callers share a single
-// refresh call (coalesced via RWMutex double-checked locking).
+// in-flight refresh via singleflight. Failed refreshes trigger exponential
+// backoff (1s, 2s, 4s, ..., 30s cap) to avoid amplifying token endpoint outages.
 func RefreshableToken(src oauth2.TokenSource, opts ...RefreshOption) (AuthProvider, error) {
 	if src == nil {
 		return nil, errNilTokenSource

openshell/v1/auth_refresh_test.go

@@ -121,29 +121,32 @@ func TestGetRequestMetadata_ConcurrentSingleFlight(t *testing.T) {
 	src := &mockTokenSource{
 		tokenFunc: func() (*oauth2.Token, error) {
 			fetchCount.Add(1)
-			time.Sleep(10 * time.Millisecond) // simulate slow token fetch
+			time.Sleep(100 * time.Millisecond)
 			return &oauth2.Token{AccessToken: "shared-token", Expiry: time.Now().Add(time.Hour)}, nil
 		},
 	}
 	provider, err := RefreshableToken(src)
 	require.NoError(t, err)
 
-	const goroutines = 1000
+	const goroutines = 20
 	var wg sync.WaitGroup
 	wg.Add(goroutines)
+	ready := make(chan struct{})
 	results := make([]string, goroutines)
 	errs := make([]error, goroutines)
 
 	for i := range goroutines {
 		go func(idx int) {
 			defer wg.Done()
+			<-ready
 			md, e := provider.GetRequestMetadata(context.Background())
 			errs[idx] = e
 			if md != nil {
 				results[idx] = md["authorization"]
 			}
 		}(i)
 	}
+	close(ready)
 	wg.Wait()
 
 	for i := range goroutines {
@@ -318,6 +321,132 @@ func TestGetRequestMetadata_ZeroExpiryNeverRefreshes(t *testing.T) {
 	assert.Equal(t, 1, src.calls(), "zero-expiry token should never be refreshed")
 }
 
+// --- Backoff tests ---
+
+func TestGetRequestMetadata_BackoffSkipsRefreshDuringWindow(t *testing.T) {
+	var callNum atomic.Int32
+	src := &mockTokenSource{
+		tokenFunc: func() (*oauth2.Token, error) {
+			n := callNum.Add(1)
+			if n == 1 {
+				return &oauth2.Token{AccessToken: "stale", Expiry: time.Now().Add(-time.Minute)}, nil
+			}
+			return nil, fmt.Errorf("idp down")
+		},
+	}
+	provider, err := RefreshableToken(src, WithLeeway(0))
+	require.NoError(t, err)
+
+	_, err = provider.GetRequestMetadata(context.Background())
+	require.NoError(t, err)
+	_, err = provider.GetRequestMetadata(context.Background())
+	require.NoError(t, err)
+	beforeCount := src.calls()
+	require.Equal(t, 2, beforeCount)
+
+	ra := provider.(*refreshableAuth)
+	ra.mu.Lock()
+	ra.nextRetry = time.Now().Add(time.Minute)
+	ra.mu.Unlock()
+
+	md, err := provider.GetRequestMetadata(context.Background())
+	require.NoError(t, err)
+	assert.Equal(t, "Bearer stale", md["authorization"])
+	assert.Equal(t, beforeCount, src.calls(), "should not call Token() during backoff window")
+}
+
+func TestGetRequestMetadata_BackoffResetsOnSuccess(t *testing.T) {
+	var callNum atomic.Int32
+	src := &mockTokenSource{
+		tokenFunc: func() (*oauth2.Token, error) {
+			n := callNum.Add(1)
+			switch n {
+			case 1:
+				return &oauth2.Token{AccessToken: "initial", Expiry: time.Now().Add(-time.Second)}, nil
+			case 2:
+				return nil, fmt.Errorf("fail once")
+			default:
+				return &oauth2.Token{AccessToken: "recovered", Expiry: time.Now().Add(time.Hour)}, nil
+			}
+		},
+	}
+	provider, err := RefreshableToken(src, WithLeeway(0))
+	require.NoError(t, err)
+
+	_, _ = provider.GetRequestMetadata(context.Background())
+	_, _ = provider.GetRequestMetadata(context.Background())
+
+	ra := provider.(*refreshableAuth)
+	ra.mu.Lock()
+	ra.nextRetry = time.Time{}
+	ra.mu.Unlock()
+
+	md, err := provider.GetRequestMetadata(context.Background())
+	require.NoError(t, err)
+	assert.Equal(t, "Bearer recovered", md["authorization"])
+
+	ra.mu.Lock()
+	assert.Equal(t, time.Duration(0), ra.backoff, "backoff should reset after success")
+	assert.True(t, ra.nextRetry.IsZero(), "nextRetry should be zero after success")
+	ra.mu.Unlock()
+}
+
+func TestGetRequestMetadata_BackoffCapsAt30s(t *testing.T) {
+	src := &mockTokenSource{
+		tokenFunc: func() (*oauth2.Token, error) {
+			return nil, fmt.Errorf("always fail")
+		},
+	}
+	provider, err := RefreshableToken(src, WithLeeway(0))
+	require.NoError(t, err)
+
+	ra := provider.(*refreshableAuth)
+
+	for range 10 {
+		ra.mu.Lock()
+		ra.nextRetry = time.Time{}
+		ra.mu.Unlock()
+
+		_, _ = provider.GetRequestMetadata(context.Background())
+	}
+
+	ra.mu.Lock()
+	assert.Equal(t, maxBackoff, ra.backoff, "backoff should cap at maxBackoff")
+	ra.mu.Unlock()
+}
+
+func TestGetRequestMetadata_ConcurrentFailureBackoffNotOverIncremented(t *testing.T) {
+	src := &mockTokenSource{
+		tokenFunc: func() (*oauth2.Token, error) {
+			time.Sleep(10 * time.Millisecond)
+			return nil, fmt.Errorf("idp down")
+		},
+	}
+	provider, err := RefreshableToken(src, WithLeeway(0))
+	require.NoError(t, err)
+
+	const goroutines = 20
+	var wg sync.WaitGroup
+	wg.Add(goroutines)
+	ready := make(chan struct{})
+	for range goroutines {
+		go func() {
+			defer wg.Done()
+			<-ready
+			_, _ = provider.GetRequestMetadata(context.Background())
+		}()
+	}
+	close(ready)
+	wg.Wait()
+
+	ra := provider.(*refreshableAuth)
+	ra.mu.Lock()
+	assert.Equal(t, initialBackoff, ra.backoff,
+		"a single coalesced failure should set backoff to initialBackoff, not escalate")
+	ra.mu.Unlock()
+	assert.Equal(t, 1, src.calls(), "singleflight should coalesce to 1 call")
+}
+
 // --- benchmarks ---
 
 func BenchmarkGetRequestMetadata_CachedToken(b *testing.B) {