Bound remote pool and rate-limiter memory
Both tables grew one entry per distinct key and never shrank, so a long-running SDK client or a public serve process accumulated state for every endpoint or identity it had ever seen. GC cannot reclaim them while the pool and limiter still reference them. Cap the pool at 64 endpoint buckets and the limiter at 4096 buckets. Both evict useless state first: connections past the idle timeout the server has likely dropped anyway, and buckets that have fully refilled, which carry no throttling information. Only then fall back to evicting the oldest entry. Evicting a limiter bucket resets throttling for that identity, which is the deliberate trade: an attacker cycling identities faster than they go idle can regain tokens, but unbounded growth would take the process down instead.
This commit is contained in:
@@ -26,6 +26,7 @@ from browser_cli.framing import frame
|
||||
# hand back one the server has just timed out and closed.
|
||||
_MAX_IDLE_SECONDS = max(5, REMOTE_SESSION_IDLE_TIMEOUT - 5)
|
||||
_MAX_PER_ENDPOINT = 8
|
||||
_MAX_ENDPOINTS = 64
|
||||
|
||||
class PooledConnection:
|
||||
__slots__ = ("sock", "secret", "last_used")
|
||||
@@ -56,10 +57,32 @@ def checkout(endpoint: str) -> PooledConnection | None:
|
||||
_close(conn.sock) # too old — assume the server has dropped it
|
||||
return None
|
||||
|
||||
def _prune_endpoints_locked(now: float) -> None:
|
||||
"""Keep the number of endpoint buckets bounded for long-running SDK users."""
|
||||
for endpoint, bucket in list(_POOL.items()):
|
||||
fresh = [conn for conn in bucket if now - conn.last_used <= _MAX_IDLE_SECONDS]
|
||||
if fresh:
|
||||
_POOL[endpoint] = fresh
|
||||
else:
|
||||
for conn in bucket:
|
||||
_close(conn.sock)
|
||||
_POOL.pop(endpoint, None)
|
||||
|
||||
while len(_POOL) >= _MAX_ENDPOINTS:
|
||||
oldest_endpoint, bucket = min(
|
||||
_POOL.items(),
|
||||
key=lambda item: min(conn.last_used for conn in item[1]) if item[1] else 0.0,
|
||||
)
|
||||
for conn in bucket:
|
||||
_close(conn.sock)
|
||||
_POOL.pop(oldest_endpoint, None)
|
||||
|
||||
def checkin(endpoint: str, conn: PooledConnection) -> None:
|
||||
"""Return a still-healthy connection to the pool for reuse."""
|
||||
conn.last_used = time.monotonic()
|
||||
with _LOCK:
|
||||
if endpoint not in _POOL and len(_POOL) >= _MAX_ENDPOINTS:
|
||||
_prune_endpoints_locked(conn.last_used)
|
||||
bucket = _POOL.setdefault(endpoint, [])
|
||||
if len(bucket) >= _MAX_PER_ENDPOINT:
|
||||
_close(conn.sock)
|
||||
|
||||
Reference in New Issue
Block a user