diff --git a/cmd/api/middleware.go b/cmd/api/middleware.go index 07f0bd5..b296e0f 100644 --- a/cmd/api/middleware.go +++ b/cmd/api/middleware.go @@ -6,6 +6,7 @@ import ( "fmt" "net" "net/http" + "strconv" "strings" "sync" "time" @@ -201,20 +202,60 @@ func (app *application) enableCORS(next http.Handler) http.Handler { }) } +type metricsResponseWriter struct { + wrapped http.ResponseWriter + statusCode int + headerWritten bool +} + +func (mw *metricsResponseWriter) Header() http.Header { + return mw.wrapped.Header() +} + +func (mw *metricsResponseWriter) WriteHeader(statusCode int) { + mw.wrapped.WriteHeader(statusCode) + + if !mw.headerWritten { + mw.statusCode = statusCode + mw.headerWritten = true + } +} + +func (mw *metricsResponseWriter) Write(b []byte) (int, error) { + if !mw.headerWritten { + mw.statusCode = http.StatusOK + mw.headerWritten = true + } + + return mw.wrapped.Write(b) +} + +func (mw *metricsResponseWriter) Unwrap() http.ResponseWriter { + return mw.wrapped +} + func (app *application) metrics(next http.Handler) http.Handler { var ( totalRequestsReceived = expvar.NewInt("total_requests_received") totalResponseSent = expvar.NewInt("total_responses_sent") totalProcessingTimeMicroseconds = expvar.NewInt("total_processing_time_μs") + + totalResponsesSentByStatus = expvar.NewMap("total_responses_sent_by_status") ) return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { start := time.Now() totalRequestsReceived.Add(1) - next.ServeHTTP(w, r) + + mw := &metricsResponseWriter{wrapped: w} + + next.ServeHTTP(mw, r) + totalResponseSent.Add(1) + totalResponsesSentByStatus.Add(strconv.Itoa(mw.statusCode), 1) + duration := time.Since(start).Microseconds() totalProcessingTimeMicroseconds.Add(duration) })