chore: add StatMiddleware for chi router (#14)

chiware/stats.go

@@ -0,0 +1,77 @@
+package chiware
+
+import (
+	"context"
+	"fmt"
+	"net/http"
+	"strconv"
+	"sync/atomic"
+	"time"
+
+	"github.com/go-chi/chi/v5"
+	"github.com/rudderlabs/rudder-go-kit/stats"
+)
+
+func StatMiddleware(ctx context.Context, router chi.Router, s stats.Stats, component string) func(http.Handler) http.Handler {
+	var concurrentRequests int32
+	activeClientCount := s.NewStat(fmt.Sprintf("%s.concurrent_requests_count", component), stats.GaugeType)
+	go func() {
+		for {
+			select {
+			case <-ctx.Done():
+				return
+			case <-time.After(10 * time.Second):
+				activeClientCount.Gauge(atomic.LoadInt32(&concurrentRequests))
+			}
+		}
+	}()
+
+	// getPath retrieves the path from the request.
+	// The matched route's template is used if a match is found,
+	// otherwise the request's URL path is used instead.
+	getPath := func(r *http.Request) string {
+		if path := chi.RouteContext(r.Context()).RoutePattern(); path != "" {
+			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)
+			start := time.Now()
+			atomic.AddInt32(&concurrentRequests, 1)
+			defer atomic.AddInt32(&concurrentRequests, -1)
+
+			next.ServeHTTP(sw, r)
+
+			s.NewSampledTaggedStat(
+				fmt.Sprintf("%s.response_time", component),
+				stats.TimerType,
+				map[string]string{
+					"reqType": getPath(r),
+					"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)
+}

chiware/stats_test.go

@@ -0,0 +1,57 @@
+package chiware_test
+
+import (
+	"context"
+	"fmt"
+	"net/http"
+	"net/http/httptest"
+	"strconv"
+	"testing"
+
+	"github.com/go-chi/chi/v5"
+	"github.com/golang/mock/gomock"
+	"github.com/rudderlabs/rudder-go-kit/chiware"
+	"github.com/rudderlabs/rudder-go-kit/stats"
+	"github.com/rudderlabs/rudder-go-kit/stats/mock_stats"
+	"github.com/stretchr/testify/require"
+)
+
+func TestStatsMiddleware(t *testing.T) {
+	component := "test"
+	testCase := func(expectedStatusCode int, pathTemplate, requestPath, 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(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": pathTemplate,
+					"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 := chi.NewRouter()
+			router.Use(
+				chiware.StatMiddleware(ctx, router, mockStats, component),
+			)
+			router.MethodFunc(expectedMethod, pathTemplate, handler)
+
+			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", "GET"))
+
+	t.Run("template without param in path", testCase(http.StatusNotFound, "/v1/some-other/key", "/v1/some-other/key", "GET"))
+}

go.mod

@@ -5,6 +5,7 @@ go 1.20
 require (
 	github.com/cenkalti/backoff/v4 v4.2.0
 	github.com/fsnotify/fsnotify v1.6.0
+	github.com/go-chi/chi/v5 v5.0.8
 	github.com/golang/mock v1.6.0
 	github.com/gorilla/mux v1.8.0
 	github.com/joho/godotenv v1.5.1

go.sum

@@ -99,6 +99,8 @@ github.com/frankban/quicktest v1.14.3 h1:FJKSZTDHjyhriyC81FLQ0LY93eSai0ZyR/ZIkd3
 github.com/fsnotify/fsnotify v1.6.0 h1:n+5WquG0fcWoWp6xPWfHdbskMCQaFnG6PfBrh1Ky4HY=
 github.com/fsnotify/fsnotify v1.6.0/go.mod h1:sl3t1tCWJFWoRz9R8WJCbQihKKwmorjAbSClcnxKAGw=
 github.com/ghodss/yaml v1.0.0/go.mod h1:4dBDuWmgqj2HViK6kFavaiC9ZROes6MMH2rRYeMEF04=
+github.com/go-chi/chi/v5 v5.0.8 h1:lD+NLqFcAi1ovnVZpsnObHGW4xb4J8lNmoYVfECH1Y0=
+github.com/go-chi/chi/v5 v5.0.8/go.mod h1:DslCQbL2OYiznFReuXYUmQ2hGd1aDpCnlMNITLSKoi8=
 github.com/go-gl/glfw v0.0.0-20190409004039-e6da0acd62b1/go.mod h1:vR7hzQXu2zJy9AVAgeJqvqgH9Q5CA+iKCZ2gyEVpxRU=
 github.com/go-gl/glfw/v3.3/glfw v0.0.0-20191125211704-12ad95a8df72/go.mod h1:tQ2UAYgL5IevRw8kRxooKSPJfGvJ9fJQFa0TUsXzTg8=
 github.com/go-gl/glfw/v3.3/glfw v0.0.0-20200222043503-6f7a984d4dc4/go.mod h1:tQ2UAYgL5IevRw8kRxooKSPJfGvJ9fJQFa0TUsXzTg8=