contrib/go-chi/chi.v5: support go-chi/chi/v5

contrib/go-chi/chi.v5/chi.go

@@ -0,0 +1,78 @@
+// 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 chi provides tracing functions for tracing the go-chi/chi/v5 package (https://github.com/go-chi/chi).
+package chi // import "gopkg.in/DataDog/dd-trace-go.v1/contrib/go-chi/chi.v5"
+
+import (
+	"fmt"
+	"math"
+	"net/http"
+	"strconv"
+
+	"gopkg.in/DataDog/dd-trace-go.v1/ddtrace"
+	"gopkg.in/DataDog/dd-trace-go.v1/ddtrace/ext"
+	"gopkg.in/DataDog/dd-trace-go.v1/ddtrace/tracer"
+	"gopkg.in/DataDog/dd-trace-go.v1/internal/log"
+
+	"github.com/go-chi/chi/v5"
+	"github.com/go-chi/chi/v5/middleware"
+)
+
+// Middleware returns middleware that will trace incoming requests.
+func Middleware(opts ...Option) func(next http.Handler) http.Handler {
+	cfg := new(config)
+	defaults(cfg)
+	for _, fn := range opts {
+		fn(cfg)
+	}
+	log.Debug("contrib/go-chi/chi.v5: Configuring Middleware: %#v", cfg)
+	return func(next http.Handler) http.Handler {
+		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			opts := []ddtrace.StartSpanOption{
+				tracer.SpanType(ext.SpanTypeWeb),
+				tracer.ServiceName(cfg.serviceName),
+				tracer.Tag(ext.HTTPMethod, r.Method),
+				tracer.Tag(ext.HTTPURL, r.URL.Path),
+				tracer.Measured(),
+			}
+			if !math.IsNaN(cfg.analyticsRate) {
+				opts = append(opts, tracer.Tag(ext.EventSampleRate, cfg.analyticsRate))
+			}
+			if spanctx, err := tracer.Extract(tracer.HTTPHeadersCarrier(r.Header)); err == nil {
+				opts = append(opts, tracer.ChildOf(spanctx))
+			}
+			opts = append(opts, cfg.spanOpts...)
+			span, ctx := tracer.StartSpanFromContext(r.Context(), "http.request", opts...)
+			defer span.Finish()
+
+			ww := middleware.NewWrapResponseWriter(w, r.ProtoMajor)
+
+			// pass the span through the request context and serve the request to the next middleware
+			next.ServeHTTP(ww, r.WithContext(ctx))
+
+			// set the resource name as we get it only once the handler is executed
+			resourceName := chi.RouteContext(r.Context()).RoutePattern()
+			if resourceName == "" {
+				resourceName = "unknown"
+			}
+			resourceName = r.Method + " " + resourceName
+			span.SetTag(ext.ResourceName, resourceName)
+
+			// set the status code
+			status := ww.Status()
+			// 0 status means one has not yet been sent in which case net/http library will write StatusOK
+			if ww.Status() == 0 {
+				status = http.StatusOK
+			}
+			span.SetTag(ext.HTTPCode, strconv.Itoa(status))
+
+			if cfg.isStatusError(status) {
+				// mark 5xx server error
+				span.SetTag(ext.Error, fmt.Errorf("%d: %s", status, http.StatusText(status)))
+			}
+		})
+	}
+}

contrib/go-chi/chi.v5/chi_test.go

@@ -0,0 +1,272 @@
+// 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 chi
+
+import (
+	"fmt"
+	"net/http"
+	"net/http/httptest"
+	"strconv"
+	"testing"
+
+	"gopkg.in/DataDog/dd-trace-go.v1/ddtrace/ext"
+	"gopkg.in/DataDog/dd-trace-go.v1/ddtrace/mocktracer"
+	"gopkg.in/DataDog/dd-trace-go.v1/ddtrace/tracer"
+	"gopkg.in/DataDog/dd-trace-go.v1/internal/globalconfig"
+
+	"github.com/go-chi/chi/v5"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestChildSpan(t *testing.T) {
+	assert := assert.New(t)
+	mt := mocktracer.Start()
+	defer mt.Stop()
+
+	router := chi.NewRouter()
+	router.Use(Middleware(WithServiceName("foobar")))
+	router.Get("/user/{id}", func(w http.ResponseWriter, r *http.Request) {
+		_, ok := tracer.SpanFromContext(r.Context())
+		assert.True(ok)
+	})
+
+	r := httptest.NewRequest("GET", "/user/123", nil)
+	w := httptest.NewRecorder()
+
+	router.ServeHTTP(w, r)
+}
+
+func TestTrace200(t *testing.T) {
+	assertDoRequest := func(assert *assert.Assertions, mt mocktracer.Tracer, router *chi.Mux) {
+		r := httptest.NewRequest("GET", "/user/123", nil)
+		w := httptest.NewRecorder()
+
+		// do and verify the request
+		router.ServeHTTP(w, r)
+		response := w.Result()
+		assert.Equal(response.StatusCode, 200)
+
+		// verify traces look good
+		spans := mt.FinishedSpans()
+		assert.Len(spans, 1)
+		if len(spans) < 1 {
+			t.Fatalf("no spans")
+		}
+		span := spans[0]
+		assert.Equal("http.request", span.OperationName())
+		assert.Equal(ext.SpanTypeWeb, span.Tag(ext.SpanType))
+		assert.Equal("foobar", span.Tag(ext.ServiceName))
+		assert.Equal("GET /user/{id}", span.Tag(ext.ResourceName))
+		assert.Equal("200", span.Tag(ext.HTTPCode))
+		assert.Equal("GET", span.Tag(ext.HTTPMethod))
+		assert.Equal("/user/123", span.Tag(ext.HTTPURL))
+	}
+
+	t.Run("response written", func(t *testing.T) {
+		assert := assert.New(t)
+		mt := mocktracer.Start()
+		defer mt.Stop()
+
+		router := chi.NewRouter()
+		router.Use(Middleware(WithServiceName("foobar")))
+		router.Get("/user/{id}", func(w http.ResponseWriter, r *http.Request) {
+			span, ok := tracer.SpanFromContext(r.Context())
+			assert.True(ok)
+			assert.Equal(span.(mocktracer.Span).Tag(ext.ServiceName), "foobar")
+			id := chi.URLParam(r, "id")
+			_, err := w.Write([]byte(id))
+			assert.NoError(err)
+		})
+		assertDoRequest(assert, mt, router)
+	})
+
+	t.Run("no response written", func(t *testing.T) {
+		assert := assert.New(t)
+		mt := mocktracer.Start()
+		defer mt.Stop()
+
+		router := chi.NewRouter()
+		router.Use(Middleware(WithServiceName("foobar")))
+		router.Get("/user/{id}", func(w http.ResponseWriter, r *http.Request) {
+			span, ok := tracer.SpanFromContext(r.Context())
+			assert.True(ok)
+			assert.Equal(span.(mocktracer.Span).Tag(ext.ServiceName), "foobar")
+		})
+		assertDoRequest(assert, mt, router)
+	})
+}
+
+func TestError(t *testing.T) {
+	assertSpan := func(assert *assert.Assertions, spans []mocktracer.Span, code int) {
+		assert.Len(spans, 1)
+		if len(spans) < 1 {
+			t.Fatalf("no spans")
+		}
+		span := spans[0]
+		assert.Equal("http.request", span.OperationName())
+		assert.Equal("foobar", span.Tag(ext.ServiceName))
+
+		assert.Equal(strconv.Itoa(code), span.Tag(ext.HTTPCode))
+
+		wantErr := fmt.Sprintf("%d: %s", code, http.StatusText(code))
+		assert.Equal(wantErr, span.Tag(ext.Error).(error).Error())
+	}
+
+	t.Run("default", func(t *testing.T) {
+		assert := assert.New(t)
+		mt := mocktracer.Start()
+		defer mt.Stop()
+
+		// setup
+		router := chi.NewRouter()
+		router.Use(Middleware(WithServiceName("foobar")))
+		code := 500
+
+		// a handler with an error and make the requests
+		router.Get("/err", func(w http.ResponseWriter, r *http.Request) {
+			http.Error(w, fmt.Sprintf("%d!", code), code)
+		})
+		r := httptest.NewRequest("GET", "/err", nil)
+		w := httptest.NewRecorder()
+		router.ServeHTTP(w, r)
+		response := w.Result()
+		assert.Equal(response.StatusCode, code)
+
+		// verify the errors and status are correct
+		spans := mt.FinishedSpans()
+		assertSpan(assert, spans, code)
+	})
+
+	t.Run("custom", func(t *testing.T) {
+		assert := assert.New(t)
+		mt := mocktracer.Start()
+		defer mt.Stop()
+
+		// setup
+		router := chi.NewRouter()
+		router.Use(Middleware(
+			WithServiceName("foobar"),
+			WithStatusCheck(func(statusCode int) bool {
+				return statusCode >= 400
+			}),
+		))
+		code := 404
+		// a handler with an error and make the requests
+		router.Get("/err", func(w http.ResponseWriter, r *http.Request) {
+			http.Error(w, fmt.Sprintf("%d!", code), code)
+		})
+		r := httptest.NewRequest("GET", "/err", nil)
+		w := httptest.NewRecorder()
+		router.ServeHTTP(w, r)
+		response := w.Result()
+		assert.Equal(response.StatusCode, code)
+
+		// verify the errors and status are correct
+		spans := mt.FinishedSpans()
+		assertSpan(assert, spans, code)
+	})
+}
+
+func TestGetSpanNotInstrumented(t *testing.T) {
+	assert := assert.New(t)
+	router := chi.NewRouter()
+	router.Get("/ping", func(w http.ResponseWriter, r *http.Request) {
+		// Assert we don't have a span on the context.
+		_, ok := tracer.SpanFromContext(r.Context())
+		assert.False(ok)
+		w.Write([]byte("ok"))
+	})
+	r := httptest.NewRequest("GET", "/ping", nil)
+	w := httptest.NewRecorder()
+	router.ServeHTTP(w, r)
+	response := w.Result()
+	assert.Equal(response.StatusCode, 200)
+}
+
+func TestPropagation(t *testing.T) {
+	assert := assert.New(t)
+	mt := mocktracer.Start()
+	defer mt.Stop()
+
+	r := httptest.NewRequest("GET", "/user/123", nil)
+	w := httptest.NewRecorder()
+
+	pspan := tracer.StartSpan("test")
+	tracer.Inject(pspan.Context(), tracer.HTTPHeadersCarrier(r.Header))
+
+	router := chi.NewRouter()
+	router.Use(Middleware(WithServiceName("foobar")))
+	router.Get("/user/{id}", func(w http.ResponseWriter, r *http.Request) {
+		span, ok := tracer.SpanFromContext(r.Context())
+		assert.True(ok)
+		assert.Equal(span.(mocktracer.Span).ParentID(), pspan.(mocktracer.Span).SpanID())
+	})
+
+	router.ServeHTTP(w, r)
+}
+
+func TestAnalyticsSettings(t *testing.T) {
+	assertRate := func(t *testing.T, mt mocktracer.Tracer, rate interface{}, opts ...Option) {
+		router := chi.NewRouter()
+		router.Use(Middleware(opts...))
+		router.Get("/user/{id}", func(w http.ResponseWriter, r *http.Request) {
+			_, ok := tracer.SpanFromContext(r.Context())
+			assert.True(t, ok)
+		})
+
+		r := httptest.NewRequest("GET", "/user/123", nil)
+		w := httptest.NewRecorder()
+
+		router.ServeHTTP(w, r)
+		spans := mt.FinishedSpans()
+		assert.Len(t, spans, 1)
+		s := spans[0]
+		assert.Equal(t, rate, s.Tag(ext.EventSampleRate))
+	}
+
+	t.Run("defaults", func(t *testing.T) {
+		mt := mocktracer.Start()
+		defer mt.Stop()
+
+		assertRate(t, mt, nil)
+	})
+
+	t.Run("global", func(t *testing.T) {
+		mt := mocktracer.Start()
+		defer mt.Stop()
+
+		rate := globalconfig.AnalyticsRate()
+		defer globalconfig.SetAnalyticsRate(rate)
+		globalconfig.SetAnalyticsRate(0.4)
+
+		assertRate(t, mt, 0.4)
+	})
+
+	t.Run("enabled", func(t *testing.T) {
+		mt := mocktracer.Start()
+		defer mt.Stop()
+
+		assertRate(t, mt, 1.0, WithAnalytics(true))
+	})
+
+	t.Run("disabled", func(t *testing.T) {
+		mt := mocktracer.Start()
+		defer mt.Stop()
+
+		assertRate(t, mt, nil, WithAnalytics(false))
+	})
+
+	t.Run("override", func(t *testing.T) {
+		mt := mocktracer.Start()
+		defer mt.Stop()
+
+		rate := globalconfig.AnalyticsRate()
+		defer globalconfig.SetAnalyticsRate(rate)
+		globalconfig.SetAnalyticsRate(0.4)
+
+		assertRate(t, mt, 0.23, WithAnalyticsRate(0.23))
+	})
+}