snapshot of https://github.com/centrifugal/centrifugo.git (master, 2026-09-16; blobs > 131072B omitted)

internal/consuming/nats_jetstream.go

@@ -0,0 +1,403 @@
+package consuming
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"strconv"
+	"strings"
+	"time"
+
+	"github.com/centrifugal/centrifugo/v6/internal/api"
+	"github.com/centrifugal/centrifugo/v6/internal/configtypes"
+	"github.com/centrifugal/centrifugo/v6/internal/logging"
+	"github.com/centrifugal/centrifugo/v6/internal/metrics"
+
+	"github.com/nats-io/nats.go"
+	"github.com/nats-io/nats.go/jetstream"
+)
+
+// NatsJetStreamConsumerConfig is an alias for our configuration type.
+type NatsJetStreamConsumerConfig = configtypes.NatsJetStreamConsumerConfig
+
+// natsConsumeBackoffReset is the minimum duration a consume session must run
+// before its end resets the recreation backoff, so rapid flapping still backs off.
+const natsConsumeBackoffReset = 30 * time.Second
+
+// NatsJetStreamConsumer consumes messages from NATS JetStream.
+type NatsJetStreamConsumer struct {
+	config         NatsJetStreamConsumerConfig
+	dispatcher     Dispatcher
+	nc             *nats.Conn
+	consumer       jetstream.Consumer
+	ctx            context.Context
+	common         *consumerCommon
+	consumeContext jetstream.ConsumeContext
+	// recreateCh signals that the consumer should be re-created (e.g. on heartbeat loss).
+	recreateCh chan struct{}
+}
+
+// NewNatsJetStreamConsumer creates a new NatsJetStreamConsumer instance.
+// On startup if creating a consumer fails then an error is returned.
+// Later at runtime, if a heartbeat error occurs or the consumer/stream is
+// deleted on the server side the consumer is re-created endlessly.
+func NewNatsJetStreamConsumer(
+	cfg NatsJetStreamConsumerConfig,
+	dispatcher Dispatcher,
+	common *consumerCommon,
+) (*NatsJetStreamConsumer, error) {
+	if err := cfg.Validate(); err != nil {
+		return nil, err
+	}
+
+	c := &NatsJetStreamConsumer{
+		config:     cfg,
+		dispatcher: dispatcher,
+		common:     common,
+		recreateCh: make(chan struct{}, 1), // buffered so duplicate signals are dropped
+	}
+
+	opts := []nats.Option{
+		nats.MaxReconnects(-1),
+		nats.ConnectHandler(func(conn *nats.Conn) {
+			c.common.log.Info().Msg("connected")
+		}),
+		nats.ReconnectHandler(func(conn *nats.Conn) {
+			c.common.log.Info().Msg("reconnected")
+		}),
+		nats.DisconnectErrHandler(func(conn *nats.Conn, err error) {
+			c.common.log.Warn().Err(err).Msg("disconnected")
+		}),
+	}
+
+	switch {
+	case cfg.CredentialsFile != "":
+		opts = append(opts, nats.UserCredentials(cfg.CredentialsFile))
+	case cfg.Username != "":
+		opts = append(opts, nats.UserInfo(cfg.Username, cfg.Password))
+	case cfg.Token != "":
+		opts = append(opts, nats.Token(cfg.Token))
+	}
+
+	if cfg.TLS.Enabled {
+		tlsConfig, err := cfg.TLS.ToGoTLSConfig("nats_jetstream")
+		if err != nil {
+			return nil, fmt.Errorf("failed to create TLS config: %w", err)
+		}
+		opts = append(opts, nats.Secure(tlsConfig))
+	}
+
+	nc, err := nats.Connect(cfg.URL, opts...)
+	if err != nil {
+		return nil, fmt.Errorf("failed to connect to NATS: %w", err)
+	}
+	c.nc = nc
+
+	createCtx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	defer cancel()
+	jsConsumer, err := createJetStreamConsumer(createCtx, nc, cfg)
+	if err != nil {
+		nc.Close()
+		return nil, fmt.Errorf("failed to create JetStream consumer: %w", err)
+	}
+	c.consumer = jsConsumer
+	return c, nil
+}
+
+// createJetStreamConsumer is a helper (used both at startup and runtime)
+// to create a new JetStream consumer based on the configuration.
+func createJetStreamConsumer(ctx context.Context, nc *nats.Conn, cfg NatsJetStreamConsumerConfig) (jetstream.Consumer, error) {
+	js, err := jetstream.New(nc)
+	if err != nil {
+		return nil, fmt.Errorf("failed to create JetStream context: %w", err)
+	}
+
+	if cfg.UseExistingConsumer {
+		// If UseExistingConsumer is enabled, just look up the existing consumer by DurableConsumerName
+		consumer, err := js.Consumer(ctx, cfg.StreamName, cfg.DurableConsumerName)
+		if err != nil {
+			return nil, fmt.Errorf("failed to lookup existing consumer %q: %w", cfg.DurableConsumerName, err)
+		}
+		return consumer, nil
+	}
+
+	// By default, create or update the consumer
+	deliverPolicy := jetstream.DeliverNewPolicy
+	if cfg.DeliverPolicy == "all" {
+		deliverPolicy = jetstream.DeliverAllPolicy
+	}
+	return js.CreateOrUpdateConsumer(ctx, cfg.StreamName, jetstream.ConsumerConfig{
+		Name:           cfg.DurableConsumerName,
+		Durable:        cfg.DurableConsumerName,
+		FilterSubjects: cfg.Subjects,
+		DeliverPolicy:  deliverPolicy,
+		AckWait:        30 * time.Second,
+		MaxAckPending:  cfg.MaxAckPending,
+	})
+}
+
+// createJetStreamConsumer creates a new consumer using the instance's connection and config.
+// This method is used for runtime re-creation.
+func (c *NatsJetStreamConsumer) createJetStreamConsumer(ctx context.Context) (jetstream.Consumer, error) {
+	ctx, cancel := context.WithTimeout(ctx, 10*time.Second)
+	defer cancel()
+	return createJetStreamConsumer(ctx, c.nc, c.config)
+}
+
+// startConsume calls Consume on the current consumer.
+func (c *NatsJetStreamConsumer) startConsume(_ context.Context) error {
+	consumeContext, err := c.consumer.Consume(
+		c.msgHandler,
+		jetstream.ConsumeErrHandler(c.errorHandler()),
+		jetstream.PullHeartbeat(5*time.Second),
+	)
+	if err != nil {
+		return err
+	}
+	c.consumeContext = consumeContext
+	return nil
+}
+
+// triggerRecreation signals that the consumer should be re-created.
+func (c *NatsJetStreamConsumer) triggerRecreation() {
+	select {
+	case c.recreateCh <- struct{}{}:
+	default:
+		// Already signaled, do nothing.
+	}
+}
+
+// msgHandler is the callback for incoming JetStream messages.
+func (c *NatsJetStreamConsumer) msgHandler(msg jetstream.Msg) {
+	if logging.Enabled(logging.DebugLevel) {
+		c.common.log.Debug().Str("subject", msg.Subject()).Msg("received message from subject")
+	}
+
+	var processErr error
+	if c.config.PublicationDataMode.Enabled {
+		processErr = c.processPublicationDataMessage(msg)
+	} else {
+		processErr = c.processCommandMessage(msg)
+	}
+
+	if processErr == nil {
+		if err := msg.Ack(); err != nil {
+			c.common.log.Error().Err(err).Msg("failed to ack message")
+			metrics.ConsumerErrorsTotal.WithLabelValues(c.common.name).Inc()
+		} else {
+			metrics.ConsumerProcessedTotal.WithLabelValues(c.common.name).Inc()
+		}
+	} else {
+		c.common.log.Error().Err(processErr).Msg("processing message failed")
+		// NakWithDelay, not Nak: a bare Nak triggers instant redelivery with no
+		// backoff and the consumer sets no MaxDeliver, so a persistently failing
+		// (poison) message or a downstream outage would spin in a tight
+		// redeliver->fail->redeliver loop, pinning CPU and blocking head-of-line
+		// progress. Delay grows with the delivery count, capped like the other
+		// consumers' backoff.
+		if err := msg.NakWithDelay(natsNakDelay(msg)); err != nil {
+			c.common.log.Error().Err(err).Msg("failed to nak message")
+		}
+	}
+}
+
+// natsNakDelay computes an exponential redelivery backoff (capped) from the
+// message's delivery count.
+func natsNakDelay(msg jetstream.Msg) time.Duration {
+	retries := 1
+	if md, err := msg.Metadata(); err == nil && md.NumDelivered > 0 {
+		retries = int(md.NumDelivered)
+		if retries > 10 { // cap to avoid 1<<retries overflow / overlong delays
+			retries = 10
+		}
+	}
+	return getNextBackoffDuration(0, retries)
+}
+
+// processPublicationDataMessage processes a message in publication data mode.
+func (c *NatsJetStreamConsumer) processPublicationDataMessage(msg jetstream.Msg) error {
+	idempotencyKey := getNatsHeaderValue(msg, c.config.PublicationDataMode.IdempotencyKeyHeader)
+	delta, err := getNatsBoolHeaderValue(msg, c.config.PublicationDataMode.DeltaHeader)
+	if err != nil {
+		c.common.log.Error().Err(err).Msg("error parsing delta header, skipping message")
+		return nil
+	}
+	channels := strings.Split(getNatsHeaderValue(msg, c.config.PublicationDataMode.ChannelsHeader), ",")
+	if len(channels) == 0 || (len(channels) == 1 && channels[0] == "") {
+		c.common.log.Info().Msg("no channels found, skipping message")
+		return nil
+	}
+	tags := publicationTagsFromNatsHeaders(msg, c.config.PublicationDataMode.TagsHeaderPrefix)
+	version, err := getNatsUint64HeaderValue(msg, c.config.PublicationDataMode.VersionHeader)
+	if err != nil {
+		c.common.log.Error().Err(err).Msg("error parsing version header, skipping message")
+		return nil
+	}
+	return c.dispatcher.DispatchPublication(c.ctx, channels, api.ConsumedPublication{
+		Data:           msg.Data(),
+		IdempotencyKey: idempotencyKey,
+		Delta:          delta,
+		Tags:           tags,
+		Version:        version,
+		VersionEpoch:   getNatsHeaderValue(msg, c.config.PublicationDataMode.VersionEpochHeader),
+	})
+}
+
+// processCommandMessage processes a command message.
+func (c *NatsJetStreamConsumer) processCommandMessage(msg jetstream.Msg) error {
+	method := getNatsHeaderValue(msg, c.config.MethodHeader)
+	return c.dispatcher.DispatchCommand(c.ctx, method, msg.Data())
+}
+
+func getNatsUint64HeaderValue(msg jetstream.Msg, key string) (uint64, error) {
+	if key == "" {
+		return 0, nil
+	}
+	val := msg.Headers().Get(key)
+	if val == "" {
+		return 0, nil
+	}
+	i, err := strconv.ParseUint(val, 10, 64)
+	if err != nil {
+		return 0, fmt.Errorf("error parsing header %q: %w", key, err)
+	}
+	return i, nil
+}
+
+func getNatsBoolHeaderValue(msg jetstream.Msg, key string) (bool, error) {
+	if key == "" {
+		return false, nil
+	}
+	val := msg.Headers().Get(key)
+	if val == "" {
+		return false, nil
+	}
+	b, err := strconv.ParseBool(val)
+	if err != nil {
+		return false, fmt.Errorf("error parsing header %q: %w", key, err)
+	}
+	return b, nil
+}
+
+// getNatsHeaderValue retrieves a header value from the NATS message.
+func getNatsHeaderValue(msg jetstream.Msg, key string) string {
+	if key == "" {
+		return ""
+	}
+	return msg.Headers().Get(key)
+}
+
+// publicationTagsFromNatsHeaders extracts tags from message headers using the given prefix.
+func publicationTagsFromNatsHeaders(msg jetstream.Msg, prefix string) map[string]string {
+	var tags map[string]string
+	if prefix == "" {
+		return tags
+	}
+	for key, vals := range msg.Headers() {
+		if strings.HasPrefix(key, prefix) && len(vals) > 0 {
+			if tags == nil {
+				tags = make(map[string]string)
+			}
+			tags[strings.TrimPrefix(key, prefix)] = vals[0]
+		}
+	}
+	return tags
+}
+
+// errorHandler returns a jetstream error handler that triggers recreation on heartbeat loss
+// or when the consumer (or its stream) was deleted on the server side. The latter happens
+// when JetStream state is lost externally – ex. stream removed and re-created by an operator,
+// or in-memory stream lost after NATS restart. Without recreation Centrifugo would keep
+// logging errors without consuming anything until the process is restarted.
+func (c *NatsJetStreamConsumer) errorHandler() jetstream.ConsumeErrHandlerFunc {
+	return func(consumeCtx jetstream.ConsumeContext, err error) {
+		if errors.Is(err, jetstream.ErrNoHeartbeat) ||
+			errors.Is(err, jetstream.ErrConsumerDeleted) ||
+			errors.Is(err, jetstream.ErrConsumerNotFound) ||
+			errors.Is(err, jetstream.ErrStreamNotFound) {
+			c.common.log.Warn().Err(err).Msg("consumer unavailable, triggering consumer recreation")
+			c.triggerRecreation()
+			return
+		}
+		c.common.log.Error().Err(err).Msg("error during consuming")
+	}
+}
+
+// Run starts the consumer and blocks until the context is canceled.
+// If a heartbeat error occurs then the consumer is re-created
+// and the consumeContext re-established with endless retries.
+func (c *NatsJetStreamConsumer) Run(ctx context.Context) error {
+	c.ctx = ctx
+	defer func() {
+		_ = c.nc.Drain()
+	}()
+
+	retries := 0
+	var backoffDuration time.Duration
+
+	firstRun := true
+	for {
+		// For subsequent runs, attempt to recreate the consumer.
+		if !firstRun {
+			for {
+				newConsumer, err := c.createJetStreamConsumer(ctx)
+				if err == nil {
+					c.consumer = newConsumer
+					c.common.log.Info().Msg("consumer recreation succeeded")
+					break
+				}
+				backoffDuration = getNextBackoffDuration(backoffDuration, retries)
+				retries++
+				c.common.log.Error().Err(err).Str("delay", backoffDuration.String()).
+					Msg("failed to recreate consumer, retrying")
+				select {
+				case <-ctx.Done():
+					return ctx.Err()
+				case <-time.After(backoffDuration):
+				}
+			}
+		}
+
+		// Start the ConsumeContext.
+		if err := c.startConsume(ctx); err != nil {
+			if firstRun {
+				// On initial startup, return error immediately.
+				return fmt.Errorf("error on start consuming: %w", err)
+			}
+			backoffDuration = getNextBackoffDuration(backoffDuration, retries)
+			retries++
+			c.common.log.Error().Err(err).Str("delay", backoffDuration.String()).
+				Msg("failed to start consumer consumeContext, retrying")
+			select {
+			case <-ctx.Done():
+				return ctx.Err()
+			case <-time.After(backoffDuration):
+			}
+			continue
+		}
+
+		firstRun = false
+		startedAt := time.Now()
+
+		c.common.log.Info().Msg("consuming started, waiting for messages")
+		// Block until context cancellation or a recreation signal is received.
+		select {
+		case <-ctx.Done():
+			c.consumeContext.Stop()
+			return ctx.Err()
+		case <-c.recreateCh:
+			c.common.log.Info().Msg("recreating consumer")
+			c.consumeContext.Stop()
+			// Recreate the consumer.
+		}
+
+		// Reset the backoff ramp only if the session ran healthily for a while. A
+		// session that starts fine but dies almost immediately (flapping) keeps the
+		// ramp so the recreation loop backs off instead of spinning, rather than
+		// resetting to zero on every short-lived start.
+		if time.Since(startedAt) > natsConsumeBackoffReset {
+			retries = 0
+			backoffDuration = 0
+		}
+	}
+}