diff --git a/api.middlewares.go b/api.middlewares.go index 4d7211f..6fe6f01 100644 --- a/api.middlewares.go +++ b/api.middlewares.go @@ -33,14 +33,15 @@ func (api *APIHandler) DurationMiddleware(next httprouter.Handle) httprouter.Han next(nw, r, ps) logger.Info( "stats", - zap.Int("request.status", nw.statusCode), + zap.Int("request.status", nw.Status()), + zap.Int("bytes.sent", nw.Bytes()), zap.Duration("request.duration", api.clock.Now().Sub(start)), ) api.stats.mu.Lock() - if num, found := api.stats.status[nw.statusCode]; !found { - api.stats.status[nw.statusCode] = 1 + if num, found := api.stats.status[nw.code]; !found { + api.stats.status[nw.code] = 1 } else { - api.stats.status[nw.statusCode] = num + 1 + api.stats.status[nw.code] = num + 1 } api.stats.mu.Unlock() } diff --git a/helpers.response.go b/helpers.response.go index 27ec5de..49ddff2 100644 --- a/helpers.response.go +++ b/helpers.response.go @@ -5,39 +5,62 @@ import ( "net/http" ) -// CustomResponseWriter is a wrapper for http.ResponseWriter. -// It is used to record response attributes like statusCode. +// CustomResponseWriter is a wrapper for http.ResponseWriter. It is +// used to record response details like status code and body size. type CustomResponseWriter struct { http.ResponseWriter - statusCode int + code int + bytes int + wrote bool } // NewCustomResponseWriter provides CustomResponseWriter with 200 as status code. func NewCustomResponseWriter(rw http.ResponseWriter) *CustomResponseWriter { return &CustomResponseWriter{ ResponseWriter: rw, - statusCode: 200, + code: 200, } } // Header implements http.Header interface. -func (rw *CustomResponseWriter) Header() http.Header { - return rw.ResponseWriter.Header() +func (cw *CustomResponseWriter) Header() http.Header { + return cw.ResponseWriter.Header() } // WriteHeader implements http.WriteHeader interface. -func (rw *CustomResponseWriter) WriteHeader(statusCode int) { - rw.ResponseWriter.WriteHeader(statusCode) - rw.statusCode = statusCode +func (cw *CustomResponseWriter) WriteHeader(code int) { + if !cw.wrote { + cw.code = code + cw.wrote = true + cw.ResponseWriter.WriteHeader(code) + } } // Write implements http.Write interface. -func (rw *CustomResponseWriter) Write(bytes []byte) (int, error) { - return rw.ResponseWriter.Write(bytes) +func (cw *CustomResponseWriter) Write(bytes []byte) (int, error) { + if !cw.wrote { + cw.WriteHeader(cw.code) + } + + n, err := cw.ResponseWriter.Write(bytes) + cw.bytes += n + return n, err +} + +// Status returns the written status code. +func (cw *CustomResponseWriter) Status() int { + return cw.code +} + +// Bytes returns bytes written as response body. +func (cw *CustomResponseWriter) Bytes() int { + return cw.bytes } -func (rw *CustomResponseWriter) Status() int { - return rw.statusCode +// Unwrap returns native response writer and used by +// the http.ResponseController during its operation. +func (cw *CustomResponseWriter) Unwrap() http.ResponseWriter { + return cw.ResponseWriter } // APIError is the data model sent when an error occurred during request processing.