feat(natsstore): serve refresh requests from the refresh queue

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"`

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)
+	})
+}

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)
+		}
+	}
+}

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