From 93d4ab36baf09dcced04a396277ea2658d554965 Mon Sep 17 00:00:00 2001 From: Mikayla Toffler <46911781+mtoffl01@users.noreply.github.com> Date: Thu, 28 Sep 2023 15:35:11 -0400 Subject: [PATCH] fasthttptrace: Add trace context propagation support for libraries built on fasthttp (#2218) --- .../fasthttptrace/fasthttpheaderscarrier.go | 41 ++++++++ .../fasthttpheaderscarrier_test.go | 96 +++++++++++++++++++ .../internal/fasthttptrace/fasthttptrace.go | 24 +++++ .../fasthttptrace/fasthttptrace_test.go | 26 +++++ ddtrace/tracer/context.go | 15 ++- ddtrace/tracer/context_test.go | 11 ++- internal/active_span_key.go | 11 +++ 7 files changed, 210 insertions(+), 14 deletions(-) create mode 100644 contrib/internal/fasthttptrace/fasthttpheaderscarrier.go create mode 100644 contrib/internal/fasthttptrace/fasthttpheaderscarrier_test.go create mode 100644 contrib/internal/fasthttptrace/fasthttptrace.go create mode 100644 contrib/internal/fasthttptrace/fasthttptrace_test.go create mode 100644 internal/active_span_key.go diff --git a/contrib/internal/fasthttptrace/fasthttpheaderscarrier.go b/contrib/internal/fasthttptrace/fasthttpheaderscarrier.go new file mode 100644 index 0000000000..f1029e2651 --- /dev/null +++ b/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) +} diff --git a/contrib/internal/fasthttptrace/fasthttpheaderscarrier_test.go b/contrib/internal/fasthttptrace/fasthttpheaderscarrier_test.go new file mode 100644 index 0000000000..a540c30c15 --- /dev/null +++ b/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()) +} diff --git a/contrib/internal/fasthttptrace/fasthttptrace.go b/contrib/internal/fasthttptrace/fasthttptrace.go new file mode 100644 index 0000000000..e03b249193 --- /dev/null +++ b/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 +} diff --git a/contrib/internal/fasthttptrace/fasthttptrace_test.go b/contrib/internal/fasthttptrace/fasthttptrace_test.go new file mode 100644 index 0000000000..9c5b474915 --- /dev/null +++ b/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) +} diff --git a/ddtrace/tracer/context.go b/ddtrace/tracer/context.go index 07b72c1ae8..5698dea685 100644 --- a/ddtrace/tracer/context.go +++ b/ddtrace/tracer/context.go @@ -9,16 +9,13 @@ 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 @@ -26,13 +23,13 @@ func ContextWithSpan(ctx context.Context, s Span) context.Context { // 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 diff --git a/ddtrace/tracer/context_test.go b/ddtrace/tracer/context_test.go index 1ee62adcd2..589343dc65 100644 --- a/ddtrace/tracer/context_test.go +++ b/ddtrace/tracer/context_test.go @@ -12,7 +12,8 @@ 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" ) @@ -20,7 +21,7 @@ import ( 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) diff --git a/internal/active_span_key.go b/internal/active_span_key.go new file mode 100644 index 0000000000..090150a587 --- /dev/null +++ b/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{}