fix: Make stats work

cmd/serve.go

@@ -40,10 +40,10 @@ import (
 	backend "github.com/dadav/gorge/internal/v3/backend"
 	"github.com/dadav/gorge/internal/v3/ui"
 	openapi "github.com/dadav/gorge/pkg/gen/v3/openapi"
+	"github.com/dadav/stampede"
 	"github.com/go-chi/chi/v5"
 	"github.com/go-chi/chi/v5/middleware"
 	"github.com/go-chi/cors"
-	"github.com/go-chi/stampede"
 	"github.com/spf13/cobra"
 	"golang.org/x/sync/errgroup"
 )
@@ -125,6 +125,16 @@ You can also enable the caching functionality to speed things up.`,
 			userService := v3.NewUserOperationsApi()
 
 			r := chi.NewRouter()
+
+			// 0. Inject statistic middleware
+			x := customMiddleware.NewStatistics()
+			r.Use(func(next http.Handler) http.Handler {
+				return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+					ctx := context.WithValue(r.Context(), "stats", x)
+					next.ServeHTTP(w, r.WithContext(ctx))
+				})
+			})
+
 			// 1. Recoverer should be first to catch panics in all other middleware
 			r.Use(middleware.Recoverer)
 			// 2. RealIP should be early to ensure all other middleware sees the correct IP
@@ -140,8 +150,6 @@ You can also enable the caching functionality to speed things up.`,
 			// 4. RequireUserAgent should be early to ensure all other middleware sees the correct user agent
 			r.Use(customMiddleware.RequireUserAgent)
 
-			x := customMiddleware.NewStatistics()
-
 			if config.UI {
 				r.Group(func(r chi.Router) {
 					r.HandleFunc("/", ui.IndexHandler)
@@ -168,10 +176,28 @@ You can also enable the caching functionality to speed things up.`,
 						return stampede.StringToHash(r.Method, requestURI, strings.ToLower(token))
 					}
 
-					cachedMiddleware := stampede.HandlerWithKey(
+					cbFunc := func(fromCache bool, w http.ResponseWriter, r *http.Request) error {
+						x.Mutex.Lock()
+						if fromCache {
+							log.Log.Debugf("Cache hit for path: %s", r.URL.Path)
+							x.TotalCacheHits++
+							x.CacheHitsPerEndpoint[r.URL.Path]++
+							w.Header().Set("X-Cache", "Hit from gorge")
+						} else {
+							log.Log.Debugf("Cache miss for path: %s", r.URL.Path)
+							x.TotalCacheMisses++
+							x.CacheMissesPerEndpoint[r.URL.Path]++
+							w.Header().Set("X-Cache", "MISS from gorge")
+						}
+						x.Mutex.Unlock()
+						return nil
+					}
+
+					cachedMiddleware := stampede.HandlerWithKeyAndCb(
 						512,
 						time.Duration(config.CacheMaxAge)*time.Second,
 						customKeyFunc,
+						cbFunc,
 					)
 					log.Log.Debugf("Cache middleware configured with prefixes: %s", config.CachePrefixes)
 
@@ -188,21 +214,7 @@ You can also enable the caching functionality to speed things up.`,
 							}
 
 							if shouldCache {
-								wrapper := customMiddleware.NewResponseWrapper(w)
-								// Set default cache status before serving
-								// w.Header().Set("X-Cache", "MISS from gorge")
-
-								cachedMiddleware(next).ServeHTTP(wrapper, r)
-
-								// Only override if it was served from cache
-								// TODO: this is not working as expected
-								if wrapper.WasWritten() {
-									log.Log.Debugf("Serving cached response for path: %s", r.URL.Path)
-									w.Header().Set("X-Cache", "HIT from gorge")
-								} else {
-									log.Log.Debugf("Cache miss for path: %s", r.URL.Path)
-									w.Header().Set("X-Cache", "MISS from gorge")
-								}
+								cachedMiddleware(next).ServeHTTP(w, r)
 							} else {
 								next.ServeHTTP(w, r)
 							}

go.mod

@@ -4,10 +4,10 @@ go 1.22.0
 
 require (
 	github.com/a-h/templ v0.2.793
+	github.com/dadav/stampede v0.0.0-20241228173147-dd16def44490
 	github.com/go-chi/chi/v5 v5.2.0
 	github.com/go-chi/cors v1.2.1
 	github.com/go-chi/jwtauth/v5 v5.3.2
-	github.com/go-chi/stampede v0.6.0
 	github.com/hashicorp/go-version v1.7.0
 	github.com/mitchellh/go-homedir v1.1.0
 	github.com/spf13/cobra v1.8.1

go.sum

@@ -3,6 +3,12 @@ github.com/a-h/templ v0.2.793/go.mod h1:lq48JXoUvuQrU0VThrK31yFwdRjTCnIE5bcPCM9I
 github.com/cespare/xxhash/v2 v2.3.0 h1:UL815xU9SqsFlibzuggzjXhog7bL6oX9BbNZnL2UFvs=
 github.com/cespare/xxhash/v2 v2.3.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs=
 github.com/cpuguy83/go-md2man/v2 v2.0.4/go.mod h1:tgQtvFlXSQOSOSIRvRPT7W67SCa46tRHOmNcaadrF8o=
+github.com/dadav/stampede v0.0.0-20241228165756-61568b6ca93b h1:t22u27sjssQ2146msFpZXX47oRgtiHDdudc4zCwwyPU=
+github.com/dadav/stampede v0.0.0-20241228165756-61568b6ca93b/go.mod h1:O/5HgfMaQjv+J8LZCVmaKdE1cD7JsUhwnCuK5BnU0XI=
+github.com/dadav/stampede v0.0.0-20241228171116-ae01e8a04d08 h1:TUK02XGb2f5GgZ5PTNFWgRr+gGSEiwhijvR/pVwSz98=
+github.com/dadav/stampede v0.0.0-20241228171116-ae01e8a04d08/go.mod h1:O/5HgfMaQjv+J8LZCVmaKdE1cD7JsUhwnCuK5BnU0XI=
+github.com/dadav/stampede v0.0.0-20241228173147-dd16def44490 h1:uSaqpXtJE+DofLTUebe7+3bxm5YVrVL8rqIDV+fpRh4=
+github.com/dadav/stampede v0.0.0-20241228173147-dd16def44490/go.mod h1:O/5HgfMaQjv+J8LZCVmaKdE1cD7JsUhwnCuK5BnU0XI=
 github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc h1:U9qPSI2PIWSS1VwoXQT9A3Wy9MM3WgvqSxFWenqJduM=
@@ -19,8 +25,6 @@ github.com/go-chi/cors v1.2.1 h1:xEC8UT3Rlp2QuWNEr4Fs/c2EAGVKBwy/1vHx3bppil4=
 github.com/go-chi/cors v1.2.1/go.mod h1:sSbTewc+6wYHBBCW7ytsFSn836hqM7JxpglAy2Vzc58=
 github.com/go-chi/jwtauth/v5 v5.3.2 h1:s+ON3ATyyMs3Me0kqyuua6Rwu+2zqIIkL0GCaMarwvs=
 github.com/go-chi/jwtauth/v5 v5.3.2/go.mod h1:O4QvPRuZLZghl9WvfVaON+ARfGzpD2PBX/QY5vUz7aQ=
-github.com/go-chi/stampede v0.6.0 h1:9YXCHtnePdj02neMOHysC93WAi3ZXZA8SygCmooNE6o=
-github.com/go-chi/stampede v0.6.0/go.mod h1:9sHbrc5N9uMJHMjw33pBvV8sGtyK3GQdn0lH5bOgjLA=
 github.com/goccy/go-json v0.10.3 h1:KZ5WoDbxAIgm2HNbYckL0se1fHD6rz5j4ywS6ebzDqA=
 github.com/goccy/go-json v0.10.3/go.mod h1:oq7eo15ShAhp70Anwd5lgX2pLfOS3QCiwU/PULtXL6M=
 github.com/google/go-cmp v0.6.0 h1:ofyhxvXcZhMsU5ulbFiLKl/XBFqE1GSq7atu8tAmTRI=

internal/middleware/proxy.go

@@ -49,7 +49,6 @@ func NewSingleHostReverseProxy(target *url.URL) *httputil.ReverseProxy {
 func ProxyFallback(upstreamHost string, forwardToProxy func(int) bool, proxiedResponseCb func(*http.Response)) func(next http.Handler) http.Handler {
 	return func(next http.Handler) http.Handler {
 		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-
 			// Store original headers before any modifications
 			originalHeaders := make(http.Header)
 			for k, v := range w.Header() {
@@ -93,6 +92,12 @@ func ProxyFallback(upstreamHost string, forwardToProxy func(int) bool, proxiedRe
 					capturedResponseWriter.sendCapturedResponse()
 				}
 
+				stats := r.Context().Value("stats").(*Statistics)
+				stats.Mutex.Lock()
+				stats.ProxiedConnections++
+				stats.ProxiedConnectionsPerEndpoint[r.URL.Path]++
+				stats.Mutex.Unlock()
+
 				proxy.ServeHTTP(w, r)
 				return
 			}

internal/middleware/stats.go

@@ -52,23 +52,8 @@ func StatisticsMiddleware(stats *Statistics) func(next http.Handler) http.Handle
 				duration := time.Since(start)
 				stats.Mutex.Lock()
 				stats.ActiveConnections--
-
-				if w.Header().Get("X-Cache") == "HIT from gorge" {
-					stats.TotalCacheHits++
-					stats.CacheHitsPerEndpoint[r.URL.Path]++
-				} else {
-					stats.TotalCacheMisses++
-					stats.CacheMissesPerEndpoint[r.URL.Path]++
-				}
-
 				stats.TotalResponseTime += duration
 				stats.ResponseTimePerEndpoint[r.URL.Path] += duration
-
-				if w.Header().Get("X-Proxied-To") != "" {
-					stats.ProxiedConnections++
-					stats.ProxiedConnectionsPerEndpoint[r.URL.Path]++
-				}
-
 				stats.Mutex.Unlock()
 			}()