fasthttptrace: Add trace context propagation support for libraries bu…

…ilt on fasthttp (#2218)

contrib/internal/fasthttptrace/fasthttpheaderscarrier.go

@@ -0,0 +1,41 @@
+// Unless explicitly stated otherwise all files in this repository are licensed
+// under the Apache License Version 2.0.
+// This product includes software developed at Datadog (https://www.datadoghq.com/).
+// Copyright 2016 Datadog, Inc.
+
+package fasthttptrace
+
+import (
+	"gopkg.in/DataDog/dd-trace-go.v1/ddtrace/tracer"
+
+	"github.com/valyala/fasthttp"
+)
+
+// FastHTTPHeadersCarrier implements tracer.TextMapWriter and tracer.TextMapReader on top
+// of fasthttp's RequestHeader object, allowing it to be used as a span context carrier for
+// distributed tracing.
+type FastHTTPHeadersCarrier struct {
+	ReqHeader *fasthttp.RequestHeader
+}
+
+var _ tracer.TextMapWriter = (*FastHTTPHeadersCarrier)(nil)
+var _ tracer.TextMapReader = (*FastHTTPHeadersCarrier)(nil)
+
+// ForeachKey iterates over fasthttp request header keys and values
+func (f *FastHTTPHeadersCarrier) ForeachKey(handler func(key, val string) error) error {
+	keys := f.ReqHeader.PeekKeys()
+	for _, key := range keys {
+		sKey := string(key)
+		v := f.ReqHeader.Peek(sKey)
+		if err := handler(sKey, string(v)); err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
+// Set adds the given value to request header for key. Key will be lowercased to match
+// the metadata implementation.
+func (f *FastHTTPHeadersCarrier) Set(key, val string) {
+	f.ReqHeader.Set(key, val)
+}

contrib/internal/fasthttptrace/fasthttpheaderscarrier_test.go

@@ -0,0 +1,96 @@
+// Unless explicitly stated otherwise all files in this repository are licensed
+// under the Apache License Version 2.0.
+// This product includes software developed at Datadog (https://www.datadoghq.com/).
+// Copyright 2016 Datadog, Inc.
+
+package fasthttptrace
+
+import (
+	"context"
+	"testing"
+
+	"gopkg.in/DataDog/dd-trace-go.v1/ddtrace/mocktracer"
+	"gopkg.in/DataDog/dd-trace-go.v1/ddtrace/tracer"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"github.com/valyala/fasthttp"
+)
+
+func TestFastHTTPHeadersCarrierSet(t *testing.T) {
+	assert := assert.New(t)
+	fcc := &FastHTTPHeadersCarrier{
+		ReqHeader: new(fasthttp.RequestHeader),
+	}
+	t.Run("key-val", func(t *testing.T) {
+		// add one item
+		fcc.Set("k1", "v1")
+		assert.Len(fcc.ReqHeader.PeekAll("k1"), 1)
+		assert.Equal("v1", string(fcc.ReqHeader.Peek("k1")))
+	})
+	t.Run("key-multival", func(t *testing.T) {
+		// add a second value, ensure the second value overwrites the first
+		fcc.Set("k1", "v1")
+		fcc.Set("k1", "v2")
+		vals := fcc.ReqHeader.PeekAll("k1")
+		assert.Len(vals, 1)
+		assert.Equal("v2", string(vals[0]))
+	})
+	t.Run("multi-key", func(t *testing.T) {
+		// // add a second key
+		fcc.Set("k1", "v1")
+		fcc.Set("k2", "v21")
+		assert.Len(fcc.ReqHeader.PeekAll("k2"), 1)
+		assert.Equal("v21", string(fcc.ReqHeader.Peek("k2")))
+	})
+	t.Run("case insensitive", func(t *testing.T) {
+		// new key
+		fcc.Set("K3", "v31")
+		assert.Equal("v31", string(fcc.ReqHeader.Peek("k3")))
+		assert.Equal("v31", string(fcc.ReqHeader.Peek("K3")))
+		// access existing, lowercase key with uppercase input
+		fcc.Set("K3", "v32")
+		vals := fcc.ReqHeader.PeekAll("k3")
+		assert.Equal("v32", string(vals[0]))
+	})
+}
+
+func TestFastHTTPHeadersCarrierForeachKey(t *testing.T) {
+	assert := assert.New(t)
+	h := new(fasthttp.RequestHeader)
+	headers := map[string][]string{
+		"K1": {"v1"},
+		"K2": {"v2", "v22"},
+	}
+	assert.Len(headers, 2)
+	for k, vs := range headers {
+		for _, v := range vs {
+			h.Add(k, v)
+		}
+	}
+	fcc := &FastHTTPHeadersCarrier{
+		ReqHeader: h,
+	}
+	err := fcc.ForeachKey(func(k, v string) error {
+		delete(headers, k)
+		return nil
+	})
+	assert.NoError(err)
+	assert.Len(headers, 0)
+}
+
+func TestInjectExtract(t *testing.T) {
+	assert := assert.New(t)
+	mt := mocktracer.Start()
+	defer mt.Stop()
+	pspan, _ := tracer.StartSpanFromContext(context.Background(), "test")
+	fcc := &FastHTTPHeadersCarrier{
+		ReqHeader: &fasthttp.RequestHeader{},
+	}
+	err := tracer.Inject(pspan.Context(), fcc)
+	require.NoError(t, err)
+	sctx, err := tracer.Extract(fcc)
+	require.NoError(t, err)
+	assert.Equal(sctx.TraceID(), pspan.Context().TraceID())
+	assert.Equal(sctx.SpanID(), pspan.Context().SpanID())
+}

contrib/internal/fasthttptrace/fasthttptrace.go

@@ -0,0 +1,24 @@
+// Unless explicitly stated otherwise all files in this repository are licensed
+// under the Apache License Version 2.0.
+// This product includes software developed at Datadog (https://www.datadoghq.com/).
+// Copyright 2016 Datadog, Inc.
+
+package fasthttptrace
+
+import (
+	"gopkg.in/DataDog/dd-trace-go.v1/ddtrace/tracer"
+	"gopkg.in/DataDog/dd-trace-go.v1/internal"
+
+	"github.com/valyala/fasthttp"
+)
+
+// StartSpanFromContext returns a new span with the given operation name and options.
+// If a span is found in the `fctx`, it will be used as the parent of the resulting span.
+// The resulting span is then set on the given `fctx`.
+// This function is similar to tracer.StartSpanFromContext, but it modifies the given fasthttp context directly.
+// If the ChildOf option is passed, it will only be used as the parent if there is no span found in `fctx`.
+func StartSpanFromContext(fctx *fasthttp.RequestCtx, operationName string, opts ...tracer.StartSpanOption) tracer.Span {
+	s, _ := tracer.StartSpanFromContext(fctx, operationName, opts...)
+	fctx.SetUserValue(internal.ActiveSpanKey, s)
+	return s
+}

contrib/internal/fasthttptrace/fasthttptrace_test.go

@@ -0,0 +1,26 @@
+// Unless explicitly stated otherwise all files in this repository are licensed
+// under the Apache License Version 2.0.
+// This product includes software developed at Datadog (https://www.datadoghq.com/).
+// Copyright 2016 Datadog, Inc.
+
+package fasthttptrace
+
+import (
+	"testing"
+
+	"gopkg.in/DataDog/dd-trace-go.v1/ddtrace/mocktracer"
+	"gopkg.in/DataDog/dd-trace-go.v1/internal"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/valyala/fasthttp"
+)
+
+func TestStartSpanFromContext(t *testing.T) {
+	assert := assert.New(t)
+	mt := mocktracer.Start()
+	defer mt.Stop()
+	fctx := &fasthttp.RequestCtx{}
+	activeSpan := StartSpanFromContext(fctx, "myOp")
+	keySpan := fctx.UserValue(internal.ActiveSpanKey)
+	assert.Equal(activeSpan, keySpan)
+}

ddtrace/tracer/context.go

@@ -9,30 +9,27 @@ import (
 	"context"
 
 	"gopkg.in/DataDog/dd-trace-go.v1/ddtrace"
-	"gopkg.in/DataDog/dd-trace-go.v1/ddtrace/internal"
+	traceinternal "gopkg.in/DataDog/dd-trace-go.v1/ddtrace/internal"
+	"gopkg.in/DataDog/dd-trace-go.v1/internal"
 )
 
-type contextKey struct{}
-
-var activeSpanKey = contextKey{}
-
 // ContextWithSpan returns a copy of the given context which includes the span s.
 func ContextWithSpan(ctx context.Context, s Span) context.Context {
-	return context.WithValue(ctx, activeSpanKey, s)
+	return context.WithValue(ctx, internal.ActiveSpanKey, s)
 }
 
 // SpanFromContext returns the span contained in the given context. A second return
 // value indicates if a span was found in the context. If no span is found, a no-op
 // span is returned.
 func SpanFromContext(ctx context.Context) (Span, bool) {
 	if ctx == nil {
-		return &internal.NoopSpan{}, false
+		return &traceinternal.NoopSpan{}, false
 	}
-	v := ctx.Value(activeSpanKey)
+	v := ctx.Value(internal.ActiveSpanKey)
 	if s, ok := v.(ddtrace.Span); ok {
 		return s, true
 	}
-	return &internal.NoopSpan{}, false
+	return &traceinternal.NoopSpan{}, false
 }
 
 // StartSpanFromContext returns a new span with the given operation name and options. If a span

ddtrace/tracer/context_test.go

@@ -12,15 +12,16 @@ import (
 	"testing"
 
 	"gopkg.in/DataDog/dd-trace-go.v1/ddtrace"
-	"gopkg.in/DataDog/dd-trace-go.v1/ddtrace/internal"
+	traceinternal "gopkg.in/DataDog/dd-trace-go.v1/ddtrace/internal"
+	"gopkg.in/DataDog/dd-trace-go.v1/internal"
 
 	"github.com/stretchr/testify/assert"
 )
 
 func TestContextWithSpan(t *testing.T) {
 	want := &span{SpanID: 123}
 	ctx := ContextWithSpan(context.Background(), want)
-	got, ok := ctx.Value(activeSpanKey).(*span)
+	got, ok := ctx.Value(internal.ActiveSpanKey).(*span)
 	assert := assert.New(t)
 	assert.True(ok)
 	assert.Equal(got, want)
@@ -39,11 +40,11 @@ func TestSpanFromContext(t *testing.T) {
 		assert := assert.New(t)
 		span, ok := SpanFromContext(context.Background())
 		assert.False(ok)
-		_, ok = span.(*internal.NoopSpan)
+		_, ok = span.(*traceinternal.NoopSpan)
 		assert.True(ok)
 		span, ok = SpanFromContext(nil)
 		assert.False(ok)
-		_, ok = span.(*internal.NoopSpan)
+		_, ok = span.(*traceinternal.NoopSpan)
 		assert.True(ok)
 	})
 }
@@ -69,7 +70,7 @@ func TestStartSpanFromContext(t *testing.T) {
 	gotctx, ok := SpanFromContext(ctx)
 	assert.True(ok)
 	assert.Equal(gotctx, got)
-	_, ok = gotctx.(*internal.NoopSpan)
+	_, ok = gotctx.(*traceinternal.NoopSpan)
 	assert.False(ok)
 
 	assert.Equal(uint64(456), got.TraceID)

internal/active_span_key.go

@@ -0,0 +1,11 @@
+// Unless explicitly stated otherwise all files in this repository are licensed
+// under the Apache License Version 2.0.
+// This product includes software developed at Datadog (https://www.datadoghq.com/).
+// Copyright 2016 Datadog, Inc.
+
+package internal
+
+type contextKey struct{}
+
+// ActiveSpanKey is used to set tracer context on a context.Context objects with a unique key
+var ActiveSpanKey = contextKey{}