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

components/metrics/builder.go

@@ -0,0 +1,140 @@
+package metrics
+
+import (
+	"github.com/ThreeDotsLabs/watermill/internal"
+	"github.com/ThreeDotsLabs/watermill/message"
+	"github.com/pkg/errors"
+	"github.com/prometheus/client_golang/prometheus"
+)
+
+type PrometheusMetricsBuilderConfig struct {
+	Namespace        string
+	Subsystem        string
+	AdditionalLabels []MetricLabel
+}
+
+func NewPrometheusMetricsBuilderWithConfig(prometheusRegistry prometheus.Registerer, config PrometheusMetricsBuilderConfig) PrometheusMetricsBuilder {
+	builder := PrometheusMetricsBuilder{
+		Namespace:          config.Namespace,
+		Subsystem:          config.Subsystem,
+		PrometheusRegistry: prometheusRegistry,
+		additionalLabels:   config.AdditionalLabels,
+	}
+	return builder
+}
+
+func NewPrometheusMetricsBuilder(prometheusRegistry prometheus.Registerer, namespace string, subsystem string) PrometheusMetricsBuilder {
+	return NewPrometheusMetricsBuilderWithConfig(prometheusRegistry, PrometheusMetricsBuilderConfig{
+		Namespace: namespace,
+		Subsystem: subsystem,
+	})
+}
+
+// PrometheusMetricsBuilder provides methods to decorate publishers, subscribers and handlers.
+type PrometheusMetricsBuilder struct {
+	// PrometheusRegistry may be filled with a pre-existing Prometheus registry, or left empty for the default registry.
+	PrometheusRegistry prometheus.Registerer
+
+	Namespace string
+	Subsystem string
+	// PublishBuckets defines the histogram buckets for publish time histogram, defaulted if nil.
+	PublishBuckets []float64
+	// HandlerBuckets defines the histogram buckets for handle execution time histogram, defaulted to watermill's default.
+	HandlerBuckets []float64
+
+	additionalLabels []MetricLabel
+}
+
+// AddPrometheusRouterMetrics is a convenience function that acts on the message router to add the metrics middleware
+// to all its handlers. The handlers' publishers and subscribers are also decorated.
+// The default buckets are used for the handler execution time histogram (use your own provisioning
+// with NewRouterMiddlewareWithConfig if needed).
+func (b PrometheusMetricsBuilder) AddPrometheusRouterMetrics(r *message.Router) {
+	r.AddPublisherDecorators(b.DecoratePublisher)
+	r.AddSubscriberDecorators(b.DecorateSubscriber)
+	r.AddMiddleware(b.NewRouterMiddleware().Middleware)
+}
+
+// DecoratePublisher wraps the underlying publisher with Prometheus metrics.
+func (b PrometheusMetricsBuilder) DecoratePublisher(pub message.Publisher) (message.Publisher, error) {
+	var err error
+	d := PublisherPrometheusMetricsDecorator{
+		pub:              pub,
+		publisherName:    internal.StructName(pub),
+		additionalLabels: b.additionalLabels,
+	}
+
+	d.publishTimeSeconds, err = b.registerHistogramVec(prometheus.NewHistogramVec(
+		prometheus.HistogramOpts{
+			Namespace: b.Namespace,
+			Subsystem: b.Subsystem,
+			Name:      "publish_time_seconds",
+			Help:      "The time that a publishing attempt (success or not) took in seconds",
+			Buckets:   b.PublishBuckets,
+		},
+		toLabelsSlice(publisherLabelKeys, b.additionalLabels),
+	))
+	if err != nil {
+		return nil, errors.Wrap(err, "could not register publish time metric")
+	}
+	return d, nil
+}
+
+// DecorateSubscriber wraps the underlying subscriber with Prometheus metrics.
+func (b PrometheusMetricsBuilder) DecorateSubscriber(sub message.Subscriber) (message.Subscriber, error) {
+	var err error
+	d := &SubscriberPrometheusMetricsDecorator{
+		closing:          make(chan struct{}),
+		subscriberName:   internal.StructName(sub),
+		additionalLabels: b.additionalLabels,
+	}
+
+	d.subscriberMessagesReceivedTotal, err = b.registerCounterVec(prometheus.NewCounterVec(
+		prometheus.CounterOpts{
+			Namespace: b.Namespace,
+			Subsystem: b.Subsystem,
+			Name:      "subscriber_messages_received_total",
+			Help:      "The total number of messages received by the subscriber",
+		},
+		toLabelsSlice(append(subscriberLabelKeys, labelAcked), b.additionalLabels),
+	))
+	if err != nil {
+		return nil, errors.Wrap(err, "could not register time to ack metric")
+	}
+
+	d.Subscriber, err = message.MessageTransformSubscriberDecorator(d.recordMetrics)(sub)
+	if err != nil {
+		return nil, errors.Wrap(err, "could not decorate subscriber with metrics decorator")
+	}
+
+	return d, nil
+}
+
+func (b PrometheusMetricsBuilder) register(c prometheus.Collector) (prometheus.Collector, error) {
+	err := b.PrometheusRegistry.Register(c)
+	if err == nil {
+		return c, nil
+	}
+
+	if are, ok := err.(prometheus.AlreadyRegisteredError); ok {
+		return are.ExistingCollector, nil
+	}
+
+	return nil, err
+}
+
+func (b PrometheusMetricsBuilder) registerCounterVec(c *prometheus.CounterVec) (*prometheus.CounterVec, error) {
+	col, err := b.register(c)
+	if err != nil {
+		return nil, err
+	}
+	return col.(*prometheus.CounterVec), nil
+}
+
+func (b PrometheusMetricsBuilder) registerHistogramVec(h *prometheus.HistogramVec) (*prometheus.HistogramVec, error) {
+	col, err := b.register(h)
+	if err != nil {
+		return nil, err
+	}
+	return col.(*prometheus.HistogramVec), nil
+}

components/metrics/ctx.go

@@ -0,0 +1,28 @@
+package metrics
+
+import "context"
+
+type contextValue int
+
+const (
+	publishObserved contextValue = iota
+	subscribeObserved
+)
+
+// setPublishObservedToCtx is used to achieve metrics idempotency in case of double applied middleware
+func setPublishObservedToCtx(ctx context.Context) context.Context {
+	return context.WithValue(ctx, publishObserved, true)
+}
+
+func publishAlreadyObserved(ctx context.Context) bool {
+	return ctx.Value(publishObserved) != nil
+}
+
+// setSubscribeObservedToCtx is used to achieve metrics idempotency in case of double applied middleware
+func setSubscribeObservedToCtx(ctx context.Context) context.Context {
+	return context.WithValue(ctx, subscribeObserved, true)
+}
+
+func subscribeAlreadyObserved(ctx context.Context) bool {
+	return ctx.Value(subscribeObserved) != nil
+}

components/metrics/handler.go

@@ -0,0 +1,92 @@
+package metrics
+
+import (
+	"time"
+
+	"github.com/pkg/errors"
+	"github.com/prometheus/client_golang/prometheus"
+
+	"github.com/ThreeDotsLabs/watermill/message"
+)
+
+var (
+	handlerLabelKeys = []string{
+		labelKeyHandlerName,
+		labelSuccess,
+	}
+
+	// defaultHandlerExecutionTimeBuckets are one order of magnitude smaller than default buckets (5ms~10s),
+	// because the handler execution times are typically shorter (µs~ms range).
+	defaultHandlerExecutionTimeBuckets = []float64{
+		0.0005,
+		0.001,
+		0.0025,
+		0.005,
+		0.01,
+		0.025,
+		0.05,
+		0.1,
+		0.25,
+		0.5,
+		1,
+	}
+)
+
+// HandlerPrometheusMetricsMiddleware is a middleware that captures Prometheus metrics.
+type HandlerPrometheusMetricsMiddleware struct {
+	handlerExecutionTimeSeconds *prometheus.HistogramVec
+	additionalLabels            []MetricLabel
+}
+
+// Middleware returns the middleware ready to be used with watermill's Router.
+func (m HandlerPrometheusMetricsMiddleware) Middleware(h message.HandlerFunc) message.HandlerFunc {
+	return func(msg *message.Message) (msgs []*message.Message, err error) {
+		now := time.Now()
+		ctx := msg.Context()
+		labels := prometheus.Labels{
+			labelKeyHandlerName: message.HandlerNameFromCtx(ctx),
+		}
+		for _, lb := range m.additionalLabels {
+			labels[lb.Label] = lb.ComputeValueFn(ctx)
+		}
+
+		defer func() {
+			if err != nil {
+				labels[labelSuccess] = "false"
+			} else {
+				labels[labelSuccess] = "true"
+			}
+			m.handlerExecutionTimeSeconds.With(labels).Observe(time.Since(now).Seconds())
+		}()
+
+		return h(msg)
+	}
+}
+
+// NewRouterMiddleware returns new middleware.
+func (b PrometheusMetricsBuilder) NewRouterMiddleware() HandlerPrometheusMetricsMiddleware {
+	var err error
+	m := HandlerPrometheusMetricsMiddleware{
+		additionalLabels: b.additionalLabels,
+	}
+
+	if b.HandlerBuckets == nil {
+		b.HandlerBuckets = defaultHandlerExecutionTimeBuckets
+	}
+
+	m.handlerExecutionTimeSeconds, err = b.registerHistogramVec(prometheus.NewHistogramVec(
+		prometheus.HistogramOpts{
+			Namespace: b.Namespace,
+			Subsystem: b.Subsystem,
+			Name:      "handler_execution_time_seconds",
+			Help:      "The total time elapsed while executing the handler function in seconds",
+			Buckets:   b.HandlerBuckets,
+		},
+		toLabelsSlice(handlerLabelKeys, b.additionalLabels),
+	))
+	if err != nil {
+		panic(errors.Wrap(err, "could not register handler execution time metric"))
+	}
+
+	return m
+}

components/metrics/http.go

@@ -0,0 +1,41 @@
+package metrics
+
+import (
+	"errors"
+	"net/http"
+
+	"github.com/go-chi/chi/v5"
+	"github.com/prometheus/client_golang/prometheus"
+	"github.com/prometheus/client_golang/prometheus/promhttp"
+)
+
+// CreateRegistryAndServeHTTP establishes an HTTP server that exposes the /metrics endpoint for Prometheus at the given address.
+// It returns a new prometheus registry (to register the metrics on) and a canceling function that ends the server.
+func CreateRegistryAndServeHTTP(addr string) (registry *prometheus.Registry, cancel func()) {
+	registry = prometheus.NewRegistry()
+	return registry, ServeHTTP(addr, registry)
+}
+
+// ServeHTTP establishes an HTTP server that exposes the /metrics endpoint for Prometheus at the given address.
+// It takes an existing Prometheus registry and returns a canceling function that ends the server.
+func ServeHTTP(addr string, registry *prometheus.Registry) (cancel func()) {
+	router := chi.NewRouter()
+
+	handler := promhttp.HandlerFor(registry, promhttp.HandlerOpts{})
+	router.Get("/metrics", func(w http.ResponseWriter, r *http.Request) {
+		handler.ServeHTTP(w, r)
+	})
+	server := http.Server{
+		Addr:    addr,
+		Handler: router,
+	}
+
+	go func() {
+		err := server.ListenAndServe()
+		if !errors.Is(err, http.ErrServerClosed) {
+			panic(err)
+		}
+	}()
+
+	return func() { _ = server.Close() }
+}

components/metrics/http_test.go

@@ -0,0 +1,66 @@
+package metrics_test
+
+import (
+	"net/http"
+	"testing"
+	"time"
+
+	"github.com/pkg/errors"
+	"github.com/prometheus/client_golang/prometheus/collectors"
+	"github.com/stretchr/testify/assert"
+
+	"github.com/ThreeDotsLabs/watermill/components/metrics"
+)
+
+func TestCreateRegistryAndServeHTTP_metrics_endpoint(t *testing.T) {
+	reg, cancel := metrics.CreateRegistryAndServeHTTP(":8090")
+	defer cancel()
+	err := reg.Register(collectors.NewBuildInfoCollector())
+	if err != nil {
+		t.Fatal(errors.Wrap(err, "registration of prometheus build info collector failed"))
+	}
+	waitServerReady(t, "http://localhost:8090")
+	resp, err := http.DefaultClient.Get("http://localhost:8090/metrics")
+	if resp != nil {
+		defer func() { _ = resp.Body.Close() }()
+	}
+
+	if err != nil {
+		t.Fatal(errors.Wrap(err, "call to metrics endpoint failed"))
+	}
+	assert.NotNil(t, resp)
+	assert.Equal(t, http.StatusOK, resp.StatusCode)
+}
+
+func TestCreateRegistryAndServeHTTP_unknown_endpoint(t *testing.T) {
+	reg, cancel := metrics.CreateRegistryAndServeHTTP(":8091")
+	defer cancel()
+	err := reg.Register(collectors.NewBuildInfoCollector())
+	if err != nil {
+		t.Error(errors.Wrap(err, "registration of prometheus build info collector failed"))
+	}
+	waitServerReady(t, "http://localhost:8091")
+	resp, err := http.DefaultClient.Get("http://localhost:8091/unknown")
+	if resp != nil {
+		defer func() { _ = resp.Body.Close() }()
+	}
+
+	if err != nil {
+		t.Fatal(errors.Wrap(err, "call to unknown endpoint failed"))
+	}
+	assert.NotNil(t, resp)
+	assert.Equal(t, http.StatusNotFound, resp.StatusCode)
+}
+
+// server might have small delay before being able to server traffic
+func waitServerReady(t *testing.T, addr string) {
+	for i := 0; i < 50; i++ {
+		_, err := http.DefaultClient.Get(addr)
+		// assume server ready when no err anymore
+		if err == nil {
+			return
+		}
+		time.Sleep(100 * time.Millisecond)
+		continue
+	}
+}

components/metrics/labels.go

@@ -0,0 +1,75 @@
+package metrics
+
+import (
+	"context"
+
+	"github.com/ThreeDotsLabs/watermill/message"
+
+	"github.com/prometheus/client_golang/prometheus"
+)
+
+const (
+	labelKeyHandlerName    = "handler_name"
+	labelKeyPublisherName  = "publisher_name"
+	labelKeySubscriberName = "subscriber_name"
+	labelSuccess           = "success"
+	labelAcked             = "acked"
+
+	labelValueNoHandler = "<no handler>"
+)
+
+var (
+	labelGetters = map[string]func(context.Context) string{
+		labelKeyHandlerName:    message.HandlerNameFromCtx,
+		labelKeyPublisherName:  message.PublisherNameFromCtx,
+		labelKeySubscriberName: message.SubscriberNameFromCtx,
+	}
+)
+
+func labelsFromCtx(ctx context.Context, labels ...string) prometheus.Labels {
+	ctxLabels := map[string]string{}
+
+	for _, l := range labels {
+		k := l
+		ctxLabels[l] = ""
+
+		getter, ok := labelGetters[k]
+		if !ok {
+			continue
+		}
+
+		v := getter(ctx)
+		if v != "" {
+			ctxLabels[l] = v
+		}
+	}
+
+	return ctxLabels
+}
+
+type LabelComputeValueFn func(msgCtx context.Context) string
+
+type MetricLabel struct {
+	Label          string
+	ComputeValueFn LabelComputeValueFn
+}
+
+func toLabelsSlice(baseLabels []string, customs []MetricLabel) []string {
+	labels := make([]string, len(baseLabels), len(baseLabels)+len(customs))
+	copy(labels, baseLabels)
+	for _, label := range customs {
+		//Check if the additional label is already in the base labels. We cannot have duplicate labels
+		//If it's in the base, just skip it as the compute function is going to overwrite the default value
+		contains := false
+		for _, baseLabel := range baseLabels {
+			if baseLabel == label.Label {
+				contains = true
+				break
+			}
+		}
+		if !contains {
+			labels = append(labels, label.Label)
+		}
+	}
+	return labels
+}

components/metrics/publisher.go

@@ -0,0 +1,70 @@
+package metrics
+
+import (
+	"time"
+
+	"github.com/ThreeDotsLabs/watermill/message"
+	"github.com/prometheus/client_golang/prometheus"
+)
+
+var (
+	publisherLabelKeys = []string{
+		labelKeyHandlerName,
+		labelKeyPublisherName,
+		labelSuccess,
+	}
+)
+
+// PublisherPrometheusMetricsDecorator decorates a publisher to capture Prometheus metrics.
+type PublisherPrometheusMetricsDecorator struct {
+	pub                message.Publisher
+	publisherName      string
+	publishTimeSeconds *prometheus.HistogramVec
+	additionalLabels   []MetricLabel
+}
+
+// Publish updates the relevant publisher metrics and calls the wrapped publisher's Publish.
+func (m PublisherPrometheusMetricsDecorator) Publish(topic string, messages ...*message.Message) (err error) {
+	if len(messages) == 0 {
+		return m.pub.Publish(topic)
+	}
+
+	// TODO: take ctx not only from first msg. Might require changing the signature of Publish, which is planned anyway.
+	ctx := messages[0].Context()
+	labels := labelsFromCtx(ctx, publisherLabelKeys...)
+	if labels[labelKeyPublisherName] == "" {
+		labels[labelKeyPublisherName] = m.publisherName
+	}
+	if labels[labelKeyHandlerName] == "" {
+		labels[labelKeyHandlerName] = labelValueNoHandler
+	}
+	for _, lb := range m.additionalLabels {
+		labels[lb.Label] = lb.ComputeValueFn(ctx)
+	}
+	start := time.Now()
+
+	defer func() {
+		if publishAlreadyObserved(ctx) {
+			// decorator idempotency when applied decorator multiple times
+			return
+		}
+
+		if err != nil {
+			labels[labelSuccess] = "false"
+		} else {
+			labels[labelSuccess] = "true"
+		}
+		m.publishTimeSeconds.With(labels).Observe(time.Since(start).Seconds())
+	}()
+
+	for _, msg := range messages {
+		msg.SetContext(setPublishObservedToCtx(msg.Context()))
+	}
+
+	return m.pub.Publish(topic, messages...)
+}
+
+// Close decreases the total publisher count, closes the Prometheus HTTP server and calls wrapped Close.
+func (m PublisherPrometheusMetricsDecorator) Close() error {
+	return m.pub.Close()
+}

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