coderd/nats/pubsub.go
Rama main · réplica de lectura del repositorio original.
package nats
import (
"context"
"errors"
"fmt"
"hash/fnv"
"net"
"net/url"
"sync"
"time"
natsserver "github.com/nats-io/nats-server/v2/server"
natsgo "github.com/nats-io/nats.go"
"golang.org/x/xerrors"
"cdr.dev/slog/v3"
"github.com/coder/coder/v2/coderd/cryptokeys"
"github.com/coder/coder/v2/coderd/database/pubsub"
"github.com/coder/quartz"
)
// DefaultServerMaxPendingBytes caps how many bytes the embedded NATS server will
// hold in memory for a single client connection while waiting to write
// them out to that connection. Each message the server needs to deliver
// to a connection is queued in that connection's outbound buffer until
// the socket can accept it; if the consumer reads slower than messages
// arrive, the buffer grows. When it exceeds this cap, NATS declares the
// connection a slow consumer and drops messages rather than buffer
// without bound.
//
// The connection that fills this buffer in practice is the subscribe
// connection: the server writes every message bound for a replica's
// local subscribers out over its subscribe connection pool, which
// defaults to a single connection. In a cluster, cross-node fan-out
// therefore concentrates all of a replica's inbound deliveries on that
// one connection's outbound buffer. Benchmarking high-fanout cluster
// workloads (10 subjects, 10 publishers, 50 subscribers) showed the 128
// MiB default overflowing and dropping 10-15% of deliveries, while 256
// MiB and above dropped none; 512 MiB is chosen for headroom.
//
// This is a ceiling, not a reservation: the buffer grows only with
// actual backlog, so a connection that keeps up holds nearly nothing and
// raising the cap costs memory only during the overload bursts it
// absorbs.
const DefaultServerMaxPendingBytes int64 = 512 << 20
// DefaultClientMaxPendingBytes is the pending byte limit applied via
// SetPendingLimits to each coalesced *natsgo.Subscription when
// PendingLimits.Bytes is zero. Unlike DefaultServerMaxPendingBytes,
// which bounds the server-side outbound buffer for a whole connection,
// this bounds the nats.go client's per-subscription pending buffer:
// messages the client has received from the server but the
// subscription's async message handler has not yet dispatched. When
// that handler falls behind and the buffer overflows, NATS marks the
// subscription a slow consumer and drops messages, which the package
// surfaces as pubsub.ErrDroppedMessages.
const DefaultClientMaxPendingBytes = 512 * 1024 * 1024
const (
defaultClusterName = "coder"
defaultClusterPort = 6222
defaultRoutePoolSize = 3
)
var errClosed = xerrors.New("nats pubsub closed")
// PendingLimits configures per-subscription NATS pending limits set
// via SetPendingLimits on each *natsgo.Subscription.
type PendingLimits struct {
// Msgs is the per-subscription pending message limit. Zero or
// negative disables the message limit, leaving the byte limit
// (PendingLimits.Bytes) as the only per-subscription bound.
Msgs int
// Bytes is the per-subscription pending byte limit.
// Zero uses the package default. Negative disables this limit.
Bytes int
}
// Options configures the embedded NATS Pubsub.
type Options struct {
// MaxPayload is the NATS max payload. Zero means server default.
MaxPayload int32
// MaxPending is the per-client outbound pending byte budget on the
// embedded server. Zero or negative means the package default,
// 128 MiB.
MaxPending int64
// PendingLimits configures per-subscription NATS pending limits.
// Positive Msgs also sets local listener queue capacity.
// Zero fields use package defaults: Msgs -1 and Bytes 512 MiB.
PendingLimits PendingLimits
// ReconnectWait controls client reconnect delay. Zero keeps the
// NATS default.
ReconnectWait time.Duration
// InProcess, when true, uses nats.InProcessServer instead of TCP
// loopback. Intended for benchmarks and tests.
InProcess bool
// PublishConns is the number of publisher connections. Each Publish
// is routed by a stable hash of the subject. Zero or negative means 1.
PublishConns int
// SubscribeConns is the number of subscriber connections. Each
// shared subscription is pinned to one connection by a stable hash
// of its subject. Zero or negative means 1.
SubscribeConns int
// ClusterHost is the embedded NATS route listener host. Empty means
// all interfaces when cluster mode is enabled.
ClusterHost string
// ClusterPort is the embedded NATS route listener port. Zero means
// 6222 when cluster mode is enabled. NATS `server.RANDOM_PORT` can be
// used to select a random port.
ClusterPort int
// ClusterAuthToken is the shared route authentication token for
// clustered embedded NATS servers. Empty disables route auth.
ClusterAuthToken string
// ClusterCA enables mutual TLS on the cluster route listener. When set
// (and cluster mode is enabled), each replica mints an ephemeral leaf
// certificate from the active nats_ca CA and verifies peers against the
// CA fetched from this cache on each handshake. Nil keeps routes
// plaintext (token auth only). cryptokeys.SigningKeycache satisfies this.
//
// The leaf's IP SAN (and the accept-side source binding) is this replica's
// ClusterHost, so ClusterHost must be an IP for mTLS to activate.
ClusterCA cryptokeys.SigningKeycache
// clock overrides the cluster TLS clock, for tests.
clock quartz.Clock
// clusterTLSTimeout overrides the cluster route TLS handshake timeout, for
// tests. Zero leaves the NATS default (2s). Tests use a longer timeout
// because handshakes are flaky under load and in CI.
clusterTLSTimeout time.Duration
// PeerFetcher provides the current set of peer route addresses.
// RefreshPeers uses it to update the configured cluster routes.
PeerFetcher PeerFetcher
// Metrics is the shared pubsub metrics. When nil, an unregistered set is
// created internally so the Pubsub still works without exporting metrics.
// The background latency measurement loop always runs (recording into no-op
// instruments when Metrics is nil) so the loop lifecycle does not depend on
// whether metrics are collected.
Metrics *pubsub.Metrics
// RoutePoolSize is the NATS route pool size. Zero means the package
// default when cluster mode is enabled.
RoutePoolSize int
// disableCluster is intended only for testing. Since we cannot reload a server
// with a cluster host/port after initialization, we start all production servers
// with clustering enabled.
disableCluster bool
}
// conn is a stripped down version of natsgo.Conn with just the methods we use, to allow us to fake it in tests.
type conn interface {
Publish(event string, message []byte) error
Close()
Flush() error
Subscribe(event string, handler natsgo.MsgHandler) (*natsgo.Subscription, error)
}
// Pubsub is an embedded NATS-backed implementation of pubsub.Pubsub.
//
// Each Pubsub owns one embedded server, a pool of publisher
// *natsgo.Conns (Options.PublishConns) and a pool of subscriber
// *natsgo.Conns (Options.SubscribeConns). Publishes and shared
// subscriptions are pinned to a connection by a stable hash of the
// subject, so same-subject traffic preserves per-subject ordering and
// every local subscriber for a subject coalesces onto one underlying
// *natsgo.Subscription.
type Pubsub struct {
mu sync.Mutex
logger slog.Logger
opts Options
Server *natsserver.Server
// publishPool and subscribePool are immutable after construction so
// the hot path can index without holding p.mu.
publishPool []conn
subscribePool []conn
// subscriptions coalesces concurrent local subscribers on the
// same subject onto a single underlying *natsgo.Subscription.
subscriptions map[string]*groupSub
closeOnce sync.Once
// ctx is canceled by Close while holding p.mu so subscriber state
// cleanup observes the canceled context.
ctx context.Context
cancel context.CancelFunc
// unsubscribeRoutines tracks outstanding unsubscribeGroup calls while closing, to ensure they all complete before
// we start tearing down connections.
unsubscribeRoutines sync.WaitGroup
clusterMu sync.Mutex
clustered bool
serverOpts *natsserver.Options
currentRoutes []*url.URL
// clusterTLS is non-nil when the cluster route listener runs mutual TLS.
clusterTLS *clusterTLS
peerFetcher PeerFetcher
peerRefresh chan struct{}
metrics *pubsub.BackendMetrics
conns *connTracker
}
// groupSub maps to one underlying *natsgo.Subscription. The first
// local subscriber creates it; later local subscribers attach to it.
// When the last local subscriber detaches, the NATS subscription is
// unsubscribed.
type groupSub struct {
// metrics records received-message metrics from handleMessage. It is
// the only part of the owning Pubsub that a groupSub needs.
metrics *pubsub.BackendMetrics
event string
// mu guards localSubs.
mu sync.Mutex
// localSubs are the local subscribers attached to this NATS subscription.
localSubs map[*localSub]struct{}
sub *subGetter
// dropMu keeps async error accounting independent from listener fan-out.
dropMu sync.Mutex
// lastDropped is the cumulative NATS dropped count last reported locally.
lastDropped uint64
}
// localSub is the local handle returned by Subscribe /
// SubscribeWithErr.
type localSub struct {
event string
queue *pubsub.MsgQueue
}
// Compile-time assertion that *Pubsub satisfies the pubsub.Pubsub interface.
var _ pubsub.Pubsub = (*Pubsub)(nil)
// newPubsub allocates a *Pubsub with initialized maps and cancel ctx. If
// opts.Metrics is non-nil it is shared; otherwise an unregistered set is used.
func newPubsub(ctx context.Context, logger slog.Logger, opts Options) *Pubsub {
ctx, cancel := context.WithCancel(ctx)
metrics := opts.Metrics
if metrics == nil {
metrics = pubsub.NewMetrics(nil)
}
m := metrics.ForBackend(logger, pubsub.BackendNATS)
return &Pubsub{
logger: logger,
opts: opts,
subscriptions: make(map[string]*groupSub),
ctx: ctx,
cancel: cancel,
peerFetcher: opts.PeerFetcher,
peerRefresh: make(chan struct{}, 1),
metrics: m,
conns: newConnTracker(m),
}
}
// defaultPendingLimits returns the effective per-subscription pending
// limits applied at Subscribe time.
func defaultPendingLimits(in PendingLimits) PendingLimits {
out := in
if out.Msgs == 0 {
out.Msgs = -1
}
if out.Bytes == 0 {
out.Bytes = DefaultClientMaxPendingBytes
}
return out
}
// buildConnHandlers returns the connHandlers stack installed on every
// owned connection. Handlers close over p so slow-consumer routing
// keeps working.
func (p *Pubsub) buildConnHandlers() connHandlers {
return connHandlers{
disconnectErr: func(conn *natsgo.Conn, err error) {
if err != nil {
p.logger.Warn(p.ctx, "nats client disconnected", slog.Error(err))
}
p.conns.onDisconnect()
p.signalSubscribersDroppedForConn(conn)
},
reconnect: func(_ *natsgo.Conn) {
p.logger.Info(p.ctx, "nats client reconnected")
p.conns.onReconnect()
},
closed: func(_ *natsgo.Conn) {
p.logger.Debug(p.ctx, "nats client closed")
},
errH: func(_ *natsgo.Conn, sub *natsgo.Subscription, err error) {
if err != nil && errors.Is(err, natsgo.ErrSlowConsumer) {
p.handleAsyncError(sub, err)
return
}
if err != nil {
p.logger.Warn(p.ctx, "nats async error", slog.Error(err))
}
},
}
}
// New creates an embedded NATS Pubsub. The returned *Pubsub owns the
// embedded server and the publisher and subscriber connection pools.
// Close shuts down all owned resources.
func New(ctx context.Context, logger slog.Logger, opts Options) (pubSub *Pubsub, retErr error) {
sopts, err := buildServerOptions(opts)
if err != nil {
return nil, err
}
// When ClusterCA is set, install the cluster TLS callbacks at boot so the
// route listener can negotiate mTLS. The callbacks read the CA cache on
// each handshake, so the default noop cache keeps routes inert (no leaf can
// be minted) until SetCACache swaps in a real cache. The leaf IP SAN is this
// replica's ClusterHost, fixed here at construction; leaf minting enforces
// that it is an IP. ClusterCA == nil keeps routes plaintext (token auth
// only).
var ct *clusterTLS
if !opts.disableCluster && opts.ClusterCA != nil {
selfIP := net.ParseIP(opts.ClusterHost)
ct = newClusterTLS(ctx, logger, opts.clock, opts.ClusterCA, selfIP)
sopts.Cluster.TLSConfig = ct.tlsConfig()
// Leave TLSTimeout unset (NATS defaults to 2s) unless a test overrides
// it; the default has not shown a need to change in production.
if opts.clusterTLSTimeout > 0 {
sopts.Cluster.TLSTimeout = opts.clusterTLSTimeout.Seconds()
}
}
ns, err := startEmbeddedServer(logger, sopts)
if err != nil {
return nil, err
}
defer func() {
if retErr != nil {
ns.Shutdown()
ns.WaitForShutdown()
}
}()
logger.Debug(context.Background(), "embedded nats server started",
slog.F("client_url", ns.ClientURL()),
)
if opts.PeerFetcher == nil {
opts.PeerFetcher = NopPeerFetcher{}
}
p := newPubsub(ctx, logger, opts)
defer func() {
if retErr != nil {
p.cancel()
}
}()
p.Server = ns
p.clustered = !opts.disableCluster
p.serverOpts = sopts.Clone()
p.currentRoutes = cloneRouteURLs(sopts.Routes)
p.clusterTLS = ct
handlers := p.buildConnHandlers()
publishPool, err := newConnPool(ns, opts, handlers, opts.PublishConns, "coder-pubsub-pub")
if err != nil {
return nil, err
}
defer func() {
if retErr != nil {
for _, c := range publishPool {
c.Close()
}
}
}()
p.publishPool = publishPool
subscribePool, err := newConnPool(ns, opts, handlers, opts.SubscribeConns, "coder-pubsub-sub")
if err != nil {
return nil, err
}
defer func() {
if retErr != nil {
for _, c := range subscribePool {
c.Close()
}
}
}()
p.subscribePool = subscribePool
// All owned connections dialed successfully above.
p.conns.markConnected(len(publishPool) + len(subscribePool))
if p.clustered {
ca := ns.ClusterAddr()
if ca == nil {
return nil, xerrors.New("no cluster address")
}
// sec checks, just to be sure
if ca.Port < 0 || ca.Port > 65535 {
return nil, xerrors.Errorf("invalid cluster port: %d", ca.Port)
}
//nolint:gosec // range checked above so conversion is safe.
opts.PeerFetcher.SetSelfNATSPort(int32(ca.Port))
go p.runPeerRefresh()
}
// Probe latency out-of-band. It records into no-op instruments when
// Prometheus is disabled, but always runs so the loop lifecycle does not
// depend on whether metrics are collected. Start it before the context
// watcher below so it is running before that watcher can trigger Close.
p.metrics.StartLatencyLoop(p)
go func() {
<-p.ctx.Done()
_ = p.Close()
}()
return p, nil
}
func newConnPool(ns *natsserver.Server, opts Options, handlers connHandlers, count int, clientName string) ([]conn, error) {
if count <= 0 {
count = 1
}
pool := make([]conn, 0, count)
for i := 0; i < count; i++ {
// Suffix names when the pool has more than one entry so server
// logs can distinguish connections.
name := clientName
if count > 1 {
name = fmt.Sprintf("%s-%d", clientName, i)
}
nc, err := connectClient(ns, opts, handlers, name)
if err != nil {
for _, c := range pool {
c.Close()
}
return nil, xerrors.Errorf("dial conn: %w", err)
}
pool = append(pool, nc)
}
return pool, nil
}
// Publish publishes a message under the given event name. The
// publisher connection is selected by a stable hash of the subject so
// same-subject publishes preserve per-subject ordering.
func (p *Pubsub) Publish(event string, message []byte) error {
if p.ctx.Err() != nil {
p.metrics.RecordPublishFailure(event)
return errClosed
}
if err := pickConn(p.publishPool, event).Publish(event, message); err != nil {
p.metrics.RecordPublishFailure(event)
return xerrors.Errorf("publish: %w", err)
}
p.metrics.RecordPublishSuccess(event, len(message))
return nil
}
// Flush blocks until every publisher connection has flushed buffered
// publishes to the embedded server. Returns the first error
// encountered; remaining connections are still flushed.
func (p *Pubsub) Flush() error {
if p.ctx.Err() != nil {
return errClosed
}
var firstErr error
for i, nc := range p.publishPool {
if err := nc.Flush(); err != nil && firstErr == nil {
firstErr = xerrors.Errorf("flush pub conn %d: %w", i, err)
}
}
return firstErr
}
// Subscribe subscribes a Listener to the given event name. Errors
// such as ErrDroppedMessages are silently ignored, mirroring the
// legacy pubsub Listener semantics.
func (p *Pubsub) Subscribe(event string, listener pubsub.Listener) (cancel func(), err error) {
return p.subscribeQueue(event, pubsub.NewMsgQueue(context.Background(), listener, nil))
}
// SubscribeWithErr subscribes a ListenerWithErr to the given event
// name. The listener also receives error deliveries such as
// pubsub.ErrDroppedMessages. Multiple local subscribers on the same
// event share a single underlying *natsgo.Subscription with
// per-listener bounded inboxes so a slow listener cannot block its
// peers.
func (p *Pubsub) SubscribeWithErr(event string, listener pubsub.ListenerWithErr) (cancel func(), err error) {
return p.subscribeQueue(event, pubsub.NewMsgQueue(context.Background(), nil, listener))
}
// subscribeQueue subscribes the given MsgQueue for the given event.
func (p *Pubsub) subscribeQueue(event string, newQ *pubsub.MsgQueue) (cancel func(), err error) {
defer func() {
if err != nil {
// If we hit an error, close the queue so we don't leak its goroutine.
newQ.Close()
}
}()
l, g := func() (*localSub, *groupSub) {
p.mu.Lock()
defer p.mu.Unlock()
if p.ctx.Err() != nil {
return nil, erroredGroupSub(p.metrics, errClosed)
}
var (
gSub *groupSub
ok bool
)
gSub, ok = p.subscriptions[event]
if !ok {
gSub = &groupSub{
metrics: p.metrics,
event: event,
localSubs: make(map[*localSub]struct{}),
sub: &subGetter{
subscribeDone: make(chan struct{}),
},
}
go p.subscribeGroup(gSub)
p.subscriptions[event] = gSub
p.metrics.AddEvent(event)
}
lSub := &localSub{
event: event,
queue: newQ,
}
gSub.mu.Lock()
defer gSub.mu.Unlock()
gSub.localSubs[lSub] = struct{}{}
return lSub, gSub
}()
if _, err := g.sub.get(); err != nil {
p.metrics.RecordSubscribeFailure(event)
// A failed subscribe was never counted (we increment only on
// success below), so there is nothing to undo here.
return nil, err
}
p.metrics.RecordSubscribeSuccess(event)
// Count the subscriber once the NATS subscription is established. The
// matching decrement is in closeLocalSubFunc when the localSub is
// removed. A mid-subscribe Close may decrement without a matching
// increment, but the gauge is irrelevant once we are shutting down.
p.metrics.AddSubscriber(event)
return p.closeLocalSubFunc(l, g), nil
}
// signalSubscribersDroppedForConn signals local subscribers assigned to conn.
func (p *Pubsub) signalSubscribersDroppedForConn(c conn) {
if c == nil || len(p.subscribePool) == 0 {
return
}
p.mu.Lock()
subs := make([]*localSub, 0)
for event, nsub := range p.subscriptions {
if pickConn(p.subscribePool, event) != c {
continue
}
nsub.mu.Lock()
for s := range nsub.localSubs {
subs = append(subs, s)
}
nsub.mu.Unlock()
}
p.mu.Unlock()
for _, s := range subs {
s.signalDrop()
}
}
// handleAsyncError routes async error callbacks. Only slow-consumer
// errors trigger drop accounting.
func (p *Pubsub) handleAsyncError(sub *natsgo.Subscription, err error) {
if sub == nil || !errors.Is(err, natsgo.ErrSlowConsumer) {
return
}
// Snapshot candidates under p.mu, then match via the blocking
// sub.get() outside the lock. Holding p.mu across get() stalls other
// Pubsub operations and can deadlock with subscribeGroup's cleanup.
p.mu.Lock()
candidates := make([]*groupSub, 0, len(p.subscriptions))
for _, candidate := range p.subscriptions {
candidates = append(candidates, candidate)
}
p.mu.Unlock()
var nsub *groupSub
for _, candidate := range candidates {
if s, _ := candidate.sub.get(); s == sub {
nsub = candidate
break
}
}
if nsub == nil {
return
}
p.handleSlowSubscriber(nsub)
}
// handleSlowSubscriber broadcasts pubsub.ErrDroppedMessages to every
// local listener on nsub when NATS reports a new drop delta. The
// slow-consumer signal is per-subscription and cannot be narrowed to a
// single local listener.
func (p *Pubsub) handleSlowSubscriber(nsub *groupSub) {
sub, err := nsub.sub.get()
if err != nil {
return
}
nsub.dropMu.Lock()
dropped, err := sub.Dropped()
if err != nil {
nsub.dropMu.Unlock()
p.logger.Warn(p.ctx, "nats: query dropped count", slog.Error(err))
return
}
if dropped < 0 {
nsub.dropMu.Unlock()
p.logger.Warn(p.ctx, "nats: negative dropped count")
return
}
// Dropped is cumulative per subscription; signal only new drops.
droppedCount := uint64(dropped)
if droppedCount < nsub.lastDropped {
nsub.lastDropped = droppedCount
nsub.dropMu.Unlock()
return
}
if droppedCount == nsub.lastDropped {
nsub.dropMu.Unlock()
return
}
nsub.lastDropped = droppedCount
nsub.dropMu.Unlock()
nsub.mu.Lock()
defer nsub.mu.Unlock()
for s := range nsub.localSubs {
s.signalDrop()
}
}
// Close stops local delivery and shuts down the Pubsub. It is idempotent.
// Close does not drain queued listener messages.
func (p *Pubsub) Close() error {
p.closeOnce.Do(func() {
// Stop the latency loop before tearing down connections. Done
// outside p.mu because an in-flight probe can block on subscriber
// state that Close holds. No-op if the loop was never started.
p.metrics.StopLatencyLoop()
// Report disconnected immediately. The owned connections are
// closed below without firing the disconnect handler, so nothing
// else resets the gauge during shutdown.
p.conns.markClosed()
p.mu.Lock()
p.logger.Debug(p.ctx, "closing pubsub")
// Cancel while holding p.mu so subscriber state cleanup below
// observes the canceled context.
p.cancel()
var closeFuncs []func()
for _, g := range p.subscriptions {
// here we don't need to hold the ss.mu lock because we are not mutating anything and holding the p.mu
// blocks any new subscriptions.
for l := range g.localSubs {
closeFuncs = append(closeFuncs, p.closeLocalSubFunc(l, g))
}
}
p.mu.Unlock()
for _, f := range closeFuncs {
f()
}
p.logger.Debug(p.ctx, "closed all local subscriptions")
// Wait for any outstanding unsubscribe routines, kicked off above or before the Close().
p.unsubscribeRoutines.Wait()
p.logger.Debug(p.ctx, "unsubscribe routines done")
for _, nc := range p.subscribePool {
if nc != nil {
nc.Close()
}
}
p.logger.Debug(p.ctx, "subscribe pool connections closed")
for _, nc := range p.publishPool {
if nc != nil {
nc.Close()
}
}
p.logger.Debug(p.ctx, "publish pool connections closed")
if p.Server != nil {
p.Server.Shutdown()
p.Server.WaitForShutdown()
p.logger.Info(p.ctx, "nats server shut down")
} else {
p.logger.Debug(p.ctx, "nats server was never started")
}
})
return nil
}
// closeLocalSubFunc returns a function that cancels local delivery without waiting for callbacks.
//
// It returns a func() rather than just closing directly because the PubSub interface wants a func() to cancel a
// subscription.
func (p *Pubsub) closeLocalSubFunc(l *localSub, g *groupSub) func() {
return func() {
p.mu.Lock()
defer p.mu.Unlock()
g.mu.Lock()
defer g.mu.Unlock()
logger := p.logger.With(slog.F("event", l.event))
logger.Debug(context.Background(), "closing local sub")
// This function must be idempotent because it is called either by the listener code, or by the pubsub itself
// while closing. If we're already removed from the group then we must have already been closed.
if _, exists := g.localSubs[l]; !exists {
return
}
l.queue.Close()
delete(g.localSubs, l)
p.metrics.RemoveSubscriber(l.event)
logger.Debug(context.Background(), "removed local sub from group", slog.F("group_size", len(g.localSubs)))
if len(g.localSubs) > 0 {
return // Not last one out
}
// Last localSub does the nats unsubscribe. Do this async so we don't hold the pubsub lock too long. Nothing is
// left listening, so no rush.
p.unsubscribeRoutines.Add(1)
go func() {
defer p.unsubscribeRoutines.Done()
p.unsubscribeGroup(g)
}()
if pSub, ok := p.subscriptions[l.event]; ok && g == pSub {
delete(p.subscriptions, l.event)
p.metrics.RemoveEvent(l.event)
}
}
}
func (p *Pubsub) subscribeGroup(g *groupSub) {
defer func() {
// Close subscribeDone before taking p.mu: a goroutine holding
// p.mu may be blocked in get(), so the reverse order deadlocks.
close(g.sub.subscribeDone)
if g.sub.err != nil {
// failed to subscribe. Kick this out of the pubsub map of subscriptions, so that we don't permanently
// fail to subscribe to this event. The subscribe that kicked this off as well as any concurrent ones will
// see an error.
p.mu.Lock()
defer p.mu.Unlock()
if psub := p.subscriptions[g.event]; psub == g {
delete(p.subscriptions, g.event)
p.metrics.RemoveEvent(g.event)
}
}
}()
logger := p.logger.With(slog.F("event", g.event))
logger.Debug(context.Background(), "subscribing on nats")
subConn := pickConn(p.subscribePool, g.event)
natsSubscription, err := subConn.Subscribe(g.event, g.handleMessage)
if err != nil {
g.sub.err = xerrors.Errorf("subscribe: %w", err)
return
}
g.sub.sub = natsSubscription
defer func() {
if g.sub.err != nil {
unsubErr := natsSubscription.Unsubscribe()
// best effort, just log if it fails
if unsubErr != nil {
// nolint: gocritic // false positive because we log two errors
logger.Error(p.ctx, "failed to unsubscribe after error subscribing",
slog.Error(unsubErr), slog.F("previous_error", g.sub.err))
}
}
}()
// Flush the SUB to the server so a publish issued immediately
// after Subscribe returns cannot race ahead of registration.
if err := subConn.Flush(); err != nil {
g.sub.err = xerrors.Errorf("flush subscribe: %w", err)
return
}
limits := defaultPendingLimits(p.opts.PendingLimits)
if err := natsSubscription.SetPendingLimits(limits.Msgs, limits.Bytes); err != nil {
g.sub.err = xerrors.Errorf("set pending limits: %w", err)
return
}
}
func (p *Pubsub) unsubscribeGroup(g *groupSub) {
logger := p.logger.With(slog.F("event", g.event))
logger.Debug(context.Background(), "unsubscribing group subscription from nats")
// wait for any pending Subscribe to complete before we attempt to unsubscribe
sub, err := g.sub.get()
if err != nil {
// subscribe failed, nothing else to do.
return
}
if err = sub.Unsubscribe(); err != nil {
logger.Error(context.Background(), "failed to unsubscribe from pubsub", slog.Error(err))
}
// TODO: should we retry?
}
// pickConn returns the connection assigned to subject. Selection uses
// a stable FNV-1a hash so same-subject traffic always targets the same
// connection within a process; pools are immutable after construction
// so the lookup is lock-free.
func pickConn(pool []conn, subject string) conn {
if len(pool) == 1 {
return pool[0]
}
h := fnv.New32a()
_, _ = h.Write([]byte(subject))
n := uint32(len(pool)) //nolint:gosec // pool size bounded by Options.{Publish,Subscribe}Conns
return pool[h.Sum32()%n]
}
// erroredGroupSub returns a groupSub that shows an error rather than an active subscription.
func erroredGroupSub(m *pubsub.BackendMetrics, err error) *groupSub {
c := make(chan struct{})
close(c)
return &groupSub{
metrics: m,
sub: &subGetter{
subscribeDone: c,
err: err,
},
}
}
// handleMessage handles messages for the shared subscription. Each
// enqueue is non-blocking and does not call user code, so one slow
// listener cannot stall the NATS delivery goroutine.
//
// Zero-copy fan-out: the same msg.Data slice is delivered to every
// local listener without cloning. Listeners on a coalesced subject MUST
// treat the delivered bytes as immutable.
func (g *groupSub) handleMessage(msg *natsgo.Msg) {
g.metrics.RecordReceived(g.event, msg.Data)
g.mu.Lock()
defer g.mu.Unlock()
for l := range g.localSubs {
l.queue.Enqueue(msg.Data)
}
}
// subGetter allows callers to asynchronously wait for the subscription to complete or error by calling the get()
// method. Routines other than the one that actually starts the natsgo.Subscription should never access sub directly.
type subGetter struct {
// closed when the initial subscribe completes
subscribeDone chan struct{}
// either sub or err are non-nil after subscribeDone is closed
sub *natsgo.Subscription
err error
}
func (s *subGetter) get() (*natsgo.Subscription, error) {
<-s.subscribeDone
return s.sub, s.err
}
// signalDrop pushes onto dropSignal without blocking. Multiple drops
// between dispatcher dequeues coalesce into a single pending signal, so
// the listener observes one ErrDroppedMessages per drop wave.
func (l *localSub) signalDrop() {
l.queue.Dropped()
}