fix: don't ignore stats middleware template path calculation (#2594)

gateway/gateway.go

@@ -1426,7 +1426,7 @@ func (gateway *HandleT) StartWebHandler(ctx context.Context) error {
 
 	srvMux := mux.NewRouter()
 	srvMux.Use(
-		middleware.StatMiddleware(ctx, srvMux),
+		middleware.StatMiddleware(ctx, srvMux, stats.Default),
 		middleware.LimitConcurrentRequests(maxConcurrentRequests),
 	)
 	srvMux.HandleFunc("/v1/batch", gateway.webBatchHandler).Methods("POST")
@@ -1502,7 +1502,7 @@ func (gateway *HandleT) StartAdminHandler(ctx context.Context) error {
 
 	srvMux := mux.NewRouter()
 	srvMux.Use(
-		middleware.StatMiddleware(ctx, srvMux),
+		middleware.StatMiddleware(ctx, srvMux, stats.Default),
 		middleware.LimitConcurrentRequests(maxConcurrentRequests),
 	)
 	srvMux.HandleFunc("/v1/pending-events", gateway.pendingEventsHandler).Methods("POST")

middleware/middleware.go

@@ -3,6 +3,7 @@ package middleware
 import (
 	"context"
 	"net/http"
+	"strconv"
 	"sync/atomic"
 	"time"
 
@@ -32,9 +33,9 @@ func LimitConcurrentRequests(maxRequests int) func(http.Handler) http.Handler {
 	}
 }
 
-func StatMiddleware(ctx context.Context, router *mux.Router) func(http.Handler) http.Handler {
+func StatMiddleware(ctx context.Context, router *mux.Router, s stats.Stats) func(http.Handler) http.Handler {
 	var concurrentRequests int32
-	activeClientCount := stats.Default.NewStat("gateway.concurrent_requests_count", stats.GaugeType)
+	activeClientCount := s.NewStat("gateway.concurrent_requests_count", stats.GaugeType)
 	go func() {
 		for {
 			select {
@@ -52,23 +53,50 @@ func StatMiddleware(ctx context.Context, router *mux.Router) func(http.Handler)
 	getPath := func(r *http.Request) string {
 		var match mux.RouteMatch
 		if router.Match(r, &match) {
-			if path, err := match.Route.GetPathTemplate(); err != nil {
+			if path, err := match.Route.GetPathTemplate(); err == nil {
 				return path
 			}
 		}
 		return r.URL.Path
 	}
 	return func(next http.Handler) http.Handler {
 		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			sw := newStatusCapturingWriter(w)
 			path := getPath(r)
-			latencyStat := stats.Default.NewSampledTaggedStat("gateway.response_time", stats.TimerType, map[string]string{"reqType": path, "method": r.Method})
-			latencyStat.Start()
-			defer latencyStat.End()
-
+			start := time.Now()
 			atomic.AddInt32(&concurrentRequests, 1)
 			defer atomic.AddInt32(&concurrentRequests, -1)
 
-			next.ServeHTTP(w, r)
+			next.ServeHTTP(sw, r)
+
+			s.NewSampledTaggedStat(
+				"gateway.response_time",
+				stats.TimerType,
+				map[string]string{
+					"reqType": path,
+					"method":  r.Method,
+					"code":    strconv.Itoa(sw.status),
+				}).Since(start)
 		})
 	}
 }
+
+// newStatusCapturingWriter returns a new, properly initialized statusCapturingWriter
+func newStatusCapturingWriter(w http.ResponseWriter) *statusCapturingWriter {
+	return &statusCapturingWriter{
+		ResponseWriter: w,
+		status:         http.StatusOK,
+	}
+}
+
+// statusCapturingWriter is a response writer decorator that captures the status code.
+type statusCapturingWriter struct {
+	http.ResponseWriter
+	status int
+}
+
+// WriteHeader override the http.ResponseWriter's `WriteHeader` method
+func (w *statusCapturingWriter) WriteHeader(status int) {
+	w.status = status
+	w.ResponseWriter.WriteHeader(status)
+}

middleware/middleware_test.go

@@ -1,14 +1,20 @@
 package middleware_test
 
 import (
+	"context"
 	"net/http"
 	"net/http/httptest"
+	"strconv"
 	"sync"
 	"sync/atomic"
 	"testing"
 	"time"
 
+	"github.com/golang/mock/gomock"
+	"github.com/gorilla/mux"
 	"github.com/rudderlabs/rudder-server/middleware"
+	mock_stats "github.com/rudderlabs/rudder-server/mocks/services/stats"
+	"github.com/rudderlabs/rudder-server/services/stats"
 	"github.com/stretchr/testify/require"
 )
 
@@ -73,3 +79,42 @@ func TestMaxConcurrentRequestsMiddleware(t *testing.T) {
 		})
 	}
 }
+
+func TestStatsMiddleware(t *testing.T) {
+	testCase := func(expectedStatusCode int, pathTemplate, requestPath, expectedReqType, expectedMethod string) func(t *testing.T) {
+		return func(t *testing.T) {
+			ctrl := gomock.NewController(t)
+			mockStats := mock_stats.NewMockStats(ctrl)
+			handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+				w.WriteHeader(expectedStatusCode)
+			})
+
+			measurement := mock_stats.NewMockMeasurement(ctrl)
+			mockStats.EXPECT().NewStat("gateway.concurrent_requests_count", stats.GaugeType).Return(measurement).Times(1)
+			mockStats.EXPECT().NewSampledTaggedStat("gateway.response_time", stats.TimerType,
+				map[string]string{
+					"reqType": expectedReqType,
+					"method":  expectedMethod,
+					"code":    strconv.Itoa(expectedStatusCode),
+				}).Return(measurement).Times(1)
+			measurement.EXPECT().Since(gomock.Any()).Times(1)
+
+			ctx, cancel := context.WithCancel(context.Background())
+			defer cancel()
+			router := mux.NewRouter()
+			router.Use(
+				middleware.StatMiddleware(ctx, router, mockStats),
+			)
+			router.HandleFunc(pathTemplate, handler).Methods(expectedMethod)
+
+			response := httptest.NewRecorder()
+			request := httptest.NewRequest("GET", "http://example.com"+requestPath, http.NoBody)
+			router.ServeHTTP(response, request)
+			require.Equal(t, expectedStatusCode, response.Code)
+		}
+	}
+
+	t.Run("template with param in path", testCase(http.StatusNotFound, "/v1/{param}", "/v1/abc", "/v1/{param}", "GET"))
+
+	t.Run("template without param in path", testCase(http.StatusNotFound, "/v1/some-other/key", "/v1/some-other/key", "/v1/some-other/key", "GET"))
+}