feat: add Prometheus metrics to NATS pubsub for parity with PGPubsub (#26441)

coderd/database/pubsub/pubsub.go

@@ -421,9 +421,9 @@ func (p *PGPubsub) listen() {
 }
 
 func (p *PGPubsub) listenReceive(notif *pq.Notification) {
-	sizeLabel := messageSizeNormal
-	if len(notif.Extra) >= colossalThreshold {
-		sizeLabel = messageSizeColossal
+	sizeLabel := MessageSizeNormal
+	if len(notif.Extra) >= ColossalThreshold {
+		sizeLabel = MessageSizeColossal
 	}
 	p.messagesTotal.WithLabelValues(sizeLabel).Inc()
 	p.receivedBytesTotal.Add(float64(len(notif.Extra)))
@@ -614,10 +614,14 @@ var (
 // notify limit. If we see a lot of colossal packets that's an indication that
 // we might be trying to send too much data over the pubsub and are in danger of
 // failing to publish.
+//
+// These are exported so other pubsub implementations (e.g. the NATS
+// pubsub) classify message size identically, keeping the messages_total
+// "size" label consistent across backends.
 const (
-	colossalThreshold   = 7600
-	messageSizeNormal   = "normal"
-	messageSizeColossal = "colossal"
+	ColossalThreshold   = 7600
+	MessageSizeNormal   = "normal"
+	MessageSizeColossal = "colossal"
 )
 
 // Describe implements, along with Collect, the prometheus.Collector interface

coderd/x/nats/metrics.go

@@ -0,0 +1,289 @@
+package nats
+
+import (
+	"context"
+	"sync"
+	"sync/atomic"
+
+	"github.com/prometheus/client_golang/prometheus"
+
+	"cdr.dev/slog/v3"
+	"github.com/coder/coder/v2/coderd/database/pubsub"
+)
+
+// Descriptors for metrics that are not backed by a stored collector:
+// current_subscribers and current_events are read from atomic counters,
+// and the latency metrics are measured during each scrape.
+var (
+	currentSubscribersDesc = prometheus.NewDesc(
+		"coder_nats_pubsub_current_subscribers",
+		"The current number of active pubsub subscribers",
+		nil, nil,
+	)
+	currentEventsDesc = prometheus.NewDesc(
+		"coder_nats_pubsub_current_events",
+		"The current number of pubsub event channels listened for",
+		nil, nil,
+	)
+	sendLatencyDesc = prometheus.NewDesc(
+		"coder_nats_pubsub_send_latency_seconds",
+		"The time taken to send a message into a pubsub event channel",
+		nil, nil,
+	)
+	recvLatencyDesc = prometheus.NewDesc(
+		"coder_nats_pubsub_receive_latency_seconds",
+		"The time taken to receive a message from a pubsub event channel",
+		nil, nil,
+	)
+	latencyMeasureCountDesc = prometheus.NewDesc(
+		"coder_nats_pubsub_latency_measures_total",
+		"The number of pubsub latency measurements",
+		nil, nil,
+	)
+	latencyMeasureErrDesc = prometheus.NewDesc(
+		"coder_nats_pubsub_latency_measure_errs_total",
+		"The number of pubsub latency measurement failures",
+		nil, nil,
+	)
+)
+
+// metrics owns all Prometheus state for the NATS Pubsub. Collaborators
+// such as groupSub depend on this narrow type rather than the whole
+// Pubsub, so the only thing they can do is record metrics.
+type metrics struct {
+	logger slog.Logger
+
+	publishesTotal      *prometheus.CounterVec
+	subscribesTotal     *prometheus.CounterVec
+	messagesTotal       *prometheus.CounterVec
+	publishedBytesTotal prometheus.Counter
+	receivedBytesTotal  prometheus.Counter
+	disconnectionsTotal prometheus.Counter
+	connected           prometheus.Gauge
+
+	latencyMeasurer       *pubsub.LatencyMeasurer
+	latencyMeasureCounter atomic.Int64
+	latencyErrCounter     atomic.Int64
+
+	// connMu guards the connection-state accounting below. Connect and
+	// disconnect callbacks are rare, so a mutex keeps the gauge update
+	// atomic with the count without meaningful contention.
+	connMu         sync.Mutex
+	totalConns     int
+	connectedConns int
+
+	// currentEvents and currentSubscribers shadow the sizes of the
+	// Pubsub's subscriptions map and per-event localSubs maps. They are
+	// maintained at the subscribe/unsubscribe sites so Collect can read
+	// the gauges without locking the Pubsub.
+	currentEvents      atomic.Int64
+	currentSubscribers atomic.Int64
+}
+
+// newMetrics builds all metric instruments up front so collaborators
+// never observe nil metric fields.
+func newMetrics(logger slog.Logger) *metrics {
+	return &metrics{
+		logger:          logger,
+		latencyMeasurer: pubsub.NewLatencyMeasurer(logger.Named("latency-measurer")),
+
+		publishesTotal: prometheus.NewCounterVec(prometheus.CounterOpts{
+			Namespace: "coder",
+			Subsystem: "nats_pubsub",
+			Name:      "publishes_total",
+			Help:      "Total number of calls to Publish",
+		}, []string{"success"}),
+		subscribesTotal: prometheus.NewCounterVec(prometheus.CounterOpts{
+			Namespace: "coder",
+			Subsystem: "nats_pubsub",
+			Name:      "subscribes_total",
+			Help:      "Total number of calls to Subscribe/SubscribeWithErr",
+		}, []string{"success"}),
+		messagesTotal: prometheus.NewCounterVec(prometheus.CounterOpts{
+			Namespace: "coder",
+			Subsystem: "nats_pubsub",
+			Name:      "messages_total",
+			Help:      "Total number of messages received from nats",
+		}, []string{"size"}),
+		publishedBytesTotal: prometheus.NewCounter(prometheus.CounterOpts{
+			Namespace: "coder",
+			Subsystem: "nats_pubsub",
+			Name:      "published_bytes_total",
+			Help:      "Total number of bytes successfully published across all publishes",
+		}),
+		receivedBytesTotal: prometheus.NewCounter(prometheus.CounterOpts{
+			Namespace: "coder",
+			Subsystem: "nats_pubsub",
+			Name:      "received_bytes_total",
+			Help:      "Total number of bytes received across all messages",
+		}),
+		disconnectionsTotal: prometheus.NewCounter(prometheus.CounterOpts{
+			Namespace: "coder",
+			Subsystem: "nats_pubsub",
+			Name:      "disconnections_total",
+			Help:      "Total number of times we disconnected unexpectedly from nats",
+		}),
+		connected: prometheus.NewGauge(prometheus.GaugeOpts{
+			Namespace: "coder",
+			Subsystem: "nats_pubsub",
+			Name:      "connected",
+			Help:      "Whether we are connected (1) or not connected (0) to nats",
+		}),
+	}
+}
+
+// recordPublishSuccess records a successful Publish of n bytes.
+func (m *metrics) recordPublishSuccess(n int) {
+	m.publishesTotal.WithLabelValues("true").Inc()
+	m.publishedBytesTotal.Add(float64(n))
+}
+
+// recordPublishFailure records a failed Publish.
+func (m *metrics) recordPublishFailure() {
+	m.publishesTotal.WithLabelValues("false").Inc()
+}
+
+// recordSubscribeSuccess records a successful Subscribe/SubscribeWithErr.
+func (m *metrics) recordSubscribeSuccess() {
+	m.subscribesTotal.WithLabelValues("true").Inc()
+}
+
+// recordSubscribeFailure records a failed Subscribe/SubscribeWithErr.
+func (m *metrics) recordSubscribeFailure() {
+	m.subscribesTotal.WithLabelValues("false").Inc()
+}
+
+// recordReceived records metrics for a single received NATS message.
+// Size is classified using the shared pubsub thresholds so the
+// messages_total "size" label matches PGPubsub.
+func (m *metrics) recordReceived(data []byte) {
+	sizeLabel := pubsub.MessageSizeNormal
+	if len(data) >= pubsub.ColossalThreshold {
+		sizeLabel = pubsub.MessageSizeColossal
+	}
+	m.messagesTotal.WithLabelValues(sizeLabel).Inc()
+	m.receivedBytesTotal.Add(float64(len(data)))
+}
+
+// markConnected records that all total owned connections have dialed
+// successfully. The connected gauge is 1 only while every owned
+// connection is up.
+func (m *metrics) markConnected(total int) {
+	m.connMu.Lock()
+	defer m.connMu.Unlock()
+	m.totalConns = total
+	m.connectedConns = total
+	m.setConnectedLocked()
+}
+
+// markClosed records that the Pubsub is shutting down and forces the
+// connected gauge to 0. Closing our own connections does not fire the
+// disconnect handler (see NoCallbacksAfterClientClose), so without this
+// the gauge would still read 1 after Close. connectedConns is zeroed so
+// a late reconnect callback cannot flip the gauge back to 1.
+func (m *metrics) markClosed() {
+	m.connMu.Lock()
+	defer m.connMu.Unlock()
+	// Zero both counters so setConnectedLocked's totalConns > 0 guard is
+	// permanently false: a late reconnect callback during the shutdown
+	// window cannot increment connectedConns back up and flip the gauge
+	// to 1.
+	m.totalConns = 0
+	m.connectedConns = 0
+	m.connected.Set(0)
+}
+
+// onDisconnect records an unexpected disconnect of one owned connection.
+func (m *metrics) onDisconnect() {
+	m.disconnectionsTotal.Inc()
+	m.connMu.Lock()
+	defer m.connMu.Unlock()
+	if m.connectedConns > 0 {
+		m.connectedConns--
+	}
+	m.setConnectedLocked()
+}
+
+// onReconnect records that one owned connection reconnected.
+func (m *metrics) onReconnect() {
+	m.connMu.Lock()
+	defer m.connMu.Unlock()
+	if m.connectedConns < m.totalConns {
+		m.connectedConns++
+	}
+	m.setConnectedLocked()
+}
+
+// setConnectedLocked sets the connected gauge to 1 only when every owned
+// connection is up. Callers must hold connMu.
+func (m *metrics) setConnectedLocked() {
+	if m.totalConns > 0 && m.connectedConns == m.totalConns {
+		m.connected.Set(1)
+		return
+	}
+	m.connected.Set(0)
+}
+
+// addEvent and removeEvent track the number of subscribed event
+// channels. addSubscriber and removeSubscriber track the number of
+// local subscribers across all events.
+func (m *metrics) addEvent()         { m.currentEvents.Add(1) }
+func (m *metrics) removeEvent()      { m.currentEvents.Add(-1) }
+func (m *metrics) addSubscriber()    { m.currentSubscribers.Add(1) }
+func (m *metrics) removeSubscriber() { m.currentSubscribers.Add(-1) }
+
+// describe sends every metric descriptor on behalf of the owning
+// Pubsub's prometheus.Collector implementation.
+func (m *metrics) describe(descs chan<- *prometheus.Desc) {
+	// explicit metrics
+	m.publishesTotal.Describe(descs)
+	m.subscribesTotal.Describe(descs)
+	m.messagesTotal.Describe(descs)
+	m.publishedBytesTotal.Describe(descs)
+	m.receivedBytesTotal.Describe(descs)
+	m.disconnectionsTotal.Describe(descs)
+	m.connected.Describe(descs)
+
+	// implicit metrics
+	descs <- currentSubscribersDesc
+	descs <- currentEventsDesc
+
+	// additional metrics
+	descs <- sendLatencyDesc
+	descs <- recvLatencyDesc
+	descs <- latencyMeasureCountDesc
+	descs <- latencyMeasureErrDesc
+}
+
+// collect emits all metrics. p is the pubsub used for the out-of-band
+// latency measurement. The current subscriber and event gauges are read
+// from atomic counters maintained at the subscribe/unsubscribe sites, so
+// Collect does not lock the Pubsub.
+func (m *metrics) collect(ch chan<- prometheus.Metric, p pubsub.Pubsub) {
+	// explicit metrics
+	m.publishesTotal.Collect(ch)
+	m.subscribesTotal.Collect(ch)
+	m.messagesTotal.Collect(ch)
+	m.publishedBytesTotal.Collect(ch)
+	m.receivedBytesTotal.Collect(ch)
+	m.disconnectionsTotal.Collect(ch)
+	m.connected.Collect(ch)
+
+	// implicit metrics
+	ch <- prometheus.MustNewConstMetric(currentSubscribersDesc, prometheus.GaugeValue, float64(m.currentSubscribers.Load()))
+	ch <- prometheus.MustNewConstMetric(currentEventsDesc, prometheus.GaugeValue, float64(m.currentEvents.Load()))
+
+	// additional metrics
+	ctx, cancel := context.WithTimeout(context.Background(), pubsub.LatencyMeasureTimeout)
+	defer cancel()
+	send, recv, err := m.latencyMeasurer.Measure(ctx, p)
+
+	ch <- prometheus.MustNewConstMetric(latencyMeasureCountDesc, prometheus.CounterValue, float64(m.latencyMeasureCounter.Add(1)))
+	if err != nil {
+		m.logger.Warn(context.Background(), "failed to measure latency", slog.Error(err))
+		ch <- prometheus.MustNewConstMetric(latencyMeasureErrDesc, prometheus.CounterValue, float64(m.latencyErrCounter.Add(1)))
+		return
+	}
+	ch <- prometheus.MustNewConstMetric(sendLatencyDesc, prometheus.GaugeValue, send.Seconds())
+	ch <- prometheus.MustNewConstMetric(recvLatencyDesc, prometheus.GaugeValue, recv.Seconds())
+}

coderd/x/nats/metrics_test.go

@@ -0,0 +1,117 @@
+package nats_test
+
+import (
+	"context"
+	"testing"
+
+	"github.com/prometheus/client_golang/prometheus"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"github.com/coder/coder/v2/coderd/database/pubsub"
+	"github.com/coder/coder/v2/coderd/x/nats"
+	"github.com/coder/coder/v2/testutil"
+)
+
+func TestPubsub_Metrics(t *testing.T) {
+	t.Parallel()
+
+	ctx := testutil.Context(t, testutil.WaitLong)
+	uut := newPubsub(t, nats.Options{})
+
+	registry := prometheus.NewRegistry()
+	err := registry.Register(uut)
+	require.NoError(t, err)
+
+	// each Gather measures pubsub latency by publishing a message & subscribing to it
+	var gatherCount float64
+
+	metrics, err := registry.Gather()
+	gatherCount++
+	require.NoError(t, err)
+	require.True(t, testutil.PromGaugeHasValue(t, metrics, 0, "coder_nats_pubsub_current_events"))
+	require.True(t, testutil.PromGaugeHasValue(t, metrics, 0, "coder_nats_pubsub_current_subscribers"))
+
+	event := "test"
+	data := "testing"
+	messageChannel := make(chan []byte)
+	unsub0, err := uut.Subscribe(event, func(_ context.Context, message []byte) {
+		messageChannel <- message
+	})
+	require.NoError(t, err)
+	go func() {
+		err := uut.Publish(event, []byte(data))
+		assert.NoError(t, err)
+	}()
+	_ = testutil.TryReceive(ctx, t, messageChannel)
+
+	require.Eventually(t, func() bool {
+		latencyBytes := gatherCount * pubsub.LatencyMessageLength
+		metrics, err = registry.Gather()
+		gatherCount++
+		assert.NoError(t, err)
+		return testutil.PromGaugeHasValue(t, metrics, 1, "coder_nats_pubsub_current_events") &&
+			testutil.PromGaugeHasValue(t, metrics, 1, "coder_nats_pubsub_current_subscribers") &&
+			testutil.PromGaugeHasValue(t, metrics, 1, "coder_nats_pubsub_connected") &&
+			testutil.PromCounterHasValue(t, metrics, gatherCount, "coder_nats_pubsub_publishes_total", "true") &&
+			testutil.PromCounterHasValue(t, metrics, gatherCount, "coder_nats_pubsub_subscribes_total", "true") &&
+			testutil.PromCounterHasValue(t, metrics, gatherCount, "coder_nats_pubsub_messages_total", "normal") &&
+			testutil.PromCounterHasValue(t, metrics, float64(len(data))+latencyBytes, "coder_nats_pubsub_received_bytes_total") &&
+			testutil.PromCounterHasValue(t, metrics, float64(len(data))+latencyBytes, "coder_nats_pubsub_published_bytes_total") &&
+			testutil.PromGaugeAssertion(t, metrics, func(in float64) bool { return in > 0 }, "coder_nats_pubsub_send_latency_seconds") &&
+			testutil.PromGaugeAssertion(t, metrics, func(in float64) bool { return in > 0 }, "coder_nats_pubsub_receive_latency_seconds") &&
+			testutil.PromCounterHasValue(t, metrics, gatherCount, "coder_nats_pubsub_latency_measures_total") &&
+			!testutil.PromCounterGathered(t, metrics, "coder_nats_pubsub_latency_measure_errs_total")
+	}, testutil.WaitShort, testutil.IntervalFast)
+
+	colossalSize := 7600
+	colossalData := make([]byte, colossalSize)
+	for i := range colossalData {
+		colossalData[i] = 'q'
+	}
+	unsub1, err := uut.Subscribe(event, func(_ context.Context, message []byte) {
+		messageChannel <- message
+	})
+	require.NoError(t, err)
+	go func() {
+		err := uut.Publish(event, colossalData)
+		assert.NoError(t, err)
+	}()
+	// should get 2 messages because we have 2 subs
+	_ = testutil.TryReceive(ctx, t, messageChannel)
+	_ = testutil.TryReceive(ctx, t, messageChannel)
+
+	require.Eventually(t, func() bool {
+		latencyBytes := gatherCount * pubsub.LatencyMessageLength
+		metrics, err = registry.Gather()
+		gatherCount++
+		assert.NoError(t, err)
+		return testutil.PromGaugeHasValue(t, metrics, 1, "coder_nats_pubsub_current_events") &&
+			testutil.PromGaugeHasValue(t, metrics, 2, "coder_nats_pubsub_current_subscribers") &&
+			testutil.PromGaugeHasValue(t, metrics, 1, "coder_nats_pubsub_connected") &&
+			testutil.PromCounterHasValue(t, metrics, 1+gatherCount, "coder_nats_pubsub_publishes_total", "true") &&
+			testutil.PromCounterHasValue(t, metrics, 1+gatherCount, "coder_nats_pubsub_subscribes_total", "true") &&
+			testutil.PromCounterHasValue(t, metrics, gatherCount, "coder_nats_pubsub_messages_total", "normal") &&
+			testutil.PromCounterHasValue(t, metrics, 1, "coder_nats_pubsub_messages_total", "colossal") &&
+			testutil.PromCounterHasValue(t, metrics, float64(colossalSize+len(data))+latencyBytes, "coder_nats_pubsub_received_bytes_total") &&
+			testutil.PromCounterHasValue(t, metrics, float64(colossalSize+len(data))+latencyBytes, "coder_nats_pubsub_published_bytes_total") &&
+			testutil.PromGaugeAssertion(t, metrics, func(in float64) bool { return in > 0 }, "coder_nats_pubsub_send_latency_seconds") &&
+			testutil.PromGaugeAssertion(t, metrics, func(in float64) bool { return in > 0 }, "coder_nats_pubsub_receive_latency_seconds") &&
+			testutil.PromCounterHasValue(t, metrics, gatherCount, "coder_nats_pubsub_latency_measures_total") &&
+			!testutil.PromCounterGathered(t, metrics, "coder_nats_pubsub_latency_measure_errs_total")
+	}, testutil.WaitShort, testutil.IntervalFast)
+
+	// Unsubscribing both local subscribers should decrement the
+	// subscriber gauge back to 0, and removing the last subscriber on the
+	// event should decrement the event gauge back to 0.
+	unsub0()
+	unsub1()
+
+	require.Eventually(t, func() bool {
+		metrics, err = registry.Gather()
+		gatherCount++
+		assert.NoError(t, err)
+		return testutil.PromGaugeHasValue(t, metrics, 0, "coder_nats_pubsub_current_events") &&
+			testutil.PromGaugeHasValue(t, metrics, 0, "coder_nats_pubsub_current_subscribers")
+	}, testutil.WaitShort, testutil.IntervalFast)
+}

coderd/x/nats/pubsub.go

@@ -11,6 +11,7 @@ import (
 
 	natsserver "github.com/nats-io/nats-server/v2/server"
 	natsgo "github.com/nats-io/nats.go"
+	"github.com/prometheus/client_golang/prometheus"
 	"golang.org/x/xerrors"
 
 	"cdr.dev/slog/v3"
@@ -183,6 +184,8 @@ type Pubsub struct {
 
 	peerFetcher PeerFetcher
 	peerRefresh chan struct{}
+
+	metrics *metrics
 }
 
 // groupSub maps to one underlying *natsgo.Subscription. The first
@@ -190,7 +193,10 @@ type Pubsub struct {
 // When the last local subscriber detaches, the NATS subscription is
 // unsubscribed.
 type groupSub struct {
-	event string
+	// metrics records received-message metrics from handleMessage. It is
+	// the only part of the owning Pubsub that a groupSub needs.
+	metrics *metrics
+	event   string
 	// mu guards localSubs.
 	mu sync.Mutex
 	// localSubs are the local subscribers attached to this NATS subscription.
@@ -214,6 +220,21 @@ type localSub struct {
 // Compile-time assertion that *Pubsub satisfies the pubsub.Pubsub interface.
 var _ pubsub.Pubsub = (*Pubsub)(nil)
 
+// Compile-time assertion that *Pubsub is a prometheus.Collector.
+var _ prometheus.Collector = (*Pubsub)(nil)
+
+// Describe implements prometheus.Collector.
+func (p *Pubsub) Describe(descs chan<- *prometheus.Desc) {
+	p.metrics.describe(descs)
+}
+
+// Collect implements prometheus.Collector. The subscriber and event
+// gauges are maintained as atomic counters by metrics, so Collect does
+// not lock the Pubsub.
+func (p *Pubsub) Collect(ch chan<- prometheus.Metric) {
+	p.metrics.collect(ch, p)
+}
+
 // newPubsub allocates a *Pubsub with initialized maps and cancel ctx.
 func newPubsub(ctx context.Context, logger slog.Logger, opts Options) *Pubsub {
 	ctx, cancel := context.WithCancel(ctx)
@@ -225,6 +246,7 @@ func newPubsub(ctx context.Context, logger slog.Logger, opts Options) *Pubsub {
 		cancel:        cancel,
 		peerFetcher:   opts.PeerFetcher,
 		peerRefresh:   make(chan struct{}, 1),
+		metrics:       newMetrics(logger),
 	}
 }
 
@@ -250,10 +272,12 @@ func (p *Pubsub) buildConnHandlers() connHandlers {
 			if err != nil {
 				p.logger.Warn(p.ctx, "nats client disconnected", slog.Error(err))
 			}
+			p.metrics.onDisconnect()
 			p.signalSubscribersDroppedForConn(conn)
 		},
 		reconnect: func(_ *natsgo.Conn) {
 			p.logger.Info(p.ctx, "nats client reconnected")
+			p.metrics.onReconnect()
 		},
 		closed: func(_ *natsgo.Conn) {
 			p.logger.Debug(p.ctx, "nats client closed")
@@ -320,6 +344,8 @@ func New(ctx context.Context, logger slog.Logger, opts Options) (*Pubsub, error)
 
 	p.publishPool = publishPool
 	p.subscribePool = subscribePool
+	// All owned connections dialed successfully above.
+	p.metrics.markConnected(len(publishPool) + len(subscribePool))
 
 	if p.clustered {
 		go p.runPeerRefresh()
@@ -361,12 +387,15 @@ func newConnPool(ns *natsserver.Server, opts Options, handlers connHandlers, cou
 // same-subject publishes preserve per-subject ordering.
 func (p *Pubsub) Publish(event string, message []byte) error {
 	if p.ctx.Err() != nil {
+		p.metrics.recordPublishFailure()
 		return errClosed
 	}
 
 	if err := pickConn(p.publishPool, event).Publish(event, message); err != nil {
+		p.metrics.recordPublishFailure()
 		return xerrors.Errorf("publish: %w", err)
 	}
+	p.metrics.recordPublishSuccess(len(message))
 	return nil
 }
 
@@ -418,7 +447,7 @@ func (p *Pubsub) subscribeQueue(event string, newQ *pubsub.MsgQueue) (cancel fun
 		defer p.mu.Unlock()
 
 		if p.ctx.Err() != nil {
-			return nil, erroredGroupSub(errClosed)
+			return nil, erroredGroupSub(p.metrics, errClosed)
 		}
 
 		var (
@@ -428,6 +457,7 @@ func (p *Pubsub) subscribeQueue(event string, newQ *pubsub.MsgQueue) (cancel fun
 		gSub, ok = p.subscriptions[event]
 		if !ok {
 			gSub = &groupSub{
+				metrics:   p.metrics,
 				event:     event,
 				localSubs: make(map[*localSub]struct{}),
 				sub: &subGetter{
@@ -436,6 +466,7 @@ func (p *Pubsub) subscribeQueue(event string, newQ *pubsub.MsgQueue) (cancel fun
 			}
 			go p.subscribeGroup(gSub)
 			p.subscriptions[event] = gSub
+			p.metrics.addEvent()
 		}
 		lSub := &localSub{
 			event: event,
@@ -448,8 +479,17 @@ func (p *Pubsub) subscribeQueue(event string, newQ *pubsub.MsgQueue) (cancel fun
 	}()
 
 	if _, err := g.sub.get(); err != nil {
+		p.metrics.recordSubscribeFailure()
+		// 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()
+	// 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()
 	return p.closeLocalSubFunc(l, g), nil
 }
 
@@ -546,6 +586,10 @@ func (p *Pubsub) handleSlowSubscriber(nsub *groupSub) {
 // Close does not drain queued listener messages.
 func (p *Pubsub) Close() error {
 	p.closeOnce.Do(func() {
+		// Report disconnected immediately. The owned connections are
+		// closed below without firing the disconnect handler, so nothing
+		// else resets the gauge during shutdown.
+		p.metrics.markClosed()
 		p.mu.Lock()
 		p.logger.Debug(p.ctx, "closing pubsub")
 		// Cancel while holding p.mu so subscriber state cleanup below
@@ -615,6 +659,7 @@ func (p *Pubsub) closeLocalSubFunc(l *localSub, g *groupSub) func() {
 		l.queue.Close()
 
 		delete(g.localSubs, l)
+		p.metrics.removeSubscriber()
 		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
@@ -628,6 +673,7 @@ func (p *Pubsub) closeLocalSubFunc(l *localSub, g *groupSub) func() {
 		}()
 		if pSub, ok := p.subscriptions[l.event]; ok && g == pSub {
 			delete(p.subscriptions, l.event)
+			p.metrics.removeEvent()
 		}
 	}
 }
@@ -642,6 +688,7 @@ func (p *Pubsub) subscribeGroup(g *groupSub) {
 			defer p.mu.Unlock()
 			if psub := p.subscriptions[g.event]; psub == g {
 				delete(p.subscriptions, g.event)
+				p.metrics.removeEvent()
 			}
 		}
 		close(g.sub.subscribeDone)
@@ -710,10 +757,11 @@ func pickConn(pool []conn, subject string) conn {
 }
 
 // erroredGroupSub returns a groupSub that shows an error rather than an active subscription.
-func erroredGroupSub(err error) *groupSub {
+func erroredGroupSub(m *metrics, err error) *groupSub {
 	c := make(chan struct{})
 	close(c)
 	return &groupSub{
+		metrics: m,
 		sub: &subGetter{
 			subscribeDone: c,
 			err:           err,
@@ -729,6 +777,7 @@ func erroredGroupSub(err error) *groupSub {
 // 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(msg.Data)
 	g.mu.Lock()
 	defer g.mu.Unlock()
 	for l := range g.localSubs {

coderd/x/nats/pubsub_internal_test.go

@@ -11,6 +11,7 @@ import (
 
 	natsserver "github.com/nats-io/nats-server/v2/server"
 	natsgo "github.com/nats-io/nats.go"
+	promtestutil "github.com/prometheus/client_golang/prometheus/testutil"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"golang.org/x/xerrors"
@@ -206,6 +207,88 @@ func Test_Pubsub_buildConnHandlers(t *testing.T) {
 	})
 }
 
+func Test_Pubsub_connectedMetric(t *testing.T) {
+	t.Parallel()
+
+	logger := slogtest.Make(t, nil)
+	ctx := testutil.Context(t, testutil.WaitShort)
+	ps := newPubsub(ctx, logger, defaultTestOptions())
+	handlers := ps.buildConnHandlers()
+
+	// Two owned connections, all up.
+	ps.metrics.markConnected(2)
+	require.Equal(t, 1.0, promtestutil.ToFloat64(ps.metrics.connected))
+	require.Equal(t, 0.0, promtestutil.ToFloat64(ps.metrics.disconnectionsTotal))
+
+	// First disconnect drops the gauge to 0 and counts a disconnection.
+	handlers.disconnectErr(nil, xerrors.New("boom"))
+	require.Equal(t, 0.0, promtestutil.ToFloat64(ps.metrics.connected))
+	require.Equal(t, 1.0, promtestutil.ToFloat64(ps.metrics.disconnectionsTotal))
+
+	// Second disconnect: still down, counts again.
+	handlers.disconnectErr(nil, xerrors.New("boom"))
+	require.Equal(t, 0.0, promtestutil.ToFloat64(ps.metrics.connected))
+	require.Equal(t, 2.0, promtestutil.ToFloat64(ps.metrics.disconnectionsTotal))
+
+	// One reconnect with the other still down keeps the gauge at 0.
+	handlers.reconnect(nil)
+	require.Equal(t, 0.0, promtestutil.ToFloat64(ps.metrics.connected))
+
+	// Once every owned connection is back the gauge returns to 1.
+	handlers.reconnect(nil)
+	require.Equal(t, 1.0, promtestutil.ToFloat64(ps.metrics.connected))
+}
+
+func Test_Pubsub_failureMetrics(t *testing.T) {
+	t.Parallel()
+
+	logger := slogtest.Make(t, nil)
+	ctx := testutil.Context(t, testutil.WaitShort)
+	ps := newPubsub(ctx, logger, defaultTestOptions())
+	// Closing makes Publish and Subscribe fail fast so we can exercise the
+	// success="false" label without needing the embedded server to error.
+	require.NoError(t, ps.Close())
+
+	require.Error(t, ps.Publish("evt", []byte("payload")))
+	_, err := ps.Subscribe("evt", func(context.Context, []byte) {})
+	require.Error(t, err)
+
+	require.Equal(t, 1.0, promtestutil.ToFloat64(ps.metrics.publishesTotal.WithLabelValues("false")))
+	require.Equal(t, 1.0, promtestutil.ToFloat64(ps.metrics.subscribesTotal.WithLabelValues("false")))
+}
+
+func Test_Pubsub_gracefulCloseDoesNotCountDisconnect(t *testing.T) {
+	t.Parallel()
+
+	ps := newTestPubsub(t, defaultTestOptions())
+	require.Equal(t, 0.0, promtestutil.ToFloat64(ps.metrics.disconnectionsTotal))
+	require.Equal(t, 1.0, promtestutil.ToFloat64(ps.metrics.connected))
+
+	handlers := ps.buildConnHandlers()
+	require.NoError(t, ps.Close())
+
+	// Close reports disconnected even though the disconnect handler is
+	// suppressed for our own connection closes.
+	require.Equal(t, 0.0, promtestutil.ToFloat64(ps.metrics.connected))
+
+	// Late reconnect callbacks during the shutdown window must not flip
+	// the gauge back to 1: markClosed zeroes totalConns so the connected
+	// guard stays false even if every owned connection reports a
+	// reconnect. Fire one per owned connection to exercise that.
+	for range len(ps.publishPool) + len(ps.subscribePool) {
+		handlers.reconnect(nil)
+	}
+	require.Equal(t, 0.0, promtestutil.ToFloat64(ps.metrics.connected))
+
+	// Closing our own connections must not invoke the disconnect handler,
+	// so disconnections_total stays 0. The async callback would fire
+	// within milliseconds if it were going to, so a short window catches a
+	// regression without making the test slow.
+	require.Never(t, func() bool {
+		return promtestutil.ToFloat64(ps.metrics.disconnectionsTotal) > 0
+	}, 2*time.Second, testutil.IntervalFast)
+}
+
 func Test_localSub(t *testing.T) {
 	t.Parallel()
 

coderd/x/nats/server.go

@@ -96,6 +96,11 @@ type connHandlers struct {
 func connectClient(ns *natsserver.Server, opts Options, handlers connHandlers, connName string) (*natsgo.Conn, error) {
 	connOpts := []natsgo.Option{
 		natsgo.Name(connName),
+		// Suppress async callbacks when we close the connection ourselves
+		// during Pubsub.Close, so a graceful shutdown does not fire the
+		// disconnect handler and inflate disconnections_total. Genuine
+		// disconnects still invoke the handler.
+		natsgo.NoCallbacksAfterClientClose(),
 	}
 	if opts.ClusterAuthToken != "" {
 		connOpts = append(connOpts, natsgo.Token(opts.ClusterAuthToken))