diff --git a/jetstream/api.go b/jetstream/api.go index 51b0b6063..0fc5d91ae 100644 --- a/jetstream/api.go +++ b/jetstream/api.go @@ -1,4 +1,4 @@ -// Copyright 2022-2023 The NATS Authors +// Copyright 2022-2025 The NATS Authors // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. // You may obtain a copy of the License at @@ -104,6 +104,9 @@ const ( // apiMsgDeleteT is the endpoint to remove a message. apiMsgDeleteT = "STREAM.MSG.DELETE.%s" + + // apiConsumerUnpinT is the endpoint to unpin a consumer. + apiConsumerUnpinT = "CONSUMER.UNPIN.%s.%s" ) func (js *jetStream) apiRequestJSON(ctx context.Context, subject string, resp any, data ...[]byte) (*jetStreamMsg, error) { diff --git a/jetstream/consumer.go b/jetstream/consumer.go index 82e9ade37..3f390627d 100644 --- a/jetstream/consumer.go +++ b/jetstream/consumer.go @@ -1,4 +1,4 @@ -// Copyright 2022-2024 The NATS Authors +// Copyright 2022-2025 The NATS Authors // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. // You may obtain a copy of the License at @@ -374,3 +374,37 @@ func validateConsumerName(dur string) error { } return nil } + +func unpinConsumer(ctx context.Context, js *jetStream, stream, consumer, group string) error { + ctx, cancel := js.wrapContextWithoutDeadline(ctx) + if cancel != nil { + defer cancel() + } + if err := validateConsumerName(consumer); err != nil { + return err + } + unpinSubject := fmt.Sprintf(apiConsumerUnpinT, stream, consumer) + + var req = consumerUnpinRequest{ + Group: group, + } + + reqJSON, err := json.Marshal(req) + if err != nil { + return err + } + + var resp apiResponse + + if _, err := js.apiRequestJSON(ctx, unpinSubject, &resp, reqJSON); err != nil { + return err + } + if resp.Error != nil { + if resp.Error.ErrorCode == JSErrCodeConsumerNotFound { + return ErrConsumerNotFound + } + return resp.Error + } + + return nil +} diff --git a/jetstream/consumer_config.go b/jetstream/consumer_config.go index 220d2aae7..e93f43ca3 100644 --- a/jetstream/consumer_config.go +++ b/jetstream/consumer_config.go @@ -76,6 +76,9 @@ type ( // TimeStamp indicates when the info was gathered by the server. TimeStamp time.Time `json:"ts"` + // PriorityGroups contains the information about the currently defined priority groups + PriorityGroups []PriorityGroupState `json:"priority_groups,omitempty"` + // Paused indicates whether the consumer is paused. Paused bool `json:"paused,omitempty"` @@ -84,6 +87,17 @@ type ( PauseRemaining time.Duration `json:"pause_remaining,omitempty"` } + PriorityGroupState struct { + // Group this status is for. + Group string `json:"group"` + + // PinnedClientID is the generated ID of the pinned client. + PinnedClientID string `json:"pinned_client_id,omitempty"` + + // PinnedTS is the timestamp when the client was pinned. + PinnedTS time.Time `json:"pinned_ts,omitempty"` + } + // ConsumerConfig is the configuration of a JetStream consumer. ConsumerConfig struct { // Name is an optional name for the consumer. If not set, one is @@ -227,6 +241,18 @@ type ( // PauseUntil is for suspending the consumer until the deadline. PauseUntil *time.Time `json:"pause_until,omitempty"` + + // PriorityPolicy represents he priority policy the consumer is set to. + // Requires nats-server v2.11.0 or later. + PriorityPolicy PriorityPolicy `json:"priority_policy,omitempty"` + + // PinnedTTL represents the time after which the client will be unpinned + // if no new pull requests are sent.Used with PriorityPolicyPinned. + // Requires nats-server v2.11.0 or later. + PinnedTTL time.Duration `json:"priority_timeout,omitempty"` + + // PriorityGroups is a list of priority groups this consumer supports. + PriorityGroups []string `json:"priority_groups,omitempty"` } // OrderedConsumerConfig is the configuration of an ordered JetStream @@ -298,8 +324,51 @@ type ( Stream uint64 `json:"stream_seq"` Last *time.Time `json:"last_active,omitempty"` } + + // PriorityPolicy determines the priority policy the consumer is set to. + PriorityPolicy int ) +const ( + // PriorityPolicyNone is the default priority policy. + PriorityPolicyNone PriorityPolicy = iota + + // PriorityPolicyPinned is the priority policy that pins a consumer to a + // specific client. + PriorityPolicyPinned + + // PriorityPolicyOverflow is the priority policy that allows for + // restricting when a consumer will receive messages based on the number of + // pending messages or acks. + PriorityPolicyOverflow +) + +func (p *PriorityPolicy) UnmarshalJSON(data []byte) error { + switch string(data) { + case jsonString(""): + *p = PriorityPolicyNone + case jsonString("pinned_client"): + *p = PriorityPolicyPinned + case jsonString("overflow"): + *p = PriorityPolicyOverflow + default: + return fmt.Errorf("nats: can not unmarshal %q", data) + } + return nil +} + +func (p PriorityPolicy) MarshalJSON() ([]byte, error) { + switch p { + case PriorityPolicyNone: + return json.Marshal("") + case PriorityPolicyPinned: + return json.Marshal("pinned_client") + case PriorityPolicyOverflow: + return json.Marshal("overflow") + } + return nil, fmt.Errorf("nats: unknown priority policy %v", p) +} + const ( // DeliverAllPolicy starts delivering messages from the very beginning of a // stream. This is the default. diff --git a/jetstream/errors.go b/jetstream/errors.go index 8d0f4d532..5e50f488e 100644 --- a/jetstream/errors.go +++ b/jetstream/errors.go @@ -195,6 +195,10 @@ var ( // consumer. ErrNoMessages JetStreamError = &jsError{message: "no messages"} + // ErrPinIDMismatch is returned when Pin ID sent in the request does not match + // the currently pinned consumer subscriber ID on the server. + ErrPinIDMismatch JetStreamError = &jsError{message: "pin ID mismatch"} + // ErrMaxBytesExceeded is returned when a message would exceed MaxBytes set // on a pull request. ErrMaxBytesExceeded JetStreamError = &jsError{message: "message size exceeds max bytes"} diff --git a/jetstream/jetstream_options.go b/jetstream/jetstream_options.go index 40cb2a4e9..1baa0e42c 100644 --- a/jetstream/jetstream_options.go +++ b/jetstream/jetstream_options.go @@ -297,6 +297,73 @@ func (t PullThresholdBytes) configureMessages(opts *consumeOpts) error { return nil } +// PullMinPending sets the minimum number of messages that should be pending for +// a consumer with PriorityPolicyOverflow to be considered for delivery. +// If provided, PullPriorityGroup must be set as well and the consumer has to have +// PriorityPolicy set to PriorityPolicyOverflow. +// +// PullMinPending implements both PullConsumeOpt and PullMessagesOpt, allowing +// it to configure Consumer.Consume and Consumer.Messages. +type PullMinPending int + +func (min PullMinPending) configureConsume(opts *consumeOpts) error { + if min < 1 { + return fmt.Errorf("%w: min pending should be more than 0", ErrInvalidOption) + } + opts.MinPending = int64(min) + return nil +} + +func (min PullMinPending) configureMessages(opts *consumeOpts) error { + if min < 1 { + return fmt.Errorf("%w: min pending should be more than 0", ErrInvalidOption) + } + opts.MinPending = int64(min) + return nil +} + +// PullMinAckPending sets the minimum number of pending acks that should be +// present for a consumer with PriorityPolicyOverflow to be considered for +// delivery. If provided, PullPriorityGroup must be set as well and the consumer +// has to have PriorityPolicy set to PriorityPolicyOverflow. +// +// PullMinAckPending implements both PullConsumeOpt and PullMessagesOpt, allowing +// it to configure Consumer.Consume and Consumer.Messages. +type PullMinAckPending int + +func (min PullMinAckPending) configureConsume(opts *consumeOpts) error { + if min < 1 { + return fmt.Errorf("%w: min pending should be more than 0", ErrInvalidOption) + } + opts.MinAckPending = int64(min) + return nil +} + +func (min PullMinAckPending) configureMessages(opts *consumeOpts) error { + if min < 1 { + return fmt.Errorf("%w: min pending should be more than 0", ErrInvalidOption) + } + opts.MinAckPending = int64(min) + return nil +} + +// PullPriorityGroup sets the priority group for a consumer. +// It has to match one of the priority groups set on the consumer. +// +// PullPriorityGroup implements both PullConsumeOpt and PullMessagesOpt, allowing +// it to configure Consumer.Consume and Consumer.Messages. +type PullPriorityGroup string + +func (g PullPriorityGroup) configureConsume(opts *consumeOpts) error { + opts.Group = string(g) + return nil +} + +func (g PullPriorityGroup) configureMessages(opts *consumeOpts) error { + opts.Group = string(g) + return nil +} + // PullHeartbeat sets the idle heartbeat duration for a pull subscription // If a client does not receive a heartbeat message from a stream for more // than the idle heartbeat setting, the subscription will be removed @@ -368,6 +435,43 @@ func WithMessagesErrOnMissingHeartbeat(hbErr bool) PullMessagesOpt { }) } +// FetchMinPending sets the minimum number of messages that should be pending for +// a consumer with PriorityPolicyOverflow to be considered for delivery. +// If provided, FetchPriorityGroup must be set as well and the consumer has to have +// PriorityPolicy set to PriorityPolicyOverflow. +func FetchMinPending(min int64) FetchOpt { + return func(req *pullRequest) error { + if min < 1 { + return fmt.Errorf("%w: min pending should be more than 0", ErrInvalidOption) + } + req.MinPending = min + return nil + } +} + +// FetchMinAckPending sets the minimum number of pending acks that should be +// present for a consumer with PriorityPolicyOverflow to be considered for +// delivery. If provided, FetchPriorityGroup must be set as well and the consumer +// has to have PriorityPolicy set to PriorityPolicyOverflow. +func FetchMinAckPending(min int64) FetchOpt { + return func(req *pullRequest) error { + if min < 1 { + return fmt.Errorf("%w: min ack pending should be more than 0", ErrInvalidOption) + } + req.MinAckPending = min + return nil + } +} + +// FetchPriorityGroup sets the priority group for a consumer. +// It has to match one of the priority groups set on the consumer. +func FetchPriorityGroup(group string) FetchOpt { + return func(req *pullRequest) error { + req.Group = group + return nil + } +} + // FetchMaxWait sets custom timeout for fetching predefined batch of messages. // // If not provided, a default of 30 seconds will be used. diff --git a/jetstream/message.go b/jetstream/message.go index c34ffde0e..9103c7457 100644 --- a/jetstream/message.go +++ b/jetstream/message.go @@ -1,4 +1,4 @@ -// Copyright 2022-2024 The NATS Authors +// Copyright 2022-2025 The NATS Authors // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. // You may obtain a copy of the License at @@ -147,6 +147,7 @@ const ( reqTimeout = "408" maxBytesExceeded = "409" noResponders = "503" + pinIdMismatch = "423" ) // Headers used when publishing messages. @@ -418,6 +419,8 @@ func checkMsg(msg *nats.Msg) (bool, error) { return false, nats.ErrTimeout case controlMsg: return false, nil + case pinIdMismatch: + return false, ErrPinIDMismatch case maxBytesExceeded: if strings.Contains(strings.ToLower(descr), "message size exceeds maxbytes") { return false, ErrMaxBytesExceeded diff --git a/jetstream/pull.go b/jetstream/pull.go index d255e61c4..2b37397d9 100644 --- a/jetstream/pull.go +++ b/jetstream/pull.go @@ -18,6 +18,7 @@ import ( "errors" "fmt" "math" + "slices" "sync" "sync/atomic" "time" @@ -87,14 +88,19 @@ type ( name string info *ConsumerInfo subs syncx.Map[string, *pullSubscription] + pinID string } pullRequest struct { - Expires time.Duration `json:"expires,omitempty"` - Batch int `json:"batch,omitempty"` - MaxBytes int `json:"max_bytes,omitempty"` - NoWait bool `json:"no_wait,omitempty"` - Heartbeat time.Duration `json:"idle_heartbeat,omitempty"` + Expires time.Duration `json:"expires,omitempty"` + Batch int `json:"batch,omitempty"` + MaxBytes int `json:"max_bytes,omitempty"` + NoWait bool `json:"no_wait,omitempty"` + Heartbeat time.Duration `json:"idle_heartbeat,omitempty"` + MinPending int64 `json:"min_pending,omitempty"` + MinAckPending int64 `json:"min_ack_pending,omitempty"` + PinID string `json:"id,omitempty"` + Group string `json:"group,omitempty"` } consumeOpts struct { @@ -102,6 +108,9 @@ type ( MaxMessages int MaxBytes int LimitSize bool + MinPending int64 + MinAckPending int64 + Group string Heartbeat time.Duration ErrHandler ConsumeErrHandlerFunc ReportMissingHeartbeats bool @@ -187,6 +196,19 @@ func (p *pullConsumer) Consume(handler MessageHandler, opts ...PullConsumeOpt) ( if err != nil { return nil, fmt.Errorf("%w: %s", ErrInvalidOption, err) } + + if len(p.info.Config.PriorityGroups) != 0 { + if consumeOpts.Group == "" { + return nil, fmt.Errorf("%w: %s", ErrInvalidOption, "priority group is required for priority consumer") + } + + if !slices.Contains(p.info.Config.PriorityGroups, consumeOpts.Group) { + return nil, fmt.Errorf("%w: %s", ErrInvalidOption, "invalid priority group") + } + } else if consumeOpts.Group != "" { + return nil, fmt.Errorf("%w: %s", ErrInvalidOption, "priority group is not supported for this consumer") + } + p.Lock() subject := p.js.apiSubject(fmt.Sprintf(apiRequestNextT, p.stream, p.name)) @@ -247,6 +269,9 @@ func (p *pullConsumer) Consume(handler MessageHandler, opts ...PullConsumeOpt) ( } return } + if pinId := msg.Header.Get("Nats-Pin-Id"); pinId != "" { + p.setPinID(pinId) + } handler(p.js.toJSMsg(msg)) sub.Lock() sub.decrementPendingMsgs(msg) @@ -283,10 +308,14 @@ func (p *pullConsumer) Consume(handler MessageHandler, opts ...PullConsumeOpt) ( batchSize = min(batchSize, sub.consumeOpts.StopAfter-sub.delivered) } if err := sub.pull(&pullRequest{ - Expires: consumeOpts.Expires, - Batch: batchSize, - MaxBytes: consumeOpts.MaxBytes, - Heartbeat: consumeOpts.Heartbeat, + Expires: consumeOpts.Expires, + Batch: batchSize, + MaxBytes: consumeOpts.MaxBytes, + Heartbeat: consumeOpts.Heartbeat, + MinPending: consumeOpts.MinPending, + MinAckPending: consumeOpts.MinAckPending, + Group: consumeOpts.Group, + PinID: p.getPinID(), }, subject); err != nil { sub.errs <- err } @@ -318,10 +347,14 @@ func (p *pullConsumer) Consume(handler MessageHandler, opts ...PullConsumeOpt) ( } sub.fetchNext <- &pullRequest{ - Expires: sub.consumeOpts.Expires, - Batch: sub.consumeOpts.MaxMessages, - MaxBytes: sub.consumeOpts.MaxBytes, - Heartbeat: sub.consumeOpts.Heartbeat, + Expires: sub.consumeOpts.Expires, + Batch: sub.consumeOpts.MaxMessages, + MaxBytes: sub.consumeOpts.MaxBytes, + Heartbeat: sub.consumeOpts.Heartbeat, + MinPending: sub.consumeOpts.MinPending, + MinAckPending: sub.consumeOpts.MinAckPending, + Group: sub.consumeOpts.Group, + PinID: p.getPinID(), } if sub.hbMonitor != nil { sub.hbMonitor.Reset(2 * sub.consumeOpts.Heartbeat) @@ -341,10 +374,14 @@ func (p *pullConsumer) Consume(handler MessageHandler, opts ...PullConsumeOpt) ( batchSize = min(batchSize, sub.consumeOpts.StopAfter-sub.delivered) } sub.fetchNext <- &pullRequest{ - Expires: sub.consumeOpts.Expires, - Batch: batchSize, - MaxBytes: sub.consumeOpts.MaxBytes, - Heartbeat: sub.consumeOpts.Heartbeat, + Expires: sub.consumeOpts.Expires, + Batch: batchSize, + MaxBytes: sub.consumeOpts.MaxBytes, + Heartbeat: sub.consumeOpts.Heartbeat, + MinPending: sub.consumeOpts.MinPending, + MinAckPending: sub.consumeOpts.MinAckPending, + Group: sub.consumeOpts.Group, + PinID: p.getPinID(), } if sub.hbMonitor != nil { sub.hbMonitor.Reset(2 * sub.consumeOpts.Heartbeat) @@ -412,11 +449,19 @@ func (s *pullSubscription) checkPending() { batchSize = min(batchSize, s.consumeOpts.StopAfter-s.delivered-s.pending.msgCount) } if batchSize > 0 { + pinID := "" + if s.consumer != nil { + pinID = s.consumer.getPinID() + } s.fetchNext <- &pullRequest{ - Expires: s.consumeOpts.Expires, - Batch: batchSize, - MaxBytes: maxBytes, - Heartbeat: s.consumeOpts.Heartbeat, + Expires: s.consumeOpts.Expires, + Batch: batchSize, + MaxBytes: maxBytes, + Heartbeat: s.consumeOpts.Heartbeat, + PinID: pinID, + Group: s.consumeOpts.Group, + MinPending: s.consumeOpts.MinPending, + MinAckPending: s.consumeOpts.MinAckPending, } s.pending.msgCount = s.consumeOpts.MaxMessages @@ -436,6 +481,18 @@ func (p *pullConsumer) Messages(opts ...PullMessagesOpt) (MessagesContext, error return nil, fmt.Errorf("%w: %s", ErrInvalidOption, err) } + if len(p.info.Config.PriorityGroups) != 0 { + if consumeOpts.Group == "" { + return nil, fmt.Errorf("%w: %s", ErrInvalidOption, "priority group is required for priority consumer") + } + + if !slices.Contains(p.info.Config.PriorityGroups, consumeOpts.Group) { + return nil, fmt.Errorf("%w: %s", ErrInvalidOption, "invalid priority group") + } + } else if consumeOpts.Group != "" { + return nil, fmt.Errorf("%w: %s", ErrInvalidOption, "priority group is not supported for this consumer") + } + p.Lock() subject := p.js.apiSubject(fmt.Sprintf(apiRequestNextT, p.stream, p.name)) @@ -552,6 +609,9 @@ func (s *pullSubscription) Next() (Msg, error) { } continue } + if pinId := msg.Header.Get("Nats-Pin-Id"); pinId != "" { + s.consumer.setPinID(pinId) + } s.decrementPendingMsgs(msg) s.incrementDeliveredMsgs() return s.consumer.js.toJSMsg(msg), nil @@ -595,6 +655,11 @@ func (s *pullSubscription) handleStatusMsg(msg *nats.Msg, msgErr error) error { if errors.Is(msgErr, ErrConsumerDeleted) || errors.Is(msgErr, ErrBadRequest) { return msgErr } + if errors.Is(msgErr, ErrPinIDMismatch) { + s.consumer.setPinID("") + s.pending.msgCount = 0 + s.pending.byteCount = 0 + } if s.consumeOpts.ErrHandler != nil { s.consumeOpts.ErrHandler(s, msgErr) } @@ -776,6 +841,7 @@ func (p *pullConsumer) fetch(req *pullRequest) (MessageBatch, error) { if err != nil { return nil, err } + req.PinID = p.getPinID() if err := sub.pull(req, subject); err != nil { return nil, err } @@ -798,6 +864,9 @@ func (p *pullConsumer) fetch(req *pullRequest) (MessageBatch, error) { if errNotTimeoutOrNoMsgs && !errors.Is(err, ErrMaxBytesExceeded) { res.err = err } + if errors.Is(err, ErrPinIDMismatch) { + p.setPinID("") + } res.done = true res.Unlock() return @@ -806,6 +875,9 @@ func (p *pullConsumer) fetch(req *pullRequest) (MessageBatch, error) { res.Unlock() continue } + if pinId := msg.Header.Get("Nats-Pin-Id"); pinId != "" { + p.setPinID(pinId) + } res.msgs <- p.js.toJSMsg(msg) meta, err := msg.Metadata() if err != nil { @@ -946,8 +1018,8 @@ func (s *pullSubscription) pull(req *pullRequest, subject string) error { if err != nil { return err } - reply := s.subscription.Subject + if err := s.consumer.js.conn.PublishRequest(subject, reply, reqJSON); err != nil { return err } @@ -1041,3 +1113,15 @@ func (consumeOpts *consumeOpts) setDefaults(ordered bool) error { } return nil } + +func (c *pullConsumer) getPinID() string { + c.Lock() + defer c.Unlock() + return c.pinID +} + +func (c *pullConsumer) setPinID(pinID string) { + c.Lock() + defer c.Unlock() + c.pinID = pinID +} diff --git a/jetstream/stream.go b/jetstream/stream.go index ce64cf5b6..42a5f11c1 100644 --- a/jetstream/stream.go +++ b/jetstream/stream.go @@ -1,4 +1,4 @@ -// Copyright 2022-2024 The NATS Authors +// Copyright 2022-2025 The NATS Authors // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. // You may obtain a copy of the License at @@ -114,6 +114,10 @@ type ( // ConsumerNames returns a ConsumerNameLister enabling iterating over a // channel of consumer names. ConsumerNames(context.Context) ConsumerNameLister + + // UnpinConsumer unpins the currently pinned client for a consumer for the given group name. + // If consumer does not exist, ErrConsumerNotFound is returned. + UnpinConsumer(ctx context.Context, consumer string, group string) error } RawStreamMsg struct { @@ -258,6 +262,10 @@ type ( apiPaged Consumers []string `json:"consumers"` } + + consumerUnpinRequest struct { + Group string `json:"group"` + } ) // CreateOrUpdateConsumer creates a consumer on a given stream with @@ -750,3 +758,9 @@ func (s *consumerLister) consumerNames(ctx context.Context, stream string) ([]st s.offset += len(resp.Consumers) return resp.Consumers, nil } + +// UnpinConsumer unpins the currently pinned client for a consumer for the given group name. +// If consumer does not exist, ErrConsumerNotFound is returned. +func (s *stream) UnpinConsumer(ctx context.Context, consumer string, group string) error { + return unpinConsumer(ctx, s.js, s.name, consumer, group) +} diff --git a/jetstream/test/consumer_test.go b/jetstream/test/consumer_test.go index 3ad59bb9d..c1ea42980 100644 --- a/jetstream/test/consumer_test.go +++ b/jetstream/test/consumer_test.go @@ -1,4 +1,4 @@ -// Copyright 2023 The NATS Authors +// Copyright 2023-2025 The NATS Authors // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. // You may obtain a copy of the License at @@ -16,6 +16,8 @@ package test import ( "context" "errors" + "fmt" + "sync/atomic" "testing" "time" @@ -129,6 +131,1013 @@ func TestConsumerInfo(t *testing.T) { }) } +func TestConsumerOverflow(t *testing.T) { + t.Run("fetch", func(t *testing.T) { + srv := RunBasicJetStreamServer() + defer shutdownJSServerAndRemoveStorage(t, srv) + + nc, err := nats.Connect(srv.ClientURL()) + if err != nil { + t.Fatalf("Unexpected error: %v", err) + } + defer nc.Close() + + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + + js, err := jetstream.New(nc) + if err != nil { + t.Fatalf("Unexpected error: %v", err) + } + + s, err := js.CreateStream(ctx, jetstream.StreamConfig{Name: "foo", Subjects: []string{"FOO.*"}}) + if err != nil { + t.Fatalf("Unexpected error: %v", err) + } + + c, err := s.CreateOrUpdateConsumer(ctx, jetstream.ConsumerConfig{ + PriorityPolicy: jetstream.PriorityPolicyOverflow, + PriorityGroups: []string{"A"}, + }) + if err != nil { + t.Fatalf("Unexpected error: %v", err) + } + + // Check that consumer got proper priority policy and TTL + info := c.CachedInfo() + if info.Config.PriorityPolicy != jetstream.PriorityPolicyOverflow { + t.Fatalf("Invalid priority policy; expected: %v; got: %v", jetstream.PriorityPolicyOverflow, info.Config.PriorityPolicy) + } + + for range 100 { + _, err = js.Publish(ctx, "FOO.bar", []byte("hello")) + if err != nil { + t.Fatalf("Unexpected error: %v", err) + } + } + + // We are below overflow, so we should not get any messages. + msgs, err := c.Fetch(10, jetstream.FetchMinPending(110), jetstream.FetchMaxWait(500*time.Millisecond), jetstream.FetchPriorityGroup("A")) + if err != nil { + t.Fatalf("Unexpected error: %v", err) + } + count := 0 + for msg := range msgs.Messages() { + msg.Ack() + count++ + } + if count != 0 { + t.Fatalf("Expected 0 messages, got %d", count) + } + + // Add more messages + for range 100 { + _, err = js.Publish(ctx, "FOO.bar", []byte("hello")) + if err != nil { + t.Fatalf("Unexpected error: %v", err) + } + } + + msgs, err = c.Fetch(10, jetstream.FetchMinPending(110), jetstream.FetchPriorityGroup("A")) + if err != nil { + t.Fatalf("Unexpected error: %v", err) + } + count = 0 + for msg := range msgs.Messages() { + if err := msg.DoubleAck(context.Background()); err != nil { + t.Fatalf("Unexpected error: %v", err) + } + count++ + } + if count != 10 { + t.Fatalf("Expected 10 messages, got %d", count) + } + + // try fetching messages with min ack pending + msgs, err = c.Fetch(10, jetstream.FetchMaxWait(500*time.Millisecond), jetstream.FetchMinAckPending(10), jetstream.FetchPriorityGroup("A")) + if err != nil { + t.Fatalf("Unexpected error: %v", err) + } + count = 0 + for msg := range msgs.Messages() { + msg.Ack() + count++ + } + if count != 0 { + t.Fatalf("Expected 0 messages, got %d", count) + } + + // now fetch some more messages but do not ack them, + // we will test min ack pending + msgs, err = c.Fetch(10, jetstream.FetchPriorityGroup("A")) + if err != nil { + t.Fatalf("Unexpected error: %v", err) + } + count = 0 + for range msgs.Messages() { + count++ + } + if msgs.Error() != nil { + t.Fatalf("Unexpected error: %v", msgs.Error()) + } + if count != 10 { + t.Fatalf("Expected 10 messages, got %d", count) + } + + // now fetch messages with min ack pending + msgs, err = c.Fetch(10, jetstream.FetchMinAckPending(10), jetstream.FetchPriorityGroup("A")) + if err != nil { + t.Fatalf("Unexpected error: %v", err) + } + count = 0 + for msg := range msgs.Messages() { + msg.Ack() + count++ + } + if msgs.Error() != nil { + t.Fatalf("Unexpected error: %v", msgs.Error()) + } + if count != 10 { + t.Fatalf("Expected 10 messages, got %d", count) + } + }) + + t.Run("consume", func(t *testing.T) { + srv := RunBasicJetStreamServer() + defer shutdownJSServerAndRemoveStorage(t, srv) + + nc, err := nats.Connect(srv.ClientURL()) + if err != nil { + t.Fatalf("Unexpected error: %v", err) + } + defer nc.Close() + + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + + js, err := jetstream.New(nc) + if err != nil { + t.Fatalf("Unexpected error: %v", err) + } + + s, err := js.CreateStream(ctx, jetstream.StreamConfig{Name: "foo", Subjects: []string{"FOO.*"}}) + if err != nil { + t.Fatalf("Unexpected error: %v", err) + } + + c, err := s.CreateOrUpdateConsumer(ctx, jetstream.ConsumerConfig{ + PriorityPolicy: jetstream.PriorityPolicyOverflow, + PriorityGroups: []string{"A"}, + }) + if err != nil { + t.Fatalf("Unexpected error: %v", err) + } + + // Check that consumer got proper priority policy and TTL + info := c.CachedInfo() + if info.Config.PriorityPolicy != jetstream.PriorityPolicyOverflow { + t.Fatalf("Invalid priority policy; expected: %v; got: %v", jetstream.PriorityPolicyOverflow, info.Config.PriorityPolicy) + } + + for i := range 100 { + _, err = js.Publish(ctx, "FOO.bar", []byte(fmt.Sprintf("hello %d", i))) + if err != nil { + t.Fatalf("Unexpected error: %v", err) + } + } + + count := atomic.Uint32{} + + handler := func(m jetstream.Msg) { + if err := m.Ack(); err != nil { + t.Fatalf("Unexpected error: %v", err) + } + count.Add(1) + } + cc, err := c.Consume(handler, + jetstream.PullPriorityGroup("A"), + jetstream.PullMinPending(110), + jetstream.PullMaxMessages(20), + jetstream.PullExpiry(time.Second)) + if err != nil { + t.Fatalf("Unexpected error: %v", err) + } + defer cc.Stop() + + time.Sleep(3 * time.Second) + // there are 100 messages on the stream, and min pending is 110 + // so we should not get any messages + if count.Load() != 0 { + t.Fatalf("Expected 0 messages, got %d", count.Load()) + } + + // Add more messages + for i := 100; i < 200; i++ { + _, err = js.Publish(ctx, "FOO.bar", []byte(fmt.Sprintf("hello %d", i))) + if err != nil { + t.Fatalf("Unexpected error: %v", err) + } + } + + time.Sleep(100 * time.Millisecond) + // now we should get 91 messages, because `Consume` will + // keep getting messages until it drops below min pending + if count.Load() != 91 { + t.Fatalf("Expected 10 messages, got %d", count.Load()) + } + cc.Stop() + + // now test min ack pending + count.Store(0) + + // consume with min ack pending + cc, err = c.Consume(handler, + jetstream.PullPriorityGroup("A"), + jetstream.PullMinAckPending(5), + jetstream.PullMaxMessages(20)) + if err != nil { + t.Fatalf("Unexpected error: %v", err) + } + defer cc.Stop() + + time.Sleep(100 * time.Millisecond) + // no messages should be received, there are no pending acks + if count.Load() != 0 { + t.Fatalf("Expected 0 messages, got %d", count.Load()) + } + // fetch some messages, do not ack them + msgs, err := c.Fetch(10, jetstream.FetchPriorityGroup("A")) + if err != nil { + t.Fatalf("Unexpected error: %v", err) + } + i := 0 + for range msgs.Messages() { + i++ + } + if msgs.Error() != nil { + t.Fatalf("Unexpected error: %v", msgs.Error()) + } + if i != 10 { + t.Fatalf("Expected 10 messages, got %d", i) + } + + time.Sleep(100 * time.Millisecond) + // we should process the rest of the stream minus the 10 unacked messages + // 200 - 91 - 10 = 99 + if count.Load() != 99 { + t.Fatalf("Expected 5 messages, got %d", count.Load()) + } + }) + + t.Run("messages", func(t *testing.T) { + srv := RunBasicJetStreamServer() + defer shutdownJSServerAndRemoveStorage(t, srv) + + nc, err := nats.Connect(srv.ClientURL()) + if err != nil { + t.Fatalf("Unexpected error: %v", err) + } + defer nc.Close() + + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + + js, err := jetstream.New(nc) + if err != nil { + t.Fatalf("Unexpected error: %v", err) + } + + s, err := js.CreateStream(ctx, jetstream.StreamConfig{Name: "foo", Subjects: []string{"FOO.*"}}) + if err != nil { + t.Fatalf("Unexpected error: %v", err) + } + + c, err := s.CreateOrUpdateConsumer(ctx, jetstream.ConsumerConfig{ + PriorityPolicy: jetstream.PriorityPolicyOverflow, + PriorityGroups: []string{"A"}, + }) + if err != nil { + t.Fatalf("Unexpected error: %v", err) + } + + // Check that consumer got proper priority policy and TTL + info := c.CachedInfo() + if info.Config.PriorityPolicy != jetstream.PriorityPolicyOverflow { + t.Fatalf("Invalid priority policy; expected: %v; got: %v", jetstream.PriorityPolicyOverflow, info.Config.PriorityPolicy) + } + + for i := range 100 { + _, err = js.Publish(ctx, "FOO.bar", []byte(fmt.Sprintf("hello %d", i))) + if err != nil { + t.Fatalf("Unexpected error: %v", err) + } + } + + count := atomic.Uint32{} + errs := make(chan error, 10) + done := make(chan struct{}, 1) + handler := func(it jetstream.MessagesContext) { + for { + msg, err := it.Next() + if err != nil { + if !errors.Is(err, jetstream.ErrMsgIteratorClosed) { + errs <- err + } + break + } + if err := msg.Ack(); err != nil { + errs <- err + break + } + count.Add(1) + } + done <- struct{}{} + } + + it, err := c.Messages(jetstream.PullPriorityGroup("A"), + jetstream.PullMinPending(110), + jetstream.PullMaxMessages(20)) + if err != nil { + t.Fatalf("Unexpected error: %v", err) + } + go handler(it) + + time.Sleep(100 * time.Millisecond) + // there are 100 messages on the stream, and min pending is 110 + // so we should not get any messages + if count.Load() != 0 { + t.Fatalf("Expected 0 messages, got %d", count.Load()) + } + + // Add more messages + for i := 100; i < 200; i++ { + _, err = js.Publish(ctx, "FOO.bar", []byte(fmt.Sprintf("hello %d", i))) + if err != nil { + t.Fatalf("Unexpected error: %v", err) + } + } + + time.Sleep(100 * time.Millisecond) + // now we should get 91 messages, because `Consume` will + // keep getting messages until it drops below min pending + if count.Load() != 91 { + t.Fatalf("Expected 10 messages, got %d", count.Load()) + } + it.Stop() + <-done + + // now test min ack pending + count.Store(0) + + it, err = c.Messages(jetstream.PullPriorityGroup("A"), + jetstream.PullMinAckPending(5), + jetstream.PullMaxMessages(20)) + if err != nil { + t.Fatalf("Unexpected error: %v", err) + } + go handler(it) + + time.Sleep(100 * time.Millisecond) + // no messages should be received, there are no pending acks + if count.Load() != 0 { + t.Fatalf("Expected 0 messages, got %d", count.Load()) + } + // fetch some messages, do not ack them + msgs, err := c.Fetch(10, jetstream.FetchPriorityGroup("A")) + if err != nil { + t.Fatalf("Unexpected error: %v", err) + } + i := 0 + for range msgs.Messages() { + i++ + } + if msgs.Error() != nil { + t.Fatalf("Unexpected error: %v", msgs.Error()) + } + if i != 10 { + t.Fatalf("Expected 10 messages, got %d", i) + } + + time.Sleep(100 * time.Millisecond) + // we should process the rest of the stream minus the 10 unacked messages + // 200 - 91 - 10 = 99 + if count.Load() != 99 { + t.Fatalf("Expected 5 messages, got %d", count.Load()) + } + it.Stop() + <-done + }) +} + +func TestConsumerPinned(t *testing.T) { + t.Run("messages", func(t *testing.T) { + srv := RunBasicJetStreamServer() + defer shutdownJSServerAndRemoveStorage(t, srv) + + nc, err := nats.Connect(srv.ClientURL()) + if err != nil { + t.Fatalf("Unexpected error: %v", err) + } + defer nc.Close() + + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + + js, err := jetstream.New(nc) + if err != nil { + t.Fatalf("Unexpected error: %v", err) + } + + s, err := js.CreateStream(ctx, jetstream.StreamConfig{Name: "foo", Subjects: []string{"FOO.*"}}) + if err != nil { + t.Fatalf("Unexpected error: %v", err) + } + c, err := s.CreateOrUpdateConsumer(ctx, jetstream.ConsumerConfig{ + Durable: "cons", + AckPolicy: jetstream.AckExplicitPolicy, + Description: "test consumer", + PriorityPolicy: jetstream.PriorityPolicyPinned, + PinnedTTL: time.Second, + PriorityGroups: []string{"A"}, + }) + if err != nil { + t.Fatalf("Unexpected error: %v", err) + } + + for range 1000 { + _, err = js.Publish(ctx, "FOO.bar", []byte("hello")) + if err != nil { + t.Fatalf("Unexpected error: %v", err) + } + } + + gcount := make(chan struct{}, 1000) + count := atomic.Uint32{} + + // Initially pinned consumer instance + initiallyPinned, err := s.Consumer(ctx, "cons") + if err != nil { + t.Fatalf("Unexpected error: %v", err) + } + + handler := func(it jetstream.MessagesContext, counter *atomic.Uint32, doneCh chan struct{}) { + for { + msg, err := it.Next() + if err != nil { + break + } + if err := msg.Ack(); err != nil { + break + } + counter.Add(1) + gcount <- struct{}{} + } + doneCh <- struct{}{} + } + // test priority group validation + // invalid priority group + _, err = initiallyPinned.Messages(jetstream.PullPriorityGroup("BAD")) + if err == nil || err.Error() != "nats: invalid jetstream option: invalid priority group" { + t.Fatalf("Expected invalid priority group error, got %v", err) + } + + // no priority group + _, err = initiallyPinned.Messages() + if err == nil || err.Error() != "nats: invalid jetstream option: priority group is required for priority consumer" { + t.Fatalf("Expected invalid priority group error") + } + + ipDoneCh := make(chan struct{}) + ip, err := initiallyPinned.Messages(jetstream.PullPriorityGroup("A"), jetstream.PullHeartbeat(500*time.Millisecond)) + if err != nil { + t.Fatalf("Unexpected error: %v", err) + } + defer ip.Stop() + + _, err = ip.Next() + if err != nil { + t.Fatalf("Unexpected error: %v", err) + } + count.Store(1) + go handler(ip, &count, ipDoneCh) + + time.Sleep(100 * time.Millisecond) + + // Second consume instance that should remain passive. + notPinnedC := atomic.Uint32{} + npDoneCh := make(chan struct{}) + np, err := c.Messages(jetstream.PullPriorityGroup("A"), jetstream.PullHeartbeat(500*time.Millisecond)) + if err != nil { + t.Fatalf("Unexpected error: %v", err) + } + defer np.Stop() + go handler(np, ¬PinnedC, npDoneCh) + + waitForCounter := func(t *testing.T, c *atomic.Uint32, expected int) { + t.Helper() + + outer: + for { + select { + case <-gcount: + if c.Load() == uint32(expected) { + break outer + } + case <-time.After(10 * time.Second): + t.Fatalf("Did not get all messages in time; expected %d, got %d", expected, c.Load()) + } + } + } + + waitForCounter(t, &count, 1000) + if notPinnedC.Load() != 0 { + t.Fatalf("Expected 0 messages for not pinned, got %d", notPinnedC.Load()) + } + + count.Store(0) + ip.Stop() + for range 100 { + _, err = js.Publish(ctx, "FOO.bar", []byte("hello")) + if err != nil { + t.Fatalf("Unexpected error: %v", err) + } + } + time.Sleep(100 * time.Millisecond) + if notPinnedC.Load() != 0 { + t.Fatalf("Expected 0 messages for not pinned, got %d", notPinnedC.Load()) + } + + //wait for pinned ttl to expire and messages to be consumed by the second consumer + waitForCounter(t, ¬PinnedC, 100) + if count.Load() != 0 { + t.Fatalf("Expected 0 messages for pinned, got %d", count.Load()) + } + np.Stop() + select { + case <-ipDoneCh: + case <-time.After(5 * time.Second): + t.Fatalf("Expected pinned consumer to be done") + } + select { + case <-npDoneCh: + case <-time.After(5 * time.Second): + t.Fatalf("Expected not pinned consumer to be done") + } + }) + + t.Run("consume", func(t *testing.T) { + srv := RunBasicJetStreamServer() + defer shutdownJSServerAndRemoveStorage(t, srv) + + nc, err := nats.Connect(srv.ClientURL()) + if err != nil { + t.Fatalf("Unexpected error: %v", err) + } + defer nc.Close() + + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + + js, err := jetstream.New(nc) + if err != nil { + t.Fatalf("Unexpected error: %v", err) + } + + s, err := js.CreateStream(ctx, jetstream.StreamConfig{Name: "foo", Subjects: []string{"FOO.*"}}) + if err != nil { + t.Fatalf("Unexpected error: %v", err) + } + c, err := s.CreateOrUpdateConsumer(ctx, jetstream.ConsumerConfig{ + Durable: "cons", + AckPolicy: jetstream.AckExplicitPolicy, + Description: "test consumer", + PriorityPolicy: jetstream.PriorityPolicyPinned, + PinnedTTL: 1 * time.Second, + PriorityGroups: []string{"A"}, + }) + if err != nil { + t.Fatalf("Unexpected error: %v", err) + } + + for range 1000 { + _, err = js.Publish(ctx, "FOO.bar", []byte("hello")) + if err != nil { + t.Fatalf("Unexpected error: %v", err) + } + } + + // Initially pinned consumer instance + initiallyPinned, err := s.Consumer(ctx, "cons") + if err != nil { + t.Fatalf("Unexpected error: %v", err) + } + + // test priority group validation + // invalid priority group + _, err = initiallyPinned.Consume(func(m jetstream.Msg) { + }, jetstream.PullPriorityGroup("BAD")) + if err == nil || err.Error() != "nats: invalid jetstream option: invalid priority group" { + t.Fatalf("Expected invalid priority group error") + } + + // no priority group + _, err = initiallyPinned.Consume(func(m jetstream.Msg) {}) + if err == nil || err.Error() != "nats: invalid jetstream option: priority group is required for priority consumer" { + t.Fatalf("Expected invalid priority group error") + } + + pinnedCount := atomic.Uint32{} + pinnedDone := make(chan struct{}) + ip, err := initiallyPinned.Consume(func(m jetstream.Msg) { + if err := m.Ack(); err != nil { + t.Fatalf("Unexpected error: %v", err) + } + if pinnedCount.Add(1) == 1000 { + close(pinnedDone) + } + }, jetstream.PullThresholdMessages(10), jetstream.PullPriorityGroup("A")) + if err != nil { + t.Fatalf("Unexpected error: %v", err) + } + defer ip.Stop() + + time.Sleep(100 * time.Millisecond) + + // Second consume instance that should remain passive. + notPinnedCount := atomic.Uint32{} + notPinnedDone := make(chan struct{}) + np, err := c.Consume(func(m jetstream.Msg) { + if err := m.Ack(); err != nil { + t.Fatalf("Unexpected error: %v", err) + } + if notPinnedCount.Add(1) == 100 { + close(notPinnedDone) + } + + }, jetstream.PullPriorityGroup("A")) + if err != nil { + t.Fatalf("Unexpected error: %v", err) + } + defer np.Stop() + + select { + case <-pinnedDone: + case <-time.After(10 * time.Second): + t.Fatalf("Expected pinned consumer to be done") + } + if notPinnedCount.Load() != 0 { + t.Fatalf("Expected 0 messages for not pinned, got %d", notPinnedCount.Load()) + } + + pinnedCount.Store(0) + ip.Stop() + for range 100 { + _, err = js.Publish(ctx, "FOO.bar", []byte("hello")) + if err != nil { + t.Fatalf("Unexpected error: %v", err) + } + } + time.Sleep(100 * time.Millisecond) + if notPinnedCount.Load() != 0 { + t.Fatalf("Expected 0 messages for not pinned, got %d", notPinnedCount.Load()) + } + + //wait for pinned ttl to expire and messages to be consumed by the second consumer + select { + case <-notPinnedDone: + case <-time.After(10 * time.Second): + t.Fatalf("Expected not pinned consumer to be done after pinned ttl expired") + } + if pinnedCount.Load() != 0 { + t.Fatalf("Expected 0 messages for pinned, got %d", pinnedCount.Load()) + } + }) + + t.Run("fetch", func(t *testing.T) { + srv := RunBasicJetStreamServer() + defer shutdownJSServerAndRemoveStorage(t, srv) + + nc, err := nats.Connect(srv.ClientURL()) + if err != nil { + t.Fatalf("Unexpected error: %v", err) + } + defer nc.Close() + + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + + js, err := jetstream.New(nc) + if err != nil { + t.Fatalf("Unexpected error: %v", err) + } + + s, err := js.CreateStream(ctx, jetstream.StreamConfig{Name: "foo", Subjects: []string{"FOO.*"}}) + if err != nil { + t.Fatalf("Unexpected error: %v", err) + } + c, err := s.CreateOrUpdateConsumer(ctx, jetstream.ConsumerConfig{ + Durable: "cons", + AckPolicy: jetstream.AckExplicitPolicy, + Description: "test consumer", + PriorityPolicy: jetstream.PriorityPolicyPinned, + PinnedTTL: time.Second, + PriorityGroups: []string{"A"}, + }) + if err != nil { + t.Fatalf("Unexpected error: %v", err) + } + + // Check that consumer got proper priority policy and TTL + info := c.CachedInfo() + if info.Config.PriorityPolicy != jetstream.PriorityPolicyPinned { + t.Fatalf("Invalid priority policy; expected: %v; got: %v", jetstream.PriorityPolicyPinned, info.Config.PriorityPolicy) + } + if info.Config.PinnedTTL != time.Second { + t.Fatalf("Invalid pinned TTL; expected: %v; got: %v", time.Second, info.Config.PinnedTTL) + } + + for range 100 { + _, err = js.Publish(ctx, "FOO.bar", []byte("hello")) + if err != nil { + t.Fatalf("Unexpected error: %v", err) + } + } + + // Initial fetch. + // Should get all messages and get a Pin ID. + msgs, err := c.Fetch(10, jetstream.FetchPriorityGroup("A")) + if err != nil { + t.Fatalf("Unexpected error: %v", err) + } + count := 0 + id := "" + for msg := range msgs.Messages() { + msg.Ack() + count++ + natsMsgId := msg.Headers().Get("Nats-Pin-Id") + if id == "" { + id = natsMsgId + } else { + if id != natsMsgId { + t.Fatalf("Expected Nats-Msg-Id to be the same for all messages") + } + } + } + if count != 10 { + t.Fatalf("Expected 10 messages, got %d", count) + + } + + // Different consumer instance. + cdiff, err := js.Consumer(ctx, "foo", "cons") + if err != nil { + t.Fatalf("Unexpected error: %v", err) + } + msgs, err = cdiff.Fetch(10, jetstream.FetchMaxWait(200*time.Millisecond), jetstream.FetchPriorityGroup("A")) + if err != nil { + t.Fatalf("Unexpected error: %v", err) + } + count = 0 + for msg := range msgs.Messages() { + msg.Ack() + count++ + } + if count != 0 { + t.Fatalf("Expected 0 messages, got %d", count) + } + if msgs.Error() != nil { + t.Fatalf("Unexpected error: %v", msgs.Error()) + } + + count = 0 + + // Now lets fetch from the pinned one, which should be fine. + msgs, err = c.Fetch(10, jetstream.FetchMaxWait(3*time.Second), jetstream.FetchPriorityGroup("A")) + if err != nil { + t.Fatalf("Unexpected error: %v", err) + } + for msg := range msgs.Messages() { + if pinId := msg.Headers().Get("Nats-Pin-Id"); pinId == "" { + t.Fatalf("missing Nats-Pin-Id header") + } + msg.Ack() + count++ + } + if count != 10 { + t.Fatalf("Expected 10 messages, got %d", count) + } + if msgs.Error() != nil { + t.Fatalf("Unexpected error: %v", msgs.Error()) + } + + // Wait for the TTL to expire, expect different ID + count = 0 + time.Sleep(1500 * time.Millisecond) + // The same instance, should work fine. + + msgs, err = c.Fetch(10, jetstream.FetchMaxWait(500*time.Millisecond), jetstream.FetchPriorityGroup("A")) + if err != nil { + t.Fatalf("Unexpected error: %v", err) + } + for msg := range msgs.Messages() { + msg.Ack() + count++ + } + if !errors.Is(msgs.Error(), jetstream.ErrPinIDMismatch) { + t.Fatalf("Expected error: %v, got: %v", jetstream.ErrPinIDMismatch, msgs.Error()) + } + if count != 0 { + t.Fatalf("Expected 0 messages, got %d", count) + } + + msgs, err = c.Fetch(10, jetstream.FetchMaxWait(500*time.Millisecond), jetstream.FetchPriorityGroup("A")) + if err != nil { + t.Fatalf("Unexpected error: %v", err) + } + for msg := range msgs.Messages() { + if msg == nil { + break + } + newId := msg.Headers().Get("Nats-Pin-Id") + if newId == id { + t.Fatalf("Expected new pull to have different ID. old: %s, new: %s", id, newId) + } + msg.Ack() + count++ + } + if msgs.Error() != nil { + t.Fatalf("Unexpected error: %v", msgs.Error()) + } + if count != 10 { + t.Fatalf("Expected 10 messages, got %d", count) + } + }) +} + +func TestConsumerUnpin(t *testing.T) { + t.Run("unpin consumer", func(t *testing.T) { + srv := RunBasicJetStreamServer() + defer shutdownJSServerAndRemoveStorage(t, srv) + + nc, err := nats.Connect(srv.ClientURL()) + if err != nil { + t.Fatalf("Unexpected error: %v", err) + } + defer nc.Close() + + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + + js, err := jetstream.New(nc) + if err != nil { + t.Fatalf("Unexpected error: %v", err) + } + + s, err := js.CreateStream(ctx, jetstream.StreamConfig{Name: "foo", Subjects: []string{"FOO.*"}}) + if err != nil { + t.Fatalf("Unexpected error: %v", err) + } + c, err := s.CreateOrUpdateConsumer(ctx, jetstream.ConsumerConfig{ + Durable: "cons", + AckPolicy: jetstream.AckExplicitPolicy, + Description: "test consumer", + PriorityPolicy: jetstream.PriorityPolicyPinned, + PinnedTTL: 50 * time.Second, + PriorityGroups: []string{"A"}, + }) + if err != nil { + t.Fatalf("Unexpected error: %v", err) + } + + for range 1000 { + _, err = js.Publish(ctx, "FOO.bar", []byte("hello")) + if err != nil { + t.Fatalf("Unexpected error: %v", err) + } + } + + msgs, err := c.Messages(jetstream.PullPriorityGroup("A")) + if err != nil { + t.Fatalf("Unexpected error: %v", err) + } + defer msgs.Stop() + + msg, err := msgs.Next() + if err != nil { + t.Fatalf("Unexpected error: %v", err) + } + + firstPinID := msg.Headers().Get("Nats-Pin-Id") + if firstPinID == "" { + t.Fatalf("Expected pinned message") + } + + second, err := s.Consumer(ctx, "cons") + if err != nil { + t.Fatalf("Unexpected error: %v", err) + } + + noMsgs, err := second.Messages(jetstream.PullPriorityGroup("A")) + if err != nil { + t.Fatalf("Unexpected error: %v", err) + } + defer noMsgs.Stop() + + done := make(chan struct{}) + errC := make(chan error) + go func() { + _, err := noMsgs.Next() + if err != nil { + errC <- err + return + } + done <- struct{}{} + }() + + select { + case <-done: + t.Fatalf("Expected no message") + case <-time.After(500 * time.Millisecond): + noMsgs.Stop() + } + select { + case <-time.After(5 * time.Second): + t.Fatalf("Expected error") + case err := <-errC: + if !errors.Is(err, jetstream.ErrMsgIteratorClosed) { + t.Fatalf("Expected error: %v, got: %v", jetstream.ErrMsgIteratorClosed, err) + } + } + + third, err := s.Consumer(ctx, "cons") + if err != nil { + t.Fatalf("Unexpected error: %v", err) + } + + yesMsgs, err := third.Messages(jetstream.PullPriorityGroup("A"), jetstream.PullExpiry(time.Second)) + if err != nil { + t.Fatalf("Unexpected error: %v", err) + } + defer yesMsgs.Stop() + + go func() { + msg, err := yesMsgs.Next() + newPinID := msg.Headers().Get("Nats-Pin-Id") + if newPinID == firstPinID || newPinID == "" { + errC <- fmt.Errorf("Expected new pin ID, got %s", newPinID) + return + } + if err != nil { + errC <- err + return + } + done <- struct{}{} + }() + + err = s.UnpinConsumer(ctx, "cons", "A") + if err != nil { + t.Fatalf("Unexpected error: %v", err) + } + + select { + case <-done: + case err := <-errC: + t.Fatalf("Unexpected error: %v", err) + case <-time.After(4 * time.Second): + t.Fatalf("Should not time out") + } + }) + t.Run("consumer not found", func(t *testing.T) { + srv := RunBasicJetStreamServer() + defer shutdownJSServerAndRemoveStorage(t, srv) + + nc, err := nats.Connect(srv.ClientURL()) + if err != nil { + t.Fatalf("Unexpected error: %v", err) + } + defer nc.Close() + + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + + js, err := jetstream.New(nc) + if err != nil { + t.Fatalf("Unexpected error: %v", err) + } + + s, err := js.CreateStream(ctx, jetstream.StreamConfig{Name: "foo", Subjects: []string{"FOO.*"}}) + if err != nil { + t.Fatalf("Unexpected error: %v", err) + } + + // try unpinning consumer with invalid name + err = s.UnpinConsumer(ctx, "cons", "A") + if !errors.Is(err, jetstream.ErrConsumerNotFound) { + t.Fatalf("Expected error: %v, got: %v", jetstream.ErrConsumerNotFound, err) + } + }) +} + func TestConsumerCachedInfo(t *testing.T) { srv := RunBasicJetStreamServer() defer shutdownJSServerAndRemoveStorage(t, srv) diff --git a/jetstream/test/stream_test.go b/jetstream/test/stream_test.go index 7221aff07..030014e2d 100644 --- a/jetstream/test/stream_test.go +++ b/jetstream/test/stream_test.go @@ -388,6 +388,11 @@ func TestConsumer(t *testing.T) { durable: "dur.123", withError: jetstream.ErrInvalidConsumerName, }, + { + name: "empty consumer name", + durable: "", + withError: jetstream.ErrInvalidConsumerName, + }, } srv := RunBasicJetStreamServer()