From c5735ac6e8f739d9824e40f05ebc64cae122e1a0 Mon Sep 17 00:00:00 2001 From: William Johansson Date: Thu, 5 Sep 2024 20:52:27 +0200 Subject: [PATCH] feat(pubsub): allow trace extraction from protobuf message 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. --- pubsub/trace.go | 24 +++++++++++++++++------- 1 file changed, 17 insertions(+), 7 deletions(-) diff --git a/pubsub/trace.go b/pubsub/trace.go index 1d41e9d89422..51112bb50f5f 100644 --- 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++ }