Use httpsnoop.Wrap in datadog package

go.mod

@@ -12,6 +12,7 @@ require (
 	github.com/aws/aws-sdk-go v1.44.81
 	github.com/benesch/cgosymbolizer v0.0.0-20190515212042-bec6fe6e597b
 	github.com/bugsnag/bugsnag-go/v2 v2.1.2
+	github.com/felixge/httpsnoop v1.0.3
 	github.com/fsouza/fake-gcs-server v1.38.4
 	github.com/getsentry/sentry-go v0.13.0
 	github.com/go-ole/go-ole v1.2.6 // indirect

go.sum

@@ -304,8 +304,9 @@ github.com/fatih/color v1.9.0/go.mod h1:eQcE1qtQxscV5RaZvpXrrb8Drkc3/DdQ+uUYCNjL
 github.com/fatih/color v1.13.0/go.mod h1:kLAiJbzzSOZDVNGyDpeOxJ47H46qBXwg5ILebYFFOfk=
 github.com/fatih/structs v1.1.0/go.mod h1:9NiDSp5zOcgEDl+j00MP/WkGVPOlPRLejGD8Ga6PJ7M=
 github.com/felixge/httpsnoop v1.0.1/go.mod h1:m8KPJKqk1gH5J9DgRY2ASl2lWCfGKXixSwevea8zH2U=
-github.com/felixge/httpsnoop v1.0.2 h1:+nS9g82KMXccJ/wp0zyRW9ZBHFETmMGtkk+2CTTrW4o=
 github.com/felixge/httpsnoop v1.0.2/go.mod h1:m8KPJKqk1gH5J9DgRY2ASl2lWCfGKXixSwevea8zH2U=
+github.com/felixge/httpsnoop v1.0.3 h1:s/nj+GCswXYzN5v2DpNMuMQYe+0DDwt5WVCU6CWBdXk=
+github.com/felixge/httpsnoop v1.0.3/go.mod h1:m8KPJKqk1gH5J9DgRY2ASl2lWCfGKXixSwevea8zH2U=
 github.com/fogleman/gg v1.2.1-0.20190220221249-0403632d5b90/go.mod h1:R/bRT+9gY/C5z7JzPU0zXsXHKM4/ayA+zqcVNZzPa1k=
 github.com/form3tech-oss/jwt-go v3.2.2+incompatible h1:TcekIExNqud5crz4xD2pavyTgWiPvpYe4Xau31I0PRk=
 github.com/form3tech-oss/jwt-go v3.2.2+incompatible/go.mod h1:pbq4aXjuKjdthFRnoDwaVPLA+WlJuPGy+QneDUgJi2k=

metrics/datadog/datadog.go

@@ -9,6 +9,7 @@ import (
 	"time"
 
 	"github.com/DataDog/datadog-go/v5/statsd"
+	"github.com/felixge/httpsnoop"
 	log "github.com/sirupsen/logrus"
 	"gopkg.in/DataDog/dd-trace-go.v1/ddtrace/ext"
 	"gopkg.in/DataDog/dd-trace-go.v1/ddtrace/tracer"
@@ -109,7 +110,15 @@ func StartRootSpan(ctx context.Context, rw http.ResponseWriter, r *http.Request)
 		tracer.Tag(ext.HTTPURL, r.RequestURI),
 	)
 	cancel := func() { span.Finish() }
-	newRw := dataDogResponseWriter{rw, span}
+
+	newRw := httpsnoop.Wrap(rw, httpsnoop.Hooks{
+		WriteHeader: func(next httpsnoop.WriteHeaderFunc) httpsnoop.WriteHeaderFunc {
+			return func(statusCode int) {
+				span.SetTag(ext.HTTPCode, statusCode)
+				next(statusCode)
+			}
+		},
+	})
 
 	return context.WithValue(ctx, spanCtxKey{}, span), cancel, newRw
 }
@@ -192,19 +201,3 @@ type dataDogLogger struct {
 func (l dataDogLogger) Log(msg string) {
 	log.Info(msg)
 }
-
-type dataDogResponseWriter struct {
-	rw   http.ResponseWriter
-	span tracer.Span
-}
-
-func (ddrw dataDogResponseWriter) Header() http.Header {
-	return ddrw.rw.Header()
-}
-func (ddrw dataDogResponseWriter) Write(data []byte) (int, error) {
-	return ddrw.rw.Write(data)
-}
-func (ddrw dataDogResponseWriter) WriteHeader(statusCode int) {
-	ddrw.span.SetTag(ext.HTTPCode, statusCode)
-	ddrw.rw.WriteHeader(statusCode)
-}