Documentation
¶
Overview ¶
Package ratelimit provides a comprehensive rate limiting system for Hector v2.
Features:
- Multi-layer time windows (minute, hour, day, week, month)
- Dual tracking (token count AND request count)
- Flexible scopes (per-session or per-user)
- Multiple storage backends (in-memory and SQL)
- Atomic check-and-record operations
- Detailed usage statistics
Basic Usage ¶
// Create store (memory or SQL)
store := ratelimit.NewMemoryStore()
// Create limiter with config
limiter, err := ratelimit.NewRateLimiter(config, store)
// Check and record usage
result, err := limiter.CheckAndRecord(ctx, ratelimit.ScopeSession, "session-123", 1000, 1)
if !result.Allowed {
// Handle rate limit exceeded
}
Configuration ¶
rate_limiting:
enabled: true
scope: "session" # or "user"
backend: "memory" # or "sql"
limits:
- type: token
window: day
limit: 100000
- type: count
window: minute
limit: 60
Time Windows ¶
- minute: 60 seconds (burst protection)
- hour: 60 minutes (short-term limits)
- day: 24 hours (daily quotas)
- week: 7 days (weekly budgets)
- month: 30 days (monthly billing)
Limit Types ¶
- token: Track token usage (LLM API tokens, cost control)
- count: Track request count (rate throttling, DDoS protection)
Scopes ¶
- session: Each session has independent quotas
- user: All sessions for a user share quotas
Index ¶
- Variables
- func IsRateLimitError(err error) bool
- func Middleware(cfg MiddlewareConfig) func(http.Handler) http.Handler
- func SimpleMiddleware(limiter RateLimiter, excludedPaths ...string) func(http.Handler) http.Handler
- type CheckResult
- type Config
- type DefaultRateLimiter
- func (rl *DefaultRateLimiter) Check(ctx context.Context, scope Scope, identifier string) (*CheckResult, error)
- func (rl *DefaultRateLimiter) CheckAndRecord(ctx context.Context, scope Scope, identifier string, tokenCount int64, ...) (*CheckResult, error)
- func (rl *DefaultRateLimiter) GetUsage(ctx context.Context, scope Scope, identifier string) ([]Usage, error)
- func (rl *DefaultRateLimiter) IsEnabled() bool
- func (rl *DefaultRateLimiter) Record(ctx context.Context, scope Scope, identifier string, tokenCount int64, ...) error
- func (rl *DefaultRateLimiter) Reset(ctx context.Context, scope Scope, identifier string) error
- func (rl *DefaultRateLimiter) ResetExpired(ctx context.Context, before time.Time) error
- func (rl *DefaultRateLimiter) Store() Store
- type IdentifierFunc
- type LimitRule
- type LimitType
- type MemoryStore
- func (s *MemoryStore) Close() error
- func (s *MemoryStore) DeleteExpired(ctx context.Context, before time.Time) error
- func (s *MemoryStore) DeleteUsage(ctx context.Context, scope Scope, identifier string) error
- func (s *MemoryStore) Dump() map[string]interface{}
- func (s *MemoryStore) GetUsage(ctx context.Context, scope Scope, identifier string, limitType LimitType, ...) (int64, time.Time, error)
- func (s *MemoryStore) IncrementUsage(ctx context.Context, scope Scope, identifier string, limitType LimitType, ...) (int64, time.Time, error)
- func (s *MemoryStore) SetUsage(ctx context.Context, scope Scope, identifier string, limitType LimitType, ...) error
- func (s *MemoryStore) Size() int
- type MiddlewareConfig
- type RateLimitError
- type RateLimiter
- type SQLStore
- func (s *SQLStore) Close() error
- func (s *SQLStore) DeleteExpired(ctx context.Context, before time.Time) error
- func (s *SQLStore) DeleteUsage(ctx context.Context, scope Scope, identifier string) error
- func (s *SQLStore) Dialect() string
- func (s *SQLStore) GetUsage(ctx context.Context, scope Scope, identifier string, limitType LimitType, ...) (int64, time.Time, error)
- func (s *SQLStore) IncrementUsage(ctx context.Context, scope Scope, identifier string, limitType LimitType, ...) (int64, time.Time, error)
- func (s *SQLStore) SetUsage(ctx context.Context, scope Scope, identifier string, limitType LimitType, ...) error
- type Scope
- type Store
- type TimeWindow
- type Usage
- type ValidationError
Constants ¶
This section is empty.
Variables ¶
var ( // ErrRateLimitExceeded is returned when a rate limit is exceeded. ErrRateLimitExceeded = errors.New("rate limit exceeded") // ErrInvalidIdentifier is returned when an identifier is invalid. ErrInvalidIdentifier = errors.New("invalid identifier") ErrStoreUnavailable = errors.New("store unavailable") )
Common errors.
Functions ¶
func IsRateLimitError ¶
IsRateLimitError checks if an error is a rate limit error.
func Middleware ¶
func Middleware(cfg MiddlewareConfig) func(http.Handler) http.Handler
Middleware creates an HTTP middleware that enforces rate limits.
func SimpleMiddleware ¶
SimpleMiddleware creates a simple rate limiting middleware. This is a convenience function for common use cases.
Types ¶
type CheckResult ¶
type CheckResult struct {
// Allowed indicates whether the operation is allowed.
Allowed bool `json:"allowed"`
// Reason provides a human-readable reason if denied.
Reason string `json:"reason,omitempty"`
// Usages contains current usage for all limits.
Usages []Usage `json:"usages"`
// RetryAfter indicates how long to wait before retrying (if denied).
RetryAfter *time.Duration `json:"retry_after,omitempty"`
}
CheckResult represents the result of a rate limit check.
func GetRateLimitResult ¶
func GetRateLimitResult(err error) *CheckResult
GetRateLimitResult extracts the CheckResult from a rate limit error. Returns nil if the error is not a RateLimitError.
func UsageFromContext ¶
func UsageFromContext(ctx context.Context) *CheckResult
UsageFromContext extracts rate limit usage from the request context.
func (*CheckResult) GetHighestUsagePercentage ¶
func (r *CheckResult) GetHighestUsagePercentage() float64
GetHighestUsagePercentage returns the highest usage percentage across all limits.
func (*CheckResult) GetUsage ¶
func (r *CheckResult) GetUsage(limitType LimitType, window TimeWindow) *Usage
GetUsage returns usage for a specific limit type and window.
func (*CheckResult) IsExceeded ¶
func (r *CheckResult) IsExceeded() bool
IsExceeded returns true if any limit is exceeded.
type Config ¶
type Config struct {
// Enabled controls whether rate limiting is active.
Enabled bool
// Limits defines the rate limit rules.
Limits []LimitRule
}
Config holds rate limiting configuration.
type DefaultRateLimiter ¶
type DefaultRateLimiter struct {
// contains filtered or unexported fields
}
DefaultRateLimiter implements the RateLimiter interface.
func NewRateLimiter ¶
func NewRateLimiter(cfg *Config, store Store) (*DefaultRateLimiter, error)
NewRateLimiter creates a new rate limiter with the given configuration and store.
func (*DefaultRateLimiter) Check ¶
func (rl *DefaultRateLimiter) Check(ctx context.Context, scope Scope, identifier string) (*CheckResult, error)
Check verifies if the operation is allowed without recording usage.
func (*DefaultRateLimiter) CheckAndRecord ¶
func (rl *DefaultRateLimiter) CheckAndRecord(ctx context.Context, scope Scope, identifier string, tokenCount int64, requestCount int64) (*CheckResult, error)
CheckAndRecord checks limits and records usage in a single atomic operation.
func (*DefaultRateLimiter) GetUsage ¶
func (rl *DefaultRateLimiter) GetUsage(ctx context.Context, scope Scope, identifier string) ([]Usage, error)
GetUsage returns current usage statistics for an identifier.
func (*DefaultRateLimiter) IsEnabled ¶
func (rl *DefaultRateLimiter) IsEnabled() bool
IsEnabled returns whether rate limiting is enabled.
func (*DefaultRateLimiter) Record ¶
func (rl *DefaultRateLimiter) Record(ctx context.Context, scope Scope, identifier string, tokenCount int64, requestCount int64) error
Record records actual usage (tokens and/or count).
func (*DefaultRateLimiter) ResetExpired ¶
ResetExpired removes expired usage records.
func (*DefaultRateLimiter) Store ¶
func (rl *DefaultRateLimiter) Store() Store
Store returns the underlying store (for testing).
type IdentifierFunc ¶
IdentifierFunc extracts the rate limit identifier from an HTTP request. This function should return the identifier and scope for the request.
type LimitRule ¶
type LimitRule struct {
// Type is the limit type (token or count).
Type LimitType
// Window is the time window for this limit.
Window TimeWindow
// Limit is the maximum allowed in the window.
Limit int64
}
LimitRule defines a single rate limit rule.
type LimitType ¶
type LimitType string
LimitType represents the type of rate limit.
func ParseLimitType ¶
ParseLimitType converts a config string to LimitType.
type MemoryStore ¶
type MemoryStore struct {
// contains filtered or unexported fields
}
MemoryStore is an in-memory implementation of Store. It is thread-safe and suitable for development, testing, and single-instance deployments.
func NewMemoryStore ¶
func NewMemoryStore() *MemoryStore
NewMemoryStore creates a new in-memory store.
func (*MemoryStore) DeleteExpired ¶
DeleteExpired deletes expired usage records.
func (*MemoryStore) DeleteUsage ¶
DeleteUsage deletes usage records for an identifier.
func (*MemoryStore) Dump ¶
func (s *MemoryStore) Dump() map[string]interface{}
Dump returns all records as a map (for debugging).
func (*MemoryStore) GetUsage ¶
func (s *MemoryStore) GetUsage(ctx context.Context, scope Scope, identifier string, limitType LimitType, window TimeWindow) (int64, time.Time, error)
GetUsage gets current usage for a specific limit.
func (*MemoryStore) IncrementUsage ¶
func (s *MemoryStore) IncrementUsage(ctx context.Context, scope Scope, identifier string, limitType LimitType, window TimeWindow, amount int64) (int64, time.Time, error)
IncrementUsage increments usage for a specific limit.
func (*MemoryStore) SetUsage ¶
func (s *MemoryStore) SetUsage(ctx context.Context, scope Scope, identifier string, limitType LimitType, window TimeWindow, amount int64, windowEnd time.Time) error
SetUsage sets usage for a specific limit.
func (*MemoryStore) Size ¶
func (s *MemoryStore) Size() int
Size returns the number of records in the store (for testing).
type MiddlewareConfig ¶
type MiddlewareConfig struct {
// Limiter is the rate limiter to use.
Limiter RateLimiter
// IdentifierFunc extracts the identifier and scope from requests.
// If nil, DefaultIdentifierFunc is used.
IdentifierFunc IdentifierFunc
// TokenEstimator estimates token count for a request.
// If nil, token count is set to 0 (only count-based limiting).
TokenEstimator func(r *http.Request) int64
// ExcludedPaths are paths that bypass rate limiting.
ExcludedPaths []string
// OnLimited is called when a request is rate limited.
// If nil, a default JSON error response is sent.
OnLimited func(w http.ResponseWriter, r *http.Request, result *CheckResult)
}
MiddlewareConfig configures the rate limiting middleware.
type RateLimitError ¶
type RateLimitError struct {
// Message is a human-readable error message.
Message string
// Result contains the detailed rate limit check result.
Result *CheckResult
}
RateLimitError represents a rate limit error with detailed information.
func NewRateLimitError ¶
func NewRateLimitError(result *CheckResult) *RateLimitError
NewRateLimitError creates a new RateLimitError from a CheckResult.
func (*RateLimitError) Error ¶
func (e *RateLimitError) Error() string
Error returns the error message.
func (*RateLimitError) Unwrap ¶
func (e *RateLimitError) Unwrap() error
Unwrap returns the underlying error.
type RateLimiter ¶
type RateLimiter interface {
// Check verifies if the operation is allowed without recording usage.
// Use this when you want to check limits before potentially expensive operations.
Check(ctx context.Context, scope Scope, identifier string) (*CheckResult, error)
// Record records actual usage (tokens and/or count).
// Use this after an operation completes to record the actual usage.
Record(ctx context.Context, scope Scope, identifier string, tokenCount int64, requestCount int64) error
// CheckAndRecord checks limits and records usage in a single atomic operation.
// This is the recommended method for most use cases as it prevents race conditions.
CheckAndRecord(ctx context.Context, scope Scope, identifier string, tokenCount int64, requestCount int64) (*CheckResult, error)
// GetUsage returns current usage statistics for an identifier.
// Returns usage for all configured limits.
GetUsage(ctx context.Context, scope Scope, identifier string) ([]Usage, error)
// Reset resets usage for an identifier.
// Useful for testing or manual quota resets.
Reset(ctx context.Context, scope Scope, identifier string) error
// ResetExpired removes expired usage records.
// Should be called periodically for cleanup.
ResetExpired(ctx context.Context, before time.Time) error
}
RateLimiter is the main interface for rate limiting.
Implementations must be thread-safe and support concurrent access.
func NewRateLimiterFromConfig ¶
NewRateLimiterFromConfig creates a RateLimiter using the primary database. If rate limiting is disabled, returns nil.
func NewRateLimiterFromConfigWithStore ¶
func NewRateLimiterFromConfigWithStore(cfg *config.RateLimitConfig, store Store) (RateLimiter, error)
NewRateLimiterFromConfigWithStore creates a RateLimiter with a custom store. Useful for testing or when you need to share a store across multiple limiters.
type SQLStore ¶
type SQLStore struct {
// contains filtered or unexported fields
}
SQLStore is a SQL-based implementation of Store. It supports Postgres, MySQL, and SQLite.
func NewSQLStore ¶
NewSQLStore creates a new SQL-based store. Supported dialects: "postgres", "mysql", "sqlite".
func (*SQLStore) Close ¶
Close closes the store. Note: This does NOT close the underlying database connection, as that connection may be shared with other components.
func (*SQLStore) DeleteExpired ¶
DeleteExpired deletes expired usage records.
func (*SQLStore) DeleteUsage ¶
DeleteUsage deletes usage records for an identifier.
func (*SQLStore) GetUsage ¶
func (s *SQLStore) GetUsage(ctx context.Context, scope Scope, identifier string, limitType LimitType, window TimeWindow) (int64, time.Time, error)
GetUsage gets current usage for a specific limit.
type Scope ¶
type Scope string
Scope represents the scope of rate limiting.
func DefaultIdentifierFunc ¶
DefaultIdentifierFunc extracts the identifier from the request. It uses the session ID header if present, otherwise falls back to remote address.
func ScopeFromConfig ¶
func ScopeFromConfig(cfg *config.RateLimitConfig) Scope
ScopeFromConfig returns the rate limiting scope from configuration.
type Store ¶
type Store interface {
// GetUsage gets current usage for a specific limit.
// Returns the current amount, window end time, and any error.
// If no usage exists, returns 0 with a new window end time.
GetUsage(ctx context.Context, scope Scope, identifier string, limitType LimitType, window TimeWindow) (int64, time.Time, error)
// IncrementUsage increments usage for a specific limit.
// Returns the new amount, window end time, and any error.
// If the window has expired, it resets and starts a new window.
IncrementUsage(ctx context.Context, scope Scope, identifier string, limitType LimitType, window TimeWindow, amount int64) (int64, time.Time, error)
// SetUsage sets usage for a specific limit.
// Used for explicit resets or window rollovers.
SetUsage(ctx context.Context, scope Scope, identifier string, limitType LimitType, window TimeWindow, amount int64, windowEnd time.Time) error
// DeleteUsage deletes all usage records for an identifier.
DeleteUsage(ctx context.Context, scope Scope, identifier string) error
// DeleteExpired deletes all expired usage records.
// Records with windowEnd before the specified time are deleted.
DeleteExpired(ctx context.Context, before time.Time) error
// Close closes the store and releases resources.
Close() error
}
Store is the persistence layer for rate limit data.
Implementations must be thread-safe and support concurrent access.
type TimeWindow ¶
type TimeWindow string
TimeWindow represents a rate limiting time window.
const ( // WindowMinute represents a 60-second window (burst protection). WindowMinute TimeWindow = "minute" // WindowHour represents a 60-minute window (short-term limits). WindowHour TimeWindow = "hour" // WindowDay represents a 24-hour window (daily quotas). WindowDay TimeWindow = "day" // WindowWeek represents a 7-day window (weekly budgets). WindowWeek TimeWindow = "week" // WindowMonth represents a 30-day window (monthly billing). WindowMonth TimeWindow = "month" )
func ParseTimeWindow ¶
func ParseTimeWindow(s string) TimeWindow
ParseTimeWindow converts a config string to TimeWindow.
func (TimeWindow) Duration ¶
func (w TimeWindow) Duration() time.Duration
Duration returns the duration for the time window.
func (TimeWindow) String ¶
func (w TimeWindow) String() string
String returns the string representation of the time window.
type Usage ¶
type Usage struct {
// LimitType is the type of limit (token or count).
LimitType LimitType `json:"limit_type"`
// Window is the time window for this usage.
Window TimeWindow `json:"window"`
// Current is the current usage in the window.
Current int64 `json:"current"`
// Limit is the maximum allowed in the window.
Limit int64 `json:"limit"`
// WindowEnd is when the current window ends.
WindowEnd time.Time `json:"window_end"`
// Remaining is the remaining quota in the window.
Remaining int64 `json:"remaining"`
// Percentage is the usage percentage (0-100+).
Percentage float64 `json:"percentage"`
}
Usage represents current usage for a specific limit.
type ValidationError ¶
ValidationError represents a configuration validation error.
func NewValidationError ¶
func NewValidationError(field, message string) *ValidationError
NewValidationError creates a new ValidationError.
func (*ValidationError) Error ¶
func (e *ValidationError) Error() string
Error returns the validation error message.