Курнат Андрей
1225 lines
33 KiB
Go
1225 lines
33 KiB
Go
// 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
|
|
//
|
|
// http://www.apache.org/licenses/LICENSE-2.0
|
|
//
|
|
// Unless required by applicable law or agreed to in writing, software
|
|
// distributed under the License is distributed on an "AS IS" BASIS,
|
|
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
// See the License for the specific language governing permissions and
|
|
// limitations under the License.
|
|
|
|
package jetstream
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"errors"
|
|
"fmt"
|
|
"math"
|
|
"slices"
|
|
"sync"
|
|
"sync/atomic"
|
|
"time"
|
|
|
|
"github.com/nats-io/nats.go"
|
|
"github.com/nats-io/nats.go/internal/syncx"
|
|
"github.com/nats-io/nuid"
|
|
)
|
|
|
|
type (
|
|
// MessagesContext supports iterating over a messages on a stream.
|
|
// It is returned by [Consumer.Messages] method.
|
|
MessagesContext interface {
|
|
// Next retrieves next message on a stream. If MessagesContext is closed
|
|
// (either stopped or drained), Next will return ErrMsgIteratorClosed
|
|
// error. An optional timeout or context can be provided using NextOpt
|
|
// options. If none are provided, Next will block indefinitely until a
|
|
// message is available, iterator is closed or a heartbeat error occurs.
|
|
Next(opts ...NextOpt) (Msg, error)
|
|
|
|
// Stop unsubscribes from the stream and cancels subscription. Calling
|
|
// Next after calling Stop will return ErrMsgIteratorClosed error.
|
|
// All messages that are already in the buffer are discarded.
|
|
Stop()
|
|
|
|
// Drain unsubscribes from the stream and cancels subscription. All
|
|
// messages that are already in the buffer will be available on
|
|
// subsequent calls to Next. After the buffer is drained, Next will
|
|
// return ErrMsgIteratorClosed error.
|
|
Drain()
|
|
}
|
|
|
|
// ConsumeContext supports processing incoming messages from a stream.
|
|
// It is returned by [Consumer.Consume] method.
|
|
ConsumeContext interface {
|
|
// Stop unsubscribes from the stream and cancels subscription.
|
|
// No more messages will be received after calling this method.
|
|
// All messages that are already in the buffer are discarded.
|
|
Stop()
|
|
|
|
// Drain unsubscribes from the stream and cancels subscription.
|
|
// All messages that are already in the buffer will be processed in callback function.
|
|
Drain()
|
|
|
|
// Closed returns a channel that is closed when the consuming is
|
|
// fully stopped/drained. When the channel is closed, no more messages
|
|
// will be received and processing is complete.
|
|
Closed() <-chan struct{}
|
|
}
|
|
|
|
// MessageHandler is a handler function used as callback in [Consume].
|
|
MessageHandler func(msg Msg)
|
|
|
|
// PullConsumeOpt represent additional options used in [Consume] for pull consumers.
|
|
PullConsumeOpt interface {
|
|
configureConsume(*consumeOpts) error
|
|
}
|
|
|
|
// PullMessagesOpt represent additional options used in [Messages] for pull consumers.
|
|
PullMessagesOpt interface {
|
|
configureMessages(*consumeOpts) error
|
|
}
|
|
|
|
pullConsumer struct {
|
|
sync.Mutex
|
|
js *jetStream
|
|
stream string
|
|
durable bool
|
|
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"`
|
|
MinPending int64 `json:"min_pending,omitempty"`
|
|
MinAckPending int64 `json:"min_ack_pending,omitempty"`
|
|
PinID string `json:"id,omitempty"`
|
|
Group string `json:"group,omitempty"`
|
|
Priority uint8 `json:"priority,omitempty"`
|
|
ctx context.Context `json:"-"`
|
|
maxWaitSet bool `json:"-"`
|
|
}
|
|
|
|
consumeOpts struct {
|
|
Expires time.Duration
|
|
MaxMessages int
|
|
MaxBytes int
|
|
LimitSize bool
|
|
MinPending int64
|
|
MinAckPending int64
|
|
Priority uint8
|
|
Group string
|
|
Heartbeat time.Duration
|
|
ErrHandler ConsumeErrHandler
|
|
ReportMissingHeartbeats bool
|
|
ThresholdMessages int
|
|
ThresholdBytes int
|
|
StopAfter int
|
|
stopAfterMsgsLeft chan int
|
|
notifyOnReconnect bool
|
|
}
|
|
|
|
ConsumeErrHandlerFunc func(consumeCtx ConsumeContext, err error)
|
|
|
|
pullSubscription struct {
|
|
sync.Mutex
|
|
id string
|
|
consumer *pullConsumer
|
|
subscription *nats.Subscription
|
|
msgs chan *nats.Msg
|
|
msgsClosed atomic.Uint32
|
|
errs chan error
|
|
pending pendingMsgs
|
|
hbMonitor *hbMonitor
|
|
fetchInProgress atomic.Uint32
|
|
closed atomic.Uint32
|
|
draining atomic.Uint32
|
|
done chan struct{}
|
|
connStatusChanged chan nats.Status
|
|
fetchNext chan *pullRequest
|
|
consumeOpts *consumeOpts
|
|
delivered int
|
|
closedCh chan struct{}
|
|
}
|
|
|
|
pendingMsgs struct {
|
|
msgCount int
|
|
byteCount int
|
|
}
|
|
|
|
MessageBatch interface {
|
|
Messages() <-chan Msg
|
|
Error() error
|
|
}
|
|
|
|
fetchResult struct {
|
|
sync.Mutex
|
|
msgs chan Msg
|
|
err error
|
|
done bool
|
|
sseq uint64
|
|
}
|
|
|
|
FetchOpt func(*pullRequest) error
|
|
|
|
hbMonitor struct {
|
|
timer *time.Timer
|
|
sync.Mutex
|
|
}
|
|
|
|
// NextOpt is an option for configuring the behavior of MessagesContext.Next.
|
|
NextOpt interface {
|
|
configureNext(*nextOpts)
|
|
}
|
|
|
|
nextOpts struct {
|
|
timeout time.Duration
|
|
ctx context.Context
|
|
}
|
|
)
|
|
|
|
const (
|
|
DefaultMaxMessages = 500
|
|
DefaultExpires = 30 * time.Second
|
|
defaultBatchMaxBytesOnly = 1_000_000
|
|
unset = -1
|
|
)
|
|
|
|
// Consume can be used to continuously receive messages and handle them
|
|
// with the provided callback function. Consume cannot be used concurrently
|
|
// when using ordered consumer.
|
|
//
|
|
// See [Consumer.Consume] for more details.
|
|
func (p *pullConsumer) Consume(handler MessageHandler, opts ...PullConsumeOpt) (ConsumeContext, error) {
|
|
if handler == nil {
|
|
return nil, ErrHandlerRequired
|
|
}
|
|
consumeOpts, err := parseConsumeOpts(false, opts...)
|
|
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))
|
|
|
|
consumeID := nuid.Next()
|
|
sub := &pullSubscription{
|
|
id: consumeID,
|
|
consumer: p,
|
|
errs: make(chan error, 10),
|
|
done: make(chan struct{}, 1),
|
|
fetchNext: make(chan *pullRequest, 1),
|
|
consumeOpts: consumeOpts,
|
|
}
|
|
sub.connStatusChanged = p.js.conn.StatusChanged(nats.CONNECTED, nats.RECONNECTING, nats.CLOSED)
|
|
|
|
sub.hbMonitor = sub.scheduleHeartbeatCheck(consumeOpts.Heartbeat)
|
|
|
|
p.subs.Store(sub.id, sub)
|
|
p.Unlock()
|
|
|
|
internalHandler := func(msg *nats.Msg) {
|
|
if sub.hbMonitor != nil {
|
|
sub.hbMonitor.Stop()
|
|
}
|
|
userMsg, msgErr := checkMsg(msg)
|
|
if !userMsg && msgErr == nil {
|
|
if sub.hbMonitor != nil {
|
|
sub.hbMonitor.Reset(2 * consumeOpts.Heartbeat)
|
|
}
|
|
return
|
|
}
|
|
if !userMsg {
|
|
// heartbeat message
|
|
if msgErr == nil {
|
|
return
|
|
}
|
|
|
|
sub.Lock()
|
|
termErr, notifyErr := sub.handleStatusMsg(msg, msgErr)
|
|
if termErr == nil {
|
|
sub.checkPending()
|
|
if sub.hbMonitor != nil {
|
|
sub.hbMonitor.Reset(2 * consumeOpts.Heartbeat)
|
|
}
|
|
}
|
|
sub.Unlock()
|
|
|
|
if sub.consumeOpts.ErrHandler != nil && notifyErr != nil {
|
|
sub.consumeOpts.ErrHandler(sub, notifyErr)
|
|
}
|
|
if termErr != nil {
|
|
if sub.closed.Load() == 1 {
|
|
return
|
|
}
|
|
if sub.consumeOpts.ErrHandler != nil {
|
|
sub.consumeOpts.ErrHandler(sub, termErr)
|
|
}
|
|
sub.Stop()
|
|
}
|
|
return
|
|
}
|
|
if pinId := msg.Header.Get("Nats-Pin-Id"); pinId != "" {
|
|
p.setPinID(pinId)
|
|
}
|
|
handler(p.js.toJSMsg(msg))
|
|
sub.Lock()
|
|
sub.decrementPendingMsgs(msg)
|
|
sub.incrementDeliveredMsgs()
|
|
sub.checkPending()
|
|
if sub.hbMonitor != nil {
|
|
sub.hbMonitor.Reset(2 * consumeOpts.Heartbeat)
|
|
}
|
|
sub.Unlock()
|
|
|
|
if sub.consumeOpts.StopAfter > 0 && sub.consumeOpts.StopAfter == sub.delivered {
|
|
sub.Stop()
|
|
}
|
|
}
|
|
inbox := p.js.conn.NewInbox()
|
|
sub.subscription, err = p.js.conn.Subscribe(inbox, internalHandler)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
sub.subscription.SetClosedHandler(func(sid string) func(string) {
|
|
return func(subject string) {
|
|
p.subs.Delete(sid)
|
|
sub.draining.CompareAndSwap(1, 0)
|
|
sub.Lock()
|
|
if sub.closedCh != nil {
|
|
close(sub.closedCh)
|
|
sub.closedCh = nil
|
|
}
|
|
sub.Unlock()
|
|
}
|
|
}(sub.id))
|
|
|
|
sub.Lock()
|
|
// initial pull
|
|
sub.resetPendingMsgs()
|
|
batchSize := sub.consumeOpts.MaxMessages
|
|
if sub.consumeOpts.StopAfter > 0 {
|
|
batchSize = min(batchSize, sub.consumeOpts.StopAfter-sub.delivered)
|
|
}
|
|
if err := sub.pull(&pullRequest{
|
|
Expires: consumeOpts.Expires,
|
|
Batch: batchSize,
|
|
MaxBytes: consumeOpts.MaxBytes,
|
|
Heartbeat: consumeOpts.Heartbeat,
|
|
MinPending: consumeOpts.MinPending,
|
|
MinAckPending: consumeOpts.MinAckPending,
|
|
Priority: consumeOpts.Priority,
|
|
Group: consumeOpts.Group,
|
|
PinID: p.getPinID(),
|
|
}, subject); err != nil {
|
|
sub.errs <- err
|
|
}
|
|
sub.Unlock()
|
|
|
|
go func() {
|
|
isConnected := true
|
|
for {
|
|
if sub.closed.Load() == 1 {
|
|
return
|
|
}
|
|
select {
|
|
case status, ok := <-sub.connStatusChanged:
|
|
if !ok {
|
|
continue
|
|
}
|
|
switch status {
|
|
case nats.RECONNECTING:
|
|
if sub.hbMonitor != nil {
|
|
sub.hbMonitor.Stop()
|
|
}
|
|
isConnected = false
|
|
case nats.CONNECTED:
|
|
sub.Lock()
|
|
if !isConnected {
|
|
isConnected = true
|
|
if sub.consumeOpts.notifyOnReconnect {
|
|
sub.errs <- errConnected
|
|
}
|
|
|
|
sub.fetchNext <- &pullRequest{
|
|
Expires: sub.consumeOpts.Expires,
|
|
Batch: sub.consumeOpts.MaxMessages,
|
|
MaxBytes: sub.consumeOpts.MaxBytes,
|
|
Heartbeat: sub.consumeOpts.Heartbeat,
|
|
MinPending: sub.consumeOpts.MinPending,
|
|
MinAckPending: sub.consumeOpts.MinAckPending,
|
|
Priority: sub.consumeOpts.Priority,
|
|
Group: sub.consumeOpts.Group,
|
|
PinID: p.getPinID(),
|
|
}
|
|
if sub.hbMonitor != nil {
|
|
sub.hbMonitor.Reset(2 * sub.consumeOpts.Heartbeat)
|
|
}
|
|
sub.resetPendingMsgs()
|
|
}
|
|
sub.Unlock()
|
|
|
|
case nats.CLOSED:
|
|
sub.errs <- ErrConnectionClosed
|
|
}
|
|
case err := <-sub.errs:
|
|
sub.Lock()
|
|
if errors.Is(err, ErrNoHeartbeat) {
|
|
batchSize := sub.consumeOpts.MaxMessages
|
|
if sub.consumeOpts.StopAfter > 0 {
|
|
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,
|
|
MinPending: sub.consumeOpts.MinPending,
|
|
MinAckPending: sub.consumeOpts.MinAckPending,
|
|
Priority: sub.consumeOpts.Priority,
|
|
Group: sub.consumeOpts.Group,
|
|
PinID: p.getPinID(),
|
|
}
|
|
if sub.hbMonitor != nil {
|
|
sub.hbMonitor.Reset(2 * sub.consumeOpts.Heartbeat)
|
|
}
|
|
sub.resetPendingMsgs()
|
|
}
|
|
sub.Unlock()
|
|
if sub.consumeOpts.ErrHandler != nil {
|
|
sub.consumeOpts.ErrHandler(sub, err)
|
|
}
|
|
if errors.Is(err, ErrConnectionClosed) {
|
|
sub.Stop()
|
|
}
|
|
case <-sub.done:
|
|
return
|
|
}
|
|
}
|
|
}()
|
|
|
|
go sub.pullMessages(subject)
|
|
|
|
return sub, nil
|
|
}
|
|
|
|
// resetPendingMsgs resets pending message count and byte count
|
|
// to the values set in consumeOpts
|
|
// lock should be held before calling this method
|
|
func (s *pullSubscription) resetPendingMsgs() {
|
|
s.pending.msgCount = s.consumeOpts.MaxMessages
|
|
s.pending.byteCount = s.consumeOpts.MaxBytes
|
|
}
|
|
|
|
// decrementPendingMsgs decrements pending message count and byte count
|
|
// lock should be held before calling this method
|
|
func (s *pullSubscription) decrementPendingMsgs(msg *nats.Msg) {
|
|
s.pending.msgCount--
|
|
if s.consumeOpts.MaxBytes != 0 && !s.consumeOpts.LimitSize {
|
|
s.pending.byteCount -= msg.Size()
|
|
}
|
|
}
|
|
|
|
// incrementDeliveredMsgs increments delivered message count
|
|
// lock should be held before calling this method
|
|
func (s *pullSubscription) incrementDeliveredMsgs() {
|
|
s.delivered++
|
|
}
|
|
|
|
// checkPending verifies whether there are enough messages in
|
|
// the buffer to trigger a new pull request.
|
|
// lock should be held before calling this method
|
|
func (s *pullSubscription) checkPending() {
|
|
// check if we went below any threshold
|
|
// we don't want to track bytes threshold if either it's not set or we used
|
|
// PullMaxMessagesWithBytesLimit
|
|
if (s.pending.msgCount < s.consumeOpts.ThresholdMessages ||
|
|
(s.pending.byteCount < s.consumeOpts.ThresholdBytes && s.consumeOpts.MaxBytes != 0 && !s.consumeOpts.LimitSize)) &&
|
|
s.fetchInProgress.Load() == 0 {
|
|
|
|
var batchSize, maxBytes int
|
|
batchSize = s.consumeOpts.MaxMessages - s.pending.msgCount
|
|
if s.consumeOpts.MaxBytes != 0 {
|
|
if s.consumeOpts.LimitSize {
|
|
maxBytes = s.consumeOpts.MaxBytes
|
|
} else {
|
|
maxBytes = s.consumeOpts.MaxBytes - s.pending.byteCount
|
|
// when working with max bytes only, always ask for full batch
|
|
batchSize = s.consumeOpts.MaxMessages
|
|
}
|
|
}
|
|
if s.consumeOpts.StopAfter > 0 {
|
|
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,
|
|
PinID: pinID,
|
|
Group: s.consumeOpts.Group,
|
|
MinPending: s.consumeOpts.MinPending,
|
|
MinAckPending: s.consumeOpts.MinAckPending,
|
|
Priority: s.consumeOpts.Priority,
|
|
}
|
|
|
|
s.pending.msgCount = s.consumeOpts.MaxMessages
|
|
s.pending.byteCount = s.consumeOpts.MaxBytes
|
|
}
|
|
}
|
|
}
|
|
|
|
// Messages returns MessagesContext, allowing continuously iterating
|
|
// over messages on a stream. Messages cannot be used concurrently
|
|
// when using ordered consumer.
|
|
//
|
|
// See [Consumer.Messages] for more details.
|
|
func (p *pullConsumer) Messages(opts ...PullMessagesOpt) (MessagesContext, error) {
|
|
consumeOpts, err := parseMessagesOpts(false, opts...)
|
|
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))
|
|
|
|
msgs := make(chan *nats.Msg, consumeOpts.MaxMessages)
|
|
|
|
consumeID := nuid.Next()
|
|
sub := &pullSubscription{
|
|
id: consumeID,
|
|
consumer: p,
|
|
done: make(chan struct{}, 1),
|
|
msgs: msgs,
|
|
errs: make(chan error, 10),
|
|
fetchNext: make(chan *pullRequest, 1),
|
|
consumeOpts: consumeOpts,
|
|
}
|
|
sub.connStatusChanged = p.js.conn.StatusChanged(nats.CONNECTED, nats.RECONNECTING)
|
|
inbox := p.js.conn.NewInbox()
|
|
sub.subscription, err = p.js.conn.ChanSubscribe(inbox, sub.msgs)
|
|
if err != nil {
|
|
p.Unlock()
|
|
return nil, err
|
|
}
|
|
sub.subscription.SetClosedHandler(func(sid string) func(string) {
|
|
return func(subject string) {
|
|
if sub.draining.Load() != 1 {
|
|
// if we're not draining, subscription can be closed as soon
|
|
// as closed handler is called
|
|
// otherwise, we need to wait until all messages are drained
|
|
// in Next
|
|
p.subs.Delete(sid)
|
|
}
|
|
sub.closeMsgs()
|
|
}
|
|
}(sub.id))
|
|
|
|
p.subs.Store(sub.id, sub)
|
|
p.Unlock()
|
|
|
|
go sub.pullMessages(subject)
|
|
|
|
go func() {
|
|
for {
|
|
select {
|
|
case status, ok := <-sub.connStatusChanged:
|
|
if !ok {
|
|
return
|
|
}
|
|
if status == nats.CONNECTED {
|
|
sub.errs <- errConnected
|
|
}
|
|
if status == nats.RECONNECTING {
|
|
sub.errs <- errDisconnected
|
|
}
|
|
case <-sub.done:
|
|
return
|
|
}
|
|
}
|
|
}()
|
|
|
|
return sub, nil
|
|
}
|
|
|
|
var (
|
|
errConnected = errors.New("connected")
|
|
errDisconnected = errors.New("disconnected")
|
|
)
|
|
|
|
// Next retrieves next message on a stream. If MessagesContext is closed
|
|
// (either stopped or drained), Next will return ErrMsgIteratorClosed
|
|
// error. An optional timeout or context can be provided using NextOpt
|
|
// options. If none are provided, Next will block indefinitely until a
|
|
// message is available, iterator is closed or a heartbeat error occurs.
|
|
func (s *pullSubscription) Next(opts ...NextOpt) (Msg, error) {
|
|
var nextOpts nextOpts
|
|
for _, opt := range opts {
|
|
opt.configureNext(&nextOpts)
|
|
}
|
|
|
|
if nextOpts.timeout > 0 && nextOpts.ctx != nil {
|
|
return nil, fmt.Errorf("%w: cannot specify both NextMaxWait and NextContext", ErrInvalidOption)
|
|
}
|
|
|
|
// Create timeout channel if needed
|
|
var timeoutCh <-chan time.Time
|
|
if nextOpts.timeout > 0 {
|
|
timer := time.NewTimer(nextOpts.timeout)
|
|
defer timer.Stop()
|
|
timeoutCh = timer.C
|
|
}
|
|
|
|
// Use context if provided
|
|
var ctxDone <-chan struct{}
|
|
if nextOpts.ctx != nil {
|
|
ctxDone = nextOpts.ctx.Done()
|
|
}
|
|
|
|
s.Lock()
|
|
defer s.Unlock()
|
|
drainMode := s.draining.Load() == 1
|
|
closed := s.closed.Load() == 1
|
|
if closed && !drainMode {
|
|
// Check if iterator was closed due to connection closure
|
|
if s.consumer.js.conn.IsClosed() {
|
|
return nil, fmt.Errorf("%w: %w", ErrMsgIteratorClosed, ErrConnectionClosed)
|
|
}
|
|
return nil, ErrMsgIteratorClosed
|
|
}
|
|
hbMonitor := s.scheduleHeartbeatCheck(s.consumeOpts.Heartbeat)
|
|
defer func() {
|
|
if hbMonitor != nil {
|
|
hbMonitor.Stop()
|
|
}
|
|
}()
|
|
|
|
isConnected := true
|
|
if s.consumeOpts.StopAfter > 0 && s.delivered >= s.consumeOpts.StopAfter {
|
|
s.Stop()
|
|
return nil, ErrMsgIteratorClosed
|
|
}
|
|
|
|
for {
|
|
s.checkPending()
|
|
select {
|
|
case msg, ok := <-s.msgs:
|
|
if !ok {
|
|
// if msgs channel is closed, it means that subscription was either drained or stopped
|
|
s.consumer.subs.Delete(s.id)
|
|
s.draining.CompareAndSwap(1, 0)
|
|
// Check if iterator was closed due to connection closure
|
|
if s.consumer.js.conn.IsClosed() {
|
|
return nil, fmt.Errorf("%w: %w", ErrMsgIteratorClosed, ErrConnectionClosed)
|
|
}
|
|
return nil, ErrMsgIteratorClosed
|
|
}
|
|
if hbMonitor != nil {
|
|
hbMonitor.Reset(2 * s.consumeOpts.Heartbeat)
|
|
}
|
|
userMsg, msgErr := checkMsg(msg)
|
|
if !userMsg {
|
|
// heartbeat message
|
|
if msgErr == nil {
|
|
continue
|
|
}
|
|
if termErr, _ := s.handleStatusMsg(msg, msgErr); termErr != nil {
|
|
s.Stop()
|
|
return nil, termErr
|
|
}
|
|
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
|
|
case err := <-s.errs:
|
|
if errors.Is(err, ErrNoHeartbeat) {
|
|
s.pending.msgCount = 0
|
|
s.pending.byteCount = 0
|
|
if s.consumeOpts.ReportMissingHeartbeats {
|
|
return nil, err
|
|
}
|
|
if hbMonitor != nil {
|
|
hbMonitor.Reset(2 * s.consumeOpts.Heartbeat)
|
|
}
|
|
}
|
|
if errors.Is(err, errConnected) {
|
|
if !isConnected {
|
|
isConnected = true
|
|
|
|
if s.consumeOpts.notifyOnReconnect {
|
|
return nil, errConnected
|
|
}
|
|
s.pending.msgCount = 0
|
|
s.pending.byteCount = 0
|
|
if hbMonitor != nil {
|
|
hbMonitor.Reset(2 * s.consumeOpts.Heartbeat)
|
|
}
|
|
}
|
|
}
|
|
if errors.Is(err, errDisconnected) {
|
|
if hbMonitor != nil {
|
|
hbMonitor.Stop()
|
|
}
|
|
isConnected = false
|
|
}
|
|
case <-timeoutCh:
|
|
return nil, nats.ErrTimeout
|
|
case <-ctxDone:
|
|
return nil, nextOpts.ctx.Err()
|
|
}
|
|
}
|
|
}
|
|
|
|
// handleStatusMsg processes a status message from the server.
|
|
// It returns a terminal error (caller should stop) and a non-terminal
|
|
// error to notify the user about via ErrHandler. The caller should invoke
|
|
// ErrHandler outside the lock to avoid deadlocks.
|
|
func (s *pullSubscription) handleStatusMsg(msg *nats.Msg, msgErr error) (error, error) {
|
|
if !errors.Is(msgErr, nats.ErrTimeout) && !errors.Is(msgErr, ErrMaxBytesExceeded) && !errors.Is(msgErr, ErrBatchCompleted) {
|
|
if errors.Is(msgErr, ErrConsumerDeleted) || errors.Is(msgErr, ErrBadRequest) {
|
|
return msgErr, nil
|
|
}
|
|
if errors.Is(msgErr, ErrPinIDMismatch) {
|
|
s.consumer.setPinID("")
|
|
s.pending.msgCount = 0
|
|
s.pending.byteCount = 0
|
|
}
|
|
if errors.Is(msgErr, ErrConsumerLeadershipChanged) {
|
|
s.pending.msgCount = 0
|
|
s.pending.byteCount = 0
|
|
}
|
|
return nil, msgErr
|
|
}
|
|
msgsLeft, bytesLeft, err := parsePending(msg)
|
|
if err != nil {
|
|
return err, nil
|
|
}
|
|
s.pending.msgCount -= msgsLeft
|
|
if s.pending.msgCount < 0 {
|
|
s.pending.msgCount = 0
|
|
}
|
|
if s.consumeOpts.MaxBytes > 0 && !s.consumeOpts.LimitSize {
|
|
s.pending.byteCount -= bytesLeft
|
|
if s.pending.byteCount < 0 {
|
|
s.pending.byteCount = 0
|
|
}
|
|
}
|
|
return nil, nil
|
|
}
|
|
|
|
func (hb *hbMonitor) Stop() {
|
|
hb.Lock()
|
|
hb.timer.Stop()
|
|
hb.Unlock()
|
|
}
|
|
|
|
func (hb *hbMonitor) Reset(dur time.Duration) {
|
|
hb.Lock()
|
|
hb.timer.Reset(dur)
|
|
hb.Unlock()
|
|
}
|
|
|
|
// Stop unsubscribes from the stream and cancels subscription. Calling
|
|
// Next after calling Stop will return ErrMsgIteratorClosed error.
|
|
// All messages that are already in the buffer are discarded.
|
|
func (s *pullSubscription) Stop() {
|
|
if !s.closed.CompareAndSwap(0, 1) {
|
|
return
|
|
}
|
|
close(s.done)
|
|
if s.consumeOpts.stopAfterMsgsLeft != nil {
|
|
if s.delivered >= s.consumeOpts.StopAfter {
|
|
close(s.consumeOpts.stopAfterMsgsLeft)
|
|
} else {
|
|
s.consumeOpts.stopAfterMsgsLeft <- s.consumeOpts.StopAfter - s.delivered
|
|
}
|
|
}
|
|
}
|
|
|
|
// Drain unsubscribes from the stream and cancels subscription. All
|
|
// messages that are already in the buffer will be available on
|
|
// subsequent calls to Next. After the buffer is drained, Next will
|
|
// return ErrMsgIteratorClosed error.
|
|
func (s *pullSubscription) Drain() {
|
|
if !s.closed.CompareAndSwap(0, 1) {
|
|
return
|
|
}
|
|
s.draining.Store(1)
|
|
close(s.done)
|
|
if s.consumeOpts.stopAfterMsgsLeft != nil {
|
|
if s.delivered >= s.consumeOpts.StopAfter {
|
|
close(s.consumeOpts.stopAfterMsgsLeft)
|
|
} else {
|
|
s.consumeOpts.stopAfterMsgsLeft <- s.consumeOpts.StopAfter - s.delivered
|
|
}
|
|
}
|
|
}
|
|
|
|
// Closed returns a channel that is closed when consuming is
|
|
// fully stopped/drained. When the channel is closed, no more messages
|
|
// will be received and processing is complete.
|
|
func (s *pullSubscription) Closed() <-chan struct{} {
|
|
s.Lock()
|
|
defer s.Unlock()
|
|
closedCh := s.closedCh
|
|
if closedCh == nil {
|
|
closedCh = make(chan struct{})
|
|
s.closedCh = closedCh
|
|
}
|
|
if !s.subscription.IsValid() {
|
|
close(s.closedCh)
|
|
s.closedCh = nil
|
|
}
|
|
return closedCh
|
|
}
|
|
|
|
// Fetch sends a single request to retrieve given number of messages.
|
|
// It will wait up to provided expiry time if not all messages are available.
|
|
func (p *pullConsumer) Fetch(batch int, opts ...FetchOpt) (MessageBatch, error) {
|
|
req := &pullRequest{
|
|
Batch: batch,
|
|
Expires: DefaultExpires,
|
|
Heartbeat: unset,
|
|
}
|
|
for _, opt := range opts {
|
|
if err := opt(req); err != nil {
|
|
return nil, err
|
|
}
|
|
}
|
|
|
|
if req.ctx != nil && req.maxWaitSet {
|
|
return nil, fmt.Errorf("%w: cannot specify both FetchContext and FetchMaxWait", ErrInvalidOption)
|
|
}
|
|
|
|
// if heartbeat was not explicitly set, set it to 5 seconds for longer pulls
|
|
// and disable it for shorter pulls
|
|
if req.Heartbeat == unset {
|
|
if req.Expires >= 10*time.Second {
|
|
req.Heartbeat = 5 * time.Second
|
|
} else {
|
|
req.Heartbeat = 0
|
|
}
|
|
}
|
|
if req.Expires < 2*req.Heartbeat {
|
|
return nil, fmt.Errorf("%w: expiry time should be at least 2 times the heartbeat", ErrInvalidOption)
|
|
}
|
|
|
|
return p.fetch(req)
|
|
}
|
|
|
|
// FetchBytes is used to retrieve up to a provided bytes from the stream.
|
|
func (p *pullConsumer) FetchBytes(maxBytes int, opts ...FetchOpt) (MessageBatch, error) {
|
|
req := &pullRequest{
|
|
Batch: defaultBatchMaxBytesOnly,
|
|
MaxBytes: maxBytes,
|
|
Expires: DefaultExpires,
|
|
Heartbeat: unset,
|
|
}
|
|
for _, opt := range opts {
|
|
if err := opt(req); err != nil {
|
|
return nil, err
|
|
}
|
|
}
|
|
|
|
if req.ctx != nil && req.maxWaitSet {
|
|
return nil, fmt.Errorf("%w: cannot specify both FetchContext and FetchMaxWait", ErrInvalidOption)
|
|
}
|
|
|
|
// if heartbeat was not explicitly set, set it to 5 seconds for longer pulls
|
|
// and disable it for shorter pulls
|
|
if req.Heartbeat == unset {
|
|
if req.Expires >= 10*time.Second {
|
|
req.Heartbeat = 5 * time.Second
|
|
} else {
|
|
req.Heartbeat = 0
|
|
}
|
|
}
|
|
if req.Expires < 2*req.Heartbeat {
|
|
return nil, fmt.Errorf("%w: expiry time should be at least 2 times the heartbeat", ErrInvalidOption)
|
|
}
|
|
|
|
return p.fetch(req)
|
|
}
|
|
|
|
// FetchNoWait sends a single request to retrieve given number of messages.
|
|
// FetchNoWait will only return messages that are available at the time of the
|
|
// request. It will not wait for more messages to arrive.
|
|
func (p *pullConsumer) FetchNoWait(batch int) (MessageBatch, error) {
|
|
req := &pullRequest{
|
|
Batch: batch,
|
|
NoWait: true,
|
|
}
|
|
|
|
return p.fetch(req)
|
|
}
|
|
|
|
func (p *pullConsumer) fetch(req *pullRequest) (MessageBatch, error) {
|
|
res := &fetchResult{
|
|
msgs: make(chan Msg, req.Batch),
|
|
}
|
|
msgs := make(chan *nats.Msg, 2*req.Batch)
|
|
subject := p.js.apiSubject(fmt.Sprintf(apiRequestNextT, p.stream, p.name))
|
|
|
|
sub := &pullSubscription{
|
|
consumer: p,
|
|
done: make(chan struct{}, 1),
|
|
msgs: msgs,
|
|
errs: make(chan error, 10),
|
|
}
|
|
inbox := p.js.conn.NewInbox()
|
|
var err error
|
|
sub.subscription, err = p.js.conn.ChanSubscribe(inbox, sub.msgs)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
req.PinID = p.getPinID()
|
|
if err := sub.pull(req, subject); err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
var receivedMsgs, receivedBytes int
|
|
hbTimer := sub.scheduleHeartbeatCheck(req.Heartbeat)
|
|
|
|
// Use context if provided
|
|
var ctxDone <-chan struct{}
|
|
if req.ctx != nil {
|
|
ctxDone = req.ctx.Done()
|
|
}
|
|
|
|
go func(res *fetchResult) {
|
|
defer sub.subscription.Unsubscribe()
|
|
defer close(res.msgs)
|
|
for {
|
|
select {
|
|
case msg := <-msgs:
|
|
res.Lock()
|
|
if hbTimer != nil {
|
|
hbTimer.Reset(2 * req.Heartbeat)
|
|
}
|
|
userMsg, err := checkMsg(msg)
|
|
if err != nil {
|
|
errNotTimeoutOrNoMsgs := !errors.Is(err, nats.ErrTimeout) && !errors.Is(err, ErrNoMessages)
|
|
if errNotTimeoutOrNoMsgs && !errors.Is(err, ErrMaxBytesExceeded) {
|
|
res.err = err
|
|
}
|
|
if errors.Is(err, ErrPinIDMismatch) {
|
|
p.setPinID("")
|
|
}
|
|
res.done = true
|
|
res.Unlock()
|
|
return
|
|
}
|
|
if !userMsg {
|
|
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 {
|
|
res.err = fmt.Errorf("parsing message metadata: %s", err)
|
|
res.done = true
|
|
res.Unlock()
|
|
return
|
|
}
|
|
res.sseq = meta.Sequence.Stream
|
|
receivedMsgs++
|
|
if req.MaxBytes != 0 {
|
|
receivedBytes += msg.Size()
|
|
}
|
|
if receivedMsgs == req.Batch || (req.MaxBytes != 0 && receivedBytes >= req.MaxBytes) {
|
|
res.done = true
|
|
res.Unlock()
|
|
return
|
|
}
|
|
res.Unlock()
|
|
case err := <-sub.errs:
|
|
res.Lock()
|
|
res.err = err
|
|
res.done = true
|
|
res.Unlock()
|
|
return
|
|
case <-time.After(req.Expires + 1*time.Second):
|
|
res.Lock()
|
|
res.done = true
|
|
res.Unlock()
|
|
return
|
|
case <-ctxDone:
|
|
res.Lock()
|
|
res.err = req.ctx.Err()
|
|
res.done = true
|
|
res.Unlock()
|
|
return
|
|
}
|
|
}
|
|
}(res)
|
|
return res, nil
|
|
}
|
|
|
|
func (fr *fetchResult) Messages() <-chan Msg {
|
|
fr.Lock()
|
|
defer fr.Unlock()
|
|
return fr.msgs
|
|
}
|
|
|
|
func (fr *fetchResult) Error() error {
|
|
fr.Lock()
|
|
defer fr.Unlock()
|
|
return fr.err
|
|
}
|
|
|
|
func (fr *fetchResult) closed() bool {
|
|
fr.Lock()
|
|
defer fr.Unlock()
|
|
return fr.done
|
|
}
|
|
|
|
// Next is used to retrieve the next message from the stream. This
|
|
// method will block until the message is retrieved or timeout is
|
|
// reached.
|
|
func (p *pullConsumer) Next(opts ...FetchOpt) (Msg, error) {
|
|
res, err := p.Fetch(1, opts...)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
msg := <-res.Messages()
|
|
if msg != nil {
|
|
return msg, nil
|
|
}
|
|
if res.Error() == nil {
|
|
return nil, nats.ErrTimeout
|
|
}
|
|
return nil, res.Error()
|
|
}
|
|
|
|
func (s *pullSubscription) pullMessages(subject string) {
|
|
for {
|
|
select {
|
|
case req := <-s.fetchNext:
|
|
s.fetchInProgress.Store(1)
|
|
|
|
if err := s.pull(req, subject); err != nil {
|
|
if errors.Is(err, ErrMsgIteratorClosed) {
|
|
s.cleanup()
|
|
return
|
|
}
|
|
s.errs <- err
|
|
}
|
|
s.fetchInProgress.Store(0)
|
|
case <-s.done:
|
|
s.cleanup()
|
|
return
|
|
}
|
|
}
|
|
}
|
|
|
|
func (s *pullSubscription) closeMsgs() {
|
|
if s.msgsClosed.CompareAndSwap(0, 1) {
|
|
close(s.msgs)
|
|
}
|
|
}
|
|
|
|
func (s *pullSubscription) scheduleHeartbeatCheck(dur time.Duration) *hbMonitor {
|
|
if dur == 0 {
|
|
return nil
|
|
}
|
|
return &hbMonitor{
|
|
timer: time.AfterFunc(2*dur, func() {
|
|
s.errs <- ErrNoHeartbeat
|
|
}),
|
|
}
|
|
}
|
|
|
|
func (s *pullSubscription) cleanup() {
|
|
// For now this function does not need to hold the lock.
|
|
// Holding the lock here might cause a deadlock if Next()
|
|
// is already holding the lock and waiting.
|
|
// The fields that are read (subscription, hbMonitor)
|
|
// are read only (Only written on creation of pullSubscription).
|
|
if s.subscription == nil || !s.subscription.IsValid() {
|
|
return
|
|
}
|
|
if s.consumer != nil {
|
|
nc := s.consumer.js.conn
|
|
nc.RemoveStatusListener(s.connStatusChanged)
|
|
}
|
|
if s.hbMonitor != nil {
|
|
s.hbMonitor.Stop()
|
|
}
|
|
drainMode := s.draining.Load() == 1
|
|
if drainMode {
|
|
s.subscription.Drain()
|
|
} else {
|
|
s.subscription.Unsubscribe()
|
|
}
|
|
s.closed.Store(1)
|
|
}
|
|
|
|
// pull sends a pull request to the server and waits for messages using a subscription from [pullSubscription].
|
|
// Messages will be fetched up to given batch_size or until there are no more messages or timeout is returned
|
|
func (s *pullSubscription) pull(req *pullRequest, subject string) error {
|
|
s.consumer.Lock()
|
|
defer s.consumer.Unlock()
|
|
if s.closed.Load() == 1 {
|
|
return ErrMsgIteratorClosed
|
|
}
|
|
if req.Batch < 1 {
|
|
return fmt.Errorf("%w: batch size must be at least 1", nats.ErrInvalidArg)
|
|
}
|
|
reqJSON, err := json.Marshal(req)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
reply := s.subscription.Subject
|
|
|
|
if err := s.consumer.js.conn.PublishRequest(subject, reply, reqJSON); err != nil {
|
|
return err
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func parseConsumeOpts(ordered bool, opts ...PullConsumeOpt) (*consumeOpts, error) {
|
|
consumeOpts := &consumeOpts{
|
|
MaxMessages: unset,
|
|
MaxBytes: unset,
|
|
Expires: DefaultExpires,
|
|
Heartbeat: unset,
|
|
ReportMissingHeartbeats: true,
|
|
StopAfter: unset,
|
|
}
|
|
for _, opt := range opts {
|
|
if err := opt.configureConsume(consumeOpts); err != nil {
|
|
return nil, err
|
|
}
|
|
}
|
|
if err := consumeOpts.setDefaults(ordered); err != nil {
|
|
return nil, err
|
|
}
|
|
return consumeOpts, nil
|
|
}
|
|
|
|
func parseMessagesOpts(ordered bool, opts ...PullMessagesOpt) (*consumeOpts, error) {
|
|
consumeOpts := &consumeOpts{
|
|
MaxMessages: unset,
|
|
MaxBytes: unset,
|
|
Expires: DefaultExpires,
|
|
Heartbeat: unset,
|
|
ReportMissingHeartbeats: true,
|
|
StopAfter: unset,
|
|
}
|
|
for _, opt := range opts {
|
|
if err := opt.configureMessages(consumeOpts); err != nil {
|
|
return nil, err
|
|
}
|
|
}
|
|
if err := consumeOpts.setDefaults(ordered); err != nil {
|
|
return nil, err
|
|
}
|
|
return consumeOpts, nil
|
|
}
|
|
|
|
func (consumeOpts *consumeOpts) setDefaults(ordered bool) error {
|
|
// we cannot use both max messages and max bytes unless we're using max bytes as fetch size limiter
|
|
if consumeOpts.MaxBytes != unset && consumeOpts.MaxMessages != unset && !consumeOpts.LimitSize {
|
|
return errors.New("only one of MaxMessages and MaxBytes can be specified")
|
|
}
|
|
if consumeOpts.MaxBytes != unset && !consumeOpts.LimitSize {
|
|
// we used PullMaxBytes setting, set MaxMessages to a high value
|
|
consumeOpts.MaxMessages = defaultBatchMaxBytesOnly
|
|
} else if consumeOpts.MaxMessages == unset {
|
|
// otherwise, if max messages is not set, set it to default value
|
|
consumeOpts.MaxMessages = DefaultMaxMessages
|
|
}
|
|
// if user did not set max bytes, set it to 0
|
|
if consumeOpts.MaxBytes == unset {
|
|
consumeOpts.MaxBytes = 0
|
|
}
|
|
|
|
if consumeOpts.ThresholdMessages == 0 {
|
|
// half of the max messages, rounded up
|
|
consumeOpts.ThresholdMessages = int(math.Ceil(float64(consumeOpts.MaxMessages) / 2))
|
|
}
|
|
if consumeOpts.ThresholdBytes == 0 {
|
|
// half of the max bytes, rounded up
|
|
consumeOpts.ThresholdBytes = int(math.Ceil(float64(consumeOpts.MaxBytes) / 2))
|
|
}
|
|
|
|
// set default heartbeats
|
|
if consumeOpts.Heartbeat == unset {
|
|
// by default, use 50% of expiry time
|
|
consumeOpts.Heartbeat = consumeOpts.Expires / 2
|
|
if ordered {
|
|
// for ordered consumers, the default heartbeat is 5 seconds
|
|
if consumeOpts.Expires < 10*time.Second {
|
|
consumeOpts.Heartbeat = consumeOpts.Expires / 2
|
|
} else {
|
|
consumeOpts.Heartbeat = 5 * time.Second
|
|
}
|
|
} else if consumeOpts.Heartbeat > 30*time.Second {
|
|
// cap the heartbeat to 30 seconds
|
|
consumeOpts.Heartbeat = 30 * time.Second
|
|
}
|
|
}
|
|
if consumeOpts.Heartbeat > consumeOpts.Expires/2 {
|
|
return fmt.Errorf("%w: the value of Heartbeat must be less than 50%% of expiry", ErrInvalidOption)
|
|
}
|
|
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
|
|
}
|