Pythonrate_limiter.py
Copy
1
import time
2
import redis
3
 
4
# docker run -d -p 6379:6379 redis:alpine
5
r = redis.Redis(host='localhost', port=6379, decode_responses=True)
6
 
7
# ── Token Bucket ──────────────────────────────────────────────────────
8
class TokenBucket:
9
    """
10
    Token Bucket rate limiter.
11
    Tokens refill at `rate` tokens/second up to `capacity`.
12
    """
13
    def __init__(self, capacity: int, rate: float):
14
        self.capacity = capacity
15
        self.rate = rate  # tokens per second
16
    
17
    def _key(self, identifier: str) -> tuple[str, str]:
18
        return f"bucket:{identifier}:tokens", f"bucket:{identifier}:last_refill"
19
    
20
    def consume(self, identifier: str, tokens: int = 1) -> tuple[bool, dict]:
21
        """
22
        Try to consume tokens. Returns (allowed, info).
23
        Uses Lua script for atomicity.
24
        """
25
        tokens_key, time_key = self._key(identifier)
26
        now = time.time()
27
        
28
        script = """
29
        local tokens_key = KEYS[1]
30
        local time_key = KEYS[2]
31
        local capacity = tonumber(ARGV[1])
32
        local rate = tonumber(ARGV[2])
33
        local now = tonumber(ARGV[3])
34
        local requested = tonumber(ARGV[4])
35
        
36
        local last_time = tonumber(redis.call('GET', time_key) or now)
37
        local current_tokens = tonumber(redis.call('GET', tokens_key) or capacity)
38
        
39
        -- Refill tokens based on elapsed time
40
        local elapsed = now - last_time
41
        local new_tokens = math.min(capacity, current_tokens + (elapsed * rate))
42
        
43
        redis.call('SET', time_key, now, 'EX', 3600)
44
        
45
        if new_tokens >= requested then
46
            redis.call('SET', tokens_key, new_tokens - requested, 'EX', 3600)
47
            return {1, math.floor(new_tokens - requested)}
48
        else
49
            redis.call('SET', tokens_key, new_tokens, 'EX', 3600)
50
            return {0, math.floor(new_tokens)}
51
        end
52
        """
53
        
54
        result = r.eval(script, 2, tokens_key, time_key,
55
                        self.capacity, self.rate, now, tokens)
56
        allowed = bool(result[0])
57
        remaining = int(result[1])
58
        return allowed, {"remaining_tokens": remaining, "capacity": self.capacity}
59
 
60
# ── Sliding Window ────────────────────────────────────────────────────
61
class SlidingWindowCounter:
62
    """
63
    Sliding window rate limiter using Redis ZSET.
64
    Each request is a member with score = timestamp.
65
    Window = [now - window_size, now].
66
    """
67
    def __init__(self, limit: int, window_seconds: int):
68
        self.limit = limit
69
        self.window = window_seconds
70
    
71
    def check(self, identifier: str) -> tuple[bool, dict]:
72
        key = f"sliding:{identifier}"
73
        now = time.time()
74
        window_start = now - self.window
75
        
76
        pipe = r.pipeline()
77
        # Remove old requests outside window
78
        pipe.zremrangebyscore(key, 0, window_start)
79
        # Add current request
80
        pipe.zadd(key, {str(now): now})
81
        # Count requests in window
82
        pipe.zcard(key)
83
        # Set TTL
84
        pipe.expire(key, self.window + 1)
85
        results = pipe.execute()
86
        
87
        count = results[2]
88
        allowed = count <= self.limit
89
        
90
        if not allowed:
91
            # Remove the just-added request
92
            r.zrem(key, str(now))
93
        
94
        return allowed, {
95
            "count": count if allowed else count - 1,
96
            "limit": self.limit,
97
            "window_seconds": self.window,
98
            "remaining": max(0, self.limit - count),
99
        }
100
 
101
# ── Test ─────────────────────────────────────────────────────────────
102
# Clear test keys
103
for k in r.scan_iter("bucket:*"): r.delete(k)
104
for k in r.scan_iter("sliding:*"): r.delete(k)
105
 
106
bucket = TokenBucket(capacity=10, rate=2.0)   # 10 tokens, refill 2/second
107
window = SlidingWindowCounter(limit=5, window_seconds=10)
108
 
109
print("=== Token Bucket (capacity=10, rate=2/s) ===")
110
user = "user:alice"
111
allowed_count = 0
112
for i in range(15):
113
    allowed, info = bucket.consume(user)
114
    status = "✓" if allowed else "✗"
115
    print(f"  Request {i+1:2d}: {status} (tokens remaining: {info['remaining_tokens']})")
116
    if allowed:
117
        allowed_count += 1
118
 
119
print(f"Allowed: {allowed_count}/15 (burst of 10 then throttled)
120
")
121
 
122
# Refill test
123
time.sleep(1)
124
allowed, info = bucket.consume(user)
125
print(f"After 1s sleep: {'✓' if allowed else '✗'} (tokens: {info['remaining_tokens']})
126
")
127
 
128
print("=== Sliding Window (limit=5 per 10s) ===")
129
user2 = "user:bob"
130
for i in range(8):
131
    allowed, info = window.check(user2)
132
    status = "✓" if allowed else "✗"
133
    print(f"  Request {i+1}: {status} (count: {info['count']}/{info['limit']}, remaining: {info['remaining']})")
134