Merge pull request #3 from Vinelab/pubsub

Add support for Google Cloud PubSub

Author: adiachenko

README.md

@@ -296,6 +296,7 @@ import "github.com/Vinelab/tracing-go/formats"
 spanCtx, err := Trace.Extract(&carrier, formats.TextMap)
 spanCtx, err := Trace.Extract(&carrier, formats.HTTP)
 spanCtx, err := Trace.Extract(&carrier, formats.AMQP)
+spanCtx, err := Trace.Extract(&carrier, formats.GooglePubSub)
 ```
 
 You may also add your own format using `RegisterExtractionFormat` method:
@@ -328,6 +329,7 @@ import "github.com/Vinelab/tracing-go/formats"
 err := Trace.Inject(&carrier, formats.TextMap)
 err := Trace.Inject(&carrier, formats.HTTP)
 err := Trace.Inject(&carrier, formats.AMQP)
+err := Trace.Inject(&carrier, formats.GooglePubSub)
 ```
 
 You may also add your own format using `RegisterInjectionFormat` method.

drivers/zipkin/extract_google_pubsub.go

@@ -0,0 +1,31 @@
+package zipkin
+
+import (
+	"log"
+
+	"cloud.google.com/go/pubsub"
+	"github.com/Vinelab/tracing-go"
+	"github.com/Vinelab/tracing-go/drivers/zipkin/propagation"
+)
+
+// GooglePubSubExtractor manages trace extraction from Google PubSub carrier
+type GooglePubSubExtractor struct {
+	TracerSetter
+}
+
+// NewGooglePubSubExtractor returns the instance of AMQPExtractor
+func NewGooglePubSubExtractor() *GooglePubSubExtractor {
+	return &GooglePubSubExtractor{}
+}
+
+// Extract deserializes SpanContext from amqp.Delivery object
+func (extractor *GooglePubSubExtractor) Extract(carrier interface{}) (tracing.SpanContext, error) {
+	msg, ok := carrier.(*pubsub.Message)
+
+	if !ok {
+		log.Fatalf("Expected *pubsub.Message, got %T", carrier)
+	}
+
+	rawCtx := extractor.Tracing.Extract(propagation.ExtractGooglePubSub(msg))
+	return NewSpanContext(rawCtx), nil
+}

drivers/zipkin/inject_google_pubsub.go

@@ -0,0 +1,38 @@
+package zipkin
+
+import (
+	"log"
+
+	"cloud.google.com/go/pubsub"
+	"github.com/Vinelab/tracing-go"
+	"github.com/Vinelab/tracing-go/drivers/zipkin/propagation"
+	"github.com/openzipkin/zipkin-go/model"
+)
+
+// GooglePubSubInjector manages trace injection into Google Cloud PubSub carrier
+type GooglePubSubInjector struct {
+	//
+}
+
+// NewGooglePubSubInjector returns the instance of GooglePubSubInjector
+func NewGooglePubSubInjector() *GooglePubSubInjector {
+	return &GooglePubSubInjector{}
+}
+
+// Inject serialises given SpanContext into given amqp.Publishing object
+func (extractor *GooglePubSubInjector) Inject(spanCtx tracing.SpanContext, carrier interface{}) error {
+	msg, ok := carrier.(*pubsub.Message)
+	if !ok {
+		log.Fatalf("Expected *pubsub.Message, got %T", carrier)
+	}
+
+	rawCtx := spanCtx.RawContext()
+
+	zipkinCtx, ok := rawCtx.(model.SpanContext)
+	if !ok {
+		log.Fatalf("Expected %T, got %T", model.SpanContext{}, rawCtx)
+	}
+
+	inject := propagation.InjectGooglePubSub(msg)
+	return inject(zipkinCtx)
+}

drivers/zipkin/propagation/google_pubsub.go

@@ -0,0 +1,57 @@
+package propagation
+
+import (
+	"cloud.google.com/go/pubsub"
+	"github.com/openzipkin/zipkin-go/model"
+	"github.com/openzipkin/zipkin-go/propagation"
+	"github.com/openzipkin/zipkin-go/propagation/b3"
+)
+
+// ExtractGooglePubSub will extract a span.Context from the Google Cloud PubSub message if found in B3 header format.
+func ExtractGooglePubSub(msg *pubsub.Message) propagation.Extractor {
+	return func() (*model.SpanContext, error) {
+		var (
+			traceIDHeader      = msg.Attributes[b3.TraceID]
+			spanIDHeader       = msg.Attributes[b3.SpanID]
+			parentSpanIDHeader = msg.Attributes[b3.ParentSpanID]
+			sampledHeader      = msg.Attributes[b3.Sampled]
+			flagsHeader        = msg.Attributes[b3.Flags]
+		)
+
+		return b3.ParseHeaders(
+			traceIDHeader, spanIDHeader, parentSpanIDHeader, sampledHeader,
+			flagsHeader,
+		)
+	}
+}
+
+// InjectGooglePubSub will inject a span.Context into a Google Cloud PubSub message
+func InjectGooglePubSub(msg *pubsub.Message) propagation.Injector {
+	return func(sc model.SpanContext) error {
+		if (model.SpanContext{}) == sc {
+			return b3.ErrEmptyContext
+		}
+
+		if sc.Debug {
+			msg.Attributes[b3.Flags] = "1"
+		} else if sc.Sampled != nil {
+			// Debug is encoded as X-B3-Flags: 1. Since Debug implies Sampled,
+			// so don't also send "X-B3-Sampled: 1".
+			if *sc.Sampled {
+				msg.Attributes[b3.Sampled] = "1"
+			} else {
+				msg.Attributes[b3.Sampled] = "0"
+			}
+		}
+
+		if !sc.TraceID.Empty() && sc.ID > 0 {
+			msg.Attributes[b3.TraceID] = sc.TraceID.String()
+			msg.Attributes[b3.SpanID] = sc.ID.String()
+			if sc.ParentID != nil {
+				msg.Attributes[b3.ParentSpanID] = sc.ParentID.String()
+			}
+		}
+
+		return nil
+	}
+}

drivers/zipkin/tracer.go

@@ -253,6 +253,7 @@ func registerDefaultExtractionFormats() map[string]tracing.Extractor {
 	extractionFormats[formats.TextMap] = NewTextMapExtractor()
 	extractionFormats[formats.HTTP] = NewHTTPExtractor()
 	extractionFormats[formats.AMQP] = NewAMQPExtractor()
+	extractionFormats[formats.GooglePubSub] = NewGooglePubSubExtractor()
 
 	return extractionFormats
 }
@@ -263,6 +264,7 @@ func registerDefaultInjectionFormats() map[string]tracing.Injector {
 	injectionFormats[formats.TextMap] = NewTextMapInjector()
 	injectionFormats[formats.HTTP] = NewHTTPInjector()
 	injectionFormats[formats.AMQP] = NewAMQPInjector()
+	injectionFormats[formats.GooglePubSub] = NewGooglePubSubInjector()
 
 	return injectionFormats
 }

formats/formats.go

@@ -9,4 +9,7 @@ var (
 
 	// AMQP is a format descriptor for propagating trace context via AMQP message
 	AMQP = "amqp"
+
+	// GooglePubSub is a format descriptor for propagating trace context via Google Cloud PubSub message
+	GooglePubSub = "google_pubsub"
 )