Add Redis-backed chat queue

ai_knowledge_base_v2/apps/backend/app/services/chat_queue_redis_service.py

@@ -0,0 +1,155 @@
+from __future__ import annotations
+
+import asyncio
+import time
+from uuid import uuid4
+
+from redis.asyncio import Redis
+
+from app.services.chat_queue_service import ChatQueueConfig, QueueRequest
+from app.services.redis_client import get_redis_client
+
+
+class RedisChatQueueManager:
+    _active_key = "ai_kb:chat_queue:active"
+    _waiting_key = "ai_kb:chat_queue:waiting"
+    _request_script = """
+local active_key = KEYS[1]
+local waiting_key = KEYS[2]
+local now = tonumber(ARGV[1])
+local token = ARGV[2]
+local max_active = tonumber(ARGV[3])
+local max_waiting = tonumber(ARGV[4])
+local active_expire_at = tonumber(ARGV[5])
+local key_ttl = tonumber(ARGV[6])
+
+redis.call('ZREMRANGEBYSCORE', active_key, '-inf', now)
+local active_count = redis.call('ZCARD', active_key)
+local waiting_count = redis.call('LLEN', waiting_key)
+
+if active_count < max_active and waiting_count == 0 then
+  redis.call('ZADD', active_key, active_expire_at, token)
+  redis.call('EXPIRE', active_key, key_ttl)
+  return {'acquired', token, 0, active_count + 1, 0}
+end
+
+if waiting_count >= max_waiting then
+  redis.call('EXPIRE', active_key, key_ttl)
+  redis.call('EXPIRE', waiting_key, key_ttl)
+  return {'rejected', token, 0, active_count, waiting_count}
+end
+
+redis.call('RPUSH', waiting_key, token)
+redis.call('EXPIRE', waiting_key, key_ttl)
+redis.call('EXPIRE', active_key, key_ttl)
+return {'queued', token, waiting_count + 1, active_count, waiting_count + 1}
+"""
+    _promote_script = """
+local active_key = KEYS[1]
+local waiting_key = KEYS[2]
+local now = tonumber(ARGV[1])
+local token = ARGV[2]
+local max_active = tonumber(ARGV[3])
+local active_expire_at = tonumber(ARGV[4])
+local key_ttl = tonumber(ARGV[5])
+
+redis.call('ZREMRANGEBYSCORE', active_key, '-inf', now)
+local active_count = redis.call('ZCARD', active_key)
+local waiting_count = redis.call('LLEN', waiting_key)
+local head = redis.call('LINDEX', waiting_key, 0)
+
+if head == token and active_count < max_active then
+  redis.call('LPOP', waiting_key)
+  redis.call('ZADD', active_key, active_expire_at, token)
+  redis.call('EXPIRE', active_key, key_ttl)
+  redis.call('EXPIRE', waiting_key, key_ttl)
+  return {'acquired', token, 0, active_count + 1, waiting_count - 1}
+end
+
+local position = redis.call('LPOS', waiting_key, token)
+if not position then
+  return {'canceled', token, 0, active_count, waiting_count}
+end
+
+redis.call('EXPIRE', active_key, key_ttl)
+redis.call('EXPIRE', waiting_key, key_ttl)
+return {'queued', token, position + 1, active_count, waiting_count}
+"""
+
+    async def request_slot(self, config: ChatQueueConfig) -> QueueRequest:
+        client = _require_client()
+        now = time.time()
+        token = uuid4().hex
+        result = await client.eval(
+            self._request_script,
+            2,
+            self._active_key,
+            self._waiting_key,
+            now,
+            token,
+            config.max_active_requests,
+            config.max_queue_size,
+            now + config.active_lease_seconds,
+            _key_ttl(config),
+        )
+        return _queue_request(result)
+
+    async def wait_for_slot(self, token: str, config: ChatQueueConfig) -> QueueRequest:
+        client = _require_client()
+        deadline = time.monotonic() + config.queue_timeout_seconds
+        while True:
+            now = time.time()
+            result = await client.eval(
+                self._promote_script,
+                2,
+                self._active_key,
+                self._waiting_key,
+                now,
+                token,
+                config.max_active_requests,
+                now + config.active_lease_seconds,
+                _key_ttl(config),
+            )
+            request = _queue_request(result)
+            if request.status in {"acquired", "canceled"}:
+                return request
+
+            remaining = deadline - time.monotonic()
+            if remaining <= 0:
+                await self.cancel_waiter(token)
+                return QueueRequest(status="timeout", token=token)
+            await asyncio.sleep(min(0.25, remaining))
+
+    async def cancel_waiter(self, token: str) -> None:
+        client = _require_client()
+        await client.lrem(self._waiting_key, 0, token)
+
+    async def release_slot(self, token: str) -> None:
+        client = _require_client()
+        await client.zrem(self._active_key, token)
+
+
+redis_chat_queue_manager = RedisChatQueueManager()
+
+
+def _require_client() -> Redis:
+    client = get_redis_client()
+    if client is None:
+        raise RuntimeError("Redis is not configured")
+    return client
+
+
+def _queue_request(result: list) -> QueueRequest:
+    status = str(result[0])
+    token = str(result[1]) if len(result) > 1 and result[1] else None
+    return QueueRequest(
+        status=status,
+        token=token,
+        position=int(result[2]) if len(result) > 2 else 0,
+        active_count=int(result[3]) if len(result) > 3 else 0,
+        waiting_count=int(result[4]) if len(result) > 4 else 0,
+    )
+
+
+def _key_ttl(config: ChatQueueConfig) -> int:
+    return max(config.active_lease_seconds, config.queue_timeout_seconds) + 300

ai_knowledge_base_v2/apps/backend/app/services/chat_queue_runtime.py

@@ -0,0 +1,65 @@
+from __future__ import annotations
+
+import asyncio
+from dataclasses import dataclass
+
+from app.services.chat_queue_redis_service import redis_chat_queue_manager
+from app.services.chat_queue_service import ChatQueueConfig, QueueRequest, chat_queue_manager
+from app.services.redis_client import ping_redis
+
+
+@dataclass
+class ChatQueueLease:
+    backend: str
+    request: QueueRequest
+
+
+async def request_chat_slot(config: ChatQueueConfig) -> ChatQueueLease:
+    if await ping_redis():
+        try:
+            return ChatQueueLease("redis", await redis_chat_queue_manager.request_slot(config))
+        except Exception:
+            pass
+    request = await asyncio.to_thread(chat_queue_manager.request_slot, config)
+    return ChatQueueLease("memory", request)
+
+
+async def wait_for_chat_slot(lease: ChatQueueLease, config: ChatQueueConfig) -> ChatQueueLease:
+    token = lease.request.token
+    if lease.backend == "redis" and isinstance(token, str):
+        try:
+            lease.request = await redis_chat_queue_manager.wait_for_slot(token, config)
+            return lease
+        except Exception:
+            await cancel_chat_waiter(lease)
+            lease.backend = "memory"
+            lease.request = await asyncio.to_thread(chat_queue_manager.request_slot, config)
+            if lease.request.status == "queued":
+                lease.request = await asyncio.to_thread(chat_queue_manager.wait_for_slot, lease.request.token, config)
+            return lease
+
+    lease.request = await asyncio.to_thread(chat_queue_manager.wait_for_slot, token, config)
+    return lease
+
+
+async def cancel_chat_waiter(lease: ChatQueueLease) -> None:
+    token = lease.request.token
+    if lease.backend == "redis" and isinstance(token, str):
+        try:
+            await redis_chat_queue_manager.cancel_waiter(token)
+        except Exception:
+            pass
+        return
+    if token is not None:
+        await asyncio.to_thread(chat_queue_manager.cancel_waiter, token)
+
+
+async def release_chat_slot(lease: ChatQueueLease) -> None:
+    token = lease.request.token
+    if lease.backend == "redis" and isinstance(token, str):
+        try:
+            await redis_chat_queue_manager.release_slot(token)
+        except Exception:
+            pass
+        return
+    await asyncio.to_thread(chat_queue_manager.release_slot)

ai_knowledge_base_v2/apps/backend/app/services/chat_queue_service.py

@@ -17,6 +17,7 @@ class ChatQueueConfig:
     max_active_requests: int
     max_queue_size: int
     queue_timeout_seconds: int
+    active_lease_seconds: int
 
 
 @dataclass(frozen=True)
@@ -132,6 +133,13 @@ def load_chat_queue_config(db: Session) -> ChatQueueConfig:
             minimum=1,
             maximum=3600,
         ),
+        active_lease_seconds=_config_int(
+            db,
+            "chat_active_lease_seconds",
+            settings.chat_active_lease_seconds,
+            minimum=60,
+            maximum=86400,
+        ),
     )
 
 

ai_knowledge_base_v2/apps/backend/app/services/redis_client.py

@@ -0,0 +1,25 @@
+from __future__ import annotations
+
+from functools import lru_cache
+
+from redis.asyncio import Redis
+
+from app.core.config import get_settings
+
+
+@lru_cache(maxsize=1)
+def get_redis_client() -> Redis | None:
+    redis_url = get_settings().redis_url.strip()
+    if not redis_url:
+        return None
+    return Redis.from_url(redis_url, decode_responses=True)
+
+
+async def ping_redis() -> bool:
+    client = get_redis_client()
+    if client is None:
+        return False
+    try:
+        return bool(await client.ping())
+    except Exception:
+        return False