fix(#5543): fire-and-forget commit_session_memory + drain-safe managed stop
Move the 1-5s commit_session_memory() extraction off the POST /api/session/new request thread (was blocking + New Chat) into a fire-and-forget daemon thread, guarded by the existing per-session in-flight serialization. Data-loss prevention on managed stop: - server.py installs a SIGTERM handler that requests an orderly stop so serve_forever()'s finally block runs drain_all_on_shutdown() (the default SIGTERM disposition terminated without unwinding, losing in-flight commits). - session_lifecycle: bounded self-unregistering background-commit registry; a _draining flag so no NEW worker starts once the drain begins (the inline generation-drain commits any late arrival instead); a 30s overall drain deadline so a stuck worker cannot hang the stop path forever. Co-authored-by: luperrypf <luperrypf@users.noreply.github.com>
This commit is contained in:
1 parent
04717ad305
commit
19d00e3586
4 files changed
+539
-17
No files matched your search
+43
-15
@@ -13067,21 +13067,49 @@ def handle_post(handler, parsed) -> bool:
|
|||||||
# the new session (#5420).
|
# the new session (#5420).
|
||||||
prev_session_id = None
|
prev_session_id = None
|
||||||
if prev_session_id:
|
if prev_session_id:
|
||||||
try:
|
# Fire-and-forget: commit_memory_session() can take 1-5+ seconds
|
||||||
from api.session_lifecycle import commit_session_memory
|
# (extraction call to the memory provider), and blocking the
|
||||||
from api.config import SESSION_AGENT_CACHE, SESSION_AGENT_CACHE_LOCK
|
# response here made "+ New Chat" feel slow/unresponsive.
|
||||||
prev_agent = None
|
# commit_session_memory() already serialises overlapping commits
|
||||||
with SESSION_AGENT_CACHE_LOCK:
|
# for a session via its own in-flight guard, so running it off
|
||||||
_cached = SESSION_AGENT_CACHE.get(prev_session_id)
|
# the request thread is safe.
|
||||||
if _cached:
|
def _commit_prev_session_memory(_sid=prev_session_id):
|
||||||
prev_agent = _cached[0]
|
try:
|
||||||
commit_session_memory(prev_session_id, agent=prev_agent)
|
from api.session_lifecycle import commit_session_memory
|
||||||
except Exception:
|
from api.config import SESSION_AGENT_CACHE, SESSION_AGENT_CACHE_LOCK
|
||||||
logger.debug(
|
prev_agent = None
|
||||||
"Lifecycle commit for prev_session %s failed",
|
with SESSION_AGENT_CACHE_LOCK:
|
||||||
prev_session_id,
|
_cached = SESSION_AGENT_CACHE.get(_sid)
|
||||||
exc_info=True,
|
if _cached:
|
||||||
)
|
prev_agent = _cached[0]
|
||||||
|
commit_session_memory(_sid, agent=prev_agent)
|
||||||
|
except Exception:
|
||||||
|
logger.warning(
|
||||||
|
"Lifecycle commit for prev_session %s failed",
|
||||||
|
_sid,
|
||||||
|
exc_info=True,
|
||||||
|
)
|
||||||
|
finally:
|
||||||
|
# Self-unregister so the background-commit registry does
|
||||||
|
# not leak completed threads; drain only tracks live ones.
|
||||||
|
try:
|
||||||
|
from api.session_lifecycle import _unregister_background_commit_thread
|
||||||
|
_unregister_background_commit_thread(threading.current_thread())
|
||||||
|
except Exception:
|
||||||
|
pass
|
||||||
|
|
||||||
|
t = threading.Thread(
|
||||||
|
target=_commit_prev_session_memory,
|
||||||
|
daemon=True,
|
||||||
|
name=f"commit-memory-{prev_session_id}",
|
||||||
|
)
|
||||||
|
from api.session_lifecycle import _register_background_commit_thread
|
||||||
|
# Refused only if shutdown draining has already begun; in that
|
||||||
|
# window the inline drain commits the pending generation instead,
|
||||||
|
# so skipping the worker start is safe (avoids a late daemon
|
||||||
|
# thread the drain snapshot already missed).
|
||||||
|
if _register_background_commit_thread(t):
|
||||||
|
t.start()
|
||||||
s = new_session(
|
s = new_session(
|
||||||
workspace=workspace,
|
workspace=workspace,
|
||||||
model=model,
|
model=model,
|
||||||
|
|||||||
+108
-2
@@ -45,6 +45,56 @@ _condition = threading.Condition(_lock)
|
|||||||
|
|
||||||
_sessions: dict[str, dict] = {}
|
_sessions: dict[str, dict] = {}
|
||||||
|
|
||||||
|
# Background commit threads spawned by fire-and-forget paths (e.g. POST /api/session/new).
|
||||||
|
# Tracked so drain_all_on_shutdown() can join them before interpreter teardown,
|
||||||
|
# preventing silent data loss when daemon threads are abandoned mid-commit.
|
||||||
|
_background_commit_threads: set[threading.Thread] = set()
|
||||||
|
_background_commit_threads_lock = threading.Lock()
|
||||||
|
# Set once drain_all_on_shutdown() begins so no NEW background worker is started
|
||||||
|
# after the drain has snapshotted the registry. A late arrival is not lost: its
|
||||||
|
# uncommitted generation is still committed inline by the generation-drain loop
|
||||||
|
# below (which re-snapshots _sessions each pass and commits with wait=True).
|
||||||
|
_draining = False
|
||||||
|
|
||||||
|
|
||||||
|
def _register_background_commit_thread(t: threading.Thread) -> bool:
|
||||||
|
"""Register a fire-and-forget commit thread. Returns True if the caller
|
||||||
|
should start it. Returns False when shutdown draining has already begun —
|
||||||
|
the caller must NOT start a new worker in that window; the inline
|
||||||
|
generation-drain will commit the pending work instead."""
|
||||||
|
with _background_commit_threads_lock:
|
||||||
|
if _draining:
|
||||||
|
return False
|
||||||
|
_background_commit_threads.add(t)
|
||||||
|
return True
|
||||||
|
|
||||||
|
|
||||||
|
def _unregister_background_commit_thread(t: threading.Thread) -> None:
|
||||||
|
# A completed worker removes itself so the registry does not grow without
|
||||||
|
# bound over the process lifetime; drain only needs threads still running.
|
||||||
|
with _background_commit_threads_lock:
|
||||||
|
_background_commit_threads.discard(t)
|
||||||
|
|
||||||
|
|
||||||
|
def _drain_background_commit_threads(timeout: float = 5.0, deadline_fn=None) -> None:
|
||||||
|
"""Join tracked fire-and-forget commit threads. When ``deadline_fn`` is given
|
||||||
|
(returns the remaining overall shutdown budget in seconds), each join is
|
||||||
|
clamped to that budget and the phase stops early once it is exhausted, so the
|
||||||
|
join phase cannot overrun ``drain_all_on_shutdown``'s deadline."""
|
||||||
|
with _background_commit_threads_lock:
|
||||||
|
threads = list(_background_commit_threads)
|
||||||
|
_background_commit_threads.clear()
|
||||||
|
for t in threads:
|
||||||
|
if not t.is_alive():
|
||||||
|
continue
|
||||||
|
if deadline_fn is None:
|
||||||
|
t.join(timeout)
|
||||||
|
continue
|
||||||
|
remaining = deadline_fn()
|
||||||
|
if remaining <= 0:
|
||||||
|
return
|
||||||
|
t.join(min(timeout, remaining))
|
||||||
|
|
||||||
|
|
||||||
def _new_entry() -> dict:
|
def _new_entry() -> dict:
|
||||||
return {
|
return {
|
||||||
@@ -225,16 +275,72 @@ def commit_session_memory(session_id: str, agent=None, *, wait: bool = False, ti
|
|||||||
return True
|
return True
|
||||||
|
|
||||||
|
|
||||||
def drain_all_on_shutdown() -> None:
|
def drain_all_on_shutdown(deadline: float = 30.0) -> None:
|
||||||
|
# Signal that draining has begun so no NEW background commit worker is
|
||||||
|
# started after this point (see _register_background_commit_thread). Any
|
||||||
|
# already-registered worker is joined below; a caller that loses the race
|
||||||
|
# (registered after our snapshot, or refused a start) still has its pending
|
||||||
|
# generation committed by the inline loop.
|
||||||
|
global _draining
|
||||||
|
with _background_commit_threads_lock:
|
||||||
|
_draining = True
|
||||||
|
|
||||||
|
# Overall wall-clock budget so a stuck worker holding in_flight cannot hang
|
||||||
|
# the shutdown path forever. Start the clock BEFORE the join phase so the
|
||||||
|
# join and the inline drain share one budget. The inline commit runs in a
|
||||||
|
# daemon worker joined only for the remaining budget — otherwise a provider
|
||||||
|
# wedged inside commit_memory_session() (whether we are waiting on an
|
||||||
|
# existing in_flight commit, OR the inline drain owns the segment and calls
|
||||||
|
# the provider directly) would block indefinitely and the deadline check
|
||||||
|
# would never run again.
|
||||||
|
started = time.monotonic()
|
||||||
|
|
||||||
|
def _remaining() -> float:
|
||||||
|
return deadline - (time.monotonic() - started)
|
||||||
|
|
||||||
|
def _commit_within_budget(session_id: str, budget: float) -> bool:
|
||||||
|
"""Run one commit in a daemon worker, join at most `budget` seconds.
|
||||||
|
Returns True only if the commit completed within budget AND made
|
||||||
|
progress. On timeout the daemon worker is left running (it will finish
|
||||||
|
or die at interpreter exit; either way it mutates generation state only
|
||||||
|
under _condition, so nothing is corrupted) and we report no progress so
|
||||||
|
the drain can honor its deadline instead of hanging on a wedged
|
||||||
|
provider."""
|
||||||
|
holder = {"ok": False}
|
||||||
|
|
||||||
|
def _run():
|
||||||
|
holder["ok"] = commit_session_memory(session_id, wait=True)
|
||||||
|
|
||||||
|
worker = threading.Thread(target=_run, daemon=True, name=f"drain-commit-{session_id}")
|
||||||
|
worker.start()
|
||||||
|
worker.join(budget)
|
||||||
|
return holder["ok"] if not worker.is_alive() else False
|
||||||
|
|
||||||
|
# Drain in-flight background commit threads first (fire-and-forget from
|
||||||
|
# POST /api/session/new) so their commit work completes before we drain any
|
||||||
|
# remaining uncommitted generations inline. Each join consumes the shared
|
||||||
|
# budget so it cannot overrun the overall deadline.
|
||||||
|
_drain_background_commit_threads(deadline_fn=_remaining)
|
||||||
|
|
||||||
while True:
|
while True:
|
||||||
with _lock:
|
with _lock:
|
||||||
snapshot = [sid for sid, entry in _sessions.items() if entry["generation"] > entry["committed_generation"]]
|
snapshot = [sid for sid, entry in _sessions.items() if entry["generation"] > entry["committed_generation"]]
|
||||||
if not snapshot:
|
if not snapshot:
|
||||||
return
|
return
|
||||||
|
remaining = _remaining()
|
||||||
|
if remaining <= 0:
|
||||||
|
logger.warning("drain_all_on_shutdown: deadline hit with uncommitted sessions: %s", sorted(snapshot))
|
||||||
|
return
|
||||||
|
|
||||||
made_progress = False
|
made_progress = False
|
||||||
for sid in snapshot:
|
for sid in snapshot:
|
||||||
if commit_session_memory(sid, wait=True):
|
budget = _remaining()
|
||||||
|
if budget <= 0:
|
||||||
|
logger.warning("drain_all_on_shutdown: deadline hit with uncommitted sessions: %s", sorted(snapshot))
|
||||||
|
return
|
||||||
|
# Budget-bounded (covers both the wait-for-in_flight path AND the
|
||||||
|
# unbounded provider call when the inline drain owns the segment).
|
||||||
|
if _commit_within_budget(sid, budget):
|
||||||
made_progress = True
|
made_progress = True
|
||||||
if not made_progress:
|
if not made_progress:
|
||||||
logger.debug("drain_all_on_shutdown: stopped with uncommitted sessions: %s", sorted(snapshot))
|
logger.debug("drain_all_on_shutdown: stopped with uncommitted sessions: %s", sorted(snapshot))
|
||||||
|
|||||||
@@ -688,6 +688,36 @@ def main() -> None:
|
|||||||
print(f' Remote access: ssh -N -L {PORT}:127.0.0.1:{PORT} <user>@<your-server>', flush=True)
|
print(f' Remote access: ssh -N -L {PORT}:127.0.0.1:{PORT} <user>@<your-server>', flush=True)
|
||||||
print(f' Then open: {scheme}://localhost:{PORT}', flush=True)
|
print(f' Then open: {scheme}://localhost:{PORT}', flush=True)
|
||||||
print('', flush=True)
|
print('', flush=True)
|
||||||
|
|
||||||
|
# ctl.sh stops the WebUI with SIGTERM. Python's default SIGTERM handler
|
||||||
|
# terminates the process WITHOUT unwinding the try/finally around
|
||||||
|
# serve_forever(), so drain_all_on_shutdown() (which flushes in-flight
|
||||||
|
# fire-and-forget memory commits) would never run on the normal managed
|
||||||
|
# stop. Install a handler that requests an orderly shutdown so
|
||||||
|
# serve_forever() returns and the existing `finally` block drains cleanly.
|
||||||
|
#
|
||||||
|
# httpd.shutdown() blocks until serve_forever() has exited and MUST NOT be
|
||||||
|
# called from the thread running serve_forever() (it would deadlock), so we
|
||||||
|
# dispatch it from a short-lived helper thread. The handler is idempotent
|
||||||
|
# and guards against double-shutdown (e.g. repeated SIGTERM/SIGINT).
|
||||||
|
_shutdown_requested = threading.Event()
|
||||||
|
|
||||||
|
def _request_shutdown(signum, _frame):
|
||||||
|
if _shutdown_requested.is_set():
|
||||||
|
return
|
||||||
|
_shutdown_requested.set()
|
||||||
|
threading.Thread(
|
||||||
|
target=httpd.shutdown,
|
||||||
|
name="webui-sigterm-shutdown",
|
||||||
|
daemon=True,
|
||||||
|
).start()
|
||||||
|
|
||||||
|
try:
|
||||||
|
signal.signal(signal.SIGTERM, _request_shutdown)
|
||||||
|
except (ValueError, OSError):
|
||||||
|
# Not on the main thread (e.g. embedded/test harness); skip handler.
|
||||||
|
logger.debug("Could not install SIGTERM handler", exc_info=True)
|
||||||
|
|
||||||
try:
|
try:
|
||||||
httpd.serve_forever()
|
httpd.serve_forever()
|
||||||
finally:
|
finally:
|
||||||
|
|||||||
@@ -0,0 +1,358 @@
|
|||||||
|
"""Tests for issue #5543: fire-and-forget commit_session_memory + shutdown drain.
|
||||||
|
|
||||||
|
PR #5543 (@luperrypf) moves commit_session_memory() off the POST
|
||||||
|
/api/session/new request thread into a daemon "background commit" thread so
|
||||||
|
"+ New Chat" returns immediately instead of blocking 1-5s on memory-provider
|
||||||
|
extraction. The gate-certifier required three guarantees this file locks in:
|
||||||
|
|
||||||
|
1. Fire-and-forget: spawning the background commit does not block the caller,
|
||||||
|
even while the underlying commit is still running.
|
||||||
|
2. Registry hygiene: a completed background-commit worker unregisters itself so
|
||||||
|
``_background_commit_threads`` does not grow without bound (Finding 2).
|
||||||
|
3. Shutdown drain: drain_all_on_shutdown() joins in-flight background commit
|
||||||
|
threads AND flushes the pending generation, and the server's SIGTERM handler
|
||||||
|
routes the managed ``ctl.sh stop`` through the serve_forever() ``finally`` so
|
||||||
|
the drain actually runs (Finding 1).
|
||||||
|
"""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import ast
|
||||||
|
import importlib
|
||||||
|
import socket
|
||||||
|
import threading
|
||||||
|
import time
|
||||||
|
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
|
||||||
|
REPO = Path(__file__).resolve().parents[1]
|
||||||
|
|
||||||
|
|
||||||
|
def _fresh_lifecycle():
|
||||||
|
"""Import/reload lifecycle module and clear process-global test state."""
|
||||||
|
lifecycle = importlib.import_module("api.session_lifecycle")
|
||||||
|
lifecycle = importlib.reload(lifecycle)
|
||||||
|
reset = getattr(lifecycle, "_reset_for_tests", None)
|
||||||
|
if callable(reset):
|
||||||
|
reset()
|
||||||
|
# Ensure the background-commit registry starts empty for each test.
|
||||||
|
with lifecycle._background_commit_threads_lock:
|
||||||
|
lifecycle._background_commit_threads.clear()
|
||||||
|
return lifecycle
|
||||||
|
|
||||||
|
|
||||||
|
class RecordingAgent:
|
||||||
|
"""Agent whose commit blocks on an event so we can hold it "in flight"."""
|
||||||
|
|
||||||
|
def __init__(self):
|
||||||
|
self.calls = 0
|
||||||
|
self.entered = threading.Event()
|
||||||
|
self.release = threading.Event()
|
||||||
|
|
||||||
|
def commit_memory_session(self):
|
||||||
|
self.calls += 1
|
||||||
|
self.entered.set()
|
||||||
|
# Block until the test releases us (bounded so a buggy test can't hang).
|
||||||
|
self.release.wait(timeout=5)
|
||||||
|
|
||||||
|
|
||||||
|
def _make_background_worker(lifecycle, sid, agent):
|
||||||
|
"""Mirror the routes.py fire-and-forget worker: commit then self-unregister."""
|
||||||
|
|
||||||
|
def _worker():
|
||||||
|
try:
|
||||||
|
lifecycle.commit_session_memory(sid, agent=agent)
|
||||||
|
finally:
|
||||||
|
lifecycle._unregister_background_commit_thread(threading.current_thread())
|
||||||
|
|
||||||
|
t = threading.Thread(target=_worker, daemon=True, name=f"commit-memory-{sid}")
|
||||||
|
lifecycle._register_background_commit_thread(t)
|
||||||
|
return t
|
||||||
|
|
||||||
|
|
||||||
|
# --------------------------------------------------------------------------- #
|
||||||
|
# Finding 2: registry hygiene #
|
||||||
|
# --------------------------------------------------------------------------- #
|
||||||
|
|
||||||
|
|
||||||
|
def test_register_and_unregister_roundtrip():
|
||||||
|
lifecycle = _fresh_lifecycle()
|
||||||
|
t = threading.Thread(target=lambda: None)
|
||||||
|
|
||||||
|
# Registration returns True (caller should start) when not draining.
|
||||||
|
assert lifecycle._register_background_commit_thread(t) is True
|
||||||
|
with lifecycle._background_commit_threads_lock:
|
||||||
|
assert t in lifecycle._background_commit_threads
|
||||||
|
|
||||||
|
lifecycle._unregister_background_commit_thread(t)
|
||||||
|
with lifecycle._background_commit_threads_lock:
|
||||||
|
assert t not in lifecycle._background_commit_threads
|
||||||
|
|
||||||
|
# Idempotent: unregistering an absent thread is a harmless no-op.
|
||||||
|
lifecycle._unregister_background_commit_thread(t)
|
||||||
|
|
||||||
|
|
||||||
|
def test_register_refused_once_draining():
|
||||||
|
"""After drain begins, registration is refused (returns False) and the
|
||||||
|
thread is NOT added — the caller must skip start(); the inline generation
|
||||||
|
drain commits the pending work instead. This closes Codex's
|
||||||
|
register-before-start shutdown window."""
|
||||||
|
lifecycle = _fresh_lifecycle()
|
||||||
|
# Nothing pending, so drain_all_on_shutdown returns quickly but still sets
|
||||||
|
# the _draining flag.
|
||||||
|
lifecycle.drain_all_on_shutdown()
|
||||||
|
|
||||||
|
t = threading.Thread(target=lambda: None)
|
||||||
|
assert lifecycle._register_background_commit_thread(t) is False
|
||||||
|
with lifecycle._background_commit_threads_lock:
|
||||||
|
assert t not in lifecycle._background_commit_threads
|
||||||
|
|
||||||
|
|
||||||
|
def test_completed_worker_self_unregisters():
|
||||||
|
"""A finished background-commit worker must drop itself from the registry."""
|
||||||
|
lifecycle = _fresh_lifecycle()
|
||||||
|
agent = RecordingAgent()
|
||||||
|
agent.release.set() # commit completes immediately
|
||||||
|
sid = "self-unregister"
|
||||||
|
lifecycle.register_agent(sid, agent)
|
||||||
|
lifecycle.mark_turn_completed(sid, agent=agent)
|
||||||
|
|
||||||
|
t = _make_background_worker(lifecycle, sid, agent)
|
||||||
|
t.start()
|
||||||
|
t.join(timeout=5)
|
||||||
|
assert not t.is_alive()
|
||||||
|
|
||||||
|
with lifecycle._background_commit_threads_lock:
|
||||||
|
assert t not in lifecycle._background_commit_threads
|
||||||
|
assert len(lifecycle._background_commit_threads) == 0
|
||||||
|
assert agent.calls == 1
|
||||||
|
|
||||||
|
|
||||||
|
# --------------------------------------------------------------------------- #
|
||||||
|
# Fire-and-forget: the caller does not block #
|
||||||
|
# --------------------------------------------------------------------------- #
|
||||||
|
|
||||||
|
|
||||||
|
def test_fire_and_forget_does_not_block_caller():
|
||||||
|
"""Starting the background commit returns immediately, even mid-commit."""
|
||||||
|
lifecycle = _fresh_lifecycle()
|
||||||
|
agent = RecordingAgent() # blocks in commit until released
|
||||||
|
sid = "non-blocking"
|
||||||
|
lifecycle.register_agent(sid, agent)
|
||||||
|
lifecycle.mark_turn_completed(sid, agent=agent)
|
||||||
|
|
||||||
|
t = _make_background_worker(lifecycle, sid, agent)
|
||||||
|
|
||||||
|
start = time.monotonic()
|
||||||
|
t.start()
|
||||||
|
# Caller "returns" as soon as the thread is started; measure that dispatch.
|
||||||
|
elapsed = time.monotonic() - start
|
||||||
|
|
||||||
|
# The commit is genuinely in flight (proves we didn't run it inline).
|
||||||
|
assert agent.entered.wait(timeout=2)
|
||||||
|
assert elapsed < 1.0, f"caller blocked for {elapsed:.3f}s"
|
||||||
|
assert lifecycle.has_uncommitted_work(sid) is True
|
||||||
|
|
||||||
|
# Let the in-flight commit finish and clean up.
|
||||||
|
agent.release.set()
|
||||||
|
t.join(timeout=5)
|
||||||
|
|
||||||
|
|
||||||
|
# --------------------------------------------------------------------------- #
|
||||||
|
# Finding 1: shutdown drain flushes the pending commit #
|
||||||
|
# --------------------------------------------------------------------------- #
|
||||||
|
|
||||||
|
|
||||||
|
def test_drain_joins_inflight_background_commit_and_flushes():
|
||||||
|
"""drain_all_on_shutdown() waits for the in-flight worker, then commits."""
|
||||||
|
lifecycle = _fresh_lifecycle()
|
||||||
|
agent = RecordingAgent() # held in flight until we release it
|
||||||
|
sid = "drain-inflight"
|
||||||
|
lifecycle.register_agent(sid, agent)
|
||||||
|
lifecycle.mark_turn_completed(sid, agent=agent)
|
||||||
|
|
||||||
|
t = _make_background_worker(lifecycle, sid, agent)
|
||||||
|
t.start()
|
||||||
|
assert agent.entered.wait(timeout=2) # commit is in flight
|
||||||
|
|
||||||
|
drained = threading.Event()
|
||||||
|
|
||||||
|
def _do_drain():
|
||||||
|
lifecycle.drain_all_on_shutdown()
|
||||||
|
drained.set()
|
||||||
|
|
||||||
|
dt = threading.Thread(target=_do_drain, daemon=True)
|
||||||
|
dt.start()
|
||||||
|
|
||||||
|
# While the background commit is still in flight, drain must be blocked
|
||||||
|
# joining it — it cannot have finished yet.
|
||||||
|
time.sleep(0.1)
|
||||||
|
assert not drained.is_set()
|
||||||
|
|
||||||
|
# Release the commit; drain should now join the worker and finish.
|
||||||
|
agent.release.set()
|
||||||
|
assert drained.wait(timeout=5)
|
||||||
|
dt.join(timeout=2)
|
||||||
|
|
||||||
|
assert agent.calls == 1
|
||||||
|
assert lifecycle.has_uncommitted_work(sid) is False
|
||||||
|
with lifecycle._background_commit_threads_lock:
|
||||||
|
assert t not in lifecycle._background_commit_threads
|
||||||
|
|
||||||
|
|
||||||
|
def test_sigterm_style_shutdown_runs_finally_drain():
|
||||||
|
"""The SIGTERM handler mechanism drives serve_forever() through its finally.
|
||||||
|
|
||||||
|
Replicates server.py's approach: a signal handler requests shutdown by
|
||||||
|
dispatching httpd.shutdown() from a short-lived helper thread (you cannot
|
||||||
|
call it from serve_forever()'s own thread). serve_forever() returns, the
|
||||||
|
``finally`` runs, and drain_all_on_shutdown() flushes the pending commit.
|
||||||
|
"""
|
||||||
|
lifecycle = _fresh_lifecycle()
|
||||||
|
agent = RecordingAgent()
|
||||||
|
agent.release.set() # commit completes immediately during drain
|
||||||
|
sid = "sigterm-drain"
|
||||||
|
lifecycle.register_agent(sid, agent)
|
||||||
|
lifecycle.mark_turn_completed(sid, agent=agent)
|
||||||
|
|
||||||
|
class _Handler(BaseHTTPRequestHandler):
|
||||||
|
def log_message(self, *args): # silence test noise
|
||||||
|
pass
|
||||||
|
|
||||||
|
httpd = ThreadingHTTPServer(("127.0.0.1", 0), _Handler)
|
||||||
|
finally_ran = threading.Event()
|
||||||
|
|
||||||
|
def _serve():
|
||||||
|
try:
|
||||||
|
httpd.serve_forever()
|
||||||
|
finally:
|
||||||
|
httpd.server_close()
|
||||||
|
lifecycle.drain_all_on_shutdown()
|
||||||
|
finally_ran.set()
|
||||||
|
|
||||||
|
serve_thread = threading.Thread(target=_serve, daemon=True)
|
||||||
|
serve_thread.start()
|
||||||
|
|
||||||
|
# Wait until the server is actually accepting connections before shutdown.
|
||||||
|
addr = httpd.server_address
|
||||||
|
for _ in range(100):
|
||||||
|
try:
|
||||||
|
with socket.create_connection(addr, timeout=0.1):
|
||||||
|
break
|
||||||
|
except OSError:
|
||||||
|
time.sleep(0.02)
|
||||||
|
|
||||||
|
# --- exact handler mechanism from server.py ---
|
||||||
|
shutdown_requested = threading.Event()
|
||||||
|
|
||||||
|
def _request_shutdown():
|
||||||
|
if shutdown_requested.is_set():
|
||||||
|
return
|
||||||
|
shutdown_requested.set()
|
||||||
|
threading.Thread(target=httpd.shutdown, daemon=True).start()
|
||||||
|
|
||||||
|
_request_shutdown()
|
||||||
|
_request_shutdown() # idempotent: a repeated signal must not double-shutdown
|
||||||
|
|
||||||
|
serve_thread.join(timeout=5)
|
||||||
|
assert not serve_thread.is_alive(), "serve_forever() did not exit on shutdown"
|
||||||
|
assert finally_ran.wait(timeout=5), "finally block / drain did not run"
|
||||||
|
assert agent.calls == 1
|
||||||
|
assert lifecycle.has_uncommitted_work(sid) is False
|
||||||
|
|
||||||
|
|
||||||
|
# --------------------------------------------------------------------------- #
|
||||||
|
# Static wiring guards (the inline handler in main() cannot be imported) #
|
||||||
|
# --------------------------------------------------------------------------- #
|
||||||
|
|
||||||
|
|
||||||
|
def test_server_installs_sigterm_handler_and_drains_in_finally():
|
||||||
|
"""server.py must wire a SIGTERM handler and drain in serve_forever's finally."""
|
||||||
|
src = (REPO / "server.py").read_text(encoding="utf-8")
|
||||||
|
tree = ast.parse(src)
|
||||||
|
|
||||||
|
# A SIGTERM handler is installed via signal.signal(signal.SIGTERM, ...).
|
||||||
|
installs_sigterm = any(
|
||||||
|
isinstance(node, ast.Call)
|
||||||
|
and isinstance(node.func, ast.Attribute)
|
||||||
|
and node.func.attr == "signal"
|
||||||
|
and node.args
|
||||||
|
and ast.dump(node.args[0]).find("SIGTERM") != -1
|
||||||
|
for node in ast.walk(tree)
|
||||||
|
)
|
||||||
|
assert installs_sigterm, "server.py does not install a SIGTERM handler"
|
||||||
|
|
||||||
|
# httpd.shutdown() is dispatched (from a helper thread) rather than the
|
||||||
|
# default SIGTERM behavior that would skip the finally.
|
||||||
|
assert "httpd.shutdown" in src
|
||||||
|
# The existing shutdown drain is still wired.
|
||||||
|
assert "drain_all_on_shutdown" in src
|
||||||
|
# SIGPIPE handling must be preserved.
|
||||||
|
assert "SIGPIPE" in src
|
||||||
|
|
||||||
|
|
||||||
|
# --------------------------------------------------------------------------- #
|
||||||
|
# Deadline enforcement: a wedged worker must not hang the managed-stop drain #
|
||||||
|
# --------------------------------------------------------------------------- #
|
||||||
|
|
||||||
|
|
||||||
|
def test_drain_deadline_returns_despite_stuck_worker():
|
||||||
|
"""A background commit wedged in commit_memory_session() (e.g. a hung memory
|
||||||
|
provider) must NOT hang drain_all_on_shutdown() past its deadline. The inline
|
||||||
|
drain runs each commit in a budget-bounded daemon worker, so a stuck in_flight
|
||||||
|
session cannot block the drain. Regression for the dual-gate finding that the
|
||||||
|
top-of-loop deadline check never fired while blocked inside a stuck commit."""
|
||||||
|
lifecycle = _fresh_lifecycle()
|
||||||
|
agent = RecordingAgent() # NEVER released → stays wedged in commit
|
||||||
|
sid = "stuck-drain"
|
||||||
|
lifecycle.register_agent(sid, agent)
|
||||||
|
lifecycle.mark_turn_completed(sid, agent=agent)
|
||||||
|
|
||||||
|
# Put the worker in flight, but keep it OUT of the registry so the 5s join
|
||||||
|
# phase does not dominate the timing — we are isolating the inline-drain
|
||||||
|
# deadline path.
|
||||||
|
worker = threading.Thread(
|
||||||
|
target=lambda: lifecycle.commit_session_memory(sid, agent=agent),
|
||||||
|
daemon=True,
|
||||||
|
name="stuck-commit",
|
||||||
|
)
|
||||||
|
worker.start()
|
||||||
|
assert agent.entered.wait(timeout=2) # worker is now wedged holding in_flight
|
||||||
|
|
||||||
|
started = time.monotonic()
|
||||||
|
lifecycle.drain_all_on_shutdown(deadline=0.5)
|
||||||
|
elapsed = time.monotonic() - started
|
||||||
|
|
||||||
|
# Without the timeout fix this blocks until the worker is released (never);
|
||||||
|
# with it, the drain returns within roughly the deadline budget.
|
||||||
|
assert elapsed < 3.0, f"drain hung {elapsed:.1f}s past its 0.5s deadline"
|
||||||
|
|
||||||
|
# Let the wedged worker unwind so the test process exits cleanly.
|
||||||
|
agent.release.set()
|
||||||
|
worker.join(timeout=5)
|
||||||
|
|
||||||
|
|
||||||
|
def test_drain_deadline_returns_when_inline_provider_hangs_no_prior_worker():
|
||||||
|
"""Deeper case (2nd-round gate finding): a DIRTY session with NO pre-existing
|
||||||
|
in_flight worker, whose provider call hangs when the inline drain owns the
|
||||||
|
segment and calls commit_memory_session() DIRECTLY. The timeout= param only
|
||||||
|
bounds waiting on an already-in_flight commit; the direct provider call must
|
||||||
|
also be budget-bounded (run in a joined daemon worker), or the drain hangs."""
|
||||||
|
lifecycle = _fresh_lifecycle()
|
||||||
|
agent = RecordingAgent() # NEVER released → the DIRECT provider call wedges
|
||||||
|
sid = "inline-hang"
|
||||||
|
lifecycle.register_agent(sid, agent)
|
||||||
|
lifecycle.mark_turn_completed(sid, agent=agent)
|
||||||
|
# No background worker started: the session is simply dirty
|
||||||
|
# (generation > committed_generation) with in_flight=False, so the inline
|
||||||
|
# drain itself acquires the segment and calls the hanging provider directly.
|
||||||
|
|
||||||
|
started = time.monotonic()
|
||||||
|
lifecycle.drain_all_on_shutdown(deadline=0.5)
|
||||||
|
elapsed = time.monotonic() - started
|
||||||
|
|
||||||
|
assert elapsed < 3.0, f"inline-drain hung {elapsed:.1f}s past its 0.5s deadline"
|
||||||
|
|
||||||
|
# Let the wedged inline commit worker unwind cleanly.
|
||||||
|
agent.release.set()
|
||||||
Reference in new issue
Block a user