zrouter: Add logging middleware and request-id (#66)

* Add logging middleware

* add request-id middleware

* fix applyMiddlewares and add request-id

* fix go mod

Author: lucaslopezf

go.mod

@@ -10,6 +10,7 @@ require (
 	github.com/go-chi/cors v1.2.1
 	github.com/go-redis/redis/v8 v8.11.5
 	github.com/go-redsync/redsync/v4 v4.12.0
+	github.com/google/uuid v1.6.0
 	github.com/prometheus/client_golang v1.18.0
 	github.com/spf13/cobra v1.8.0
 	github.com/spf13/viper v1.18.2
@@ -33,7 +34,6 @@ require (
 	github.com/fsnotify/fsnotify v1.7.0 // indirect
 	github.com/go-faster/city v1.0.1 // indirect
 	github.com/go-faster/errors v0.6.1 // indirect
-	github.com/google/uuid v1.6.0 // indirect
 	github.com/hashicorp/errwrap v1.1.0 // indirect
 	github.com/hashicorp/go-multierror v1.1.1 // indirect
 	github.com/hashicorp/go-version v1.6.0 // indirect

pkg/zrouter/zmiddlewares/cache.go

@@ -27,9 +27,9 @@ func CacheMiddleware(cache zcache.ZCache, config domain.CacheConfig, logger *zap
 					return
 				}
 
-				mrw := &metricsResponseWriter{ResponseWriter: w}
-				next.ServeHTTP(mrw, r) // Important: This line needs to be BEFORE setting the cache.
-				cacheResponseIfNeeded(mrw, r, cache, key, ttl, logger)
+				rw := &responseWriter{ResponseWriter: w}
+				next.ServeHTTP(rw, r) // Important: This line needs to be BEFORE setting the cache.
+				cacheResponseIfNeeded(rw, r, cache, key, ttl, logger)
 			}
 		})
 	}
@@ -53,17 +53,21 @@ func tryServeFromCache(w http.ResponseWriter, r *http.Request, cache zcache.ZCac
 	if err == nil && cachedResponse != nil {
 		w.Header().Set(domain.ContentTypeHeader, domain.ContentTypeApplicationJSON)
 		_, _ = w.Write(cachedResponse)
+		requestID := r.Header.Get(RequestIDHeader)
+
+		zap.S().Debugf("[Cache] request_id: %s - Method: %s - URL: %s | Status: %d - Response Body: %s",
+			requestID, r.Method, r.URL.String(), http.StatusOK, string(cachedResponse))
 		return true
 	}
 	return false
 }
 
-func cacheResponseIfNeeded(mrw *metricsResponseWriter, r *http.Request, cache zcache.ZCache, key string, ttl time.Duration, logger *zap.SugaredLogger) {
-	if mrw.status != http.StatusOK {
+func cacheResponseIfNeeded(rw *responseWriter, r *http.Request, cache zcache.ZCache, key string, ttl time.Duration, logger *zap.SugaredLogger) {
+	if rw.status != http.StatusOK {
 		return
 	}
 
-	responseBody := mrw.Body()
+	responseBody := rw.Body()
 	if err := cache.Set(r.Context(), key, responseBody, ttl); err != nil {
 		logger.Errorf("Internal error when setting cache response: %v\n%s", err, debug.Stack())
 	}

pkg/zrouter/zmiddlewares/http.go

@@ -1,13 +1,32 @@
 package zmiddlewares
 
 import (
-	"github.com/go-chi/chi/v5/middleware"
+	"github.com/google/uuid"
+	"net/http"
+)
+
+const (
+	RequestIDHeader = "X-Request-ID"
 )
 
 func RequestID() Middleware {
-	return middleware.RequestID
+	return requestIDMiddleware
+}
+
+func requestIDMiddleware(next http.Handler) http.Handler {
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		requestID := r.Header.Get(RequestIDHeader)
+		if requestID == "" {
+			requestID = uuid.New().String()
+			r.Header.Set(RequestIDHeader, requestID)
+		}
+
+		w.Header().Set(RequestIDHeader, requestID)
+		rw := &responseWriter{ResponseWriter: w}
+		next.ServeHTTP(rw, r)
+	})
 }
 
 func Logger() Middleware {
-	return middleware.Logger
+	return LoggingMiddleware
 }

pkg/zrouter/zmiddlewares/logging.go

@@ -0,0 +1,27 @@
+package zmiddlewares
+
+import (
+	"bytes"
+	"go.uber.org/zap"
+	"net/http"
+	"time"
+)
+
+func LoggingMiddleware(next http.Handler) http.Handler {
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		buffer := &bytes.Buffer{}
+
+		rw := &responseWriter{
+			ResponseWriter: w,
+			body:           buffer,
+		}
+
+		start := time.Now()
+		next.ServeHTTP(rw, r)
+		duration := time.Since(start)
+		requestID := r.Header.Get(RequestIDHeader)
+
+		zap.S().Debugf("request_id: %s - Method: %s - URL: %s | Status: %d - Duration: %s - Response Body: %s",
+			requestID, r.Method, r.URL.String(), rw.status, duration, rw.Body())
+	})
+}

pkg/zrouter/zmiddlewares/metrics.go

@@ -60,7 +60,7 @@ func RequestMetrics(appName string, metricsServer metrics.TaskMetrics) Middlewar
 			_ = metricsServer.UpdateMetric(activeConnectionsMetricName, float64(activeConnections))
 			mu.Unlock()
 
-			mrw := &metricsResponseWriter{ResponseWriter: w}
+			mrw := &responseWriter{ResponseWriter: w}
 			next.ServeHTTP(mrw, r)
 
 			mu.Lock()

pkg/zrouter/zmiddlewares/middleware.go

@@ -7,31 +7,31 @@ import (
 
 type Middleware func(next http.Handler) http.Handler
 
-type metricsResponseWriter struct {
+type responseWriter struct {
 	http.ResponseWriter
 	status  int
 	written int64
 	body    *bytes.Buffer
 }
 
-func (mrw *metricsResponseWriter) WriteHeader(statusCode int) {
-	mrw.status = statusCode
-	mrw.ResponseWriter.WriteHeader(statusCode)
+func (rw *responseWriter) WriteHeader(statusCode int) {
+	rw.status = statusCode
+	rw.ResponseWriter.WriteHeader(statusCode)
 }
 
-func (mrw *metricsResponseWriter) Write(p []byte) (int, error) {
-	if mrw.body == nil {
-		mrw.body = new(bytes.Buffer)
+func (rw *responseWriter) Write(p []byte) (int, error) {
+	if rw.body == nil {
+		rw.body = new(bytes.Buffer)
 	}
-	mrw.body.Write(p)
-	n, err := mrw.ResponseWriter.Write(p)
-	mrw.written += int64(n)
+	rw.body.Write(p)
+	n, err := rw.ResponseWriter.Write(p)
+	rw.written += int64(n)
 	return n, err
 }
 
-func (mrw *metricsResponseWriter) Body() []byte {
-	if mrw.body != nil {
-		return mrw.body.Bytes()
+func (rw *responseWriter) Body() []byte {
+	if rw.body != nil {
+		return rw.body.Bytes()
 	}
 	return nil
 }

pkg/zrouter/zrouter.go

@@ -99,6 +99,7 @@ func (r *zrouter) SetDefaultMiddlewares() {
 	}
 
 	r.Use(zmiddlewares.RequestMetrics(r.appName, r.metricsServer))
+	r.Use(zmiddlewares.Logger())
 }
 
 func (r *zrouter) Group(prefix string) Routes {
@@ -132,14 +133,20 @@ func (r *zrouter) Run(addr ...string) error {
 
 func (r *zrouter) applyMiddlewares(handler http.HandlerFunc, middlewares ...zmiddlewares.Middleware) http.Handler {
 	var wrappedHandler http.Handler = handler
+	// The order of middleware application is crucial: Route-specific middlewares are applied first,
+	// followed by router-level general middlewares. This ensures that general middlewares, which often
+	// handle logging, security, etc... are executed first. This sequence is
+	// important to maintain consistency in logging and to apply security measures before route-specific
+	// logic is executed.
 
-	for _, mw := range r.middlewares {
+	for _, mw := range middlewares {
 		wrappedHandler = mw(wrappedHandler)
 	}
 
-	for _, mw := range middlewares {
+	for _, mw := range r.middlewares {
 		wrappedHandler = mw(wrappedHandler)
 	}
+
 	return wrappedHandler
 }