contrib/go-chi/chi: Added option to set custom error status check

contrib/go-chi/chi/chi.go

@@ -67,7 +67,7 @@ func Middleware(opts ...Option) func(next http.Handler) http.Handler {
 			}
 			span.SetTag(ext.HTTPCode, strconv.Itoa(status))
 
-			if status >= 500 && status < 600 {
+			if cfg.statusCheck(status) {
 				// mark 5xx server error
 				span.SetTag(ext.Error, fmt.Errorf("%d: %s", status, http.StatusText(status)))
 			}

contrib/go-chi/chi/chi_test.go

@@ -9,6 +9,7 @@ import (
 	"fmt"
 	"net/http"
 	"net/http/httptest"
+	"strconv"
 	"testing"
 
 	"gopkg.in/DataDog/dd-trace-go.v1/ddtrace/ext"
@@ -99,37 +100,74 @@ func TestTrace200(t *testing.T) {
 }
 
 func TestError(t *testing.T) {
-	assert := assert.New(t)
-	mt := mocktracer.Start()
-	defer mt.Stop()
+	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))
 
-	// setup
-	router := chi.NewRouter()
-	router.Use(Middleware(WithServiceName("foobar")))
-	code := 500
-	wantErr := fmt.Sprintf("%d: %s", code, http.StatusText(code))
+		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("with default isError function", 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)
 
-	// 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)
+		// verify the errors and status are correct
+		spans := mt.FinishedSpans()
+		assertSpan(assert, spans, code)
 	})
-	r := httptest.NewRequest("GET", "/err", nil)
-	w := httptest.NewRecorder()
-	router.ServeHTTP(w, r)
-	response := w.Result()
-	assert.Equal(response.StatusCode, 500)
 
-	// verify the errors and status are correct
-	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("foobar", span.Tag(ext.ServiceName))
-	assert.Equal("500", span.Tag(ext.HTTPCode))
-	assert.Equal(wantErr, span.Tag(ext.Error).(error).Error())
+	t.Run("with custome isError function", 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) {

contrib/go-chi/chi/option.go

@@ -17,6 +17,7 @@ type config struct {
 	serviceName   string
 	spanOpts      []ddtrace.StartSpanOption // additional span options to be applied
 	analyticsRate float64
+	statusCheck   func(statusCode int) bool
 }
 
 // Option represents an option that can be passed to NewRouter.
@@ -32,6 +33,8 @@ func defaults(cfg *config) {
 	} else {
 		cfg.analyticsRate = globalconfig.AnalyticsRate()
 	}
+
+	cfg.statusCheck = isServerError
 }
 
 // WithServiceName sets the given service name for the router.
@@ -71,3 +74,15 @@ func WithAnalyticsRate(rate float64) Option {
 		}
 	}
 }
+
+// WithStatusCheck sets the function which is used to decide whther the
+// status should be marked as an error
+func WithStatusCheck(statusCheck func(statusCode int) bool) Option {
+	return func(cfg *config) {
+		cfg.statusCheck = statusCheck
+	}
+}
+
+func isServerError(statusCode int) bool {
+	return statusCode >= 500 && statusCode < 600
+}