chore: refactor NATS pubsub to use MsgQueue (#26197)

Closes https://github.com/coder/scaletest/issues/151  
Closes GRU-71  
  
Use the existing MsgQueue from the original PGPubsub instead of the 2-channel solution originally built here.

Renames `natsSub` to `groupSub`, since conceptually, a "NATS Subscription" already refers to the underlying subscription on the NATS server.

This PR also simplifies the closing of the PubSub to just close each `localSub`. When the last `localSub` for an event is closed, it unsubscribes and remove the `groupSub`. This ensures we go through the same code paths closing normally and at end of day.

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()
+}

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
+}