mirror of
https://github.com/open-webui/open-webui.git
synced 2026-09-19 00:03:21 +00:00
d2e62db69b
A slow Redis freezes the whole worker during sign-in, not just the user signing in. RateLimiter held a synchronous redis-py client and signin called is_limited inline from a coroutine, so every attempt did blocking round trips on the event-loop thread, with REDIS_SOCKET_TIMEOUT defaulting to None so nothing bounded the wait. Its Redis methods are now async and take the handle as their first argument, and both handlers pass request.app.state.redis, the async client the lifespan already creates. Building one in the limiter instead would pin its pooled connection to the first event loop that used it. Without Redis, which is the default single-instance setup, the fallback store leaked. It was keyed by the rate-limit key and pruned a key's expired buckets only when that same key was checked again, so a login email never seen again was never reclaimed, and that email comes straight from an unauthenticated request body. It is now keyed by bucket, so one prune drops every key an expired bucket held, and it lives on the instance: pruning uses the per-instance num_buckets, so a shared store would let a limiter with a short window delete buckets a longer-windowed one still needs. With a Redis costing a second per call, the widest event-loop tick gap drops from 2.010s to 0.010s and a concurrent request is answered at 0.05s instead of 2.05s, at no cost to the caller's own latency. Across 20,000 distinct keys the store goes from 40,000 entries and 6.4 MB, growing linearly, to a flat 1,004 entries and 100 KB. Rate-limiting decisions are unchanged across 700,000 randomised calls over 14 window, bucket and limit combinations, against a real Redis and the in-memory fallback alike, and sign-in still returns its first 429 on attempt 16. Two behaviour changes worth naming. Pruning is now global rather than per key, so a wall clock that jumps forward past a full window and back forgets a hit it previously kept. The two limiters also stop sharing a store, which previously let a sign-in attempt with an IP-shaped email touch the token-exchange limiter's counters.
114 lines
3.9 KiB
Python
114 lines
3.9 KiB
Python
import time
|
|
from typing import Optional
|
|
|
|
from open_webui.env import REDIS_KEY_PREFIX
|
|
from redis.asyncio import Redis
|
|
|
|
|
|
class RateLimiter:
|
|
"""
|
|
General-purpose rate limiter using Redis with a rolling window strategy.
|
|
Falls back to in-memory storage if Redis is not available.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
limit: int,
|
|
window: int,
|
|
bucket_size: int = 60,
|
|
enabled: bool = True,
|
|
):
|
|
"""
|
|
:param limit: Max allowed events in the window
|
|
:param window: Time window in seconds
|
|
:param bucket_size: Bucket resolution
|
|
:param enabled: Turn on/off rate limiting globally
|
|
"""
|
|
self.limit = limit
|
|
self.window = window
|
|
self.bucket_size = bucket_size
|
|
self.num_buckets = window // bucket_size
|
|
self.enabled = enabled
|
|
# bucket index -> rate-limit key -> hits
|
|
self._memory_store: dict[int, dict[str, int]] = {}
|
|
|
|
def _bucket_key(self, key: str, bucket_index: int) -> str:
|
|
return f'{REDIS_KEY_PREFIX}:ratelimit:{key.lower()}:{bucket_index}'
|
|
|
|
def _current_bucket(self) -> int:
|
|
return int(time.time()) // self.bucket_size
|
|
|
|
def _prune_memory_store(self, now_bucket: int) -> None:
|
|
min_bucket = now_bucket - self.num_buckets
|
|
expired = [bucket_index for bucket_index in self._memory_store if bucket_index < min_bucket]
|
|
for bucket_index in expired:
|
|
del self._memory_store[bucket_index]
|
|
|
|
async def is_limited(self, redis: Redis | None, key: str) -> bool:
|
|
"""
|
|
Main rate-limit check.
|
|
Gracefully handles missing or failing Redis.
|
|
"""
|
|
if not self.enabled:
|
|
return False
|
|
|
|
if redis is not None:
|
|
try:
|
|
return await self._is_limited_redis(redis, key)
|
|
except Exception:
|
|
return self._is_limited_memory(key)
|
|
else:
|
|
return self._is_limited_memory(key)
|
|
|
|
async def get_count(self, redis: Redis | None, key: str) -> int:
|
|
if not self.enabled:
|
|
return 0
|
|
|
|
if redis is not None:
|
|
try:
|
|
return await self._get_count_redis(redis, key)
|
|
except Exception:
|
|
return self._get_count_memory(key)
|
|
else:
|
|
return self._get_count_memory(key)
|
|
|
|
async def remaining(self, redis: Redis | None, key: str) -> int:
|
|
used = await self.get_count(redis, key)
|
|
return max(0, self.limit - used)
|
|
|
|
async def _is_limited_redis(self, redis: Redis, key: str) -> bool:
|
|
now_bucket = self._current_bucket()
|
|
bucket_key = self._bucket_key(key, now_bucket)
|
|
|
|
attempts = await redis.incr(bucket_key)
|
|
if attempts == 1:
|
|
await redis.expire(bucket_key, self.window + self.bucket_size)
|
|
|
|
# Collect buckets
|
|
buckets = [self._bucket_key(key, now_bucket - i) for i in range(self.num_buckets + 1)]
|
|
|
|
counts = await redis.mget(buckets)
|
|
total = sum(int(c) for c in counts if c)
|
|
|
|
return total > self.limit
|
|
|
|
async def _get_count_redis(self, redis: Redis, key: str) -> int:
|
|
now_bucket = self._current_bucket()
|
|
buckets = [self._bucket_key(key, now_bucket - i) for i in range(self.num_buckets + 1)]
|
|
counts = await redis.mget(buckets)
|
|
return sum(int(c) for c in counts if c)
|
|
|
|
def _is_limited_memory(self, key: str) -> bool:
|
|
now_bucket = self._current_bucket()
|
|
self._prune_memory_store(now_bucket)
|
|
|
|
current_bucket_counts = self._memory_store.setdefault(now_bucket, {})
|
|
current_bucket_counts[key] = current_bucket_counts.get(key, 0) + 1
|
|
|
|
total = sum(bucket_counts.get(key, 0) for bucket_counts in self._memory_store.values())
|
|
return total > self.limit
|
|
|
|
def _get_count_memory(self, key: str) -> int:
|
|
now_bucket = self._current_bucket()
|
|
self._prune_memory_store(now_bucket)
|
|
return sum(bucket_counts.get(key, 0) for bucket_counts in self._memory_store.values())
|