add ability to serialize/deserialize b3 to Go map[string]string (#145)

Add ability to serialize/deserialize b3 to Go map[string]string

Author: basvanbeek

propagation/b3/http_test.go

@@ -24,6 +24,10 @@ import (
 	"github.com/openzipkin/zipkin-go/reporter/recorder"
 )
 
+const (
+	invalidID = "invalid_data"
+)
+
 func TestHTTPExtractFlagsOnly(t *testing.T) {
 	r := newHTTPRequest(t)
 
@@ -188,7 +192,7 @@ func TestHTTPExtractScope(t *testing.T) {
 func TestHTTPExtractTraceIDError(t *testing.T) {
 	r := newHTTPRequest(t)
 
-	r.Header.Set(b3.TraceID, "invalid_data")
+	r.Header.Set(b3.TraceID, invalidID)
 
 	_, err := b3.ExtractHTTP(r)()
 
@@ -200,7 +204,7 @@ func TestHTTPExtractTraceIDError(t *testing.T) {
 func TestHTTPExtractSpanIDError(t *testing.T) {
 	r := newHTTPRequest(t)
 
-	r.Header.Set(b3.SpanID, "invalid_data")
+	r.Header.Set(b3.SpanID, invalidID)
 
 	_, err := b3.ExtractHTTP(r)()
 
@@ -250,7 +254,7 @@ func TestHTTPExtractInvalidParentIDError(t *testing.T) {
 
 	r.Header.Set(b3.TraceID, "1")
 	r.Header.Set(b3.SpanID, "2")
-	r.Header.Set(b3.ParentSpanID, "invalid_data")
+	r.Header.Set(b3.ParentSpanID, invalidID)
 
 	_, err := b3.ExtractHTTP(r)()
 
@@ -279,7 +283,7 @@ func TestHTTPExtractSingleFailsAndMultipleFallsbackFailing(t *testing.T) {
 	r.Header.Set(b3.Context, "0000000000000001-0000000000000002-x")
 	r.Header.Set(b3.TraceID, "1")
 	r.Header.Set(b3.SpanID, "2")
-	r.Header.Set(b3.ParentSpanID, "invalid_data")
+	r.Header.Set(b3.ParentSpanID, invalidID)
 
 	_, err := b3.ExtractHTTP(r)()
 

propagation/b3/map.go

@@ -0,0 +1,101 @@
+// Copyright 2019 The OpenZipkin Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package b3
+
+import (
+	"github.com/openzipkin/zipkin-go/model"
+	"github.com/openzipkin/zipkin-go/propagation"
+)
+
+// Map allows serialization and deserialization of SpanContext into a standard Go map.
+type Map map[string]string
+
+// Extract implements Extractor
+func (m *Map) Extract() (*model.SpanContext, error) {
+	var (
+		traceIDHeader      = (*m)[TraceID]
+		spanIDHeader       = (*m)[SpanID]
+		parentSpanIDHeader = (*m)[ParentSpanID]
+		sampledHeader      = (*m)[Sampled]
+		flagsHeader        = (*m)[Flags]
+		singleHeader       = (*m)[Context]
+	)
+
+	var (
+		sc   *model.SpanContext
+		sErr error
+		mErr error
+	)
+	if singleHeader != "" {
+		sc, sErr = ParseSingleHeader(singleHeader)
+		if sErr == nil {
+			return sc, nil
+		}
+	}
+
+	sc, mErr = ParseHeaders(
+		traceIDHeader, spanIDHeader, parentSpanIDHeader,
+		sampledHeader, flagsHeader,
+	)
+
+	if mErr != nil && sErr != nil {
+		return nil, sErr
+	}
+
+	return sc, mErr
+
+}
+
+// Inject implements Injector
+func (m *Map) Inject(opts ...InjectOption) propagation.Injector {
+	options := InjectOptions{shouldInjectMultiHeader: true}
+	for _, opt := range opts {
+		opt(&options)
+	}
+
+	return func(sc model.SpanContext) error {
+		if (model.SpanContext{}) == sc {
+			return ErrEmptyContext
+		}
+
+		if options.shouldInjectMultiHeader {
+			if sc.Debug {
+				(*m)[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 {
+					(*m)[Sampled] = "1"
+				} else {
+					(*m)[Sampled] = "0"
+				}
+			}
+
+			if !sc.TraceID.Empty() && sc.ID > 0 {
+				(*m)[TraceID] = sc.TraceID.String()
+				(*m)[SpanID] = sc.ID.String()
+				if sc.ParentID != nil {
+					(*m)[ParentSpanID] = sc.ParentID.String()
+				}
+			}
+		}
+
+		if options.shouldInjectSingleHeader {
+			(*m)[Context] = BuildSingleHeader(sc)
+		}
+
+		return nil
+	}
+}