84 lines
1.6 KiB
Go
84 lines
1.6 KiB
Go
package track1
|
|
|
|
import (
|
|
"sync"
|
|
"time"
|
|
)
|
|
|
|
// InMemoryRateLimiter is a simple in-memory rate limiter
|
|
// In production, use Redis for distributed rate limiting
|
|
type InMemoryRateLimiter struct {
|
|
limits map[string]*limitEntry
|
|
mu sync.RWMutex
|
|
config RateLimitConfig
|
|
}
|
|
|
|
// RateLimitConfig defines rate limit configuration
|
|
type RateLimitConfig struct {
|
|
RequestsPerSecond int
|
|
RequestsPerMinute int
|
|
BurstSize int
|
|
}
|
|
|
|
// limitEntry tracks rate limit state for a key
|
|
type limitEntry struct {
|
|
count int
|
|
resetAt time.Time
|
|
lastReset time.Time
|
|
}
|
|
|
|
// NewInMemoryRateLimiter creates a new in-memory rate limiter
|
|
func NewInMemoryRateLimiter(config RateLimitConfig) *InMemoryRateLimiter {
|
|
return &InMemoryRateLimiter{
|
|
limits: make(map[string]*limitEntry),
|
|
config: config,
|
|
}
|
|
}
|
|
|
|
// Allow checks if a request is allowed for the given key
|
|
func (rl *InMemoryRateLimiter) Allow(key string) bool {
|
|
rl.mu.Lock()
|
|
defer rl.mu.Unlock()
|
|
|
|
now := time.Now()
|
|
entry, exists := rl.limits[key]
|
|
|
|
if !exists {
|
|
rl.limits[key] = &limitEntry{
|
|
count: 1,
|
|
resetAt: now.Add(time.Minute),
|
|
lastReset: now,
|
|
}
|
|
return true
|
|
}
|
|
|
|
// Reset if minute has passed
|
|
if now.After(entry.resetAt) {
|
|
entry.count = 1
|
|
entry.resetAt = now.Add(time.Minute)
|
|
entry.lastReset = now
|
|
return true
|
|
}
|
|
|
|
// Check limits
|
|
if entry.count >= rl.config.RequestsPerMinute {
|
|
return false
|
|
}
|
|
|
|
entry.count++
|
|
return true
|
|
}
|
|
|
|
// Cleanup removes old entries (call periodically)
|
|
func (rl *InMemoryRateLimiter) Cleanup() {
|
|
rl.mu.Lock()
|
|
defer rl.mu.Unlock()
|
|
|
|
now := time.Now()
|
|
for key, entry := range rl.limits {
|
|
if now.After(entry.resetAt.Add(5 * time.Minute)) {
|
|
delete(rl.limits, key)
|
|
}
|
|
}
|
|
}
|