fix(proxy): kill retirements scheduled after the shutdown flush synchronously

An in-flight refresh triggered by the __getattr__ stale-token fallback
can complete its rotation after flush_engine_retirements has taken its
snapshot, scheduling a retirement task that no longer has an owner and
gets abandoned at event-loop teardown. Once the flush has run,
_schedule_engine_retirement now SIGKILLs the old engine inline instead
of creating a task. The event loop is single threaded and both the
flag-set-plus-snapshot and the schedule are synchronous, so every
schedule either lands in the flush snapshot or sees the flag; no
retirement can escape both
This commit is contained in:
ryan-crabbe-berri 2026-07-28 14:33:21 -07:00
parent 1219fde8a1
commit a637b32528
2 changed files with 20 additions and 0 deletions

View file

@ -229,6 +229,7 @@ class PrismaWrapper:
self._last_refresh_time: datetime | None = None
self._active_drain_tracker = self._instrument_prisma_client(original_prisma)
self._retirement_tasks: frozenset[_EngineRetirement] = frozenset()
self._retirements_flushed = False
# Coordination for planned engine restarts (issue #29176). Every
# `recreate_prisma_client` SIGTERMs the running query-engine on
@ -321,6 +322,7 @@ class PrismaWrapper:
)
async def flush_engine_retirements(self) -> None:
self._retirements_flushed = True
pending = self._retirement_tasks
if not pending:
return
@ -333,6 +335,9 @@ class PrismaWrapper:
def _schedule_engine_retirement(self, pid: int, tracker: _PrismaDrainTracker | None) -> None:
if pid <= 0:
return
if self._retirements_flushed:
self._kill_engine_process_now(pid)
return
retirement_task = asyncio.create_task(self._retire_engine_when_drained(pid, tracker))
self._retirement_tasks = self._retirement_tasks.union((_EngineRetirement(pid=pid, task=retirement_task),))
retirement_task.add_done_callback(self._retirement_finished)

View file

@ -443,6 +443,21 @@ async def test_stop_token_refresh_flushes_pending_engine_retirement(mock_prisma_
assert wrapper._retirement_tasks == frozenset()
@pytest.mark.asyncio
async def test_retirement_scheduled_after_flush_kills_engine_synchronously(mock_prisma_binary):
"""A rotation that completes during shutdown (e.g. an in-flight __getattr__
refresh) must not create an unowned retirement task after the flush; the
old engine is killed inline instead."""
wrapper = _make_wrapper(engine_pid=111, iam=True)
await wrapper.stop_token_refresh_task()
with patch("os.kill") as mock_kill:
wrapper._schedule_engine_retirement(333, wrapper._active_drain_tracker)
assert mock_kill.call_args_list == [call(333, signal.SIGKILL)]
assert wrapper._retirement_tasks == frozenset()
@pytest.mark.asyncio
async def test_safe_refresh_cancellation_restores_token_and_cleans_replacement(
mock_prisma_binary, monkeypatch