diff --git a/go.mod b/go.mod index e1669cafc..0427a5124 100644 --- a/go.mod +++ b/go.mod @@ -17,6 +17,7 @@ require ( github.com/davecgh/go-spew v1.1.1 github.com/dgrijalva/jwt-go v3.2.0+incompatible github.com/edsrzf/mmap-go v1.0.0 // indirect + github.com/felixge/httpsnoop v1.0.2 github.com/franela/goblin v0.0.0-20200105215937-c9ffbefa60db // indirect github.com/franela/goreq v0.0.0-20171204163338-bcd34c9993f8 // indirect github.com/go-logfmt/logfmt v0.5.0 diff --git a/go.sum b/go.sum index b41a5f30d..da63e50c1 100644 --- a/go.sum +++ b/go.sum @@ -72,6 +72,8 @@ github.com/envoyproxy/go-control-plane v0.6.9/go.mod h1:SBwIajubJHhxtWwsL9s8ss4s github.com/envoyproxy/go-control-plane v0.9.1-0.20191026205805-5f8ba28d4473/go.mod h1:YTl/9mNaCwkRvm6d1a2C3ymFceY/DCBVvsKhRF0iEA4= github.com/envoyproxy/protoc-gen-validate v0.1.0/go.mod h1:iSmxcyjqTsJpI2R4NaDN7+kN2VEUnK/pcBlmesArF7c= github.com/fatih/color v1.7.0/go.mod h1:Zm6kSWBoL9eyXnKyktHP6abPY2pDugNf5KwzbycvMj4= +github.com/felixge/httpsnoop v1.0.2 h1:+nS9g82KMXccJ/wp0zyRW9ZBHFETmMGtkk+2CTTrW4o= +github.com/felixge/httpsnoop v1.0.2/go.mod h1:m8KPJKqk1gH5J9DgRY2ASl2lWCfGKXixSwevea8zH2U= github.com/franela/goblin v0.0.0-20200105215937-c9ffbefa60db h1:gb2Z18BhTPJPpLQWj4T+rfKHYCHxRHCtRxhKKjRidVw= github.com/franela/goblin v0.0.0-20200105215937-c9ffbefa60db/go.mod h1:7dvUGVsVBjqR7JHJk0brhHOZYGmfBYOrK0ZhYMEtBr4= github.com/franela/goreq v0.0.0-20171204163338-bcd34c9993f8 h1:a9ENSRDFBUPkJ5lCgVZh26+ZbGyoVJG7yb5SSzF5H54= diff --git a/transport/http/server.go b/transport/http/server.go index 34912d072..ee40a0f64 100644 --- a/transport/http/server.go +++ b/transport/http/server.go @@ -5,6 +5,7 @@ import ( "encoding/json" "net/http" + "github.com/felixge/httpsnoop" "github.com/go-kit/kit/endpoint" "github.com/go-kit/kit/log" "github.com/go-kit/kit/transport" @@ -96,15 +97,31 @@ func (s Server) ServeHTTP(w http.ResponseWriter, r *http.Request) { ctx := r.Context() if len(s.finalizer) > 0 { - iw := &interceptingWriter{w, http.StatusOK, 0} + metrics := httpsnoop.Metrics{Code: http.StatusOK} + + w = httpsnoop.Wrap(w, httpsnoop.Hooks{ + WriteHeader: func(whf httpsnoop.WriteHeaderFunc) httpsnoop.WriteHeaderFunc { + return func(code int) { + metrics.Code = code + whf(code) + } + }, + Write: func(wf httpsnoop.WriteFunc) httpsnoop.WriteFunc { + return func(b []byte) (int, error) { + n, err := wf(b) + metrics.Written += int64(n) + return n, err + } + }, + }) + defer func() { - ctx = context.WithValue(ctx, ContextKeyResponseHeaders, iw.Header()) - ctx = context.WithValue(ctx, ContextKeyResponseSize, iw.written) + ctx = context.WithValue(ctx, ContextKeyResponseHeaders, w.Header()) + ctx = context.WithValue(ctx, ContextKeyResponseSize, metrics.Written) for _, f := range s.finalizer { - f(ctx, iw.code, r) + f(ctx, metrics.Code, r) } }() - w = iw } for _, f := range s.before { @@ -223,22 +240,3 @@ type StatusCoder interface { type Headerer interface { Headers() http.Header } - -type interceptingWriter struct { - http.ResponseWriter - code int - written int64 -} - -// WriteHeader may not be explicitly called, so care must be taken to -// initialize w.code to its default value of http.StatusOK. -func (w *interceptingWriter) WriteHeader(code int) { - w.code = code - w.ResponseWriter.WriteHeader(code) -} - -func (w *interceptingWriter) Write(p []byte) (int, error) { - n, err := w.ResponseWriter.Write(p) - w.written += int64(n) - return n, err -}