feat: added retryMax middleware

middleware/retry_max.go

@@ -0,0 +1,324 @@
+// SPDX-License-Identifier: MIT
+// SPDX-FileCopyrightText: © 2015 LabStack LLC and Echo contributors
+
+package middleware
+
+import (
+	"math"
+	"net/http"
+	"sync"
+	"time"
+
+	"github.com/labstack/echo/v4"
+)
+
+// RetryMaxStore is the interface to be implemented by custom retry stores.
+type RetryMaxStore interface {
+	// AllowAttempt returns true if another attempt is allowed for the identifier.
+	AllowAttempt(identifier string) (bool, error)
+	// ResetAttempts resets the attempt counter for the identifier.
+	ResetAttempts(identifier string) error
+}
+
+// RetryableChecker determines if an error should trigger a retry
+type RetryableChecker func(err error) bool
+
+// BackoffStrategy defines how to calculate retry delays
+type BackoffStrategy func(attempt int, minTimeout, maxTimeout time.Duration) time.Duration
+
+// RetryMaxConfig defines the configuration for the RetryMax middleware.
+type RetryMaxConfig struct {
+	Skipper    Skipper
+	BeforeFunc BeforeFunc
+	// IdentifierExtractor extracts an identifier (e.g., IP, user ID) for tracking retries.
+	IdentifierExtractor Extractor
+	// Store defines a store for retry counts (optional - if nil, retries per request).
+	Store RetryMaxStore
+	// MaxAttempts is the total number of attempts before giving up.
+	MaxAttempts int
+	// MinTimeout is the minimum wait between attempts.
+	MinTimeout time.Duration
+	// MaxTimeout is the maximum wait between attempts.
+	MaxTimeout time.Duration
+	// BackoffStrategy determines how retry delays are calculated.
+	BackoffStrategy BackoffStrategy
+	// RetryableChecker determines if an error should trigger a retry.
+	RetryableChecker RetryableChecker
+	// ErrorHandler is called if IdentifierExtractor fails.
+	ErrorHandler func(c echo.Context, err error) error
+	// DenyHandler is called when the max attempts are exhausted.
+	DenyHandler func(c echo.Context, identifier string, err error) error
+	// OnRetry is called before each retry attempt.
+	OnRetry func(c echo.Context, identifier string, attempt int, err error)
+}
+
+// ExponentialBackoff implements exponential backoff with jitter
+func ExponentialBackoff(attempt int, minTimeout, maxTimeout time.Duration) time.Duration {
+	backoff := min(time.Duration(math.Pow(2, float64(attempt-1)))*minTimeout, maxTimeout)
+	// Add up to 10% jitter to prevent thundering herd
+	jitter := time.Duration(float64(backoff) * 0.1 * (0.5 - float64(time.Now().UnixNano()%1000)/1000))
+	return backoff + jitter
+}
+
+// LinearBackoff implements linear backoff
+func LinearBackoff(attempt int, minTimeout, maxTimeout time.Duration) time.Duration {
+	backoff := time.Duration(attempt) * minTimeout
+	if backoff > maxTimeout {
+		backoff = maxTimeout
+	}
+	return backoff
+}
+
+// DefaultRetryableChecker determines if an error is retryable
+func DefaultRetryableChecker(err error) bool {
+	if err == nil {
+		return false
+	}
+
+	// Check for HTTP errors that shouldn't be retried
+	if httpErr, ok := err.(*echo.HTTPError); ok {
+		code := httpErr.Code
+		// Don't retry client errors (4xx) except specific ones
+		if code >= 400 && code < 500 {
+			return code == http.StatusRequestTimeout || code == http.StatusTooManyRequests
+		}
+	}
+
+	// Retry on server errors and other errors by default
+	return true
+}
+
+// DefaultRetryMaxConfig provides default values for RetryMaxConfig.
+var DefaultRetryMaxConfig = RetryMaxConfig{
+	Skipper: DefaultSkipper,
+	IdentifierExtractor: func(c echo.Context) (string, error) {
+		return c.RealIP(), nil
+	},
+	MaxAttempts:      5,
+	MinTimeout:       1 * time.Second,
+	MaxTimeout:       5 * time.Second,
+	BackoffStrategy:  ExponentialBackoff,
+	RetryableChecker: DefaultRetryableChecker,
+	ErrorHandler: func(c echo.Context, err error) error {
+		return echo.NewHTTPError(http.StatusInternalServerError, "error extracting identifier")
+	},
+	DenyHandler: func(c echo.Context, identifier string, err error) error {
+		return echo.NewHTTPError(http.StatusTooManyRequests, "max retry attempts exceeded")
+	},
+}
+
+// RetryMax returns a RetryMax middleware with the given store.
+func RetryMax(store RetryMaxStore) echo.MiddlewareFunc {
+	config := DefaultRetryMaxConfig
+	config.Store = store
+	return RetryMaxWithConfig(config)
+}
+
+// RetryMaxWithConfig returns a RetryMax middleware with custom config.
+func RetryMaxWithConfig(config RetryMaxConfig) echo.MiddlewareFunc {
+	// Apply defaults for nil fields
+	if config.Skipper == nil {
+		config.Skipper = DefaultRetryMaxConfig.Skipper
+	}
+	if config.IdentifierExtractor == nil {
+		config.IdentifierExtractor = DefaultRetryMaxConfig.IdentifierExtractor
+	}
+	if config.ErrorHandler == nil {
+		config.ErrorHandler = DefaultRetryMaxConfig.ErrorHandler
+	}
+	if config.DenyHandler == nil {
+		config.DenyHandler = DefaultRetryMaxConfig.DenyHandler
+	}
+	if config.BackoffStrategy == nil {
+		config.BackoffStrategy = DefaultRetryMaxConfig.BackoffStrategy
+	}
+	if config.RetryableChecker == nil {
+		config.RetryableChecker = DefaultRetryMaxConfig.RetryableChecker
+	}
+	if config.MaxAttempts <= 0 {
+		config.MaxAttempts = DefaultRetryMaxConfig.MaxAttempts
+	}
+	if config.MinTimeout <= 0 {
+		config.MinTimeout = DefaultRetryMaxConfig.MinTimeout
+	}
+	if config.MaxTimeout <= 0 {
+		config.MaxTimeout = DefaultRetryMaxConfig.MaxTimeout
+	}
+
+	return func(next echo.HandlerFunc) echo.HandlerFunc {
+		return func(c echo.Context) error {
+			if config.Skipper(c) {
+				return next(c)
+			}
+
+			if config.BeforeFunc != nil {
+				config.BeforeFunc(c)
+			}
+
+			// If using a store, check identifier extraction
+			var identifier string
+			var err error
+			if config.Store != nil {
+				identifier, err = config.IdentifierExtractor(c)
+				if err != nil {
+					return config.ErrorHandler(c, err)
+				}
+			}
+
+			var lastErr error
+			ctx := c.Request().Context()
+
+			for attempt := 1; attempt <= config.MaxAttempts; attempt++ {
+				// If using a store, check if we're allowed to attempt
+				if config.Store != nil {
+					allowed, err := config.Store.AllowAttempt(identifier)
+					if err != nil {
+						return config.ErrorHandler(c, err)
+					}
+					if !allowed {
+						return config.DenyHandler(c, identifier, lastErr)
+					}
+				}
+
+				// Execute the handler
+				lastErr = next(c)
+
+				// Success case
+				if lastErr == nil {
+					if config.Store != nil {
+						_ = config.Store.ResetAttempts(identifier)
+					}
+					return nil
+				}
+
+				// Check if error is retryable
+				if !config.RetryableChecker(lastErr) {
+					return lastErr
+				}
+
+				// If this was the last attempt, return the error
+				if attempt >= config.MaxAttempts {
+					return lastErr
+				}
+
+				// Call retry callback if provided
+				if config.OnRetry != nil {
+					config.OnRetry(c, identifier, attempt, lastErr)
+				}
+
+				// Calculate and apply backoff
+				backoff := config.BackoffStrategy(attempt, config.MinTimeout, config.MaxTimeout)
+
+				select {
+				case <-ctx.Done():
+					return ctx.Err()
+				case <-time.After(backoff):
+					// Continue to next attempt
+				}
+			}
+
+			return lastErr
+		}
+	}
+}
+
+// RetryMaxMemoryStore is an improved in-memory store for RetryMax.
+type RetryMaxMemoryStore struct {
+	attempts    map[string]int
+	lastAccess  map[string]time.Time
+	mutex       sync.RWMutex
+	expiresIn   time.Duration
+	maxAttempts int
+}
+
+// NewRetryMaxMemoryStore creates a new RetryMaxMemoryStore.
+func NewRetryMaxMemoryStore(expiresIn time.Duration) *RetryMaxMemoryStore {
+	return NewRetryMaxMemoryStoreWithAttempts(expiresIn, 5)
+}
+
+// NewRetryMaxMemoryStoreWithAttempts creates a new RetryMaxMemoryStore with custom max attempts.
+func NewRetryMaxMemoryStoreWithAttempts(expiresIn time.Duration, maxAttempts int) *RetryMaxMemoryStore {
+	store := &RetryMaxMemoryStore{
+		attempts:    make(map[string]int),
+		lastAccess:  make(map[string]time.Time),
+		expiresIn:   expiresIn,
+		maxAttempts: maxAttempts,
+	}
+
+	// Start cleanup goroutine
+	go store.startCleanup()
+
+	return store
+}
+
+// AllowAttempt implements RetryMaxStore.AllowAttempt.
+func (s *RetryMaxMemoryStore) AllowAttempt(identifier string) (bool, error) {
+	s.mutex.Lock()
+	defer s.mutex.Unlock()
+
+	// Check if expired and reset if so
+	if s.isExpiredUnsafe(identifier) {
+		s.attempts[identifier] = 0
+	}
+
+	// Increment attempt count
+	s.attempts[identifier]++
+	s.lastAccess[identifier] = time.Now()
+
+	// Check if we've exceeded max attempts
+	return s.attempts[identifier] <= s.maxAttempts, nil
+}
+
+// ResetAttempts implements RetryMaxStore.ResetAttempts.
+func (s *RetryMaxMemoryStore) ResetAttempts(identifier string) error {
+	s.mutex.Lock()
+	defer s.mutex.Unlock()
+
+	s.attempts[identifier] = 0
+	s.lastAccess[identifier] = time.Now()
+	return nil
+}
+
+// GetAttempts returns the current attempt count for an identifier.
+func (s *RetryMaxMemoryStore) GetAttempts(identifier string) int {
+	s.mutex.RLock()
+	defer s.mutex.RUnlock()
+
+	if s.isExpiredUnsafe(identifier) {
+		return 0
+	}
+
+	return s.attempts[identifier]
+}
+
+// Cleanup removes expired entries.
+func (s *RetryMaxMemoryStore) Cleanup() error {
+	s.mutex.Lock()
+	defer s.mutex.Unlock()
+
+	now := time.Now()
+	for identifier, lastAccess := range s.lastAccess {
+		if now.Sub(lastAccess) > s.expiresIn {
+			delete(s.attempts, identifier)
+			delete(s.lastAccess, identifier)
+		}
+	}
+
+	return nil
+}
+
+// isExpiredUnsafe checks if an entry has expired (no locking, caller must hold lock).
+func (s *RetryMaxMemoryStore) isExpiredUnsafe(identifier string) bool {
+	lastAccess, exists := s.lastAccess[identifier]
+	return !exists || time.Since(lastAccess) > s.expiresIn
+}
+
+// startCleanup runs periodic cleanup in a goroutine.
+func (s *RetryMaxMemoryStore) startCleanup() {
+	ticker := time.NewTicker(s.expiresIn / 2) // Clean up twice per expiry period
+	defer ticker.Stop()
+
+	for range ticker.C {
+		_ = s.Cleanup()
+	}
+}