diff --git a/coderd/x/nats/pubsub.go b/coderd/x/nats/pubsub.go index 4c6d902fd2..b1fea0b5b8 100644 --- a/coderd/x/nats/pubsub.go +++ b/coderd/x/nats/pubsub.go @@ -99,6 +99,14 @@ type Options struct { 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 @@ -117,12 +125,12 @@ type Pubsub struct { Server *natsserver.Server // publishPool and subscribePool are immutable after construction so // the hot path can index without holding p.mu. - publishPool []*natsgo.Conn - subscribePool []*natsgo.Conn + publishPool []conn + subscribePool []conn // subscriptions coalesces concurrent local subscribers on the // same subject onto a single underlying *natsgo.Subscription. - subscriptions map[string]*natsSub + subscriptions map[string]*groupSub closeOnce sync.Once // ctx is canceled by Close while holding p.mu so subscriber state @@ -130,6 +138,10 @@ type Pubsub struct { 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 @@ -139,20 +151,19 @@ type Pubsub struct { peerRefresh chan struct{} } -// natsSub maps to one underlying *natsgo.Subscription. The first +// 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 natsSub struct { - // sub is set before this natsSub is published in Pubsub.subscriptions - // and is immutable after that. - sub *natsgo.Subscription - +type groupSub struct { + 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. @@ -160,26 +171,10 @@ type natsSub struct { } // localSub is the local handle returned by Subscribe / -// SubscribeWithErr. Each local subscriber gets its own bounded inbox -// and dispatcher goroutine so one slow listener cannot block peers on -// the same subject. +// SubscribeWithErr. type localSub struct { - cancelOnce sync.Once - - ctx context.Context - - event string - listener pubsub.ListenerWithErr - - // queue is the per-listener data fan-out inbox. The shared NATS - // callback enqueues non-blockingly; on overflow the message is - // dropped and a drop signal is raised. - queue chan []byte - // dropSignal is a size-1 buffered channel that coalesces drop - // notifications from local overflow and NATS slow-consumer - // broadcasts onto a single pending wake. - dropSignal chan struct{} - cancel context.CancelFunc + event string + queue *pubsub.MsgQueue } // Compile-time assertion that *Pubsub satisfies the pubsub.Pubsub interface. @@ -191,7 +186,7 @@ func newPubsub(ctx context.Context, logger slog.Logger, opts Options) *Pubsub { return &Pubsub{ logger: logger, opts: opts, - subscriptions: make(map[string]*natsSub), + subscriptions: make(map[string]*groupSub), ctx: ctx, cancel: cancel, peerFetcher: opts.PeerFetcher, @@ -303,11 +298,11 @@ func New(ctx context.Context, logger slog.Logger, opts Options) (*Pubsub, error) return p, nil } -func newConnPool(ns *natsserver.Server, opts Options, handlers connHandlers, count int, clientName string) ([]*natsgo.Conn, error) { +func newConnPool(ns *natsserver.Server, opts Options, handlers connHandlers, count int, clientName string) ([]conn, error) { if count <= 0 { count = 1 } - pool := make([]*natsgo.Conn, 0, count) + 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. @@ -362,12 +357,7 @@ func (p *Pubsub) Flush() error { // 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.SubscribeWithErr(event, func(ctx context.Context, msg []byte, err error) { - if err != nil { - return - } - listener(ctx, msg) - }) + return p.subscribeQueue(event, pubsub.NewMsgQueue(context.Background(), listener, nil)) } // SubscribeWithErr subscribes a ListenerWithErr to the given event @@ -377,203 +367,68 @@ func (p *Pubsub) Subscribe(event string, listener pubsub.Listener) (cancel func( // 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) { - s, err := p.addSubscriber(event, listener) - if err != nil { - return nil, err - } - - cancelFn := func() { - s.close() - p.unsubscribeLocal(s) - } - return cancelFn, nil + return p.subscribeQueue(event, pubsub.NewMsgQueue(context.Background(), nil, listener)) } -// listenerQueueSize returns the per-listener inbox capacity. A -// positive PendingLimits.Msgs sets the cap (giving callers a knob to -// trigger local-overflow drops since coalescing makes NATS-level -// slow-consumer signals rare). Otherwise the default is used. -func listenerQueueSize(in PendingLimits) int { - if in.Msgs > 0 { - return in.Msgs - } - return defaultListenerQueueSize -} +// 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() + } + }() -const defaultListenerQueueSize = 1024 - -// addSubscriber creates a local subscriber and attaches it to the natsSub -// for event. New natsSub entries are published only after NATS setup succeeds. -func (p *Pubsub) addSubscriber(event string, listener pubsub.ListenerWithErr) (*localSub, error) { - ctx, cancel := context.WithCancel(p.ctx) - s := &localSub{ - ctx: ctx, - cancel: cancel, - event: event, - listener: listener, - queue: make(chan []byte, listenerQueueSize(p.opts.PendingLimits)), - dropSignal: make(chan struct{}, 1), - } - s.init() - - cleanupSub, err := func() (*natsgo.Subscription, error) { + l, g := func() (*localSub, *groupSub) { p.mu.Lock() defer p.mu.Unlock() if p.ctx.Err() != nil { - return nil, errClosed + return nil, erroredGroupSub(errClosed) } - nsub, ok := p.subscriptions[event] - if ok { - nsub.mu.Lock() - nsub.localSubs[s] = struct{}{} - nsub.mu.Unlock() - return nsub.sub, nil - } - - nsub = &natsSub{ - localSubs: map[*localSub]struct{}{ - s: {}, - }, - } - - subConn := pickConn(p.subscribePool, event) - natsSubscription, err := subConn.Subscribe(event, nsub.handleMessage) - if err != nil { - return nil, xerrors.Errorf("subscribe: %w", err) - } - nsub.sub = natsSubscription - - // 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 { - return natsSubscription, xerrors.Errorf("flush subscribe: %w", err) - } - limits := defaultPendingLimits(p.opts.PendingLimits) - if err := natsSubscription.SetPendingLimits(limits.Msgs, limits.Bytes); err != nil { - return natsSubscription, xerrors.Errorf("set pending limits: %w", err) - } - - p.subscriptions[event] = nsub - return natsSubscription, nil - }() - if err != nil { - s.close() - if cleanupSub != nil { - if unsubscribeErr := cleanupSub.Unsubscribe(); unsubscribeErr != nil { - err = errors.Join(err, xerrors.Errorf("unsubscribe: %w", unsubscribeErr)) + var ( + gSub *groupSub + ok bool + ) + gSub, ok = p.subscriptions[event] + if !ok { + gSub = &groupSub{ + event: event, + localSubs: make(map[*localSub]struct{}), + sub: &subGetter{ + subscribeDone: make(chan struct{}), + }, } + go p.subscribeGroup(gSub) + p.subscriptions[event] = gSub } + 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 { return nil, err } - return s, nil -} - -// unsubscribeLocal removes s from its natsSub. If s was the last -// listener, it also removes and unsubscribes the underlying NATS -// subscription. -func (p *Pubsub) unsubscribeLocal(s *localSub) { - natsSub := func() *natsgo.Subscription { - p.mu.Lock() - defer p.mu.Unlock() - - nsub := p.subscriptions[s.event] - if nsub == nil { - return nil - } - - nsub.mu.Lock() - defer nsub.mu.Unlock() - if _, tracked := nsub.localSubs[s]; !tracked { - return nil - } - delete(nsub.localSubs, s) - if len(nsub.localSubs) > 0 { - return nil - } - // Last listener: remove the nsub entry so a new Subscribe to this - // subject creates a fresh underlying subscription. - delete(p.subscriptions, s.event) - return nsub.sub - }() - if natsSub != nil { - _ = natsSub.Unsubscribe() - } -} - -// 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 (nsub *natsSub) handleMessage(msg *natsgo.Msg) { - nsub.mu.Lock() - defer nsub.mu.Unlock() - - for s := range nsub.localSubs { - s.enqueue(msg.Data) - } -} - -// init starts the per-listener delivery goroutine. -func (s *localSub) init() { - go func() { - for { - select { - case <-s.ctx.Done(): - return - case data := <-s.queue: - s.listener(s.ctx, data, nil) - case <-s.dropSignal: - s.listener(s.ctx, nil, pubsub.ErrDroppedMessages) - } - } - }() -} - -// close cancels local delivery without waiting for callbacks. -func (s *localSub) close() { - s.cancelOnce.Do(func() { - if s.cancel != nil { - s.cancel() - } - }) -} - -// enqueue non-blockingly sends data onto s.queue. On overflow it drops the -// message and raises a drop signal so pubsub.ErrDroppedMessages is surfaced. -// If s is canceled the message is silently dropped. -func (s *localSub) enqueue(data []byte) { - select { - case s.queue <- data: - default: - s.signalDrop() - } -} - -// 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 (s *localSub) signalDrop() { - select { - case s.dropSignal <- struct{}{}: - default: - } + return p.closeLocalSubFunc(l, g), nil } // signalSubscribersDroppedForConn signals local subscribers assigned to conn. -func (p *Pubsub) signalSubscribersDroppedForConn(conn *natsgo.Conn) { - if conn == nil || len(p.subscribePool) == 0 { +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) != conn { + if pickConn(p.subscribePool, event) != c { continue } nsub.mu.Lock() @@ -596,9 +451,9 @@ func (p *Pubsub) handleAsyncError(sub *natsgo.Subscription, err error) { return } p.mu.Lock() - var nsub *natsSub + var nsub *groupSub for _, candidate := range p.subscriptions { - if candidate.sub == sub { + if s, _ := candidate.sub.get(); s == sub { nsub = candidate break } @@ -614,9 +469,13 @@ func (p *Pubsub) handleAsyncError(sub *natsgo.Subscription, err error) { // 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 *natsSub) { +func (p *Pubsub) handleSlowSubscriber(nsub *groupSub) { + sub, err := nsub.sub.get() + if err != nil { + return + } nsub.dropMu.Lock() - dropped, err := nsub.sub.Dropped() + dropped, err := sub.Dropped() if err != nil { nsub.dropMu.Unlock() p.logger.Warn(p.ctx, "nats: query dropped count", slog.Error(err)) @@ -654,59 +513,159 @@ func (p *Pubsub) handleSlowSubscriber(nsub *natsSub) { func (p *Pubsub) Close() error { p.closeOnce.Do(func() { 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 subs []*localSub - shareds := make([]*natsSub, 0, len(p.subscriptions)) - for _, ss := range p.subscriptions { - shareds = append(shareds, ss) - ss.mu.Lock() - for s := range ss.localSubs { - subs = append(subs, s) - delete(ss.localSubs, s) + 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)) } - ss.mu.Unlock() } - clear(p.subscriptions) p.mu.Unlock() - // Unsubscribe shared subscriptions before closing connections. - for _, ss := range shareds { - if ss.sub != nil { - _ = ss.sub.Unsubscribe() - } - } - - // Signal per-listener goroutines without waiting for callbacks. - for _, s := range subs { - s.close() + 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) + 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) + } + } +} + +func (p *Pubsub) subscribeGroup(g *groupSub) { + defer func() { + 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) + } + } + close(g.sub.subscribeDone) + }() + 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 []*natsgo.Conn, subject string) *natsgo.Conn { +func pickConn(pool []conn, subject string) conn { if len(pool) == 1 { return pool[0] } @@ -715,3 +674,52 @@ func pickConn(pool []*natsgo.Conn, subject string) *natsgo.Conn { 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(err error) *groupSub { + c := make(chan struct{}) + close(c) + return &groupSub{ + 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.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() +} diff --git a/coderd/x/nats/pubsub_internal_test.go b/coderd/x/nats/pubsub_internal_test.go index 678db23a4b..38384fc137 100644 --- a/coderd/x/nats/pubsub_internal_test.go +++ b/coderd/x/nats/pubsub_internal_test.go @@ -2,17 +2,16 @@ package nats import ( "context" - "errors" "fmt" "net/url" "slices" - "sync" "sync/atomic" "testing" "time" natsserver "github.com/nats-io/nats-server/v2/server" natsgo "github.com/nats-io/nats.go" + "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "golang.org/x/xerrors" @@ -75,19 +74,21 @@ func Test_pickConn(t *testing.T) { t.Run("DifferentSubjects", func(t *testing.T) { t.Parallel() - var a, b natsgo.Conn - pool := []*natsgo.Conn{&a, &b} - - require.NotSame(t, pickConn(pool, "a"), pickConn(pool, "b")) + a := new(fakeConn) + b := new(fakeConn) + pool := []conn{a, b} + ca := pickConn(pool, "a") + cb := pickConn(pool, "b") + require.NotSame(t, ca, cb) }) } -func subjectForConn(t *testing.T, pool []*natsgo.Conn, conn *natsgo.Conn, prefix string) string { +func subjectForConn(t *testing.T, pool []conn, c conn, prefix string) string { t.Helper() for i := range 10_000 { subject := fmt.Sprintf("%s_%d", prefix, i) - if pickConn(pool, subject) == conn { + if pickConn(pool, subject) == c { return subject } } @@ -159,124 +160,58 @@ func Test_Pubsub_buildConnHandlers(t *testing.T) { ps := newPubsub(ctx, logger, defaultTestOptions()) var subConnA, subConnB, pubConn natsgo.Conn - ps.subscribePool = []*natsgo.Conn{&subConnA, &subConnB} + ps.subscribePool = []conn{&subConnA, &subConnB} matchingEvent := subjectForConn(t, ps.subscribePool, &subConnA, "disconnect_match") otherEvent := subjectForConn(t, ps.subscribePool, &subConnB, "disconnect_other") - newLocal := func(event string) *localSub { + newLocal := func(event string, errCh chan error) *localSub { + queue := pubsub.NewMsgQueue(ctx, nil, func(_ context.Context, _ []byte, err error) { + testutil.RequireSend(ctx, t, errCh, err) + }) + // normally, closing the pubsub would clean this, but we don't actually close pubsub in this test because + // it uses fake connections. So, we need to close these to avoid leaking goroutines. + t.Cleanup(func() { + queue.Close() + }) return &localSub{ - event: event, - dropSignal: make(chan struct{}, 1), + event: event, + queue: queue, } } - matchingSub := newLocal(matchingEvent) - otherSub := newLocal(otherEvent) - ps.subscriptions[matchingSub.event] = &natsSub{localSubs: map[*localSub]struct{}{matchingSub: {}}} - ps.subscriptions[otherSub.event] = &natsSub{localSubs: map[*localSub]struct{}{otherSub: {}}} + matchErr := make(chan error) + matchingSub := newLocal(matchingEvent, matchErr) + otherErr := make(chan error) + otherSub := newLocal(otherEvent, otherErr) + ps.subscriptions[matchingSub.event] = &groupSub{localSubs: map[*localSub]struct{}{matchingSub: {}}} + ps.subscriptions[otherSub.event] = &groupSub{localSubs: map[*localSub]struct{}{otherSub: {}}} handlers := ps.buildConnHandlers() handlers.disconnectErr(&subConnA, xerrors.New("disconnect")) + err := testutil.RequireReceive(ctx, t, matchErr) + require.ErrorIs(t, err, pubsub.ErrDroppedMessages) select { - case <-matchingSub.dropSignal: - default: - require.Fail(t, "matching subscriber did not receive drop signal") - } - select { - case <-otherSub.dropSignal: + case <-otherErr: require.Fail(t, "non-matching subscriber received drop signal") default: } handlers.disconnectErr(&pubConn, xerrors.New("publisher disconnect")) select { - case <-otherSub.dropSignal: + case <-otherErr: require.Fail(t, "publisher connection disconnect signaled subscriber") default: } }) } -func Test_localSub_init(t *testing.T) { +func Test_localSub(t *testing.T) { t.Parallel() - t.Run("SerializesCallbacks", func(t *testing.T) { - t.Parallel() - ctx := testutil.Context(t, testutil.WaitShort) - - dataStarted := make(chan struct{}) - dropDelivered := make(chan struct{}) - release := make(chan struct{}) - var dataOnce sync.Once - var dropOnce sync.Once - var releaseOnce sync.Once - var active atomic.Int64 - var concurrent atomic.Bool - - s := &localSub{ - ctx: ctx, - cancel: func() {}, - listener: func(_ context.Context, _ []byte, ferr error) { - if active.Add(1) != 1 { - concurrent.Store(true) - } - defer active.Add(-1) - - if errors.Is(ferr, pubsub.ErrDroppedMessages) { - dropOnce.Do(func() { close(dropDelivered) }) - return - } - - dataOnce.Do(func() { close(dataStarted) }) - <-release - }, - queue: make(chan []byte, 1), - dropSignal: make(chan struct{}, 1), - } - s.init() - t.Cleanup(func() { - releaseOnce.Do(func() { close(release) }) - s.close() - }) - - s.enqueue([]byte("data")) - require.Eventually(t, func() bool { - select { - case <-dataStarted: - return true - default: - return false - } - }, testutil.WaitShort, testutil.IntervalFast) - - s.signalDrop() - require.Never(t, func() bool { - select { - case <-dropDelivered: - return true - default: - return false - } - }, testutil.IntervalMedium, testutil.IntervalFast, - "drop callback must wait for the blocked data callback") - require.False(t, concurrent.Load(), "listener callback ran concurrently") - - releaseOnce.Do(func() { close(release) }) - require.Eventually(t, func() bool { - select { - case <-dropDelivered: - return true - default: - return false - } - }, testutil.WaitShort, testutil.IntervalFast) - require.False(t, concurrent.Load(), "listener callback ran concurrently") - }) - t.Run("SameSubjectSlowListenerDoesNotBlockPeer", func(t *testing.T) { t.Parallel() - logger := slogtest.Make(t, nil) + logger := testutil.Logger(t) ctx := testutil.Context(t, testutil.WaitLong) ps, err := New(ctx, logger, defaultTestOptions()) require.NoError(t, err) @@ -327,8 +262,14 @@ func Test_localSub_init(t *testing.T) { // One coalesced subscription on one subConn; the slow consumer must // not tear it down. require.Len(t, ps.subscribePool, 1) - require.False(t, ps.subscribePool[0].IsClosed(), "subConn must not be closed by slow consumer") - require.True(t, ps.subscribePool[0].IsConnected(), "subConn must stay connected") + natsConn, ok := ps.subscribePool[0].(*natsgo.Conn) + require.True(t, ok) + require.False(t, natsConn.IsClosed(), "subConn must not be closed by slow consumer") + require.True(t, natsConn.IsConnected(), "subConn must stay connected") + + err = ps.Close() + require.NoError(t, err) + require.Empty(t, ps.subscriptions) }) } @@ -468,6 +409,48 @@ func TestPubsubCluster(t *testing.T) { }) } +func TestSubscribeError(t *testing.T) { + t.Parallel() + + testCases := []struct { + name string + fConn *fakeConn + }{ + { + name: "Subscribe", + fConn: &fakeConn{ + subError: assert.AnError, + }, + }, + { + name: "Flush", + fConn: &fakeConn{ + flushError: assert.AnError, + }, + }, + } + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + logger := slogtest.Make(t, &slogtest.Options{ + IgnoredErrorIs: []error{natsgo.ErrConnectionClosed, assert.AnError}, + }) + ctx := testutil.Context(t, testutil.WaitShort) + ps := newPubsub(ctx, logger, defaultTestOptions()) + ps.subscribePool = []conn{tc.fConn} + cancel, err := ps.SubscribeWithErr("foo", func(ctx context.Context, message []byte, err error) { + t.Error("should not get any events") + }) + require.ErrorIs(t, err, assert.AnError) + require.Nil(t, cancel) + ps.mu.Lock() + defer ps.mu.Unlock() + require.Empty(t, ps.subscriptions) + }) + } +} + func defaultTestOptions() Options { return Options{disableCluster: true} } @@ -562,3 +545,26 @@ func routeStrings(routes []*url.URL) []string { } return out } + +type fakeConn struct { + subError error + flushError error +} + +func (*fakeConn) Publish(string, []byte) error { + // TODO implement me + panic("implement me") +} + +func (*fakeConn) Close() { + // TODO implement me + panic("implement me") +} + +func (f *fakeConn) Flush() error { + return f.flushError +} + +func (f *fakeConn) Subscribe(string, natsgo.MsgHandler) (*natsgo.Subscription, error) { + return &natsgo.Subscription{}, f.subError +}