chore: modified statsMiddleware to be generic (#2626)

* chore: modified statsMiddleware to be generic * minor fix * Update gateway/gateway.go Co-authored-by: Francesco Casula <fracasula@users.noreply.github.com> * Update gateway/gateway.go Co-authored-by: Francesco Casula <fracasula@users.noreply.github.com> * lint fix * test fix Co-authored-by: chandumlg <54652834+chandumlg@users.noreply.github.com> Co-authored-by: Francesco Casula <fracasula@users.noreply.github.com>
rudderlabs · Oct 28, 2022 · b6ce9b9 · b6ce9b9
1 parent 46b5b3f
commit b6ce9b9
Show file tree

Hide file tree

Showing 3 changed files with 13 additions and 10 deletions.
diff --git a/gateway/gateway.go b/gateway/gateway.go
@@ -1423,10 +1423,10 @@ func (gateway *HandleT) StartWebHandler(ctx context.Context) error {
 	gateway.logger.Infof("WebHandler waiting for BackendConfig before starting on %d", webPort)
 	gateway.backendConfig.WaitForConfig(ctx)
 	gateway.logger.Infof("WebHandler Starting on %d", webPort)
-
+	component := "gateway"
 	srvMux := mux.NewRouter()
 	srvMux.Use(
-		middleware.StatMiddleware(ctx, srvMux, stats.Default),
+		middleware.StatMiddleware(ctx, srvMux, stats.Default, component),
 		middleware.LimitConcurrentRequests(maxConcurrentRequests),
 	)
 	srvMux.HandleFunc("/v1/batch", gateway.webBatchHandler).Methods("POST")
@@ -1499,10 +1499,10 @@ func (gateway *HandleT) StartAdminHandler(ctx context.Context) error {
 	gateway.logger.Infof("AdminHandler waiting for BackendConfig before starting on %d", adminWebPort)
 	gateway.backendConfig.WaitForConfig(ctx)
 	gateway.logger.Infof("AdminHandler starting on %d", adminWebPort)
-
+	component := "gateway"
 	srvMux := mux.NewRouter()
 	srvMux.Use(
-		middleware.StatMiddleware(ctx, srvMux, stats.Default),
+		middleware.StatMiddleware(ctx, srvMux, stats.Default, component),
 		middleware.LimitConcurrentRequests(maxConcurrentRequests),
 	)
 	srvMux.HandleFunc("/v1/pending-events", gateway.pendingEventsHandler).Methods("POST")

diff --git a/middleware/middleware.go b/middleware/middleware.go
@@ -2,6 +2,7 @@ package middleware
 
 import (
 	"context"
+	"fmt"
 	"net/http"
 	"strconv"
 	"sync/atomic"
@@ -33,9 +34,9 @@ func LimitConcurrentRequests(maxRequests int) func(http.Handler) http.Handler {
 	}
 }
 
-func StatMiddleware(ctx context.Context, router *mux.Router, s stats.Stats) func(http.Handler) http.Handler {
+func StatMiddleware(ctx context.Context, router *mux.Router, s stats.Stats, component string) func(http.Handler) http.Handler {
 	var concurrentRequests int32
-	activeClientCount := s.NewStat("gateway.concurrent_requests_count", stats.GaugeType)
+	activeClientCount := s.NewStat(fmt.Sprintf("%s.concurrent_requests_count", component), stats.GaugeType)
 	go func() {
 		for {
 			select {
@@ -70,7 +71,7 @@ func StatMiddleware(ctx context.Context, router *mux.Router, s stats.Stats) func
 			next.ServeHTTP(sw, r)
 
 			s.NewSampledTaggedStat(
-				"gateway.response_time",
+				fmt.Sprintf("%s.response_time", component),
 				stats.TimerType,
 				map[string]string{
 					"reqType": path,

diff --git a/middleware/middleware_test.go b/middleware/middleware_test.go
@@ -2,6 +2,7 @@ package middleware_test
 
 import (
 	"context"
+	"fmt"
 	"net/http"
 	"net/http/httptest"
 	"strconv"
@@ -81,6 +82,7 @@ func TestMaxConcurrentRequestsMiddleware(t *testing.T) {
 }
 
 func TestStatsMiddleware(t *testing.T) {
+	component := "test"
 	testCase := func(expectedStatusCode int, pathTemplate, requestPath, expectedReqType, expectedMethod string) func(t *testing.T) {
 		return func(t *testing.T) {
 			ctrl := gomock.NewController(t)
@@ -90,8 +92,8 @@ func TestStatsMiddleware(t *testing.T) {
 			})
 
 			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,
+			mockStats.EXPECT().NewStat(fmt.Sprintf("%s.concurrent_requests_count", component), stats.GaugeType).Return(measurement).Times(1)
+			mockStats.EXPECT().NewSampledTaggedStat(fmt.Sprintf("%s.response_time", component), stats.TimerType,
 				map[string]string{
 					"reqType": expectedReqType,
 					"method":  expectedMethod,
@@ -103,7 +105,7 @@ func TestStatsMiddleware(t *testing.T) {
 			defer cancel()
 			router := mux.NewRouter()
 			router.Use(
-				middleware.StatMiddleware(ctx, router, mockStats),
+				middleware.StatMiddleware(ctx, router, mockStats, component),
 			)
 			router.HandleFunc(pathTemplate, handler).Methods(expectedMethod)