maint: add instrumentation for GoRedisPubSub

internal/peer/peers_test.go

@@ -18,6 +18,7 @@ import (
 	"github.com/sirupsen/logrus"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+	"go.opentelemetry.io/otel/trace/noop"
 )
 
 func newPeers(c config.Config) (Peers, error) {
@@ -46,7 +47,10 @@ func newPeers(c config.Config) (Peers, error) {
 			pubsubber = &pubsub.GoRedisPubSub{}
 		}
 	case "redis":
-		pubsubber = &pubsub.GoRedisPubSub{}
+		pubsubber = &pubsub.GoRedisPubSub{
+			Metrics: &metrics.NullMetrics{},
+			Tracer:  noop.NewTracerProvider().Tracer("test"),
+		}
 		peers = &RedisPubsubPeers{}
 	default:
 		// this should have been caught by validation

internal/peer/pubsub_redis.go

@@ -96,7 +96,7 @@ func (p *RedisPubsubPeers) checkHash() {
 	p.Metrics.Gauge("peer_hash", float64(p.hash))
 }
 
-func (p *RedisPubsubPeers) listen(msg string) {
+func (p *RedisPubsubPeers) listen(ctx context.Context, msg string) {
 	cmd := &peerCommand{}
 	if !cmd.unmarshal(msg) {
 		return

pubsub/pubsub.go

@@ -26,7 +26,7 @@ type PubSub interface {
 	// There is no unsubscribe method; close the subscription to stop receiving messages.
 	// The subscription only exists to provide a way to stop receiving messages; if you don't need to stop,
 	// you can ignore the return value.
-	Subscribe(ctx context.Context, topic string, callback func(msg string)) Subscription
+	Subscribe(ctx context.Context, topic string, callback SubscriptionCallback) Subscription
 	// Close shuts down all topics and the pubsub connection.
 	Close()
 
@@ -36,6 +36,9 @@ type PubSub interface {
 	startstop.Stopper
 }
 
+// SubscriptionCallback is the function signature for a subscription callback.
+type SubscriptionCallback func(context.Context, string)
+
 type Subscription interface {
 	// Close stops the subscription which means the callback will no longer be called.
 	// Optional; the topic will be closed when the pubsub connection is closed.

pubsub/pubsub_goredis.go

@@ -5,7 +5,10 @@ import (
 	"strings"
 	"sync"
 
+	"go.opentelemetry.io/otel/trace"
+
 	"github.com/honeycombio/refinery/config"
+	"github.com/honeycombio/refinery/internal/otelutil"
 	"github.com/honeycombio/refinery/logger"
 	"github.com/honeycombio/refinery/metrics"
 	"github.com/redis/go-redis/v9"
@@ -25,6 +28,7 @@ type GoRedisPubSub struct {
 	Config  config.Config   `inject:""`
 	Logger  logger.Logger   `inject:""`
 	Metrics metrics.Metrics `inject:"metrics"`
+	Tracer  trace.Tracer    `inject:"tracer"`
 	client  redis.UniversalClient
 	subs    []*GoRedisSubscription
 	mut     sync.RWMutex
@@ -36,7 +40,7 @@ var _ PubSub = (*GoRedisPubSub)(nil)
 type GoRedisSubscription struct {
 	topic  string
 	pubsub *redis.PubSub
-	cb     func(msg string)
+	cb     SubscriptionCallback
 	done   chan struct{}
 	once   sync.Once
 }
@@ -111,6 +115,13 @@ func (ps *GoRedisPubSub) Close() {
 }
 
 func (ps *GoRedisPubSub) Publish(ctx context.Context, topic, message string) error {
+	ctx, span := otelutil.StartSpanMulti(ctx, ps.Tracer, "GoRedisPubSub.Publish", map[string]interface{}{
+		"topic":   topic,
+		"message": message,
+	})
+
+	defer span.End()
+
 	ps.Metrics.Count("redis_pubsub_published", 1)
 	return ps.client.Publish(ctx, topic, message).Err()
 }
@@ -119,7 +130,10 @@ func (ps *GoRedisPubSub) Publish(ctx context.Context, topic, message string) err
 // whenever a message is received on that topic.
 // Note that the same topic is Subscribed to multiple times, this will incur a separate
 // connection to Redis for each Subscription.
-func (ps *GoRedisPubSub) Subscribe(ctx context.Context, topic string, callback func(string)) Subscription {
+func (ps *GoRedisPubSub) Subscribe(ctx context.Context, topic string, callback SubscriptionCallback) Subscription {
+	ctx, span := otelutil.StartSpanWith(ctx, ps.Tracer, "GoRedisPubSub.Subscribe", "topic", topic)
+	defer span.End()
+
 	sub := &GoRedisSubscription{
 		topic:  topic,
 		pubsub: ps.client.Subscribe(ctx, topic),
@@ -130,6 +144,7 @@ func (ps *GoRedisPubSub) Subscribe(ctx context.Context, topic string, callback f
 	ps.subs = append(ps.subs, sub)
 	ps.mut.Unlock()
 	go func() {
+		receiveRootCtx := context.Background()
 		redisch := sub.pubsub.Channel()
 		for {
 			select {
@@ -139,8 +154,18 @@ func (ps *GoRedisPubSub) Subscribe(ctx context.Context, topic string, callback f
 				if msg == nil {
 					continue
 				}
+				receiveCtx, span := otelutil.StartSpanMulti(receiveRootCtx, ps.Tracer, "GoRedisPubSub.Receive", map[string]interface{}{
+					"topic":              topic,
+					"message_queue_size": len(redisch),
+					"message":            msg.Payload,
+				})
 				ps.Metrics.Count("redis_pubsub_received", 1)
-				go sub.cb(msg.Payload)
+
+				go func(cbCtx context.Context, span trace.Span) {
+					defer span.End()
+
+					sub.cb(cbCtx, msg.Payload)
+				}(receiveCtx, span)
 			}
 		}
 	}()

pubsub/pubsub_local.go

@@ -24,7 +24,7 @@ var _ PubSub = (*LocalPubSub)(nil)
 type LocalSubscription struct {
 	ps    *LocalPubSub
 	topic string
-	cb    func(string)
+	cb    SubscriptionCallback
 	mut   sync.RWMutex
 }
 
@@ -71,13 +71,13 @@ func (ps *LocalPubSub) Publish(ctx context.Context, topic, message string) error
 	for _, sub := range ps.topics[topic] {
 		// don't wait around for slow consumers
 		if sub.cb != nil {
-			go sub.cb(message)
+			go sub.cb(ctx, message)
 		}
 	}
 	return nil
 }
 
-func (ps *LocalPubSub) Subscribe(ctx context.Context, topic string, callback func(msg string)) Subscription {
+func (ps *LocalPubSub) Subscribe(ctx context.Context, topic string, callback SubscriptionCallback) Subscription {
 	ps.mut.Lock()
 	ps.ensureTopic(topic)
 	sub := &LocalSubscription{ps: ps, topic: topic, cb: callback}

pubsub/pubsub_test.go

@@ -12,6 +12,7 @@ import (
 	"github.com/honeycombio/refinery/pubsub"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+	"go.opentelemetry.io/otel/trace/noop"
 )
 
 var types = []string{
@@ -23,10 +24,12 @@ func newPubSub(typ string) pubsub.PubSub {
 	var ps pubsub.PubSub
 	m := &metrics.NullMetrics{}
 	m.Start()
+	tracer := noop.NewTracerProvider().Tracer("test")
 	switch typ {
 	case "goredis":
 		ps = &pubsub.GoRedisPubSub{
 			Metrics: m,
+			Tracer:  tracer,
 		}
 	case "local":
 		ps = &pubsub.LocalPubSub{
@@ -44,7 +47,7 @@ type pubsubListener struct {
 	msgs []string
 }
 
-func (l *pubsubListener) Listen(msg string) {
+func (l *pubsubListener) Listen(ctx context.Context, msg string) {
 	l.lock.Lock()
 	defer l.lock.Unlock()
 	l.msgs = append(l.msgs, msg)
@@ -196,7 +199,7 @@ func TestPubSubLatency(t *testing.T) {
 				wg.Done()
 			}()
 
-			ps.Subscribe(ctx, "topic", func(msg string) {
+			ps.Subscribe(ctx, "topic", func(ctx context.Context, msg string) {
 				sent, err := strconv.Atoi(msg)
 				require.NoError(t, err)
 				rcvd := time.Now().UnixNano()