contrib/go-chi: add option to ignore requests. (#1124)

Adds an option allowing a user-provided function to disable tracing for certain requests. Fixes: #1121
DataDog · Jan 14, 2022 · 307fc73 · 307fc73
1 parent a0853da
commit 307fc73
Show file tree

Hide file tree

Showing 6 changed files with 92 additions and 0 deletions.
diff --git a/contrib/go-chi/chi.v5/chi.go b/contrib/go-chi/chi.v5/chi.go
@@ -31,6 +31,10 @@ func Middleware(opts ...Option) func(next http.Handler) http.Handler {
 	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) {
+			if cfg.ignoreRequest(r) {
+				next.ServeHTTP(w, r)
+				return
+			}
 			opts := []ddtrace.StartSpanOption{
 				tracer.SpanType(ext.SpanTypeWeb),
 				tracer.ServiceName(cfg.serviceName),

diff --git a/contrib/go-chi/chi.v5/chi_test.go b/contrib/go-chi/chi.v5/chi_test.go
@@ -10,6 +10,7 @@ import (
 	"net/http"
 	"net/http/httptest"
 	"strconv"
+	"strings"
 	"testing"
 
 	"gopkg.in/DataDog/dd-trace-go.v1/ddtrace/ext"
@@ -270,3 +271,33 @@ func TestAnalyticsSettings(t *testing.T) {
 		assertRate(t, mt, 0.23, WithAnalyticsRate(0.23))
 	})
 }
+
+func TestIgnoreRequest(t *testing.T) {
+	router := chi.NewRouter()
+	router.Use(Middleware(
+		WithIgnoreRequest(func(r *http.Request) bool {
+			return strings.HasPrefix(r.URL.Path, "/skip")
+		}),
+	))
+
+	router.Get("/ok", func(w http.ResponseWriter, r *http.Request) {
+		w.Write([]byte("ok"))
+	})
+
+	router.Get("/skip", func(w http.ResponseWriter, r *http.Request) {
+		w.Write([]byte("skip"))
+	})
+
+	for path, shouldSkip := range map[string]bool{
+		"/ok":      false,
+		"/skip":    true,
+		"/skipfoo": true,
+	} {
+		mt := mocktracer.Start()
+		defer mt.Reset()
+
+		r := httptest.NewRequest("GET", "http://localhost"+path, nil)
+		router.ServeHTTP(httptest.NewRecorder(), r)
+		assert.Equal(t, shouldSkip, len(mt.FinishedSpans()) == 0)
+	}
+}
diff --git a/contrib/go-chi/chi.v5/option.go b/contrib/go-chi/chi.v5/option.go
@@ -7,6 +7,7 @@ package chi
 
 import (
 	"math"
+	"net/http"
 
 	"gopkg.in/DataDog/dd-trace-go.v1/ddtrace"
 	"gopkg.in/DataDog/dd-trace-go.v1/internal"
@@ -18,6 +19,7 @@ type config struct {
 	spanOpts      []ddtrace.StartSpanOption // additional span options to be applied
 	analyticsRate float64
 	isStatusError func(statusCode int) bool
+	ignoreRequest func(r *http.Request) bool
 }
 
 // Option represents an option that can be passed to NewRouter.
@@ -34,6 +36,7 @@ func defaults(cfg *config) {
 		cfg.analyticsRate = globalconfig.AnalyticsRate()
 	}
 	cfg.isStatusError = isServerError
+	cfg.ignoreRequest = func(_ *http.Request) bool { return false }
 }
 
 // WithServiceName sets the given service name for the router.
@@ -85,3 +88,11 @@ func WithStatusCheck(fn func(statusCode int) bool) Option {
 func isServerError(statusCode int) bool {
 	return statusCode >= 500 && statusCode < 600
 }
+
+// WithIgnoreRequest specifies a function to use for determining if the
+// incoming HTTP request tracing should be skipped.
+func WithIgnoreRequest(fn func(r *http.Request) bool) Option {
+	return func(cfg *config) {
+		cfg.ignoreRequest = fn
+	}
+}
diff --git a/contrib/go-chi/chi/chi.go b/contrib/go-chi/chi/chi.go
@@ -31,6 +31,10 @@ func Middleware(opts ...Option) func(next http.Handler) http.Handler {
 	log.Debug("contrib/go-chi/chi: Configuring Middleware: %#v", cfg)
 	return func(next http.Handler) http.Handler {
 		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			if cfg.ignoreRequest(r) {
+				next.ServeHTTP(w, r)
+				return
+			}
 			opts := []ddtrace.StartSpanOption{
 				tracer.SpanType(ext.SpanTypeWeb),
 				tracer.ServiceName(cfg.serviceName),

diff --git a/contrib/go-chi/chi/chi_test.go b/contrib/go-chi/chi/chi_test.go
@@ -10,6 +10,7 @@ import (
 	"net/http"
 	"net/http/httptest"
 	"strconv"
+	"strings"
 	"testing"
 
 	"gopkg.in/DataDog/dd-trace-go.v1/ddtrace/ext"
@@ -270,3 +271,33 @@ func TestAnalyticsSettings(t *testing.T) {
 		assertRate(t, mt, 0.23, WithAnalyticsRate(0.23))
 	})
 }
+
+func TestIgnoreRequest(t *testing.T) {
+	router := chi.NewRouter()
+	router.Use(Middleware(
+		WithIgnoreRequest(func(r *http.Request) bool {
+			return strings.HasPrefix(r.URL.Path, "/skip")
+		}),
+	))
+
+	router.Get("/ok", func(w http.ResponseWriter, r *http.Request) {
+		w.Write([]byte("ok"))
+	})
+
+	router.Get("/skip", func(w http.ResponseWriter, r *http.Request) {
+		w.Write([]byte("skip"))
+	})
+
+	for path, shouldSkip := range map[string]bool{
+		"/ok":      false,
+		"/skip":    true,
+		"/skipfoo": true,
+	} {
+		mt := mocktracer.Start()
+		defer mt.Reset()
+
+		r := httptest.NewRequest("GET", "http://localhost"+path, nil)
+		router.ServeHTTP(httptest.NewRecorder(), r)
+		assert.Equal(t, shouldSkip, len(mt.FinishedSpans()) == 0)
+	}
+}
diff --git a/contrib/go-chi/chi/option.go b/contrib/go-chi/chi/option.go
@@ -7,6 +7,7 @@ package chi
 
 import (
 	"math"
+	"net/http"
 
 	"gopkg.in/DataDog/dd-trace-go.v1/ddtrace"
 	"gopkg.in/DataDog/dd-trace-go.v1/internal"
@@ -18,6 +19,7 @@ type config struct {
 	spanOpts      []ddtrace.StartSpanOption // additional span options to be applied
 	analyticsRate float64
 	isStatusError func(statusCode int) bool
+	ignoreRequest func(r *http.Request) bool
 }
 
 // Option represents an option that can be passed to NewRouter.
@@ -34,6 +36,7 @@ func defaults(cfg *config) {
 		cfg.analyticsRate = globalconfig.AnalyticsRate()
 	}
 	cfg.isStatusError = isServerError
+	cfg.ignoreRequest = func(_ *http.Request) bool { return false }
 }
 
 // WithServiceName sets the given service name for the router.
@@ -85,3 +88,11 @@ func WithStatusCheck(fn func(statusCode int) bool) Option {
 func isServerError(statusCode int) bool {
 	return statusCode >= 500 && statusCode < 600
 }
+
+// WithIgnoreRequest specifies a function to use for determining if the
+// incoming HTTP request tracing should be skipped.
+func WithIgnoreRequest(fn func(r *http.Request) bool) Option {
+	return func(cfg *config) {
+		cfg.ignoreRequest = fn
+	}
+}