diff --git a/sdk/cliproxy/auth/auto_refresh_loop.go b/sdk/cliproxy/auth/auto_refresh_loop.go index b7f87bf02d0..079723aa800 100644 --- a/sdk/cliproxy/auth/auto_refresh_loop.go +++ b/sdk/cliproxy/auth/auto_refresh_loop.go @@ -101,6 +101,10 @@ func (l *authAutoRefreshLoop) worker(ctx context.Context) { continue } if l.manager.beginRefreshJob(ctx, job) { + if retryAt, settled := l.manager.dispatchRefresh(ctx, job); settled { + l.manager.finishRefreshJob(job, retryAt, false) + continue + } _, _ = l.manager.refreshAuthForRequestAtEpoch(ctx, job.id, "", job.registrationEpoch) } l.manager.finishRefreshJob(job, time.Time{}, false) diff --git a/sdk/cliproxy/auth/conductor.go b/sdk/cliproxy/auth/conductor.go index 3f485fcf10d..ce880a79f2e 100644 --- a/sdk/cliproxy/auth/conductor.go +++ b/sdk/cliproxy/auth/conductor.go @@ -231,6 +231,9 @@ type Manager struct { // Auto refresh state refreshCancel context.CancelFunc refreshLoop *authAutoRefreshLoop + // refreshDispatcher, when set, takes the refreshes the scheduler finds due + // in place of refreshing inline, under m.mu. + refreshDispatcher RefreshDispatcher // refreshJobs retains queued and running jobs across loop restarts under m.mu. refreshJobs map[string]*authRefreshJob diff --git a/sdk/cliproxy/auth/conductor_refresh.go b/sdk/cliproxy/auth/conductor_refresh.go index aa31a893a4a..588fbb8429c 100644 --- a/sdk/cliproxy/auth/conductor_refresh.go +++ b/sdk/cliproxy/auth/conductor_refresh.go @@ -66,6 +66,7 @@ func (m *Manager) StartAutoRefresh(parent context.Context, interval time.Duratio m.refreshCancel = cancelCtx m.refreshLoop = loop m.mu.Unlock() + m.startRefreshQueue(ctx) loop.rebuild(time.Now()) go loop.run(ctx) @@ -78,6 +79,7 @@ func (m *Manager) StopAutoRefresh() { cancel := m.refreshCancel m.refreshCancel = nil m.refreshLoop = nil + m.refreshDispatcher = nil m.mu.Unlock() if cancel != nil { cancel() diff --git a/sdk/cliproxy/auth/refresh_queue.go b/sdk/cliproxy/auth/refresh_queue.go new file mode 100644 index 00000000000..cd93101947b --- /dev/null +++ b/sdk/cliproxy/auth/refresh_queue.go @@ -0,0 +1,136 @@ +package auth + +import ( + "context" + "errors" + "strings" + "time" + + log "github.com/sirupsen/logrus" +) + +// refreshDispatchTimeout bounds handing a request to the queue, so an +// unreachable queue costs one scheduler worker seconds, not its whole run. +const refreshDispatchTimeout = 10 * time.Second + +// RefreshRequest asks whichever replica takes it to refresh one credential. It +// names the credential and carries no secret: the replica that takes it reads +// the credential itself. +type RefreshRequest struct { + AuthID string + Provider string + // CredentialVersion is the version of the credential the requester saw + // expiring. It identifies a request, so asking again for the same version + // is not queued twice. + CredentialVersion uint64 +} + +// RefreshDispatcher hands a refresh to the replicas that share it, so exactly +// one of them performs it. +type RefreshDispatcher interface { + // RequestRefresh queues the request. It returns nil when the request is + // queued or an equivalent one already is. + RequestRefresh(ctx context.Context, req RefreshRequest) error +} + +// RefreshQueueHost is implemented by a store whose replicas share a refresh +// queue. The manager then asks the queue for the refreshes its scheduler finds +// due, and serves the requests it is given, instead of refreshing inline. +type RefreshQueueHost interface { + RefreshDispatcher + // ServeRefreshRequests calls handle for each request this replica wins, + // until ctx ends. A request is complete when handle returns nil; handle + // returns an error only when the request should be delivered again. + ServeRefreshRequests(ctx context.Context, handle func(context.Context, RefreshRequest) error) error +} + +// startRefreshQueue attaches the store's refresh queue to the scheduler, when +// it has one, and serves its requests until ctx ends. +func (m *Manager) startRefreshQueue(ctx context.Context) { + m.mu.Lock() + m.refreshDispatcher = nil + store := m.store + m.mu.Unlock() + host, ok := store.(RefreshQueueHost) + if !ok || host == nil || m.HomeEnabled() { + return + } + m.mu.Lock() + m.refreshDispatcher = host + m.mu.Unlock() + go func() { + if err := host.ServeRefreshRequests(ctx, m.handleRefreshRequest); err != nil && ctx.Err() == nil { + log.WithError(err).Error("refresh queue consumer stopped; this replica no longer performs queued refreshes") + } + }() +} + +// dispatchRefresh asks the queue to refresh the credential of a due job. It +// reports whether the job is settled, and when the scheduler may consider the +// credential again. It reports false when the caller must refresh inline, +// which is only ever the case for a grant that stays valid after use: a +// single-use grant waits for the queue instead, since the queue is what keeps +// two replicas from spending it together. +func (m *Manager) dispatchRefresh(ctx context.Context, job *authRefreshJob) (retryAt time.Time, settled bool) { + m.mu.RLock() + dispatcher := m.refreshDispatcher + var ( + req RefreshRequest + exec ProviderExecutor + ) + if auth := m.auths[job.id]; auth != nil { + req = RefreshRequest{AuthID: auth.ID, Provider: auth.Provider, CredentialVersion: auth.CredentialVersion} + exec, _ = m.executorLocked(executorKeyFromAuth(auth)) + } + m.mu.RUnlock() + if dispatcher == nil || req.AuthID == "" { + return time.Time{}, false + } + + ctx, cancel := context.WithTimeout(ctx, refreshDispatchTimeout) + defer cancel() + err := dispatcher.RequestRefresh(ctx, req) + if err == nil { + // Whoever takes the request refreshes it and the store tells this + // replica; asking again before then only repeats the request. + return job.pendingUntil, true + } + log.WithFields(log.Fields{"auth_id": req.AuthID, "provider": req.Provider}).Warnf("queue refresh request failed: %v", err) + if exec != nil && RefreshPolicyOf(exec) != RefreshSingleUse { + return time.Time{}, false + } + return time.Time{}, true +} + +// handleRefreshRequest refreshes the credential of a request this replica won. +// It looks at the credential again first: another replica may have refreshed it +// since the request was made. The refresh itself is the one every refresh uses, +// so a single-use grant is claimed in the journal before it is sent, and the +// result is written back with the store's revision check. A refresh that +// fails is recorded on the credential and is not delivered again; the +// scheduler asks again when the credential is next due. +func (m *Manager) handleRefreshRequest(ctx context.Context, req RefreshRequest) error { + id := strings.TrimSpace(req.AuthID) + if id == "" { + return nil + } + m.mu.RLock() + auth := m.auths[id].Clone() + m.mu.RUnlock() + if auth == nil { + return nil + } + // The requester's own pending marker must not make it skip its request. + probe := auth.Clone() + probe.NextRefreshAfter = time.Time{} + if !m.shouldRefresh(probe, time.Now()) { + return nil + } + if _, err := m.refreshAuthForRequestAtEpoch(ctx, id, "", auth.RegistrationEpoch); err != nil { + if ctx.Err() != nil && errors.Is(err, ctx.Err()) { + return err + } + log.WithFields(log.Fields{"auth_id": id, "provider": auth.Provider}).Debugf("queued refresh failed: %v", err) + } + return nil +} diff --git a/sdk/cliproxy/auth/refresh_queue_test.go b/sdk/cliproxy/auth/refresh_queue_test.go new file mode 100644 index 00000000000..1d899e2e89b --- /dev/null +++ b/sdk/cliproxy/auth/refresh_queue_test.go @@ -0,0 +1,324 @@ +package auth + +import ( + "context" + "errors" + "sync" + "testing" + "time" +) + +const queueRefreshProvider = "queue-refresh" + +// queueHostStore is a store with a refresh queue: it records the requests the +// manager makes and lets a test play the replica that serves them. +type queueHostStore struct { + mu sync.Mutex + requests []RefreshRequest + failWith error + serving chan func(context.Context, RefreshRequest) error +} + +func newQueueHostStore() *queueHostStore { + return &queueHostStore{serving: make(chan func(context.Context, RefreshRequest) error, 4)} +} + +func (s *queueHostStore) List(context.Context) ([]*Auth, error) { return nil, nil } +func (s *queueHostStore) Save(context.Context, *Auth) (string, error) { return "", nil } +func (s *queueHostStore) Delete(context.Context, string) error { return nil } + +func (s *queueHostStore) RequestRefresh(_ context.Context, req RefreshRequest) error { + s.mu.Lock() + defer s.mu.Unlock() + if s.failWith != nil { + return s.failWith + } + s.requests = append(s.requests, req) + return nil +} + +func (s *queueHostStore) ServeRefreshRequests(ctx context.Context, handle func(context.Context, RefreshRequest) error) error { + s.serving <- handle + <-ctx.Done() + return nil +} + +func (s *queueHostStore) requested() []RefreshRequest { + s.mu.Lock() + defer s.mu.Unlock() + return append([]RefreshRequest(nil), s.requests...) +} + +// reusableGrantExecutor declares its grant stays valid after use. +type reusableGrantExecutor struct{ countingRefreshExecutor } + +func (*reusableGrantExecutor) RefreshPolicy() RefreshPolicy { return RefreshNone } + +func registerQueueAuth(t *testing.T, manager *Manager, id, provider string, now time.Time) { + t.Helper() + auth := &Auth{ + ID: id, + Provider: provider, + Status: StatusActive, + LastRefreshedAt: now.Add(-time.Hour), + Metadata: map[string]any{ + "access_token": "expired-access", + "refresh_token": "refresh-" + id, + "expires_at": now.Add(-time.Hour).Format(time.RFC3339), + "refresh_interval_seconds": 1, + }, + } + if _, err := manager.Register(context.Background(), auth); err != nil { + t.Fatalf("register: %v", err) + } +} + +// dueQueueJob runs one scheduler job for the credential the way a worker does. +func dueQueueJob(t *testing.T, manager *Manager, loop *authAutoRefreshLoop, id string, now time.Time) *authRefreshJob { + t.Helper() + manager.mu.RLock() + epoch := manager.auths[id].RegistrationEpoch + manager.mu.RUnlock() + job := manager.markRefreshPending(loop, id, epoch, now) + if job == nil { + t.Fatal("credential is not due") + } + return job +} + +func newQueueManager(t *testing.T, store Store, exec ProviderExecutor) (*Manager, *authAutoRefreshLoop) { + t.Helper() + manager := NewManager(store, &RoundRobinSelector{}, nil) + manager.RegisterExecutor(exec) + loop := newAuthAutoRefreshLoop(manager, refreshCheckInterval, 1) + manager.refreshLoop = loop + return manager, loop +} + +func TestSchedulerQueuesARefreshInsteadOfRefreshingInline(t *testing.T) { + store := newQueueHostStore() + exec := &countingRefreshExecutor{id: queueRefreshProvider} + manager, loop := newQueueManager(t, store, exec) + now := time.Now() + registerQueueAuth(t, manager, "a", queueRefreshProvider, now) + manager.mu.Lock() + manager.refreshDispatcher = store + manager.mu.Unlock() + + loop.rebuild(now) + loop.handleDue(context.Background(), now) + select { + case job := <-loop.jobs: + if !manager.beginRefreshJob(context.Background(), job) { + t.Fatal("job did not begin") + } + retryAt, settled := manager.dispatchRefresh(context.Background(), job) + if !settled || !retryAt.Equal(job.pendingUntil) { + t.Fatalf("dispatch = %v, %v; want settled until %v", retryAt, settled, job.pendingUntil) + } + default: + t.Fatal("no refresh job was scheduled") + } + reqs := store.requested() + if len(reqs) != 1 || reqs[0].AuthID != "a" || reqs[0].Provider != queueRefreshProvider || reqs[0].CredentialVersion == 0 { + t.Fatalf("requests = %+v", reqs) + } + if got := exec.refreshCalls.Load(); got != 0 { + t.Fatalf("executor refreshed %d times inline, want 0", got) + } +} + +func TestWorkerQueuesWithoutRefreshingAndWaitsBeforeAsking(t *testing.T) { + store := newQueueHostStore() + exec := &countingRefreshExecutor{id: queueRefreshProvider} + manager, loop := newQueueManager(t, store, exec) + now := time.Now() + registerQueueAuth(t, manager, "a", queueRefreshProvider, now) + manager.mu.Lock() + manager.refreshDispatcher = store + manager.mu.Unlock() + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + go loop.worker(ctx) + loop.rebuild(now) + loop.handleDue(ctx, now) + + deadline := time.Now().Add(5 * time.Second) + for len(store.requested()) == 0 && time.Now().Before(deadline) { + time.Sleep(10 * time.Millisecond) + } + if len(store.requested()) != 1 { + t.Fatalf("requests = %+v", store.requested()) + } + for { + manager.mu.RLock() + pending := manager.refreshJobs["a"] + next := manager.auths["a"].NextRefreshAfter + manager.mu.RUnlock() + if pending == nil { + if !next.After(now.Add(refreshPendingBackoff - time.Second)) { + t.Fatalf("next refresh after = %v, want about a pending backoff from now", next) + } + break + } + if time.Now().After(deadline) { + t.Fatal("job never finished") + } + time.Sleep(10 * time.Millisecond) + } + if got := exec.refreshCalls.Load(); got != 0 { + t.Fatalf("executor refreshed %d times inline, want 0", got) + } +} + +func TestFailedQueueingFailsClosedForSingleUseGrants(t *testing.T) { + store := newQueueHostStore() + store.failWith = errors.New("queue unreachable") + exec := &countingRefreshExecutor{id: queueRefreshProvider} + manager, loop := newQueueManager(t, store, exec) + now := time.Now() + registerQueueAuth(t, manager, "a", queueRefreshProvider, now) + manager.mu.Lock() + manager.refreshDispatcher = store + manager.mu.Unlock() + + job := dueQueueJob(t, manager, loop, "a", now) + retryAt, settled := manager.dispatchRefresh(context.Background(), job) + if !settled || !retryAt.IsZero() { + t.Fatalf("dispatch = %v, %v; want settled and due again at once", retryAt, settled) + } +} + +func TestFailedQueueingRefreshesInlineForReusableGrants(t *testing.T) { + store := newQueueHostStore() + store.failWith = errors.New("queue unreachable") + exec := &reusableGrantExecutor{countingRefreshExecutor{id: queueRefreshProvider}} + manager, loop := newQueueManager(t, store, exec) + now := time.Now() + registerQueueAuth(t, manager, "a", queueRefreshProvider, now) + manager.mu.Lock() + manager.refreshDispatcher = store + manager.mu.Unlock() + + job := dueQueueJob(t, manager, loop, "a", now) + if _, settled := manager.dispatchRefresh(context.Background(), job); settled { + t.Fatal("a reusable grant should fall back to refreshing inline") + } +} + +func TestWithoutAQueueRefreshStaysInline(t *testing.T) { + exec := &countingRefreshExecutor{id: queueRefreshProvider} + manager, loop := newQueueManager(t, nil, exec) + now := time.Now() + registerQueueAuth(t, manager, "a", queueRefreshProvider, now) + + job := dueQueueJob(t, manager, loop, "a", now) + if _, settled := manager.dispatchRefresh(context.Background(), job); settled { + t.Fatal("a manager without a queue dispatched a refresh") + } +} + +func TestStartAutoRefreshServesTheStoresQueue(t *testing.T) { + store := newQueueHostStore() + exec := &renewingQueueExecutor{countingRefreshExecutor{id: queueRefreshProvider}} + manager, _ := newQueueManager(t, store, exec) + now := time.Now() + registerQueueAuth(t, manager, "a", queueRefreshProvider, now) + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + manager.StartAutoRefresh(ctx, time.Hour) + defer manager.StopAutoRefresh() + + var handle func(context.Context, RefreshRequest) error + select { + case handle = <-store.serving: + case <-time.After(5 * time.Second): + t.Fatal("the store's queue was not served") + } + if err := handle(ctx, RefreshRequest{AuthID: "a", Provider: queueRefreshProvider}); err != nil { + t.Fatalf("handle = %v", err) + } + if got := exec.refreshCalls.Load(); got != 1 { + t.Fatalf("a served request refreshed %d times, want 1", got) + } + if err := handle(ctx, RefreshRequest{AuthID: "a", Provider: queueRefreshProvider}); err != nil { + t.Fatalf("second handle = %v", err) + } + if got := exec.refreshCalls.Load(); got != 1 { + t.Fatalf("a request for a fresh credential refreshed again: %d calls", got) + } +} + +func TestServedRequestIgnoresTheRequestersPendingMarker(t *testing.T) { + store := newQueueHostStore() + exec := &countingRefreshExecutor{id: queueRefreshProvider} + manager, loop := newQueueManager(t, store, exec) + now := time.Now() + registerQueueAuth(t, manager, "a", queueRefreshProvider, now) + // The scheduler on this replica marked the credential pending when it queued the request. + dueQueueJob(t, manager, loop, "a", now) + + if err := manager.handleRefreshRequest(context.Background(), RefreshRequest{AuthID: "a"}); err != nil { + t.Fatal(err) + } + if got := exec.refreshCalls.Load(); got != 1 { + t.Fatalf("refreshed %d times, want 1", got) + } +} + +func TestServedRequestForAMissingCredentialIsComplete(t *testing.T) { + manager, _ := newQueueManager(t, nil, &countingRefreshExecutor{id: queueRefreshProvider}) + if err := manager.handleRefreshRequest(context.Background(), RefreshRequest{AuthID: "gone"}); err != nil { + t.Fatalf("handle = %v, want nil", err) + } + if err := manager.handleRefreshRequest(context.Background(), RefreshRequest{}); err != nil { + t.Fatalf("handle of an empty request = %v, want nil", err) + } +} + +func TestServedRequestAsksForRedeliveryOnlyWhenCanceled(t *testing.T) { + exec := &blockingQueueExecutor{countingRefreshExecutor: countingRefreshExecutor{id: queueRefreshProvider}} + manager, _ := newQueueManager(t, nil, exec) + registerQueueAuth(t, manager, "a", queueRefreshProvider, time.Now()) + + ctx, cancel := context.WithCancel(context.Background()) + cancel() + if err := manager.handleRefreshRequest(ctx, RefreshRequest{AuthID: "a"}); !errors.Is(err, context.Canceled) { + t.Fatalf("canceled handle = %v, want context.Canceled", err) + } + + exec.fail = errors.New("provider unavailable") + if err := manager.handleRefreshRequest(context.Background(), RefreshRequest{AuthID: "a"}); err != nil { + t.Fatalf("a failed refresh is recorded, not redelivered; got %v", err) + } +} + +// blockingQueueExecutor honors cancellation and can be made to fail. +type blockingQueueExecutor struct { + countingRefreshExecutor + fail error +} + +func (e *blockingQueueExecutor) Refresh(ctx context.Context, auth *Auth) (*Auth, error) { + if err := ctx.Err(); err != nil { + return nil, err + } + if e.fail != nil { + return nil, e.fail + } + return e.countingRefreshExecutor.Refresh(ctx, auth) +} + +// renewingQueueExecutor returns a credential that is fresh for an hour. +type renewingQueueExecutor struct{ countingRefreshExecutor } + +func (e *renewingQueueExecutor) Refresh(ctx context.Context, auth *Auth) (*Auth, error) { + updated, err := e.countingRefreshExecutor.Refresh(ctx, auth) + if err != nil { + return nil, err + } + updated.Metadata["expires_at"] = time.Now().Add(time.Hour).Format(time.RFC3339) + return updated, nil +}