feat: implement correlation ID middleware, enhance request logging, and prevent cache stampedes using singleflight

Author: lawalajose

internal/handlers/errors.go

@@ -35,10 +35,12 @@ const (
 
 // ErrorEnvelope represents a standardized error response
 type ErrorEnvelope struct {
-	Code    string                 `json:"code"`
-	Message string                 `json:"message"`
-	TraceID string                 `json:"trace_id"`
-	Details map[string]interface{} `json:"details,omitempty"`
+	Code          string                 `json:"code"`
+	Message       string                 `json:"message"`
+	TraceID       string                 `json:"trace_id"`
+	RequestID     string                 `json:"request_id"`
+	CorrelationID string                 `json:"correlation_id,omitempty"`
+	Details       map[string]interface{} `json:"details,omitempty"`
 }
 
 // RespondWithError sends a standardized error response
@@ -56,17 +58,22 @@ func RespondWithErrorDetails(c *gin.Context, statusCode int, code ErrorCode, mes
 		traceID = generateTraceID()
 	}
 
+	requestID := c.GetString("request_id")
+	correlationID := c.GetString("correlation_id")
+
 	// Redact message and details to prevent PII leakage
 	redactedMessage := security.MaskPII(message)
 	if details != nil {
 		details = security.RedactMap(details)
 	}
 
 	envelope := ErrorEnvelope{
-		Code:    string(code),
-		Message: redactedMessage,
-		TraceID: traceID,
-		Details: details,
+		Code:          string(code),
+		Message:       redactedMessage,
+		TraceID:       traceID,
+		RequestID:     requestID,
+		CorrelationID: correlationID,
+		Details:       details,
 	}
 
 	c.JSON(statusCode, envelope)

internal/logger/logger.go

@@ -1,32 +1,48 @@
-package logger
-
-import (
-	"os"
-
-	"github.com/sirupsen/logrus"
-	"go.opentelemetry.io/contrib/bridges/otellogrus"
-)
-
-var Log = logrus.New()
-
-func Init() {
-	Log.SetFormatter(&logrus.JSONFormatter{})
-	Log.SetOutput(os.Stdout)
-	Log.AddHook(otellogrus.NewHook("stellarbill-backend"))
-
-	level := os.Getenv("LOG_LEVEL")
-	switch level {
-	case "debug":
-		Log.SetLevel(logrus.DebugLevel)
-	case "warn":
-		Log.SetLevel(logrus.WarnLevel)
-	case "error":
-		Log.SetLevel(logrus.ErrorLevel)
-	default:
-		Log.SetLevel(logrus.InfoLevel)
-	}
-}
-
-func SafePrintf(format string, args ...interface{}) {
-	Log.Printf(format, args...)
+package logger
+
+import (
+    "os"
+
+    "github.com/gin-gonic/gin"
+    "github.com/sirupsen/logrus"
+    "go.opentelemetry.io/contrib/bridges/otellogrus"
+)
+
+var Log = logrus.New()
+
+func Init() {
+    Log.SetFormatter(&logrus.JSONFormatter{})
+    Log.SetOutput(os.Stdout)
+    Log.AddHook(otellogrus.NewHook("stellabill-backend"))
+
+    level := os.Getenv("LOG_LEVEL")
+    switch level {
+    case "debug":
+        Log.SetLevel(logrus.DebugLevel)
+    case "warn":
+        Log.SetLevel(logrus.WarnLevel)
+    case "error":
+        Log.SetLevel(logrus.ErrorLevel)
+    default:
+        Log.SetLevel(logrus.InfoLevel)
+    }
+}
+
+// WithContextFields enriches log entries with request, correlation, and trace IDs from Gin context.
+func WithContextFields(c *gin.Context) *logrus.Entry {
+    fields := logrus.Fields{}
+    if reqID := c.GetString("request_id"); reqID != "" {
+        fields["request_id"] = reqID
+    }
+    if corrID := c.GetString("correlation_id"); corrID != "" {
+        fields["correlation_id"] = corrID
+    }
+    if traceID := c.GetString("traceID"); traceID != "" {
+        fields["trace_id"] = traceID
+    }
+    return Log.WithFields(fields)
+}
+
+func SafePrintf(format string, args ...interface{}) {
+    Log.Printf(format, args...)
 }

internal/middleware/correlation.go

@@ -0,0 +1,25 @@
+package middleware
+
+import (
+    "github.com/gin-gonic/gin"
+    "stellarbill-backend/internal/correlation"
+)
+
+// CorrelationIDMiddleware injects a correlation (job) ID into the request context and Gin context.
+func CorrelationIDMiddleware() gin.HandlerFunc {
+    return func(c *gin.Context) {
+        // Retrieve existing job ID from context if present.
+        ctx := c.Request.Context()
+        jobID := correlation.JobIDFromContext(ctx)
+        if jobID == "" {
+            // Generate a new job ID and store it in the context.
+            jobID = correlation.NewID()
+            ctx = correlation.WithJobID(ctx, jobID)
+        }
+        // Update the request with the new context containing the job ID.
+        c.Request = c.Request.WithContext(ctx)
+        // Also store it in Gin context for easy access in handlers/logger.
+        c.Set("correlation_id", jobID)
+        c.Next()
+    }
+}

internal/middleware/logger.go

@@ -1,36 +1,57 @@
 package middleware
 
 import (
-	"time"
+    "time"
 
-	"stellarbill-backend/internal/logger"
-
-	"github.com/gin-gonic/gin"
-	"github.com/google/uuid"
+    "github.com/gin-gonic/gin"
+    "github.com/google/uuid"
+    "stellarbill-backend/internal/logger"
+    "stellarbill-backend/internal/correlation"
 )
 
+// RequestLogger enriches each request with IDs and logs request details.
 func RequestLogger() gin.HandlerFunc {
-	return func(c *gin.Context) {
-
-		start := time.Now()
-
-		requestID := uuid.New().String()
-		c.Set("request_id", requestID)
-
-		c.Writer.Header().Set("X-Request-ID", requestID)
-
-		c.Next()
-
-		latency := time.Since(start)
-
-		logger.Log.WithFields(map[string]interface{}{
-			"level":      "info",
-			"request_id": requestID,
-			"method":     c.Request.Method,
-			"path":       c.Request.URL.Path,
-			"status":     c.Writer.Status(),
-			"latency_ms": latency.Milliseconds(),
-			"client_ip":  c.ClientIP(),
-		}).Info("request completed")
-	}
+    return func(c *gin.Context) {
+        start := time.Now()
+
+        // Ensure request ID exists (fallback)
+        requestID := c.GetString("request_id")
+        if requestID == "" {
+            requestID = uuid.New().String()
+            c.Set("request_id", requestID)
+            c.Writer.Header().Set("X-Request-ID", requestID)
+        }
+
+        // Ensure correlation ID exists (fallback)
+        correlationID := c.GetString("correlation_id")
+        if correlationID == "" {
+            correlationID = correlation.NewID()
+            c.Set("correlation_id", correlationID)
+            c.Writer.Header().Set("X-Correlation-ID", correlationID)
+        }
+
+        // Ensure trace ID exists (fallback)
+        traceID := c.GetString("traceID")
+        if traceID == "" {
+            // TraceIDMiddleware should set this, but fallback to new UUID
+            traceID = uuid.New().String()
+            c.Set("traceID", traceID)
+            c.Writer.Header().Set("X-Trace-ID", traceID)
+        }
+
+        c.Next()
+
+        latency := time.Since(start)
+        logger.Log.WithFields(map[string]interface{}{
+            "level":          "info",
+            "request_id":     requestID,
+            "correlation_id": correlationID,
+            "trace_id":       traceID,
+            "method":         c.Request.Method,
+            "path":           c.Request.URL.Path,
+            "status":         c.Writer.Status(),
+            "latency_ms":     latency.Milliseconds(),
+            "client_ip":      c.ClientIP(),
+        }).Info("request completed")
+    }
 }

internal/repository/cached_plan_repo.go

@@ -3,11 +3,13 @@ package repository
 import (
 	"context"
 	"encoding/json"
+	"fmt"
+	"stellarbill-backend/internal/cache"
 	"sync"
 	"sync/atomic"
 	"time"
 
-	"stellarbill-backend/internal/cache"
+	"golang.org/x/sync/singleflight"
 )
 
 type cacheEnvelope struct {
@@ -25,6 +27,7 @@ type CachedPlanRepo struct {
 	misses        uint64
 	stales        uint64
 	invalidatedAt sync.Map
+	sf            singleflight.Group
 }
 
 // NewCachedPlanRepo constructs a CachedPlanRepo.
@@ -44,10 +47,12 @@ func (cpr *CachedPlanRepo) cacheKey(id string) string {
 // and updates cache on a successful backend read.
 func (cpr *CachedPlanRepo) FindByID(ctx context.Context, id string) (*PlanRow, error) {
 	key := cpr.cacheKey(id)
+	// Attempt cache fetch
 	if cpr.cache != nil {
 		if val, err := cpr.cache.Get(ctx, key); err == nil && val != nil {
 			var env cacheEnvelope
 			if err := json.Unmarshal(val, &env); err == nil {
+				// Check for staleness due to invalidation
 				stale := false
 				if invTimeVal, ok := cpr.invalidatedAt.Load(key); ok {
 					if invTime, ok := invTimeVal.(time.Time); ok && env.StoredAt.Before(invTime) {
@@ -67,29 +72,34 @@ func (cpr *CachedPlanRepo) FindByID(ctx context.Context, id string) (*PlanRow, e
 			}
 		}
 	}
+	// Cache miss, use singleflight to avoid stampede
 	atomic.AddUint64(&cpr.misses, 1)
-	pr, err := cpr.backend.FindByID(ctx, id)
-	if err != nil {
-		return nil, err
-	}
-	if cpr.cache != nil {
-		prBytes, err := json.Marshal(pr)
-		if err == nil {
-			env := cacheEnvelope{
-				Data:     prBytes,
-				StoredAt: time.Now(),
-			}
-			if envBytes, err := json.Marshal(env); err == nil {
-				_ = cpr.cache.Set(ctx, key, envBytes, cpr.ttl)
+	v, err, _ := cpr.sf.Do(key, func() (interface{}, error) {
+		pr, err := cpr.backend.FindByID(ctx, id)
+		if err != nil {
+			return nil, err
+		}
+		if cpr.cache != nil {
+			prBytes, err := json.Marshal(pr)
+			if err == nil {
+				env := cacheEnvelope{Data: prBytes, StoredAt: time.Now()}
+				if envBytes, err := json.Marshal(env); err == nil {
+					_ = cpr.cache.Set(ctx, key, envBytes, cpr.ttl)
+				}
 			}
 		}
+		return pr, nil
+	})
+	if err != nil {
+		return nil, err
 	}
-	return pr, nil
+	return v.(*PlanRow), nil
 }
 
 // List returns all plans. It caches the full list under a single key.
 func (cpr *CachedPlanRepo) List(ctx context.Context) ([]*PlanRow, error) {
 	key := "plan:list:all"
+	// Attempt cache fetch for list
 	if cpr.cache != nil {
 		if val, err := cpr.cache.Get(ctx, key); err == nil && val != nil {
 			var env cacheEnvelope
@@ -110,27 +120,34 @@ func (cpr *CachedPlanRepo) List(ctx context.Context) ([]*PlanRow, error) {
 						return out, nil
 					}
 				}
+			} else {
+				// Corrupted envelope JSON
+				return nil, fmt.Errorf("corrupted cache envelope: %w", err)
 			}
 		}
 	}
+	// Cache miss, use singleflight for list
 	atomic.AddUint64(&cpr.misses, 1)
-	out, err := cpr.backend.List(ctx)
-	if err != nil {
-		return nil, err
-	}
-	if cpr.cache != nil {
-		outBytes, err := json.Marshal(out)
-		if err == nil {
-			env := cacheEnvelope{
-				Data:     outBytes,
-				StoredAt: time.Now(),
-			}
-			if envBytes, err := json.Marshal(env); err == nil {
-				_ = cpr.cache.Set(ctx, key, envBytes, cpr.ttl)
+	v, err, _ := cpr.sf.Do(key, func() (interface{}, error) {
+		out, err := cpr.backend.List(ctx)
+		if err != nil {
+			return nil, err
+		}
+		if cpr.cache != nil {
+			outBytes, err := json.Marshal(out)
+			if err == nil {
+				env := cacheEnvelope{Data: outBytes, StoredAt: time.Now()}
+				if envBytes, err := json.Marshal(env); err == nil {
+					_ = cpr.cache.Set(ctx, key, envBytes, cpr.ttl)
+				}
 			}
 		}
+		return out, nil
+	})
+	if err != nil {
+		return nil, err
 	}
-	return out, nil
+	return v.([]*PlanRow), nil
 }
 
 // Delete invalidates a cached plan entry and records the invalidation time.