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{}