feat(pubsub): allow trace extraction from protobuf message (#10827)

radhus · hongalex · web-flow · commit caa826cea826 · 2024-09-05T19:38:57.000-07:00
Export the function NewMessageCarrierFromPB which, given a protobuf
PubsubMessage, returns a TextMapCarrier which can be used to extract
trace contexts.

This allows it to be used in e.g. HTTP push endpoints.

Co-authored-by: Alex Hong &lt;9397363+hongalex@users.noreply.github.com&gt;
diff --git a/pubsub/trace.go b/pubsub/trace.go
@@ -20,6 +20,7 @@ import (
 	"log"
 	"sync"
 
+	pb "cloud.google.com/go/pubsub/apiv1/pubsubpb"
 	"cloud.google.com/go/pubsub/internal"
 	"go.opencensus.io/stats"
 	"go.opencensus.io/stats/view"
@@ -273,33 +274,42 @@ func tracer() trace.Tracer {
 
 var _ propagation.TextMapCarrier = (*messageCarrier)(nil)
 
-// messageCarrier injects and extracts traces from a pubsub.Message.
+// messageCarrier injects and extracts traces from pubsub.Message attributes.
 type messageCarrier struct {
-	msg *Message
+	attributes map[string]string
 }
 
 const googclientPrefix string = "googclient_"
 
 // newMessageCarrier creates a new PubsubMessageCarrier.
 func newMessageCarrier(msg *Message) messageCarrier {
-	return messageCarrier{msg: msg}
+	return messageCarrier{attributes: msg.Attributes}
+}
+
+// NewMessageCarrierFromPB creates a propagation.TextMapCarrier that can be used to extract the trace
+// context from a protobuf PubsubMessage.
+//
+// Example:
+// ctx = propagation.TraceContext{}.Extract(ctx, pubsub.NewMessageCarrierFromPB(msg))
+func NewMessageCarrierFromPB(msg *pb.PubsubMessage) propagation.TextMapCarrier {
+	return messageCarrier{attributes: msg.Attributes}
 }
 
 // Get retrieves a single value for a given key.
 func (c messageCarrier) Get(key string) string {
-	return c.msg.Attributes[googclientPrefix+key]
+	return c.attributes[googclientPrefix+key]
 }
 
 // Set sets an attribute.
 func (c messageCarrier) Set(key, val string) {
-	c.msg.Attributes[googclientPrefix+key] = val
+	c.attributes[googclientPrefix+key] = val
 }
 
 // Keys returns a slice of all keys in the carrier.
 func (c messageCarrier) Keys() []string {
 	i := 0
-	out := make([]string, len(c.msg.Attributes))
-	for k := range c.msg.Attributes {
+	out := make([]string, len(c.attributes))
+	for k := range c.attributes {
 		out[i] = k
 		i++
 	}
