Skip to content
Open
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
17 changes: 4 additions & 13 deletions services/graph/pkg/metrics/middleware.go
Original file line number Diff line number Diff line change
Expand Up @@ -6,18 +6,9 @@ import (
"time"

"github.com/go-chi/chi/v5"
"github.com/go-chi/chi/v5/middleware"
)

type statusResponseWriter struct {
http.ResponseWriter
statusCode int
}

func (rw *statusResponseWriter) WriteHeader(code int) {
rw.statusCode = code
rw.ResponseWriter.WriteHeader(code)
}

// A middleware that tracks the duration of every inbound Graph API HTTP call
// and calls a function to delegate the storage of that duration into a
// histogram metric, analyzing the incoming query and deconstructing it into
Expand All @@ -34,11 +25,11 @@ func HTTPMetrics(inFlight *atomic.Int64, observe func(method, pattern string, st
return func(next http.Handler) http.Handler {
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
start := time.Now()
responseWrapper := &statusResponseWriter{ResponseWriter: w, statusCode: 200} // 200 OK is the default when it's not set
inFlight.Add(1)
defer inFlight.Add(-1)

next.ServeHTTP(responseWrapper, r)
wrapper := middleware.NewWrapResponseWriter(w, r.ProtoMajor)
next.ServeHTTP(wrapper, r)

duration := time.Since(start)

Expand All @@ -53,7 +44,7 @@ func HTTPMetrics(inFlight *atomic.Int64, observe func(method, pattern string, st
}
}

observe(r.Method, routePattern, responseWrapper.statusCode, duration)
observe(r.Method, routePattern, wrapper.Status(), duration)
})
}
}