add WithStatusCheck option

DataDog · Dec 8, 2022 · 83001e7 · 83001e7
1 parent b879761
commit 83001e7
Show file tree

Hide file tree

Showing 6 changed files with 171 additions and 13 deletions.
diff --git a/contrib/labstack/echo.v4/echotrace.go b/contrib/labstack/echo.v4/echotrace.go
@@ -7,7 +7,9 @@
 package echo
 
 import (
+	"fmt"
 	"math"
+	"net/http"
 	"strconv"
 
 	"gopkg.in/DataDog/dd-trace-go.v1/contrib/internal/httptrace"
@@ -80,16 +82,28 @@ func Middleware(opts ...Option) echo.MiddlewareFunc {
 				// This is the best we can do.
 				switch err := err.(type) {
 				case *echo.HTTPError:
-					if err.Code >= 500 {
+					if cfg.isStatusError(err.Code) {
 						finishOpts = append(finishOpts, tracer.WithError(err))
 					}
 					span.SetTag(ext.HTTPCode, strconv.Itoa(err.Code))
 				default:
 					// Any non-HTTPError errors appear as 5xx errors.
-					finishOpts = append(finishOpts, tracer.WithError(err))
+					if cfg.isStatusError(500) {
+						finishOpts = append(finishOpts, tracer.WithError(err))
+					}
 					span.SetTag(ext.HTTPCode, "500")
 				}
+			} else if status := c.Response().Status; status > 0 {
+				if cfg.isStatusError(status) {
+					// mark 5xx server error
+					finishOpts = append(finishOpts, tracer.WithError(fmt.Errorf("%d: %s", status, http.StatusText(status))))
+				}
+				span.SetTag(ext.HTTPCode, strconv.Itoa(status))
 			} else {
+				if cfg.isStatusError(200) {
+					// mark 5xx server error
+					finishOpts = append(finishOpts, tracer.WithError(fmt.Errorf("%d: %s", 200, http.StatusText(200))))
+				}
 				span.SetTag(ext.HTTPCode, "200")
 			}
 			return err

diff --git a/contrib/labstack/echo.v4/echotrace_test.go b/contrib/labstack/echo.v4/echotrace_test.go
@@ -7,6 +7,7 @@ package echo
 
 import (
 	"errors"
+	"fmt"
 	"io"
 	"net/http"
 	"net/http/httptest"
@@ -226,9 +227,10 @@ func TestErrorHandling(t *testing.T) {
 
 func TestStatusError(t *testing.T) {
 	for _, tt := range []struct {
-		err     error
-		code    string
-		handler func(c echo.Context) error
+		isStatusError func(statusCode int) bool
+		err           error
+		code          string
+		handler       func(c echo.Context) error
 	}{
 		{
 			err:  errors.New("oh no"),
@@ -251,14 +253,61 @@ func TestStatusError(t *testing.T) {
 				return echo.NewHTTPError(http.StatusBadRequest, "my error message")
 			},
 		},
+		{
+			isStatusError: func(statusCode int) bool { return statusCode >= 400 && statusCode < 500 },
+			err:           nil,
+			code:          "500",
+			handler: func(c echo.Context) error {
+				return errors.New("oh no")
+			},
+		},
+		{
+			isStatusError: func(statusCode int) bool { return statusCode >= 400 && statusCode < 500 },
+			err:           nil,
+			code:          "500",
+			handler: func(c echo.Context) error {
+				return echo.NewHTTPError(http.StatusInternalServerError, "my error message")
+			},
+		},
+		{
+			isStatusError: func(statusCode int) bool { return statusCode >= 400 },
+			err:           echo.NewHTTPError(http.StatusBadRequest, "my error message"),
+			code:          "400",
+			handler: func(c echo.Context) error {
+				return echo.NewHTTPError(http.StatusBadRequest, "my error message")
+			},
+		},
+		{
+			isStatusError: func(statusCode int) bool { return statusCode >= 200 },
+			err:           fmt.Errorf("201: Created"),
+			code:          "201",
+			handler: func(c echo.Context) error {
+				c.JSON(201, map[string]string{"status": "ok", "type": "test"})
+				return nil
+			},
+		},
+		{
+			isStatusError: func(statusCode int) bool { return statusCode >= 200 },
+			err:           fmt.Errorf("200: OK"),
+			code:          "200",
+			handler: func(c echo.Context) error {
+				// It's not clear if unset (0) status is possible naturally, but we can simulate that situation.
+				c.Response().Status = 0
+				return nil
+			},
+		},
 	} {
 		t.Run("", func(t *testing.T) {
 			assert := assert.New(t)
 			mt := mocktracer.Start()
 			defer mt.Stop()
 
 			router := echo.New()
-			router.Use(Middleware(WithServiceName("foobar")))
+			opts := []Option{WithServiceName("foobar")}
+			if tt.isStatusError != nil {
+				opts = append(opts, WithStatusCheck(tt.isStatusError))
+			}
+			router.Use(Middleware(opts...))
 			router.GET("/err", tt.handler)
 			r := httptest.NewRequest("GET", "/err", nil)
 			w := httptest.NewRecorder()
@@ -275,7 +324,9 @@ func TestStatusError(t *testing.T) {
 			assert.Equal("GET", span.Tag(ext.HTTPMethod))
 			err := span.Tag(ext.Error)
 			if tt.err != nil {
-				assert.NotNil(err)
+				if !assert.NotNil(err) {
+					return
+				}
 				assert.Equal(tt.err.Error(), err.(error).Error())
 			} else {
 				assert.Nil(err)

diff --git a/contrib/labstack/echo.v4/option.go b/contrib/labstack/echo.v4/option.go
@@ -18,6 +18,7 @@ type config struct {
 	analyticsRate     float64
 	noDebugStack      bool
 	ignoreRequestFunc IgnoreRequestFunc
+	isStatusError func(statusCode int) bool
 }
 
 // Option represents an option that can be passed to Middleware.
@@ -32,6 +33,7 @@ func defaults(cfg *config) {
 		cfg.serviceName = svc
 	}
 	cfg.analyticsRate = math.NaN()
+	cfg.isStatusError = isServerError
 }
 
 // WithServiceName sets the given service name for the system.
@@ -64,6 +66,7 @@ func WithAnalyticsRate(rate float64) Option {
 	}
 }
 
+
 // NoDebugStack prevents stack traces from being attached to spans finishing
 // with an error. This is useful in situations where errors are frequent and
 // performance is critical.
@@ -79,4 +82,15 @@ func WithIgnoreRequest(ignoreRequestFunc IgnoreRequestFunc) Option {
 	return func(cfg *config) {
 		cfg.ignoreRequestFunc = ignoreRequestFunc
 	}
+
+// WithStatusCheck specifies a function fn which reports whether the passed
+// statusCode should be considered an error.
+func WithStatusCheck(fn func(statusCode int) bool) Option {
+	return func(cfg *config) {
+		cfg.isStatusError = fn
+	}
+}
+
+func isServerError(statusCode int) bool {
+	return statusCode >= 500 && statusCode < 600
 }
diff --git a/contrib/labstack/echo/echotrace.go b/contrib/labstack/echo/echotrace.go
@@ -7,7 +7,9 @@
 package echo
 
 import (
+	"fmt"
 	"math"
+	"net/http"
 	"strconv"
 
 	"gopkg.in/DataDog/dd-trace-go.v1/contrib/internal/httptrace"
@@ -65,16 +67,28 @@ func Middleware(opts ...Option) echo.MiddlewareFunc {
 				// This is the best we can do.
 				switch err := err.(type) {
 				case *echo.HTTPError:
-					if err.Code >= 500 {
+					if cfg.isStatusError(err.Code) {
 						finishOpts = append(finishOpts, tracer.WithError(err))
 					}
 					span.SetTag(ext.HTTPCode, strconv.Itoa(err.Code))
 				default:
 					// Any non-HTTPError errors appear as 5xx errors.
-					finishOpts = append(finishOpts, tracer.WithError(err))
+					if cfg.isStatusError(500) {
+						finishOpts = append(finishOpts, tracer.WithError(err))
+					}
 					span.SetTag(ext.HTTPCode, "500")
 				}
+			} else if status := c.Response().Status; status > 0 {
+				if cfg.isStatusError(status) {
+					// mark 5xx server error
+					finishOpts = append(finishOpts, tracer.WithError(fmt.Errorf("%d: %s", status, http.StatusText(status))))
+				}
+				span.SetTag(ext.HTTPCode, strconv.Itoa(status))
 			} else {
+				if cfg.isStatusError(200) {
+					// mark 5xx server error
+					finishOpts = append(finishOpts, tracer.WithError(fmt.Errorf("%d: %s", 200, http.StatusText(200))))
+				}
 				span.SetTag(ext.HTTPCode, "200")
 			}
 			return err

diff --git a/contrib/labstack/echo/echotrace_test.go b/contrib/labstack/echo/echotrace_test.go
@@ -7,6 +7,7 @@ package echo
 
 import (
 	"errors"
+	"fmt"
 	"net/http"
 	"net/http/httptest"
 	"testing"
@@ -221,9 +222,10 @@ func TestErrorHandling(t *testing.T) {
 
 func TestStatusError(t *testing.T) {
 	for _, tt := range []struct {
-		err     error
-		code    string
-		handler func(c echo.Context) error
+		isStatusError func(statusCode int) bool
+		err           error
+		code          string
+		handler       func(c echo.Context) error
 	}{
 		{
 			err:  errors.New("oh no"),
@@ -246,14 +248,61 @@ func TestStatusError(t *testing.T) {
 				return echo.NewHTTPError(http.StatusBadRequest, "my error message")
 			},
 		},
+		{
+			isStatusError: func(statusCode int) bool { return statusCode >= 400 && statusCode < 500 },
+			err:           nil,
+			code:          "500",
+			handler: func(c echo.Context) error {
+				return errors.New("oh no")
+			},
+		},
+		{
+			isStatusError: func(statusCode int) bool { return statusCode >= 400 && statusCode < 500 },
+			err:           nil,
+			code:          "500",
+			handler: func(c echo.Context) error {
+				return echo.NewHTTPError(http.StatusInternalServerError, "my error message")
+			},
+		},
+		{
+			isStatusError: func(statusCode int) bool { return statusCode >= 400 },
+			err:           echo.NewHTTPError(http.StatusBadRequest, "my error message"),
+			code:          "400",
+			handler: func(c echo.Context) error {
+				return echo.NewHTTPError(http.StatusBadRequest, "my error message")
+			},
+		},
+		{
+			isStatusError: func(statusCode int) bool { return statusCode >= 200 },
+			err:           fmt.Errorf("201: Created"),
+			code:          "201",
+			handler: func(c echo.Context) error {
+				c.JSON(201, map[string]string{"status": "ok", "type": "test"})
+				return nil
+			},
+		},
+		{
+			isStatusError: func(statusCode int) bool { return statusCode >= 200 },
+			err:           fmt.Errorf("200: OK"),
+			code:          "200",
+			handler: func(c echo.Context) error {
+				// It's not clear if unset (0) status is possible naturally, but we can simulate that situation.
+				c.Response().Status = 0
+				return nil
+			},
+		},
 	} {
 		t.Run("", func(t *testing.T) {
 			assert := assert.New(t)
 			mt := mocktracer.Start()
 			defer mt.Stop()
 
 			router := echo.New()
-			router.Use(Middleware(WithServiceName("foobar")))
+			opts := []Option{WithServiceName("foobar")}
+			if tt.isStatusError != nil {
+				opts = append(opts, WithStatusCheck(tt.isStatusError))
+			}
+			router.Use(Middleware(opts...))
 			router.GET("/err", tt.handler)
 			r := httptest.NewRequest("GET", "/err", nil)
 			w := httptest.NewRecorder()

diff --git a/contrib/labstack/echo/option.go b/contrib/labstack/echo/option.go
@@ -16,6 +16,7 @@ type config struct {
 	serviceName   string
 	analyticsRate float64
 	noDebugStack  bool
+	isStatusError func(statusCode int) bool
 }
 
 // Option represents an option that can be passed to Middleware.
@@ -31,6 +32,7 @@ func defaults(cfg *config) {
 	} else {
 		cfg.analyticsRate = math.NaN()
 	}
+	cfg.isStatusError = isServerError
 }
 
 // WithServiceName sets the given service name for the system.
@@ -63,6 +65,7 @@ func WithAnalyticsRate(rate float64) Option {
 	}
 }
 
+<<<<<<< HEAD
 // NoDebugStack prevents stack traces from being attached to spans finishing
 // with an error. This is useful in situations where errors are frequent and
 // performance is critical.
@@ -71,3 +74,16 @@ func NoDebugStack() Option {
 		cfg.noDebugStack = true
 	}
 }
+=======
+// WithStatusCheck specifies a function fn which reports whether the passed
+// statusCode should be considered an error.
+func WithStatusCheck(fn func(statusCode int) bool) Option {
+	return func(cfg *config) {
+		cfg.isStatusError = fn
+	}
+}
+
+func isServerError(statusCode int) bool {
+	return statusCode >= 500 && statusCode < 600
+}
+>>>>>>> add WithStatusCheck option