95 lines
3.3 KiB
Python
95 lines
3.3 KiB
Python
"""
|
||
app/tenancy/engine_cache.py
|
||
---------------------------
|
||
Process-local cache of per-tenant SQLAlchemy engines, keyed by tenant id.
|
||
|
||
Each tenant has its own database, hence its own engine + connection pool.
|
||
Engines are created lazily on first use and reused across requests. Total
|
||
backend connections ≈ workers × cached-tenants × (pool_size + max_overflow),
|
||
so the cache is BOUNDED (MT-21): it holds at most
|
||
``TENANT_ENGINE_CACHE_MAX`` engines and evicts the least-recently-used one
|
||
beyond that, disposing it. Without the bound, a process that has served N
|
||
tenants holds N pools forever, and the connection count grows without limit
|
||
until MySQL's ``max_connections`` (default 151) rejects new sessions.
|
||
|
||
Worked example — 8 workers, pool_size 2 + max_overflow 3, cache cap 32:
|
||
8 × 32 × 5 = 1280 worst case, vs. unbounded before this change.
|
||
|
||
Eviction disposes the engine, which closes its *idle* pooled connections.
|
||
Connections already checked out by another thread stay valid and are closed
|
||
when returned, so eviction is safe under concurrency; the evicted tenant
|
||
simply rebuilds its engine on the next request.
|
||
|
||
`invalidate(tenant_id)` drops a cached engine (e.g. after credential rotation
|
||
or tenant suspension); the next request rebuilds it.
|
||
"""
|
||
|
||
import threading
|
||
from collections import OrderedDict
|
||
|
||
from flask import current_app
|
||
from sqlalchemy import create_engine
|
||
|
||
_engines = OrderedDict()
|
||
_lock = threading.Lock()
|
||
|
||
# Fallback cap used when no application context is available (CLI/cron paths
|
||
# that touch the cache outside a request). Mirrors the config default.
|
||
_DEFAULT_CACHE_MAX = 32
|
||
|
||
|
||
def _cache_max():
|
||
try:
|
||
return int(current_app.config.get('TENANT_ENGINE_CACHE_MAX', _DEFAULT_CACHE_MAX))
|
||
except Exception:
|
||
return _DEFAULT_CACHE_MAX
|
||
|
||
|
||
def get_tenant_engine(tenant):
|
||
"""Return (building if needed) the cached engine for a TenantContext."""
|
||
evicted = []
|
||
with _lock:
|
||
engine = _engines.get(tenant.id)
|
||
if engine is None:
|
||
engine = create_engine(
|
||
tenant.db_uri,
|
||
pool_pre_ping=True,
|
||
pool_size=current_app.config.get('TENANT_ENGINE_POOL_SIZE', 2),
|
||
max_overflow=current_app.config.get('TENANT_ENGINE_MAX_OVERFLOW', 3),
|
||
pool_recycle=current_app.config.get('TENANT_ENGINE_POOL_RECYCLE', 1800),
|
||
future=True,
|
||
)
|
||
_engines[tenant.id] = engine
|
||
# Bound the cache: drop least-recently-used engines beyond the cap.
|
||
cap = _cache_max()
|
||
while cap > 0 and len(_engines) > cap:
|
||
old_id, old_engine = _engines.popitem(last=False)
|
||
evicted.append((old_id, old_engine))
|
||
else:
|
||
_engines.move_to_end(tenant.id)
|
||
|
||
# Dispose outside the lock — dispose() can block on socket teardown.
|
||
for old_id, old_engine in evicted:
|
||
try:
|
||
old_engine.dispose()
|
||
except Exception:
|
||
pass
|
||
return engine
|
||
|
||
|
||
def invalidate(tenant_id):
|
||
"""Drop and dispose a cached tenant engine, if present."""
|
||
with _lock:
|
||
engine = _engines.pop(tenant_id, None)
|
||
if engine is not None:
|
||
engine.dispose()
|
||
|
||
|
||
def clear():
|
||
"""Dispose and drop all cached engines (test/teardown helper)."""
|
||
with _lock:
|
||
engines = list(_engines.values())
|
||
_engines.clear()
|
||
for engine in engines:
|
||
engine.dispose()
|