Implemented logic for opentracing baggage propagation

open-telemetry · Mar 3, 2022 · 3db3dc3 · 3db3dc3
1 parent 9e8d62b
commit 3db3dc3
Show file tree

Hide file tree

Showing 2 changed files with 85 additions and 17 deletions.
diff --git a/propagators/ot/ot_data_test.go b/propagators/ot/ot_data_test.go
@@ -16,6 +16,7 @@ package ot_test
 
 import (
 	"go.opentelemetry.io/otel/trace"
+	"strings"
 )
 
 const (
@@ -38,11 +39,13 @@ var (
 	traceID32    = trace.TraceID{0xa1, 0xce, 0x92, 0x9d, 0x0e, 0x0e, 0x47, 0x36, 0xa3, 0xce, 0x92, 0x9d, 0x0e, 0x0e, 0x47, 0x36}
 	spanID       = trace.SpanID{0x00, 0xf0, 0x67, 0xaa, 0x0b, 0xa9, 0x02, 0xb7}
 	emptyBaggage = map[string]string{}
-	// TODO: once baggage extraction is supported, re-enable this
-	// baggageSet   = attribute.NewSet(
-	// 	attribute.String(baggageKey, baggageValue),
-	// 	attribute.String(baggageKey2, baggageValue2),
-	// )
+	baggageSet   = map[string]string{
+		baggageKey: baggageValue,
+	}
+	baggageSet2 = map[string]string{
+		baggageKey:  baggageValue,
+		baggageKey2: baggageValue2,
+	}
 )
 
 type extractTest struct {
@@ -85,9 +88,22 @@ var extractHeaders = []extractTest{
 			SpanID:     spanID,
 			TraceFlags: trace.FlagsSampled,
 		},
-		emptyBaggage,
-		// TODO: once baggage extraction is supported, re-enable this
-		// &baggageSet,
+		baggageSet,
+	},
+	{
+		"baggage multiple values",
+		map[string]string{
+			traceIDHeader:  traceID32Str,
+			spanIDHeader:   spanIDStr,
+			sampledHeader:  "0",
+			baggageHeader:  baggageValue,
+			baggageHeader2: baggageValue2,
+		},
+		trace.SpanContextConfig{
+			TraceID: traceID32,
+			SpanID:  spanID,
+		},
+		baggageSet2,
 	},
 	{
 		"left padding 64 bit trace ID",
@@ -168,6 +184,33 @@ var invalidExtractHeaders = []extractTest{
 			sampledHeader: "wired",
 		},
 	},
+	{
+		name: "invalid baggage key",
+		headers: map[string]string{
+			traceIDHeader:     traceID32Str,
+			spanIDHeader:      spanIDStr,
+			sampledHeader:     "1",
+			"ot-baggage-d–76": "test",
+		},
+	},
+	{
+		name: "invalid baggage value",
+		headers: map[string]string{
+			traceIDHeader: traceID32Str,
+			spanIDHeader:  spanIDStr,
+			sampledHeader: "1",
+			baggageHeader: "øtel",
+		},
+	},
+	{
+		name: "invalid baggage result (too large)",
+		headers: map[string]string{
+			traceIDHeader: traceID32Str,
+			spanIDHeader:  spanIDStr,
+			sampledHeader: "1",
+			baggageHeader: strings.Repeat("s", 8188),
+		},
+	},
 	{
 		name:    "missing headers",
 		headers: map[string]string{},

diff --git a/propagators/ot/ot_propagator.go b/propagators/ot/ot_propagator.go
@@ -28,9 +28,10 @@ import (
 
 const (
 	// Default OT Header names.
-	traceIDHeader = "ot-tracer-traceid"
-	spanIDHeader  = "ot-tracer-spanid"
-	sampledHeader = "ot-tracer-sampled"
+	traceIDHeader       = "ot-tracer-traceid"
+	spanIDHeader        = "ot-tracer-spanid"
+	sampledHeader       = "ot-tracer-sampled"
+	baggageHeaderPrefix = "ot-baggage-"
 
 	otTraceIDPadding = "0000000000000000"
 
@@ -72,7 +73,7 @@ func (o OT) Inject(ctx context.Context, carrier propagation.TextMapCarrier) {
 	}
 
 	for _, m := range baggage.FromContext(ctx).Members() {
-		carrier.Set(fmt.Sprintf("ot-baggage-%s", m.Key()), m.Value())
+		carrier.Set(fmt.Sprintf("%s%s", baggageHeaderPrefix, m.Key()), m.Value())
 	}
 
 }
@@ -93,18 +94,42 @@ func (o OT) Extract(ctx context.Context, carrier propagation.TextMapCarrier) con
 	if err != nil || !sc.IsValid() {
 		return ctx
 	}
-	// TODO: implement extracting baggage
-	//
-	// this currently is not achievable without an implementation of `keys`
-	// on the carrier, see:
-	// https://github.com/open-telemetry/opentelemetry-go/issues/1493
+
+	bags, err := extractBags(carrier)
+	if err != nil {
+		return ctx
+	}
+	ctx = baggage.ContextWithBaggage(ctx, bags)
 	return trace.ContextWithRemoteSpanContext(ctx, sc)
 }
 
 func (o OT) Fields() []string {
 	return []string{traceIDHeader, spanIDHeader, sampledHeader}
 }
 
+// extractBags reconstructs the baggage information from opentracing
+func extractBags(carrier propagation.TextMapCarrier) (baggage.Baggage, error) {
+	emptyBags, _ := baggage.New()
+	var members []baggage.Member
+	for _, key := range carrier.Keys() {
+		lowerKey := strings.ToLower(key)
+		if !strings.HasPrefix(lowerKey, baggageHeaderPrefix) {
+			continue
+		}
+		strippedKey := strings.TrimPrefix(lowerKey, baggageHeaderPrefix)
+		member, err := baggage.NewMember(strippedKey, carrier.Get(key))
+		if err != nil {
+			return emptyBags, err
+		}
+		members = append(members, member)
+	}
+	bags, err := baggage.New(members...)
+	if err != nil {
+		return emptyBags, err
+	}
+	return bags, nil
+}
+
 // extract reconstructs a SpanContext from header values based on OT
 // headers.
 func extract(traceID, spanID, sampled string) (trace.SpanContext, error) {