Stop event-loop blocking that trips broker WebSocket keepalive timeouts

Under load the CoderAI broker WebSocket was being torn down with
"1011 keepalive ping timeout" and reconnecting (surfaced as 502s by the
fronting proxy). Root cause: synchronous work run directly on the asyncio
event loop starves the WebSocket keepalive so PONGs miss the 20s window.

Remove the blockers:

- analytics: dispatch the blocking token_usage / per-user INSERTs to a
  worker thread (fire-and-forget when a loop is running, inline otherwise)
  instead of calling them inline from async request handlers.
- database: add a MySQL connection pool (ping+reconnect, direct-connect
  fallback) and return connections to the pool on with-block exit instead
  of opening a fresh TCP+auth connection per query and leaking it to GC.
- utils: cache the tiktoken encoding and memoize count_messages_tokens by
  content, collapsing the several CPU-bound BPE passes per request into one.
- coderai_broker: run touch()/publish_response Redis I/O off the loop via
  asyncio.to_thread (awaited to preserve chunk order), matching the existing
  consume_request pattern.
- routes/coderai_broker: answer heartbeats immediately and process response
  frames on a dedicated FIFO worker so a slow publish can't delay the next
  heartbeat.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

aisbf/analytics.py

@@ -23,6 +23,7 @@ Why did the programmer quit his job? Because he didn't get arrays!
 Token Usage Analytics module for AISBF.
 """
 import json
+import asyncio
 from decimal import Decimal
 import csv
 import io
@@ -357,21 +358,64 @@ class Analytics:
         self._latencies[provider_id]['min_ms'] = min(self._latencies[provider_id]['min_ms'], latency_ms)
         self._latencies[provider_id]['max_ms'] = max(self._latencies[provider_id]['max_ms'], latency_ms)
         
-        # Persist to database
+        # Persist to database OFF the event loop.  record_token_usage() opens a MySQL
+        # connection and runs a blocking INSERT; calling it inline here (this method runs
+        # synchronously inside async request handlers) stalled the event loop long enough
+        # for the CoderAI broker WebSocket keepalive to time out (1011) under load and
+        # forced reconnects.  Nothing awaits the result, so fire-and-forget in a worker
+        # thread is safe.
         if tokens_used > 0:
             logger.info(f"Analytics.record_request: Recording to DB - provider={provider_id}, latency_ms={latency_ms}, success={success}, prompt={prompt_tokens}, completion={completion_tokens}, cost={actual_cost}")
+            self._persist_token_usage(
+                provider_id, model_name, tokens_used, user_id, success, latency_ms,
+                error_type, token_id, prompt_tokens, completion_tokens, actual_cost,
+                rotation_id, autoselect_id, analytics_kind,
+            )
+
+    def _persist_token_usage(self, *args):
+        """Write a token_usage row without blocking the caller's event loop.
+
+        When a running loop is present (the normal request path) the blocking DB write
+        is dispatched to a thread pool and NOT awaited — analytics is pure side-effect
+        logging that nothing downstream waits on.  With no loop (sync scripts/tests) it
+        runs inline so behaviour is unchanged there."""
+        provider_id = args[0]
+
+        # args positional order matches record_token_usage(); pull out what the
+        # per-user write also needs.
+        model_name = args[1]
+        tokens_used = args[2]
+        user_id = args[3]
+        token_id = args[7]
+
+        def _write():
             try:
-                self.db.record_token_usage(provider_id, model_name, tokens_used, user_id, success, latency_ms, error_type, token_id, prompt_tokens, completion_tokens, actual_cost, rotation_id, autoselect_id, analytics_kind)
+                self.db.record_token_usage(*args)
                 logger.info(f"✅ Analytics recording completed successfully for {provider_id}")
             except Exception as e:
                 logger.error(f"❌ Analytics recording FAILED for {provider_id}: {e}")
                 import traceback
                 logger.error(f"Traceback: {traceback.format_exc()}")
-
-            # Also record user token usage if this was an API token request
+            # Also record user token usage if this was an API token request.
             if user_id is not None and token_id is not None:
-                self.db.record_user_token_usage(user_id, token_id, provider_id, model_name, tokens_used)
-    
+                try:
+                    self.db.record_user_token_usage(user_id, token_id, provider_id, model_name, tokens_used)
+                except Exception as e:
+                    logger.error(f"❌ Per-user token usage recording FAILED for {provider_id}: {e}")
+
+        try:
+            loop = asyncio.get_running_loop()
+        except RuntimeError:
+            loop = None
+        if loop is None:
+            _write()
+            return
+        executor = getattr(self.db, 'executor', None)
+        # Fire-and-forget: the returned future is intentionally not awaited.  _write()
+        # swallows all exceptions so the future always resolves cleanly (no "exception
+        # never retrieved" warnings).
+        loop.run_in_executor(executor, _write)
+
     def get_provider_stats(
         self,
         provider_id: str,

aisbf/coderai_broker.py

@@ -255,12 +255,19 @@ class CoderAIBroker:
 
     # -------------------------------------------------- shared cache helpers
 
-    def _store_session_cache(self, session: CoderAISession) -> None:
+    def _build_session_cache_payload(self, session: CoderAISession) -> tuple[str, dict]:
+        """Build the (meta_key, payload) for the shared cache – pure in-memory work."""
         payload = self._serialize_session(session)
         payload["metadata"] = dict(payload.get("metadata") or {})
         payload["metadata"]["broker_node_id"] = self._node_id
         payload["metadata"]["connection_state"] = "connected" if not session.closed else "disconnected"
         meta_key = self._session_meta_key(session.provider_id, session.client_id)
+        return meta_key, payload
+
+    def _flush_session_cache(self, meta_key: str, payload: dict) -> None:
+        """Write the session payload + index to the shared cache.  This performs
+        synchronous Redis round-trips and MUST be run off the event loop when on a hot
+        path (see touch())."""
         self._cache.broker_set(meta_key, payload, ttl=SESSION_TTL_SECONDS)
         # Keep a global index of all live session meta keys so list_sessions
         # can enumerate them without a Redis SCAN.
@@ -271,6 +278,12 @@ class CoderAIBroker:
             index.append(meta_key)
             self._cache.broker_set(self._session_index_key(), index, ttl=SESSION_TTL_SECONDS * 10)
 
+    def _store_session_cache(self, session: CoderAISession) -> None:
+        """Synchronous store – used by low-frequency paths (register / periodic
+        heartbeat_loop) where a durable write before returning is fine."""
+        meta_key, payload = self._build_session_cache_payload(session)
+        self._flush_session_cache(meta_key, payload)
+
     def _mark_session_offline_cache(self, provider_id: str, client_id: str) -> None:
         """Write a closed=True tombstone so remote nodes see the disconnect immediately."""
         meta_key = self._session_meta_key(provider_id, client_id)
@@ -392,8 +405,12 @@ class CoderAIBroker:
                 session.metadata.update(self._normalize_hardware_metadata(metadata))
             if capabilities:
                 session.capabilities = dict(capabilities)
-            self._store_session_cache(session)
-            return session
+            meta_key, payload = self._build_session_cache_payload(session)
+        # touch() runs for EVERY inbound WebSocket frame (heartbeat / response). Flush
+        # the shared-cache TTL refresh in a worker thread so its synchronous Redis I/O
+        # can't stall the event loop / WebSocket keepalive under load.
+        await asyncio.to_thread(self._flush_session_cache, meta_key, payload)
+        return session
 
     # ------------------------------------------------------- session queries
 
@@ -720,7 +737,11 @@ class CoderAIBroker:
             return
         await self._record_request_metric(request_id, message)
         # Write to shared cache so the requesting node (possibly remote) can poll it.
-        self._cache.broker_set(self._reply_key(request_id), message, ttl=REPLY_TTL_SECONDS)
+        # Run the synchronous Redis write in a thread so it doesn't block the event loop
+        # (awaited here to preserve ordering; the loop stays free during the I/O).
+        await asyncio.to_thread(
+            self._cache.broker_set, self._reply_key(request_id), message, REPLY_TTL_SECONDS
+        )
         # Also resolve any local Future for the fast-path case.
         await self.resolve_response(message)
 
@@ -742,7 +763,11 @@ class CoderAIBroker:
             return
 
         # Slow path: push chunk to the shared cache list for the remote requester.
-        self._cache.broker_push(self._reply_key(request_id), message, ttl=REPLY_TTL_SECONDS)
+        # Offload the synchronous Redis push so a stream of chunks can't block the loop;
+        # awaited to keep chunk ordering intact.
+        await asyncio.to_thread(
+            self._cache.broker_push, self._reply_key(request_id), message, REPLY_TTL_SECONDS
+        )
 
     async def finish_stream(self, request_id: str) -> None:
         """Remove a streaming request's pending state once the consumer has drained

aisbf/database.py

@@ -26,6 +26,7 @@ import json
 import hashlib
 import time
 import copy
+import contextlib
 from pathlib import Path
 from typing import Dict, List, Optional, Tuple, Any
 from datetime import datetime, timedelta
@@ -55,10 +56,12 @@ def _verify_password(password: str, stored_hash: str) -> bool:
 
 try:
     import mysql.connector as _mysql_connector
+    from mysql.connector import pooling as _mysql_pooling
     MYSQL_AVAILABLE = True
 except ImportError:
     MYSQL_AVAILABLE = False
     _mysql_connector = None
+    _mysql_pooling = None
 
 logger = logging.getLogger(__name__)
 
@@ -90,16 +93,24 @@ class _MySQLConnectionWrapper:
 
     def __exit__(self, exc_type, exc_val, exc_tb):
         if exc_type is None:
-            self._conn.commit()
+            try:
+                self._conn.commit()
+            except Exception:
+                pass
         else:
             try:
                 self._conn.rollback()
             except Exception:
                 pass
-        # Connection is intentionally left open: cursor and conn variables in the
-        # calling function remain valid after the with-block exits (matching SQLite's
-        # context-manager behaviour).  The connection is closed by GC when the caller
-        # function returns and conn goes out of scope.
+        # Release the connection deterministically instead of waiting for GC.  For a
+        # pooled connection, close() returns it to the pool for reuse (avoiding a fresh
+        # TCP+auth handshake per query on the event loop); for a direct fallback
+        # connection it closes the socket.  No caller uses conn/cursor after the
+        # with-block, so this is safe.
+        try:
+            self._conn.close()
+        except Exception:
+            pass
         return False
 
     # Forward attribute access so the wrapper can be used directly as well
@@ -141,20 +152,67 @@ class DatabaseManager:
         self.database_type = database_type
         self.db_type = self.db_config.get('type', 'sqlite').lower()
         self.executor = get_db_executor()
-        
+        self._mysql_pool = None
+
         if self.db_type == 'mysql':
             if not MYSQL_AVAILABLE:
                 raise ImportError("MySQL connector not available. Install mysql-connector-python.")
+            self._init_mysql_pool()
 
         self._initialize_database()
         logger.info(f"Database initialized: {self.db_type}")
 
+    def _init_mysql_pool(self) -> None:
+        """Create a MySQL connection pool so queries reuse warm connections instead
+        of paying a fresh TCP+auth handshake per call on the event loop.  Pool
+        creation failures are non-fatal: _get_connection() falls back to direct
+        connections so the DB stays usable even if MySQL is briefly unreachable at
+        startup."""
+        if _mysql_pooling is None:
+            return
+        try:
+            pool_size = int(self.db_config.get('mysql_pool_size', 16) or 16)
+        except (TypeError, ValueError):
+            pool_size = 16
+        pool_size = max(1, min(pool_size, 32))  # mysql.connector hard-caps pool_size at 32
+        try:
+            self._mysql_pool = _mysql_pooling.MySQLConnectionPool(
+                pool_name=f"aisbf_{self.database_type}",
+                pool_size=pool_size,
+                pool_reset_session=True,
+                host=self.db_config['mysql_host'],
+                port=self.db_config['mysql_port'],
+                user=self.db_config['mysql_user'],
+                password=self.db_config['mysql_password'],
+                database=self.db_config['mysql_database'],
+            )
+            logger.info(f"MySQL connection pool initialized (size={pool_size})")
+        except Exception as e:
+            self._mysql_pool = None
+            logger.warning(f"MySQL pool init failed ({e}); falling back to per-call connections")
+
     def _get_connection(self):
         """Get a database connection based on the configured type."""
         if self.db_type == 'sqlite':
             db_path = Path(self.db_config['sqlite_path']).expanduser()
             return sqlite3.connect(str(db_path))
         elif self.db_type == 'mysql':
+            # Preferred path: borrow a warm connection from the pool.  ping(reconnect=True)
+            # revives a connection the server dropped (MySQL wait_timeout) so callers never
+            # see a stale socket.  Pool exhaustion or any error falls through to a direct
+            # connection so a request is never rejected just because the pool is busy.
+            if self._mysql_pool is not None:
+                try:
+                    conn = self._mysql_pool.get_connection()
+                    try:
+                        conn.ping(reconnect=True, attempts=2, delay=0)
+                    except Exception:
+                        with contextlib.suppress(Exception):
+                            conn.close()
+                        raise
+                    return _MySQLConnectionWrapper(conn)
+                except Exception as pool_err:
+                    logger.debug(f"MySQL pool unavailable ({pool_err}); using direct connection")
             try:
                 conn = _mysql_connector.connect(
                     host=self.db_config['mysql_host'],

aisbf/routes/coderai_broker.py

@@ -144,7 +144,25 @@ async def _coderai_broker_websocket_impl(websocket: WebSocket, scope_name: str):
             if queued is not None:
                 await websocket.send_text(json.dumps(queued))
 
+    # Process response frames on a dedicated worker so a slow publish_response (cache
+    # writes, metric recording) can never delay the receive loop from reading and
+    # answering the next heartbeat. A single worker draining a FIFO queue preserves
+    # streaming chunk order.
+    publish_queue: asyncio.Queue = asyncio.Queue()
+
+    async def _process_publish_queue() -> None:
+        while True:
+            msg = await publish_queue.get()
+            try:
+                await broker.touch(session.session_id)
+                await broker.publish_response(msg)
+            except Exception:
+                logger.error("CoderAI broker publish handler failed", exc_info=True)
+            finally:
+                publish_queue.task_done()
+
     queue_task = asyncio.create_task(_drain_broker_queue())
+    publish_task = asyncio.create_task(_process_publish_queue())
     try:
         while True:
             raw = await websocket.receive_text()
@@ -171,7 +189,8 @@ async def _coderai_broker_websocket_impl(websocket: WebSocket, scope_name: str):
                 asyncio.create_task(_broker_refresh_models(session.provider_id, owner_user_id))
                 continue
             if op == "heartbeat":
-                await broker.touch(session.session_id, metadata=message.get("payload") or {})
+                # Answer the keepalive immediately; the TTL touch (which does off-loop
+                # cache I/O) happens after so a slow refresh can't delay the reply.
                 await websocket.send_text(json.dumps({
                     "v": 1,
                     "request_id": message.get("request_id"),
@@ -179,17 +198,22 @@ async def _coderai_broker_websocket_impl(websocket: WebSocket, scope_name: str):
                     "event": "heartbeat",
                     "payload": {"ts": int(time.time())},
                 }))
+                await broker.touch(session.session_id, metadata=message.get("payload") or {})
                 continue
-            await broker.touch(session.session_id)
-            await broker.publish_response(message)
+            # Hand response frames to the worker and immediately loop back to reading,
+            # so heartbeats/register are never queued behind response processing.
+            publish_queue.put_nowait(message)
     except WebSocketDisconnect:
         logger.info(f"CoderAI broker disconnected provider={provider_id} client={client_id}")
     except Exception as e:
         logger.error(f"CoderAI broker websocket error: {e}", exc_info=True)
     finally:
         queue_task.cancel()
+        publish_task.cancel()
         with contextlib.suppress(asyncio.CancelledError):
             await queue_task
+        with contextlib.suppress(asyncio.CancelledError):
+            await publish_task
         await broker.unregister(session.session_id)
         await broker.fail_session_requests(session.session_id, f"CoderAI session '{session.session_id}' disconnected")
 

aisbf/utils.py

@@ -22,40 +22,107 @@ Why did the programmer quit his job? Because he didn't get arrays!
 
 Utility functions for AISBF.
 """
+import hashlib
+import logging
+import threading
+from collections import OrderedDict
 from typing import Dict, List, Optional
 from langchain_text_splitters import TokenTextSplitter
 from .config import config
 
+logger = logging.getLogger(__name__)
+
+# tiktoken's BPE tokenizer is CPU-bound and was being run several times per request
+# on the *same* large message list (each pass re-encoding tens of KB), blocking the
+# asyncio event loop.  Two caches remove that cost:
+#   * _ENCODING_BY_MODEL avoids re-resolving/loading the encoding object per call.
+#   * _TOKEN_COUNT_CACHE memoizes the total for identical (model, messages) inputs so
+#     the repeated per-request calls collapse to a single BPE pass.
+# count_messages_tokens is a pure function of its inputs, so memoization is safe.
+_ENCODING_BY_MODEL: Dict[str, object] = {}
+_TOKEN_COUNT_CACHE: "OrderedDict[str, int]" = OrderedDict()
+_TOKEN_COUNT_CACHE_MAX = 1024
+_TOKEN_COUNT_LOCK = threading.Lock()
+
+
+def _get_encoding(model: str):
+    """Return a cached tiktoken encoding for the model (loaded at most once per name)."""
+    enc = _ENCODING_BY_MODEL.get(model)
+    if enc is not None:
+        return enc
+    import tiktoken
+    # OpenAI models use cl100k_base encoding
+    if model.startswith(('gpt-', 'text-', 'davinci-', 'ada-', 'babbage-', 'curie-')):
+        try:
+            enc = tiktoken.encoding_for_model(model)
+        except KeyError:
+            # Fallback to cl100k_base if model not found
+            enc = tiktoken.get_encoding("cl100k_base")
+    else:
+        # Default encoding for other models
+        enc = tiktoken.get_encoding("cl100k_base")
+    _ENCODING_BY_MODEL[model] = enc
+    return enc
+
+
+def _messages_cache_key(messages: List[Dict], model: str) -> str:
+    """Cheap content hash of (model, messages).  Hashing is O(n) but far cheaper than a
+    BPE encode pass, so a cache hit that skips re-encoding is a clear win."""
+    h = hashlib.blake2b(digest_size=16)
+    h.update(model.encode('utf-8', 'ignore'))
+    for msg in messages:
+        content = msg.get('content', '')
+        if not content:
+            h.update(b'\x00')
+            continue
+        if isinstance(content, str):
+            h.update(b'\x01')
+            h.update(content.encode('utf-8', 'ignore'))
+        elif isinstance(content, list):
+            h.update(b'\x02')
+            for item in content:
+                if isinstance(item, dict):
+                    text = item.get('text', '')
+                    if text:
+                        h.update(text.encode('utf-8', 'ignore'))
+                elif isinstance(item, str):
+                    h.update(item.encode('utf-8', 'ignore'))
+                h.update(b'\x1f')
+        h.update(b'\x1e')
+    return h.hexdigest()
+
 
 def count_messages_tokens(messages: List[Dict], model: str) -> int:
     """
     Count the total number of tokens in a list of messages.
-    
-    This function uses tiktoken for accurate token counting.
-    
+
+    This function uses tiktoken for accurate token counting.  Results are memoized on
+    (model, messages) content so the several calls made per request collapse to one
+    BPE pass and stop starving the event loop.
+
     Args:
         messages: List of message dictionaries with 'role' and 'content' keys
         model: Model name to determine encoding
-    
+
     Returns:
         Total token count across all messages
     """
-    import tiktoken
-    import logging
-    logger = logging.getLogger(__name__)
-    
-    # Select encoding based on model
-    # OpenAI models use cl100k_base encoding
-    if model.startswith(('gpt-', 'text-', 'davinci-', 'ada-', 'babbage-', 'curie-')):
-        try:
-            encoding = tiktoken.encoding_for_model(model)
-        except KeyError:
-            # Fallback to cl100k_base if model not found
-            encoding = tiktoken.get_encoding("cl100k_base")
-    else:
-        # Default encoding for other models
-        encoding = tiktoken.get_encoding("cl100k_base")
-    
+    cache_key = None
+    try:
+        cache_key = _messages_cache_key(messages, model)
+    except Exception:
+        # Never let cache-key construction break token counting.
+        cache_key = None
+
+    if cache_key is not None:
+        with _TOKEN_COUNT_LOCK:
+            cached = _TOKEN_COUNT_CACHE.get(cache_key)
+            if cached is not None:
+                _TOKEN_COUNT_CACHE.move_to_end(cache_key)
+                return cached
+
+    encoding = _get_encoding(model)
+
     total_tokens = 0
     for msg in messages:
         content = msg.get('content', '')
@@ -71,7 +138,14 @@ def count_messages_tokens(messages: List[Dict], model: str) -> int:
                             total_tokens += len(encoding.encode(text, disallowed_special=()))
                     elif isinstance(item, str):
                         total_tokens += len(encoding.encode(item, disallowed_special=()))
-    
+
+    if cache_key is not None:
+        with _TOKEN_COUNT_LOCK:
+            _TOKEN_COUNT_CACHE[cache_key] = total_tokens
+            _TOKEN_COUNT_CACHE.move_to_end(cache_key)
+            while len(_TOKEN_COUNT_CACHE) > _TOKEN_COUNT_CACHE_MAX:
+                _TOKEN_COUNT_CACHE.popitem(last=False)
+
     logger.debug(f"Token count for model {model}: {total_tokens}")
     return total_tokens