From 373449cfa877a8e632ccb0b960ceacc94f89177c Mon Sep 17 00:00:00 2001 From: Yordis Prieto Date: Sat, 7 Nov 2026 21:51:27 -0500 Subject: [PATCH] feat(natsstore): serve refresh requests from the refresh queue The scheduler publishes a request instead of refreshing inline when the store hosts a refresh queue. Wiring the NATS store to its work-queue stream makes exactly one replica win each request and run it through the existing journal-guarded refresh, while the others learn the result from the credential bucket they already follow. Signed-off-by: Yordis Prieto --- internal/store/natsstore/refresh_queue.go | 6 +- internal/store/natsstore/refresh_worker.go | 47 +++++++++++ .../store/natsstore/refresh_worker_test.go | 77 +++++++++++++++++++ internal/store/natsstore/store.go | 2 + 4 files changed, 130 insertions(+), 2 deletions(-) create mode 100644 internal/store/natsstore/refresh_worker.go create mode 100644 internal/store/natsstore/refresh_worker_test.go diff --git a/internal/store/natsstore/refresh_queue.go b/internal/store/natsstore/refresh_queue.go index 21d40c634da..c4c27f7ad8a 100644 --- a/internal/store/natsstore/refresh_queue.go +++ b/internal/store/natsstore/refresh_queue.go @@ -209,8 +209,10 @@ type RefreshRequest struct { CredentialID string `json:"credential_id"` // Provider is informational, for logs and for handlers that route by it. Provider string `json:"provider,omitempty"` - // Version is the credential version the requester saw expiring. A request - // for a version that was already replaced is skipped by the worker. + // Version is the credential version the requester saw expiring. It names the + // request, so asking again for the same version is not queued twice. Replicas + // count versions on their own, so the worker does not compare it with the + // credential and looks at the credential's expiry instead. Version uint64 `json:"version"` // RequestedBy is the replica that published the request. RequestedBy string `json:"requested_by,omitempty"` diff --git a/internal/store/natsstore/refresh_worker.go b/internal/store/natsstore/refresh_worker.go new file mode 100644 index 00000000000..9289134dc2e --- /dev/null +++ b/internal/store/natsstore/refresh_worker.go @@ -0,0 +1,47 @@ +package natsstore + +import ( + "context" + "errors" + + cliproxyauth "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/auth" +) + +// errNoRefreshQueue is returned when the store was built without a refresh queue. +var errNoRefreshQueue = errors.New("nats store: refresh queue is not available") + +var _ cliproxyauth.RefreshQueueHost = (*Store)(nil) + +// RequestRefresh asks the replicas sharing the store to refresh a credential. +// A request that is already queued, or was made for the same credential +// version a moment ago, counts as made. +func (s *Store) RequestRefresh(ctx context.Context, req cliproxyauth.RefreshRequest) error { + if s.refresh == nil { + return errNoRefreshQueue + } + _, err := s.refresh.Publish(ctx, RefreshRequest{ + CredentialID: req.AuthID, + Provider: req.Provider, + Version: req.CredentialVersion, + RequestedBy: s.replicaID, + }) + return err +} + +// ServeRefreshRequests hands each request this replica wins to handle and +// returns when ctx ends. A request is complete once handle returns nil, and is +// delivered again, to this or another replica, when handle returns an error. +// The result of a refresh reaches the other replicas through the credential +// bucket they already follow, so nothing else is published. +func (s *Store) ServeRefreshRequests(ctx context.Context, handle func(context.Context, cliproxyauth.RefreshRequest) error) error { + if s.refresh == nil { + return errNoRefreshQueue + } + return s.refresh.Consume(ctx, func(ctx context.Context, req RefreshRequest) error { + err := handle(ctx, cliproxyauth.RefreshRequest{AuthID: req.CredentialID, Provider: req.Provider, CredentialVersion: req.Version}) + if err == nil || ctx.Err() != nil { + return err + } + return RetryRefresh(err, 0) + }) +} diff --git a/internal/store/natsstore/refresh_worker_test.go b/internal/store/natsstore/refresh_worker_test.go new file mode 100644 index 00000000000..e860c2f8c6a --- /dev/null +++ b/internal/store/natsstore/refresh_worker_test.go @@ -0,0 +1,77 @@ +package natsstore + +import ( + "context" + "errors" + "sync/atomic" + "testing" + "time" + + cliproxyauth "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/auth" +) + +func TestStoreRefreshRequestReachesExactlyOneReplica(t *testing.T) { + f := newStoreFixture(t) + a, b := f.open(), f.open() + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + var handled atomic.Int32 + got := make(chan cliproxyauth.RefreshRequest, 4) + handle := func(_ context.Context, req cliproxyauth.RefreshRequest) error { + handled.Add(1) + got <- req + return nil + } + go func() { _ = a.ServeRefreshRequests(ctx, handle) }() + go func() { _ = b.ServeRefreshRequests(ctx, handle) }() + + want := cliproxyauth.RefreshRequest{AuthID: "example/one.json", Provider: "example", CredentialVersion: 7} + if err := a.RequestRefresh(ctx, want); err != nil { + t.Fatal(err) + } + select { + case req := <-got: + if req != want { + t.Fatalf("request = %+v, want %+v", req, want) + } + case <-time.After(10 * time.Second): + t.Fatal("request was not handled") + } + time.Sleep(time.Second) + if n := handled.Load(); n != 1 { + t.Fatalf("handled %d times, want once", n) + } +} + +func TestStoreRefreshRequestIsRedeliveredAfterAHandlerError(t *testing.T) { + f := newStoreFixture(t) + s := f.open() + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + var calls atomic.Int32 + go func() { + _ = s.ServeRefreshRequests(ctx, func(context.Context, cliproxyauth.RefreshRequest) error { + if calls.Add(1) == 1 { + return errors.New("provider unavailable") + } + return nil + }) + }() + if err := s.RequestRefresh(ctx, cliproxyauth.RefreshRequest{AuthID: "example/two.json", Provider: "example", CredentialVersion: 1}); err != nil { + t.Fatal(err) + } + waitFor(t, "second delivery", func() bool { return calls.Load() >= 2 }) +} + +func TestStoreRefreshRequestRepeatIsAccepted(t *testing.T) { + f := newStoreFixture(t) + s := f.open() + req := cliproxyauth.RefreshRequest{AuthID: "example/three.json", Provider: "example", CredentialVersion: 1} + for range 2 { + if err := s.RequestRefresh(context.Background(), req); err != nil { + t.Fatalf("repeat request: %v", err) + } + } +} diff --git a/internal/store/natsstore/store.go b/internal/store/natsstore/store.go index 70ca79db73e..48e290bf161 100644 --- a/internal/store/natsstore/store.go +++ b/internal/store/natsstore/store.go @@ -30,6 +30,7 @@ type Store struct { auth *AuthStore config *ConfigStore journal *RefreshJournal + refresh *RefreshQueue spoolRoot string configPath string @@ -75,6 +76,7 @@ func NewStore(nc *nats.Conn, boot *Bootstrapped, settings StoreSettings, spoolDi auth: auth, config: cfg, journal: NewRefreshJournal(NewKV(settings.BucketName(BucketRefresh), boot.Bucket(settings, BucketRefresh), origin)), + refresh: boot.Refresh, spoolRoot: root, configPath: filepath.Join(configDir, "config.yaml"), }, nil