From d9504d0bb299b6095dedfc522db6777320c42d5f Mon Sep 17 00:00:00 2001 From: Yordis Prieto Date: Sat, 7 Nov 2026 21:42:46 -0500 Subject: [PATCH] feat(natsstore): queue refresh requests in a work-queue stream Replicas that each refresh an expiring credential race each other and waste provider calls. A work-queue stream with one request per credential and one shared consumer lets exactly one replica win each request, and duplicate requests are refused instead of queued twice. The consumer keeps a running handler alive with progress signals so a slow provider call is not redelivered to a second replica mid-refresh. Signed-off-by: Yordis Prieto --- internal/store/natsstore/bootstrap.go | 12 +- internal/store/natsstore/refresh_queue.go | 357 ++++++++++++++++++ .../store/natsstore/refresh_queue_test.go | 337 +++++++++++++++++ 3 files changed, 703 insertions(+), 3 deletions(-) create mode 100644 internal/store/natsstore/refresh_queue.go create mode 100644 internal/store/natsstore/refresh_queue_test.go diff --git a/internal/store/natsstore/bootstrap.go b/internal/store/natsstore/bootstrap.go index 2ac4cacb40b..5f074d54a04 100644 --- a/internal/store/natsstore/bootstrap.go +++ b/internal/store/natsstore/bootstrap.go @@ -49,6 +49,8 @@ type Bootstrapped struct { Features ServerFeatures // Buckets maps full bucket name to its handle. Buckets map[string]jetstream.KeyValue + // Refresh is the refresh request work queue. + Refresh *RefreshQueue } // Bucket returns the handle for a purpose such as BucketAuth. @@ -56,8 +58,8 @@ func (b *Bootstrapped) Bucket(settings StoreSettings, purpose string) jetstream. return b.Buckets[settings.BucketName(purpose)] } -// Bootstrap probes the server, ensures every bucket and records the schema -// marker. It is idempotent and safe for many replicas to run concurrently. +// Bootstrap probes the server, ensures every bucket and the refresh queue, and +// records the schema marker. It is idempotent and safe for many replicas to run concurrently. func Bootstrap(ctx context.Context, nc *nats.Conn, settings StoreSettings) (*Bootstrapped, error) { if err := settings.Validate(); err != nil { return nil, err @@ -94,7 +96,11 @@ func Bootstrap(ctx context.Context, nc *nats.Conn, settings StoreSettings) (*Boo if err := EnsureSchema(ctx, meta); err != nil { return nil, err } - return &Bootstrapped{JS: js, Features: features, Buckets: buckets}, nil + refresh, err := EnsureRefreshQueue(ctx, js, settings.RefreshQueue(), settings.Create) + if err != nil { + return nil, err + } + return &Bootstrapped{JS: js, Features: features, Buckets: buckets, Refresh: refresh}, nil } // EnsureBucket makes the bucket exist with the expected shape. With CreateAuto diff --git a/internal/store/natsstore/refresh_queue.go b/internal/store/natsstore/refresh_queue.go new file mode 100644 index 00000000000..21d40c634da --- /dev/null +++ b/internal/store/natsstore/refresh_queue.go @@ -0,0 +1,357 @@ +package natsstore + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "slices" + "strconv" + "strings" + "sync" + "time" + + "github.com/nats-io/nats.go" + "github.com/nats-io/nats.go/jetstream" + log "github.com/sirupsen/logrus" +) + +const ( + // refreshRequestMaxAge bounds how long a request waits. A request that + // reached max deliveries is never removed by the server, and with one + // message allowed per credential it would block every later request for + // that credential, so an unhandled request must age out. + refreshRequestMaxAge = 10 * time.Minute + // refreshDuplicateWindow is how long a request id stays remembered, so a + // replica asking again for the same credential version is not queued twice. + refreshDuplicateWindow = 2 * time.Minute + // refreshAckWait is how long a delivered request may go without an ack or + // a progress signal before another replica gets it. Handlers signal + // progress, so this only has to outlast a missed heartbeat, not a provider call. + refreshAckWait = time.Minute + // refreshMaxDeliver caps deliveries of one request. A redelivery after a + // crash is safe, but a request that keeps failing should not circulate. + refreshMaxDeliver = 3 + // refreshMaxInFlight caps unacknowledged requests per replica. + refreshMaxInFlight = 16 + + refreshConsumerName = "refresh-workers" + refreshNakDelay = 5 * time.Second + refreshAckTimeout = 5 * time.Second + + // subjectFullText is what the server answers when a subject already holds + // the one message the queue allows it. + subjectFullText = "maximum messages per subject exceeded" +) + +// RefreshStreamName is the stream holding refresh requests. +func (s StoreSettings) RefreshStreamName() string { return strings.ToUpper(s.Prefix) + "_REFRESH" } + +// RefreshQueue describes the refresh work queue for these settings. +func (s StoreSettings) RefreshQueue() RefreshQueueSpec { + return RefreshQueueSpec{ + Name: s.RefreshStreamName(), + SubjectPrefix: s.Prefix + ".refresh", + Replicas: s.Replicas, + MaxAge: refreshRequestMaxAge, + DuplicateWindow: refreshDuplicateWindow, + AckWait: refreshAckWait, + MaxDeliver: refreshMaxDeliver, + MaxInFlight: refreshMaxInFlight, + ConsumerName: refreshConsumerName, + Description: "CLIProxyAPI credential refresh requests", + } +} + +// RefreshQueueSpec is the shape of the refresh work queue: one stream holding +// at most one request per credential and one shared consumer, so exactly one +// replica receives each request. +type RefreshQueueSpec struct { + // Name is the stream name. + Name string + // SubjectPrefix is followed by the encoded credential id to form a subject. + SubjectPrefix string + Description string + Replicas int + MaxAge time.Duration + DuplicateWindow time.Duration + + ConsumerName string + AckWait time.Duration + MaxDeliver int + MaxInFlight int +} + +// Subject returns the subject carrying requests for the credential. +func (q RefreshQueueSpec) Subject(credentialID string) (string, error) { + token, err := EncodeKey(credentialID) + if err != nil { + return "", err + } + return q.SubjectPrefix + "." + token, nil +} + +func (q RefreshQueueSpec) streamConfig() jetstream.StreamConfig { + return jetstream.StreamConfig{ + Name: q.Name, + Description: q.Description, + Subjects: []string{q.SubjectPrefix + ".>"}, + // A request leaves the stream once it is acknowledged. + Retention: jetstream.WorkQueuePolicy, + Storage: jetstream.FileStorage, + Replicas: q.Replicas, + // One queued request per credential. With discard-new a second request + // is refused instead of replacing the first, which may be in flight. + MaxMsgsPerSubject: 1, + Discard: jetstream.DiscardNew, + DiscardNewPerSubject: true, + MaxAge: q.MaxAge, + Duplicates: q.DuplicateWindow, + } +} + +func (q RefreshQueueSpec) consumerConfig() jetstream.ConsumerConfig { + return jetstream.ConsumerConfig{ + Durable: q.ConsumerName, + Description: "CLIProxyAPI refresh workers", + AckPolicy: jetstream.AckExplicitPolicy, + AckWait: q.AckWait, + MaxDeliver: q.MaxDeliver, + MaxAckPending: q.MaxInFlight, + DeliverPolicy: jetstream.DeliverAllPolicy, + FilterSubject: q.SubjectPrefix + ".>", + } +} + +// EnsureRefreshQueue makes the stream and its shared consumer exist with the +// expected shape. With CreateAuto it creates or reconciles; with CreateNever it +// only looks. Either way the final shape is verified. A stream's retention +// cannot change once created, so a drifted one is recreated by an operator. +func EnsureRefreshQueue(ctx context.Context, js jetstream.JetStream, spec RefreshQueueSpec, mode CreateMode) (*RefreshQueue, error) { + var ( + stream jetstream.Stream + consumer jetstream.Consumer + err error + ) + switch mode { + case CreateAuto: + if stream, err = js.CreateOrUpdateStream(ctx, spec.streamConfig()); err != nil { + return nil, fmt.Errorf("failed to ensure stream %s: %w", spec.Name, err) + } + if consumer, err = js.CreateOrUpdateConsumer(ctx, spec.Name, spec.consumerConfig()); err != nil { + return nil, fmt.Errorf("failed to ensure consumer %s: %w", spec.ConsumerName, err) + } + case CreateNever: + if stream, err = js.Stream(ctx, spec.Name); err != nil { + if errors.Is(err, jetstream.ErrStreamNotFound) { + return nil, fmt.Errorf("%w: %s; provision it or set %s=%s", ErrBucketMissing, spec.Name, EnvStoreCreate, CreateAuto) + } + return nil, fmt.Errorf("failed to read stream %s: %w", spec.Name, err) + } + if consumer, err = stream.Consumer(ctx, spec.ConsumerName); err != nil { + return nil, fmt.Errorf("failed to read consumer %s: %w; provision it or set %s=%s", spec.ConsumerName, err, EnvStoreCreate, CreateAuto) + } + default: + return nil, fmt.Errorf("unknown create mode %q", mode) + } + if err := verifyRefreshQueue(ctx, stream, consumer, spec); err != nil { + return nil, err + } + return &RefreshQueue{js: js, spec: spec, consumer: consumer}, nil +} + +func verifyRefreshQueue(ctx context.Context, stream jetstream.Stream, consumer jetstream.Consumer, spec RefreshQueueSpec) error { + info, err := stream.Info(ctx) + if err != nil { + return fmt.Errorf("failed to read stream %s: %w", spec.Name, err) + } + want := spec.streamConfig() + have := info.Config + + var drift []string + check := func(name string, have, want any) { + if have != want { + drift = append(drift, fmt.Sprintf("%s: have %v, want %v", name, have, want)) + } + } + check("retention", have.Retention.String(), want.Retention.String()) + check("storage", have.Storage.String(), want.Storage.String()) + check("replicas", have.Replicas, want.Replicas) + check("discard", have.Discard.String(), want.Discard.String()) + check("discard new per subject", have.DiscardNewPerSubject, want.DiscardNewPerSubject) + check("max messages per subject", have.MaxMsgsPerSubject, want.MaxMsgsPerSubject) + check("max age", have.MaxAge, want.MaxAge) + check("duplicate window", have.Duplicates, want.Duplicates) + if !slices.Equal(have.Subjects, want.Subjects) { + drift = append(drift, fmt.Sprintf("subjects: have %v, want %v", have.Subjects, want.Subjects)) + } + + consumerInfo, err := consumer.Info(ctx) + if err != nil { + return fmt.Errorf("failed to read consumer %s: %w", spec.ConsumerName, err) + } + wantConsumer := spec.consumerConfig() + haveConsumer := consumerInfo.Config + check("consumer ack policy", haveConsumer.AckPolicy.String(), wantConsumer.AckPolicy.String()) + check("consumer ack wait", haveConsumer.AckWait, wantConsumer.AckWait) + check("consumer max deliver", haveConsumer.MaxDeliver, wantConsumer.MaxDeliver) + check("consumer max ack pending", haveConsumer.MaxAckPending, wantConsumer.MaxAckPending) + check("consumer filter", haveConsumer.FilterSubject, wantConsumer.FilterSubject) + if len(drift) > 0 { + return &DriftError{Bucket: spec.Name, Fields: drift} + } + return nil +} + +// RefreshRequest asks for one credential to be refreshed. It carries no secret: +// the worker reads the credential itself. +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 uint64 `json:"version"` + // RequestedBy is the replica that published the request. + RequestedBy string `json:"requested_by,omitempty"` +} + +// RefreshQueue publishes refresh requests and delivers each to exactly one worker. +type RefreshQueue struct { + js jetstream.JetStream + spec RefreshQueueSpec + consumer jetstream.Consumer +} + +// Spec returns the queue shape. +func (q *RefreshQueue) Spec() RefreshQueueSpec { return q.spec } + +// Publish queues a request. It reports false when an equivalent request is +// already known, either queued for the credential or published for the same +// credential version inside the duplicate window; nothing new is queued then, +// and the caller has nothing to retry. +func (q *RefreshQueue) Publish(ctx context.Context, req RefreshRequest) (bool, error) { + subject, err := q.spec.Subject(req.CredentialID) + if err != nil { + return false, err + } + body, err := json.Marshal(req) + if err != nil { + return false, fmt.Errorf("failed to encode refresh request: %w", err) + } + token, _ := EncodeKey(req.CredentialID) + msgID := token + ":" + strconv.FormatUint(req.Version, 10) + ack, err := q.js.PublishMsg(ctx, &nats.Msg{Subject: subject, Data: body}, jetstream.WithMsgID(msgID)) + if err != nil { + if isSubjectFull(err) { + return false, nil + } + return false, fmt.Errorf("failed to queue refresh of %s: %w", req.CredentialID, err) + } + return !ack.Duplicate, nil +} + +// isSubjectFull reports the refusal of a request for a credential that already has one queued. +func isSubjectFull(err error) bool { + var apiErr *jetstream.APIError + return errors.As(err, &apiErr) && strings.Contains(apiErr.Description, subjectFullText) +} + +// RefreshHandler performs one refresh request. Returning nil acknowledges it, +// RetryRefresh asks for a later redelivery, and any other error ends the +// request without redelivery, since the request was understood and failed. +type RefreshHandler func(ctx context.Context, req RefreshRequest) error + +type retryError struct { + err error + delay time.Duration +} + +func (e *retryError) Error() string { return e.err.Error() } +func (e *retryError) Unwrap() error { return e.err } + +// RetryRefresh marks a handler failure as worth another delivery after delay, +// or after a default delay when delay is not positive. +func RetryRefresh(err error, delay time.Duration) error { + if err == nil { + return nil + } + return &retryError{err: err, delay: delay} +} + +// Consume delivers requests to handle until ctx ends, then waits for running +// handlers. Every replica that consumes shares one consumer, so each request +// goes to one of them. A handler is kept alive with progress signals while it +// runs, so its request is not redelivered to a second replica mid-refresh. +func (q *RefreshQueue) Consume(ctx context.Context, handle RefreshHandler) error { + var running sync.WaitGroup + consumeCtx, err := q.consumer.Consume(func(msg jetstream.Msg) { + running.Add(1) + go func() { + defer running.Done() + q.handle(ctx, msg, handle) + }() + }, jetstream.PullMaxMessages(q.spec.MaxInFlight), jetstream.ConsumeErrHandler(func(_ jetstream.ConsumeContext, err error) { + log.WithError(err).Warn("nats store: refresh consumer error") + })) + if err != nil { + return fmt.Errorf("failed to consume refresh requests: %w", err) + } + <-ctx.Done() + consumeCtx.Stop() + running.Wait() + return nil +} + +func (q *RefreshQueue) handle(ctx context.Context, msg jetstream.Msg, handle RefreshHandler) { + var req RefreshRequest + if err := json.Unmarshal(msg.Data(), &req); err != nil || req.CredentialID == "" { + log.Warn("nats store: dropping a malformed refresh request") + _ = msg.Term() + return + } + + stopBeat := make(chan struct{}) + beatDone := make(chan struct{}) + go func() { + defer close(beatDone) + ticker := time.NewTicker(q.spec.AckWait / 3) + defer ticker.Stop() + for { + select { + case <-stopBeat: + return + case <-ticker.C: + _ = msg.InProgress() + } + } + }() + err := handle(ctx, req) + close(stopBeat) + <-beatDone + + // The acknowledgement must not depend on the consumer context, which + // has ended when the replica is shutting down. + ackCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), refreshAckTimeout) + defer cancel() + var retry *retryError + switch { + case err == nil: + if err := msg.DoubleAck(ackCtx); err != nil && !errors.Is(err, jetstream.ErrMsgAlreadyAckd) { + log.WithError(err).Warn("nats store: refresh acknowledgement was not confirmed") + } + case errors.As(err, &retry): + delay := retry.delay + if delay <= 0 { + delay = refreshNakDelay + } + _ = msg.NakWithDelay(delay) + case ctx.Err() != nil: + // Shutting down mid-request: hand it to another replica now instead of + // after the ack wait. It meets the journal's claim if this one sent a grant. + _ = msg.Nak() + default: + log.WithError(err).Warnf("nats store: refresh request for %s failed", req.CredentialID) + _ = msg.Term() + } +} diff --git a/internal/store/natsstore/refresh_queue_test.go b/internal/store/natsstore/refresh_queue_test.go new file mode 100644 index 00000000000..7845871b1f5 --- /dev/null +++ b/internal/store/natsstore/refresh_queue_test.go @@ -0,0 +1,337 @@ +package natsstore + +import ( + "context" + "errors" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/nats-io/nats.go" + "github.com/nats-io/nats.go/jetstream" + + "github.com/router-for-me/CLIProxyAPI/v8/internal/store/natsstore/natstest" +) + +// openRefreshQueues bootstraps one connection per replica against one server. +func openRefreshQueues(t *testing.T, replicas int, tune func(*StoreSettings, *RefreshQueueSpec)) (*natstest.Server, StoreSettings, []*RefreshQueue) { + t.Helper() + srv := natstest.Start(t) + settings := autoSettings("queue") + spec := settings.RefreshQueue() + if tune != nil { + tune(&settings, &spec) + } + var queues []*RefreshQueue + for range replicas { + nc, err := connect(t, map[string]string{EnvURL: srv.URL}, 5*time.Second) + if err != nil { + t.Fatal(err) + } + js, err := jetstream.New(nc) + if err != nil { + t.Fatal(err) + } + queue, err := EnsureRefreshQueue(testContext(t), js, spec, CreateAuto) + if err != nil { + t.Fatal(err) + } + queues = append(queues, queue) + } + return srv, settings, queues +} + +// consume runs the queue's consumer until the test ends. +func consume(t *testing.T, queue *RefreshQueue, handle RefreshHandler) { + t.Helper() + ctx, cancel := context.WithCancel(context.Background()) + done := make(chan struct{}) + go func() { + defer close(done) + if err := queue.Consume(ctx, handle); err != nil { + t.Error(err) + } + }() + t.Cleanup(func() { + cancel() + <-done + }) +} + +// consumeWithoutAnswering receives requests and never acknowledges them until +// ctx ends, the way a process that died mid-refresh leaves its request. +func consumeWithoutAnswering(ctx context.Context, queue *RefreshQueue, received func()) { + cc, err := queue.consumer.Consume(func(jetstream.Msg) { received() }) + if err != nil { + return + } + <-ctx.Done() + cc.Stop() +} + +func TestRefreshQueueDeliversOneRequestToOneConsumer(t *testing.T) { + _, _, queues := openRefreshQueues(t, 2, nil) + var handled atomic.Int32 + seen := make(chan RefreshRequest, 4) + handle := func(_ context.Context, req RefreshRequest) error { + handled.Add(1) + seen <- req + return nil + } + consume(t, queues[0], handle) + consume(t, queues[1], handle) + + queued, err := queues[0].Publish(testContext(t), RefreshRequest{CredentialID: "team/a@example.com.json", Provider: "codex", Version: 3, RequestedBy: "r1"}) + if err != nil || !queued { + t.Fatalf("publish = %v, %v", queued, err) + } + select { + case req := <-seen: + if req.CredentialID != "team/a@example.com.json" || req.Provider != "codex" || req.Version != 3 { + t.Fatalf("request = %+v", req) + } + case <-time.After(20 * time.Second): + t.Fatal("no consumer received the request") + } + // Give a wrongly shared delivery time to show up. + time.Sleep(time.Second) + if n := handled.Load(); n != 1 { + t.Fatalf("handlers ran %d times, want exactly 1", n) + } +} + +func TestRefreshQueueRefusesASecondRequestWhileOneIsQueued(t *testing.T) { + _, _, queues := openRefreshQueues(t, 1, nil) + queue, ctx := queues[0], testContext(t) + + if queued, err := queue.Publish(ctx, RefreshRequest{CredentialID: "a", Version: 1}); err != nil || !queued { + t.Fatalf("first publish = %v, %v", queued, err) + } + // Same version: dedupe by message id. + if queued, err := queue.Publish(ctx, RefreshRequest{CredentialID: "a", Version: 1}); err != nil || queued { + t.Fatalf("same version publish = %v, %v; want not queued", queued, err) + } + // New version while the first waits: refused by the per-credential limit. + if queued, err := queue.Publish(ctx, RefreshRequest{CredentialID: "a", Version: 2}); err != nil || queued { + t.Fatalf("second version publish = %v, %v; want not queued", queued, err) + } + // Another credential is independent. + if queued, err := queue.Publish(ctx, RefreshRequest{CredentialID: "b", Version: 1}); err != nil || !queued { + t.Fatalf("other credential publish = %v, %v", queued, err) + } + info, err := queue.consumer.Info(ctx) + if err != nil { + t.Fatal(err) + } + if got := info.NumPending; got != 2 { + t.Fatalf("pending = %d, want 2", got) + } +} + +func TestRefreshQueueAcceptsANewVersionAfterTheRequestIsHandled(t *testing.T) { + _, _, queues := openRefreshQueues(t, 1, nil) + queue, ctx := queues[0], testContext(t) + var handled atomic.Int32 + consume(t, queue, func(context.Context, RefreshRequest) error { + handled.Add(1) + return nil + }) + + if _, err := queue.Publish(ctx, RefreshRequest{CredentialID: "a", Version: 1}); err != nil { + t.Fatal(err) + } + waitFor(t, "the first request", func() bool { return handled.Load() == 1 }) + var queued bool + waitFor(t, "the slot to free", func() bool { + var err error + queued, err = queue.Publish(ctx, RefreshRequest{CredentialID: "a", Version: 2}) + return err == nil && queued + }) + waitFor(t, "the second request", func() bool { return handled.Load() == 2 }) +} + +func TestRefreshQueueRedeliversAnUnacknowledgedRequest(t *testing.T) { + _, _, queues := openRefreshQueues(t, 2, func(_ *StoreSettings, spec *RefreshQueueSpec) { + spec.AckWait = 3 * time.Second + }) + ctx := testContext(t) + crashed := make(chan struct{}) + var once sync.Once + var winner atomic.Int32 + + // The first replica takes the request and never answers, like a crashed process. + stuck, cancelStuck := context.WithCancel(context.Background()) + go func() { + consumeWithoutAnswering(stuck, queues[0], func() { once.Do(func() { close(crashed) }) }) + }() + t.Cleanup(cancelStuck) + + if _, err := queues[0].Publish(ctx, RefreshRequest{CredentialID: "a", Version: 1}); err != nil { + t.Fatal(err) + } + select { + case <-crashed: + case <-time.After(20 * time.Second): + t.Fatal("the first replica never received the request") + } + cancelStuck() + + consume(t, queues[1], func(context.Context, RefreshRequest) error { + winner.Add(1) + return nil + }) + waitFor(t, "redelivery to the second replica", func() bool { return winner.Load() == 1 }) +} + +func TestRefreshQueueStopsAtMaxDeliver(t *testing.T) { + _, _, queues := openRefreshQueues(t, 1, func(_ *StoreSettings, spec *RefreshQueueSpec) { + spec.MaxDeliver = 2 + }) + queue, ctx := queues[0], testContext(t) + var attempts atomic.Int32 + consume(t, queue, func(context.Context, RefreshRequest) error { + attempts.Add(1) + return RetryRefresh(errors.New("provider unavailable"), 100*time.Millisecond) + }) + if _, err := queue.Publish(ctx, RefreshRequest{CredentialID: "a", Version: 1}); err != nil { + t.Fatal(err) + } + waitFor(t, "both deliveries", func() bool { return attempts.Load() == 2 }) + time.Sleep(time.Second) + if n := attempts.Load(); n != 2 { + t.Fatalf("delivered %d times, want max deliver 2", n) + } +} + +func TestRefreshQueueTerminalFailureIsNotRedelivered(t *testing.T) { + _, _, queues := openRefreshQueues(t, 1, nil) + queue, ctx := queues[0], testContext(t) + var attempts atomic.Int32 + consume(t, queue, func(context.Context, RefreshRequest) error { + attempts.Add(1) + return errors.New("refresh rejected") + }) + if _, err := queue.Publish(ctx, RefreshRequest{CredentialID: "a", Version: 1}); err != nil { + t.Fatal(err) + } + waitFor(t, "the delivery", func() bool { return attempts.Load() == 1 }) + time.Sleep(time.Second) + if n := attempts.Load(); n != 1 { + t.Fatalf("delivered %d times, want 1", n) + } + // The request left the queue, so the credential can be asked for again. + waitFor(t, "the slot to free", func() bool { + queued, err := queue.Publish(ctx, RefreshRequest{CredentialID: "a", Version: 2}) + return err == nil && queued + }) +} + +func TestRefreshQueueKeepsASlowHandlerAlive(t *testing.T) { + _, _, queues := openRefreshQueues(t, 2, func(_ *StoreSettings, spec *RefreshQueueSpec) { + spec.AckWait = 2 * time.Second + }) + var handled atomic.Int32 + slow := func(context.Context, RefreshRequest) error { + handled.Add(1) + time.Sleep(5 * time.Second) + return nil + } + consume(t, queues[0], slow) + consume(t, queues[1], slow) + if _, err := queues[0].Publish(testContext(t), RefreshRequest{CredentialID: "a", Version: 1}); err != nil { + t.Fatal(err) + } + waitFor(t, "the delivery", func() bool { return handled.Load() >= 1 }) + time.Sleep(7 * time.Second) + if n := handled.Load(); n != 1 { + t.Fatalf("a handler outlasting the ack wait ran %d times, want 1", n) + } +} + +func TestRefreshQueueMaxAgeFreesAStuckRequest(t *testing.T) { + _, _, queues := openRefreshQueues(t, 1, func(_ *StoreSettings, spec *RefreshQueueSpec) { + spec.MaxAge, spec.DuplicateWindow = 2*time.Second, 2*time.Second + }) + queue, ctx := queues[0], testContext(t) + if _, err := queue.Publish(ctx, RefreshRequest{CredentialID: "a", Version: 1}); err != nil { + t.Fatal(err) + } + waitFor(t, "the request to age out", func() bool { + queued, err := queue.Publish(ctx, RefreshRequest{CredentialID: "a", Version: 2}) + return err == nil && queued + }) +} + +func TestRefreshQueueDropsMalformedRequests(t *testing.T) { + _, _, queues := openRefreshQueues(t, 1, nil) + queue, ctx := queues[0], testContext(t) + var handled atomic.Int32 + consume(t, queue, func(context.Context, RefreshRequest) error { + handled.Add(1) + return nil + }) + subject, err := queue.spec.Subject("a") + if err != nil { + t.Fatal(err) + } + if _, err := queue.js.PublishMsg(ctx, &nats.Msg{Subject: subject, Data: []byte("not json")}); err != nil { + t.Fatal(err) + } + waitFor(t, "the malformed request to leave", func() bool { + info, err := queue.consumer.Info(ctx) + return err == nil && info.NumPending == 0 && info.NumAckPending == 0 + }) + if handled.Load() != 0 { + t.Fatal("handler ran for a malformed request") + } +} + +func TestEnsureRefreshQueueDriftAndNeverMode(t *testing.T) { + srv := natstest.Start(t) + nc, err := connect(t, map[string]string{EnvURL: srv.URL}, 5*time.Second) + if err != nil { + t.Fatal(err) + } + js, err := jetstream.New(nc) + if err != nil { + t.Fatal(err) + } + ctx := testContext(t) + spec := autoSettings("drift").RefreshQueue() + + if _, err := EnsureRefreshQueue(ctx, js, spec, CreateNever); !errors.Is(err, ErrBucketMissing) { + t.Fatalf("never mode on a missing stream = %v", err) + } + if _, err := EnsureRefreshQueue(ctx, js, spec, CreateAuto); err != nil { + t.Fatal(err) + } + if _, err := EnsureRefreshQueue(ctx, js, spec, CreateNever); err != nil { + t.Fatalf("never mode on a provisioned queue: %v", err) + } + + tighter := spec + tighter.MaxDeliver = 9 + var drift *DriftError + if _, err := EnsureRefreshQueue(ctx, js, tighter, CreateNever); !errors.As(err, &drift) { + t.Fatalf("drifted consumer = %v, want DriftError", err) + } + if _, err := EnsureRefreshQueue(ctx, js, tighter, CreateAuto); err != nil { + t.Fatalf("auto mode reconciles a consumer: %v", err) + } +} + +func TestRefreshQueueIsBootstrapped(t *testing.T) { + srv := natstest.Start(t) + nc, err := connect(t, map[string]string{EnvURL: srv.URL}, 5*time.Second) + if err != nil { + t.Fatal(err) + } + boot, err := Bootstrap(testContext(t), nc, autoSettings("boot")) + if err != nil { + t.Fatal(err) + } + if boot.Refresh == nil || boot.Refresh.Spec().Name != "BOOT_REFRESH" { + t.Fatalf("refresh queue = %+v", boot.Refresh) + } +}