Add baggage to B3 codec

zipkin/propagation.go

@@ -23,13 +23,30 @@ import (
 	"github.com/uber/jaeger-client-go"
 )
 
+// Option is a function that sets an option on Propagator
+type Option func(propagator *Propagator)
+
+// BaggagePrefix is a function that sets baggage prefix on Propagator
+func BaggagePrefix(prefix string) Option {
+	return func(propagator *Propagator) {
+		propagator.baggagePrefix = prefix
+	}
+}
+
 // Propagator is an Injector and Extractor
-type Propagator struct{}
+type Propagator struct {
+	baggagePrefix string
+}
 
 // NewZipkinB3HTTPHeaderPropagator creates a Propagator for extracting and injecting
-// Zipkin HTTP B3 headers into SpanContexts.
-func NewZipkinB3HTTPHeaderPropagator() Propagator {
-	return Propagator{}
+// Zipkin HTTP B3 headers into SpanContexts. Baggage is by default enabled and uses prefix
+// 'baggage-'.
+func NewZipkinB3HTTPHeaderPropagator(opts ...Option) Propagator {
+	p := Propagator{baggagePrefix: "baggage-"}
+	for _, opt := range opts {
+		opt(&p)
+	}
+	return p
 }
 
 // Inject conforms to the Injector interface for decoding Zipkin HTTP B3 headers
@@ -53,6 +70,10 @@ func (p Propagator) Inject(
 	} else {
 		textMapWriter.Set("x-b3-sampled", "0")
 	}
+	sc.ForeachBaggageItem(func(k, v string) bool {
+		textMapWriter.Set(p.baggagePrefix+k, v)
+		return true
+	})
 	return nil
 }
 
@@ -66,6 +87,7 @@ func (p Propagator) Extract(abstractCarrier interface{}) (jaeger.SpanContext, er
 	var spanID uint64
 	var parentID uint64
 	sampled := false
+	var baggage map[string]string
 	err := textMapReader.ForeachKey(func(rawKey, value string) error {
 		key := strings.ToLower(rawKey) // TODO not necessary for plain TextMap
 		var err error
@@ -77,6 +99,11 @@ func (p Propagator) Extract(abstractCarrier interface{}) (jaeger.SpanContext, er
 			spanID, err = strconv.ParseUint(value, 16, 64)
 		} else if key == "x-b3-sampled" && value == "1" {
 			sampled = true
+		} else if strings.HasPrefix(key, p.baggagePrefix) {
+			if baggage == nil {
+				baggage = make(map[string]string)
+			}
+			baggage[key[len(p.baggagePrefix):]] = value
 		}
 		return err
 	})
@@ -91,5 +118,5 @@ func (p Propagator) Extract(abstractCarrier interface{}) (jaeger.SpanContext, er
 		jaeger.TraceID{Low: traceID},
 		jaeger.SpanID(spanID),
 		jaeger.SpanID(parentID),
-		sampled, nil), nil
+		sampled, baggage), nil
 }

zipkin/propagation_test.go

@@ -19,20 +19,22 @@ import (
 
 	opentracing "github.com/opentracing/opentracing-go"
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 	"github.com/uber/jaeger-client-go"
 )
 
 var (
-	rootSampled       = newSpanContext(1, 2, 0, true)
-	nonRootSampled    = newSpanContext(1, 2, 1, true)
-	nonRootNonSampled = newSpanContext(1, 2, 1, false)
+	rootSampled       = newSpanContext(1, 2, 0, true, map[string]string{"foo": "bar"})
+	nonRootSampled    = newSpanContext(1, 2, 1, true, nil)
+	nonRootNonSampled = newSpanContext(1, 2, 1, false, nil)
 )
 
 var (
 	rootSampledHeader = opentracing.TextMapCarrier{
 		"x-b3-traceid": "1",
 		"x-b3-spanid":  "2",
 		"x-b3-sampled": "1",
+		"baggage-foo":  "bar",
 	}
 	nonRootSampledHeader = opentracing.TextMapCarrier{
 		"x-b3-traceid":      "1",
@@ -58,13 +60,13 @@ var (
 	propagator = NewZipkinB3HTTPHeaderPropagator()
 )
 
-func newSpanContext(traceID, spanID, parentID uint64, sampled bool) jaeger.SpanContext {
+func newSpanContext(traceID, spanID, parentID uint64, sampled bool, baggage map[string]string) jaeger.SpanContext {
 	return jaeger.NewSpanContext(
 		jaeger.TraceID{Low: traceID},
 		jaeger.SpanID(spanID),
 		jaeger.SpanID(parentID),
 		sampled,
-		nil,
+		baggage,
 	)
 }
 
@@ -111,3 +113,26 @@ func TestInjectorNonRootNonSampled(t *testing.T) {
 	assert.Nil(t, err)
 	assert.EqualValues(t, nonRootNonSampledHeader, hdr)
 }
+
+func TestCustomBaggagePrefix(t *testing.T) {
+	propag := NewZipkinB3HTTPHeaderPropagator(BaggagePrefix("emoji:)"))
+	hdr := opentracing.TextMapCarrier{}
+	sc := newSpanContext(1, 2, 0, true, map[string]string{"foo": "bar"})
+	err := propag.Inject(sc, hdr)
+	assert.Nil(t, err)
+	m := opentracing.TextMapCarrier{
+		"x-b3-traceid": "1",
+		"x-b3-spanid":  "2",
+		"x-b3-sampled": "1",
+		"emoji:)foo":   "bar",
+	}
+	assert.EqualValues(t, m, hdr)
+
+	sc, err = propag.Extract(m)
+	require.NoError(t, err)
+	sc.ForeachBaggageItem(func(k, v string) bool {
+		assert.Equal(t, "foo", k)
+		assert.Equal(t, "bar", v)
+		return true
+	})
+}