134 lines
4.3 KiB
Go
134 lines
4.3 KiB
Go
package autoassignment
|
|
|
|
// RateLimiter implements Redis-backed rate limiting for auto-assignment.
|
|
//
|
|
// Reference: Chatwoot RateLimiter (Redis key-per-assignment)
|
|
// - Tracks count of assignments per agent per inbox per time window
|
|
// - Fair distribution limit (default 5) over window (default 5 min)
|
|
// - Key format: gochat:rate_limit:{inboxID}:{agentID}:{windowKey}
|
|
// - Ensures no single agent gets too many assignments in a short time
|
|
//
|
|
// In gochat, we use Redis string keys with TTL for rate counting:
|
|
// - Key: gochat:auto_assignment_rate:{inboxID}:{agentID}:{windowStart}
|
|
// - Value: count of assignments in this window
|
|
// - TTL: window duration (so keys expire automatically)
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"strconv"
|
|
"time"
|
|
|
|
"github.com/redis/go-redis/v9"
|
|
)
|
|
|
|
const (
|
|
// rateLimitKeyPrefix is the Redis key prefix for rate limit counters.
|
|
rateLimitKeyPrefix = "gochat:auto_assignment_rate:"
|
|
)
|
|
|
|
// RateLimiter tracks assignment counts per agent per inbox per time window.
|
|
type RateLimiter struct {
|
|
redis *redis.Client
|
|
}
|
|
|
|
// NewRateLimiter creates a new RateLimiter.
|
|
func NewRateLimiter(rdb *redis.Client) *RateLimiter {
|
|
return &RateLimiter{redis: rdb}
|
|
}
|
|
|
|
// key returns the Redis key for rate limiting the given agent in the given inbox.
|
|
// The key includes a time window component that changes when the window expires.
|
|
func (rl *RateLimiter) key(inboxID uint, agentID uint, windowSeconds int) string {
|
|
windowStart := time.Now().Unix() / int64(windowSeconds)
|
|
return fmt.Sprintf("%s%d:%d:%d", rateLimitKeyPrefix, inboxID, agentID, windowStart)
|
|
}
|
|
|
|
// IsAllowed checks whether the given agent is allowed to receive another
|
|
// assignment in the current time window.
|
|
//
|
|
// Returns true if the agent's assignment count in this window is below the limit.
|
|
func (rl *RateLimiter) IsAllowed(ctx context.Context, inboxID uint, agentID uint, limit int, windowSeconds int) bool {
|
|
if limit <= 0 {
|
|
return true // no rate limiting configured
|
|
}
|
|
|
|
key := rl.key(inboxID, agentID, windowSeconds)
|
|
count, err := rl.redis.Get(ctx, key).Int()
|
|
if err == redis.Nil {
|
|
return true // no assignments yet in this window
|
|
}
|
|
if err != nil {
|
|
// Redis error — allow assignment to avoid blocking service
|
|
return true
|
|
}
|
|
|
|
return count < limit
|
|
}
|
|
|
|
// Increment increments the assignment count for the given agent in the
|
|
// current time window.
|
|
//
|
|
// Sets a TTL on the key equal to the window duration, so the key
|
|
// expires automatically when the window ends.
|
|
func (rl *RateLimiter) Increment(ctx context.Context, inboxID uint, agentID uint, windowSeconds int) error {
|
|
key := rl.key(inboxID, agentID, windowSeconds)
|
|
|
|
// Increment count
|
|
if err := rl.redis.Incr(ctx, key).Err(); err != nil {
|
|
return fmt.Errorf("redis Incr: %w", err)
|
|
}
|
|
|
|
// Set TTL (only on first increment in this window)
|
|
ttl, err := rl.redis.TTL(ctx, key).Result()
|
|
if err != nil {
|
|
return fmt.Errorf("redis TTL: %w", err)
|
|
}
|
|
if ttl == -1 { // key has no TTL (first increment)
|
|
if err := rl.redis.Expire(ctx, key, time.Duration(windowSeconds)*time.Second).Err(); err != nil {
|
|
return fmt.Errorf("redis Expire: %w", err)
|
|
}
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
// GetCount returns the current assignment count for the given agent
|
|
// in the current time window.
|
|
func (rl *RateLimiter) GetCount(ctx context.Context, inboxID uint, agentID uint, windowSeconds int) (int, error) {
|
|
key := rl.key(inboxID, agentID, windowSeconds)
|
|
count, err := rl.redis.Get(ctx, key).Int()
|
|
if err == redis.Nil {
|
|
return 0, nil
|
|
}
|
|
if err != nil {
|
|
return 0, fmt.Errorf("redis Get: %w", err)
|
|
}
|
|
return count, nil
|
|
}
|
|
|
|
// Reset resets all rate limit counters for the given inbox.
|
|
// This is useful when inbox members change significantly.
|
|
func (rl *RateLimiter) Reset(ctx context.Context, inboxID uint) error {
|
|
pattern := fmt.Sprintf("%s%d:*", rateLimitKeyPrefix, inboxID)
|
|
keys, err := rl.redis.Keys(ctx, pattern).Result()
|
|
if err != nil {
|
|
return fmt.Errorf("redis Keys: %w", err)
|
|
}
|
|
if len(keys) > 0 {
|
|
if err := rl.redis.Del(ctx, keys...).Err(); err != nil {
|
|
return fmt.Errorf("redis Del: %w", err)
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// ParseAgentID extracts a uint agent ID from a string representation.
|
|
// Utility function for converting Redis string values to uint IDs.
|
|
func ParseAgentID(s string) (uint, error) {
|
|
id, err := strconv.ParseUint(s, 10, 32)
|
|
if err != nil {
|
|
return 0, err
|
|
}
|
|
return uint(id), nil
|
|
} |