snapshot of https://github.com/ThreeDotsLabs/watermill.git (master, 2026-09-15; blobs > 131072B omitted)

components/metrics/subscriber.go

@@ -0,0 +1,57 @@
+package metrics
+
+import (
+	"github.com/ThreeDotsLabs/watermill/message"
+	"github.com/prometheus/client_golang/prometheus"
+)
+
+var (
+	subscriberLabelKeys = []string{
+		labelKeyHandlerName,
+		labelKeySubscriberName,
+	}
+)
+
+// SubscriberPrometheusMetricsDecorator decorates a subscriber to capture Prometheus metrics.
+type SubscriberPrometheusMetricsDecorator struct {
+	message.Subscriber
+	subscriberName                  string
+	subscriberMessagesReceivedTotal *prometheus.CounterVec
+	closing                         chan struct{}
+	additionalLabels                []MetricLabel
+}
+
+func (s SubscriberPrometheusMetricsDecorator) recordMetrics(msg *message.Message) {
+	if msg == nil {
+		return
+	}
+
+	ctx := msg.Context()
+	labels := labelsFromCtx(ctx, subscriberLabelKeys...)
+	if labels[labelKeySubscriberName] == "" {
+		labels[labelKeySubscriberName] = s.subscriberName
+	}
+	if labels[labelKeyHandlerName] == "" {
+		labels[labelKeyHandlerName] = labelValueNoHandler
+	}
+	for _, lb := range s.additionalLabels {
+		labels[lb.Label] = lb.ComputeValueFn(ctx)
+	}
+
+	go func() {
+		if subscribeAlreadyObserved(ctx) {
+			// decorator idempotency when applied decorator multiple times
+			return
+		}
+
+		select {
+		case <-msg.Acked():
+			labels[labelAcked] = "acked"
+		case <-msg.Nacked():
+			labels[labelAcked] = "nacked"
+		}
+		s.subscriberMessagesReceivedTotal.With(labels).Inc()
+	}()
+
+	msg.SetContext(setSubscribeObservedToCtx(msg.Context()))
+}