use httpsnoop to wrap the response writer so we don't lose interfaces

open-telemetry · Oct 8, 2020 · 5a7a94c · 5a7a94c
1 parent 30ab5aa
commit 5a7a94c
Show file tree

Hide file tree

Showing 5 changed files with 33 additions and 23 deletions.
diff --git a/go.mod b/go.mod
@@ -4,8 +4,10 @@ go 1.14
 
 require (
 	github.com/davecgh/go-spew v1.1.1 // indirect
+	github.com/felixge/httpsnoop v1.0.1 // indirect
 	github.com/gorilla/mux v1.8.0 // indirect
 	github.com/kr/pretty v0.1.0 // indirect
+	go.opentelemetry.io/contrib/propagators v0.12.0 // indirect
 	go.opentelemetry.io/otel v0.12.0
 	gopkg.in/check.v1 v1.0.0-20190902080502-41f04d3bba15 // indirect
 )
diff --git a/go.sum b/go.sum
@@ -1,6 +1,8 @@
 github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
+github.com/felixge/httpsnoop v1.0.1 h1:lvB5Jl89CsZtGIWuTcDM1E/vkVs49/Ml7JJe07l8SPQ=
+github.com/felixge/httpsnoop v1.0.1/go.mod h1:m8KPJKqk1gH5J9DgRY2ASl2lWCfGKXixSwevea8zH2U=
 github.com/google/go-cmp v0.5.2 h1:X2ev0eStA3AbceY54o37/0PQ/UWqKEiiO2dKL5OPaFM=
 github.com/google/go-cmp v0.5.2/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE=
 github.com/gorilla/mux v1.8.0 h1:i40aqfkR1h2SlN9hojwV5ZA91wcXFOvkdNIeFDP5koI=
@@ -16,6 +18,10 @@ github.com/stretchr/objx v0.1.0 h1:4G4v2dO3VZwixGIRoQ5Lfboy6nUhCyYzaqnIAPPhYs4=
 github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
 github.com/stretchr/testify v1.6.1 h1:hDPOHmpOpP40lSULcqw7IrRb/u7w6RpDC9399XyoNd0=
 github.com/stretchr/testify v1.6.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
+go.opentelemetry.io v0.1.0 h1:EANZoRCOP+A3faIlw/iN6YEWoYb1vleZRKm1EvH8T48=
+go.opentelemetry.io/contrib v0.12.0/go.mod h1:onlxH6TKFRkW2Xgc5IO37kPYz3v7wMzh/FrBxsQxCt4=
+go.opentelemetry.io/contrib/propagators v0.12.0 h1:X+BFc/3zD/WGrt3rgNdz7/6zqbRdYSjH7+Mv2IAWns4=
+go.opentelemetry.io/contrib/propagators v0.12.0/go.mod h1:2cKNY+NJQjmzJkFyf8U4OyHABIxCi53ERyOjV2Tulyc=
 go.opentelemetry.io/otel v0.12.0 h1:bwWaPd/h2q+U6KdKaAiOS5GLwOMd1LDt9iNaeyIoAI8=
 go.opentelemetry.io/otel v0.12.0/go.mod h1:dlSNewoRYikTkotEnxdmuBHgzT+k/idJSfDv/FxEnOY=
 golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543 h1:E7g+9GITq07hpfrRu66IVDexMakfv52eLZ2CXBWiKr4=

diff --git a/instrumentation/github.com/gorilla/mux/otelmux/go.mod b/instrumentation/github.com/gorilla/mux/otelmux/go.mod
@@ -8,6 +8,7 @@ replace (
 )
 
 require (
+	github.com/felixge/httpsnoop v1.0.1
 	github.com/gorilla/mux v1.8.0
 	github.com/stretchr/testify v1.6.1
 	go.opentelemetry.io/contrib v0.12.0

diff --git a/instrumentation/github.com/gorilla/mux/otelmux/go.sum b/instrumentation/github.com/gorilla/mux/otelmux/go.sum
@@ -1,6 +1,8 @@
 github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
+github.com/felixge/httpsnoop v1.0.1 h1:lvB5Jl89CsZtGIWuTcDM1E/vkVs49/Ml7JJe07l8SPQ=
+github.com/felixge/httpsnoop v1.0.1/go.mod h1:m8KPJKqk1gH5J9DgRY2ASl2lWCfGKXixSwevea8zH2U=
 github.com/google/go-cmp v0.5.2 h1:X2ev0eStA3AbceY54o37/0PQ/UWqKEiiO2dKL5OPaFM=
 github.com/google/go-cmp v0.5.2/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE=
 github.com/gorilla/mux v1.8.0 h1:i40aqfkR1h2SlN9hojwV5ZA91wcXFOvkdNIeFDP5koI=

diff --git a/instrumentation/github.com/gorilla/mux/otelmux/mux.go b/instrumentation/github.com/gorilla/mux/otelmux/mux.go
@@ -19,6 +19,7 @@ import (
 	"net/http"
 	"sync"
 
+	"github.com/felixge/httpsnoop"
 	"github.com/gorilla/mux"
 
 	otelcontrib "go.opentelemetry.io/contrib"
@@ -74,27 +75,6 @@ type recordingResponseWriter struct {
 	status  int
 }
 
-func (w *recordingResponseWriter) Header() http.Header {
-	return w.writer.Header()
-}
-
-func (w *recordingResponseWriter) Write(slice []byte) (int, error) {
-	w.writeHeader(http.StatusOK)
-	return w.writer.Write(slice)
-}
-
-func (w *recordingResponseWriter) WriteHeader(statusCode int) {
-	w.writeHeader(statusCode)
-	w.writer.WriteHeader(statusCode)
-}
-
-func (w *recordingResponseWriter) writeHeader(statusCode int) {
-	if !w.written {
-		w.written = true
-		w.status = statusCode
-	}
-}
-
 var rrwPool = &sync.Pool{
 	New: func() interface{} {
 		return &recordingResponseWriter{}
@@ -105,7 +85,26 @@ func getRRW(writer http.ResponseWriter) *recordingResponseWriter {
 	rrw := rrwPool.Get().(*recordingResponseWriter)
 	rrw.written = false
 	rrw.status = 0
-	rrw.writer = writer
+	rrw.writer = httpsnoop.Wrap(writer, httpsnoop.Hooks{
+		Write: func(next httpsnoop.WriteFunc) httpsnoop.WriteFunc {
+			return func(b []byte) (int, error) {
+				if !rrw.written {
+					rrw.written = true
+					rrw.status = http.StatusOK
+				}
+				return next(b)
+			}
+		},
+		WriteHeader: func(next httpsnoop.WriteHeaderFunc) httpsnoop.WriteHeaderFunc {
+			return func(statusCode int) {
+				if !rrw.written {
+					rrw.written = true
+					rrw.status = statusCode
+				}
+				next(statusCode)
+			}
+		},
+	})
 	return rrw
 }
 
@@ -145,7 +144,7 @@ func (tw traceware) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	r2 := r.WithContext(ctx)
 	rrw := getRRW(w)
 	defer putRRW(rrw)
-	tw.handler.ServeHTTP(rrw, r2)
+	tw.handler.ServeHTTP(rrw.writer, r2)
 	attrs := semconv.HTTPAttributesFromHTTPStatusCode(rrw.status)
 	spanStatus, spanMessage := semconv.SpanStatusFromHTTPStatusCode(rrw.status)
 	span.SetAttributes(attrs...)