fix(proxy): propagate CancelledError from memory init cleanly

External cancellation while _ensure_initialized is in flight raises
CancelledError (BaseException). The existing 'except asyncio.TimeoutError:'
branch does not catch it, and caller 'except Exception:' blocks don't
either, so it propagated as-is with _backend possibly still assigned.

Add an explicit 'except asyncio.CancelledError:' handler that nulls
_backend, clears _initialized, logs at info, and re-raises. Cancellation
is a shutdown signal, not an error to swallow — but the state must be
clean so a later retry starts fresh.
This commit is contained in:
Adryan Eka Vandra 2026-04-18 01:58:05 +07:00
parent 3f1f82d7e2
commit 196d8d69a7
No known key found for this signature in database
GPG key ID: A46A577A26A97682
2 changed files with 53 additions and 0 deletions

View file

@ -198,6 +198,20 @@ class MemoryHandler:
"Subsequent requests will retry."
)
return
except asyncio.CancelledError:
# External cancellation (shutdown / task cancelled).
# CancelledError is BaseException — the TimeoutError branch
# above does NOT catch it, and caller ``except Exception``
# blocks don't either, so it propagates unconditionally.
# Reset state so any later retry starts clean, then re-raise:
# cancellation is a signal, not an error to swallow.
self._backend = None
self._initialized = False
logger.info(
"Memory: backend initialization cancelled "
f"(backend={self.config.backend})"
)
raise
async def _init_backend_locked(self) -> None:
"""Actual backend-init body. Must be called with ``_init_lock`` held."""

View file

@ -190,6 +190,45 @@ async def test_ensure_initialized_timeout_nulls_partially_initialized_backend(
assert handler._backend is None
@pytest.mark.asyncio
async def test_ensure_initialized_cancellation_propagates_and_resets_state(
tmp_path, monkeypatch
):
"""External cancellation of an in-flight ``_ensure_initialized`` must
propagate (CancelledError is BaseException not a swallowable error)
and leave the handler in a clean state."""
class HangingBackend:
def __init__(self, config):
self.config = config
async def _ensure_initialized(self) -> None:
await asyncio.Event().wait()
async def close(self) -> None:
pass
import headroom.memory.backends.local as local_mod
monkeypatch.setattr(local_mod, "LocalBackend", HangingBackend)
handler = MemoryHandler(
MemoryConfig(
enabled=True, backend="local", db_path=str(tmp_path / "mem.db")
)
)
task = asyncio.create_task(handler._ensure_initialized())
# Give the task a tick to enter _init_backend_locked and assign _backend.
await asyncio.sleep(0.05)
task.cancel()
with pytest.raises(asyncio.CancelledError):
await task
assert handler._initialized is False
assert handler._backend is None
# -------------------------------------------------------------------
# Real backend init (no monkeypatching) — integration smoke test
# -------------------------------------------------------------------