From c356752b6beb9f974889d07102b712245f8e3196 Mon Sep 17 00:00:00 2001 From: teknium1 <127238744+teknium1@users.noreply.github.com> Date: Thu, 9 Jul 2026 18:56:15 -0700 Subject: [PATCH] fix(memory): drain queued writes on shutdown --- agent/memory_manager.py | 154 ++++++++++++++++---------- tests/agent/test_memory_async_sync.py | 80 +++++++++++++ 2 files changed, 175 insertions(+), 59 deletions(-) diff --git a/agent/memory_manager.py b/agent/memory_manager.py index 37a0411ade8..3909f2ebcd3 100644 --- a/agent/memory_manager.py +++ b/agent/memory_manager.py @@ -30,7 +30,7 @@ import logging import re import inspect import threading -from concurrent.futures import ThreadPoolExecutor +from concurrent.futures import Future, ThreadPoolExecutor, wait from typing import Any, Callable, Dict, List, Optional from agent.memory_provider import MemoryProvider @@ -378,6 +378,16 @@ class MemoryManager: # _submit_background() and the sync_all/queue_prefetch_all rationale. self._sync_executor: Optional[ThreadPoolExecutor] = None self._sync_executor_lock = threading.Lock() + # Futures are tracked by durability class so shutdown can give writes + # a bounded FIFO drain, then explicitly report anything abandoned. + self._background_futures: Dict[Future, str] = {} + self._shutting_down = False + self._shutdown_drain_state: Dict[str, Any] = { + "status": "not_started", + "abandoned_writes": 0, + "abandoned_prefetches": 0, + "active_tasks": 0, + } # -- Registration -------------------------------------------------------- @@ -599,7 +609,7 @@ class MemoryManager: provider.name, e, ) - self._submit_background(_run) + self._submit_background(_run, kind="prefetch") # -- Sync ---------------------------------------------------------------- @@ -675,46 +685,57 @@ class MemoryManager: # -- Background dispatch ------------------------------------------------- - def _submit_background(self, fn) -> None: - """Run ``fn`` on the manager's background worker. - - The executor is created lazily and shared across calls. If the - executor can't be created or has already been shut down, ``fn`` - runs inline as a last-resort fallback — losing the async benefit - but never losing the write itself. ``fn`` must do its own - per-provider error handling; this wrapper only guards executor - plumbing. - """ + def _submit_background(self, fn, *, kind: str = "write") -> None: + """Queue ``fn`` on the serialized worker and track its durability class.""" executor = self._get_sync_executor() if executor is None: - # Executor unavailable (shut down / creation failed) — run - # inline rather than drop the work. Slow, but correct. + if self._shutting_down: + logger.warning("Memory manager is shutting down; rejecting late %s task", kind) + return + # Creation failure outside shutdown: preserve the historical + # fail-safe behavior and run the operation inline. try: fn() except Exception as e: # pragma: no cover - fn guards internally logger.debug("Inline memory background task failed: %s", e) return try: - executor.submit(fn) + # Make submit+tracking atomic with the shutdown snapshot. The + # callback is attached after releasing the lock because an already + # completed future invokes callbacks synchronously. + with self._sync_executor_lock: + if self._shutting_down: + logger.warning("Memory manager is shutting down; rejecting late %s task", kind) + return + future = executor.submit(fn) + self._background_futures[future] = kind + future.add_done_callback(self._forget_background_future) except RuntimeError: - # Executor was shut down between the get and the submit - # (teardown race). Fall back to inline. + if self._shutting_down: + logger.warning("Memory manager shut down during %s submission; task rejected", kind) + return try: fn() except Exception as e: # pragma: no cover - fn guards internally logger.debug("Inline memory background task failed: %s", e) + def _forget_background_future(self, future: Future) -> None: + with self._sync_executor_lock: + self._background_futures.pop(future, None) + def _get_sync_executor(self) -> Optional[ThreadPoolExecutor]: """Lazily create the single-worker background executor.""" + if self._shutting_down: + return None if self._sync_executor is not None: return self._sync_executor with self._sync_executor_lock: + if self._shutting_down: + return None if self._sync_executor is None: try: # Daemon workers (see tools.daemon_pool): a provider wedged - # on a network call must never block interpreter exit — - # stdlib ThreadPoolExecutor's atexit hook would join it - # unconditionally even after shutdown(wait=False). + # on a network call must never block interpreter exit. from tools.daemon_pool import DaemonThreadPoolExecutor self._sync_executor = DaemonThreadPoolExecutor( max_workers=1, @@ -1129,51 +1150,66 @@ class MemoryManager: provider.name, e, ) - def _drain_sync_executor(self) -> None: - """Shut down the background executor, waiting briefly for drain. - - Bounded by ``_SYNC_DRAIN_TIMEOUT_S``: a wedged provider must never - hang process/session teardown. We stop accepting new work and - cancel anything still queued, then wait at most the drain timeout - for the currently-running task on a watcher thread. The worker is - daemon, so an over-running task dies with the interpreter. - """ + @property + def shutdown_drain_state(self) -> Dict[str, Any]: + """Snapshot of the most recent bounded shutdown drain outcome.""" with self._sync_executor_lock: + return dict(self._shutdown_drain_state) + + def _drain_sync_executor(self) -> None: + """Give queued FIFO work a bounded chance, then abandon explicitly.""" + with self._sync_executor_lock: + self._shutting_down = True executor = self._sync_executor self._sync_executor = None + tracked = dict(self._background_futures) + self._shutdown_drain_state = { + "status": "draining" if executor is not None else "drained", + "abandoned_writes": 0, + "abandoned_prefetches": 0, + "active_tasks": sum(not future.done() for future in tracked), + } if executor is None: return - try: - # Stop accepting new work and drop anything still queued, but - # do NOT block here — cancel_futures cancels not-yet-started - # tasks; the in-flight one keeps running on its daemon thread. - executor.shutdown(wait=False, cancel_futures=True) - except TypeError: - # Older Python without cancel_futures kwarg. - try: - executor.shutdown(wait=False) - except Exception as e: # pragma: no cover - logger.debug("Memory sync executor shutdown failed: %s", e) - return - except Exception as e: # pragma: no cover - logger.debug("Memory sync executor shutdown failed: %s", e) - return - # Give an in-flight sync a bounded chance to finish on a watcher - # thread so we don't block the caller past the drain timeout. - drainer = threading.Thread( - target=lambda: self._bounded_executor_wait(executor), - daemon=True, - name="mem-sync-drain", - ) - drainer.start() - drainer.join(timeout=_SYNC_DRAIN_TIMEOUT_S) - @staticmethod - def _bounded_executor_wait(executor: ThreadPoolExecutor) -> None: - try: - executor.shutdown(wait=True) - except Exception as e: # pragma: no cover - logger.debug("Memory sync executor drain wait failed: %s", e) + # shutdown(wait=False) closes submission without touching the FIFO. + # Waiting on the tracked futures lets the real single-worker executor + # run every queued write/boundary task in order up to the deadline. + executor.shutdown(wait=False, cancel_futures=False) + _, pending = wait(tuple(tracked), timeout=_SYNC_DRAIN_TIMEOUT_S) + if not pending: + with self._sync_executor_lock: + self._shutdown_drain_state.update(status="drained", active_tasks=0) + return + + abandoned_writes = 0 + abandoned_prefetches = 0 + active_tasks = 0 + for future in pending: + kind = tracked[future] + if future.cancel(): + if kind == "prefetch": + abandoned_prefetches += 1 + else: + abandoned_writes += 1 + else: + active_tasks += 1 + + with self._sync_executor_lock: + self._shutdown_drain_state.update( + status="timed_out", + abandoned_writes=abandoned_writes, + abandoned_prefetches=abandoned_prefetches, + active_tasks=active_tasks, + ) + logger.warning( + "Memory shutdown drain timed out after %.2fs; abandoning %d queued " + "memory write(s) and %d queued prefetch(es); %d active task(s) remain detached", + _SYNC_DRAIN_TIMEOUT_S, + abandoned_writes, + abandoned_prefetches, + active_tasks, + ) def initialize_all(self, session_id: str, **kwargs) -> None: """Initialize all providers. diff --git a/tests/agent/test_memory_async_sync.py b/tests/agent/test_memory_async_sync.py index 7ff293e43fc..a1e4aaa23ec 100644 --- a/tests/agent/test_memory_async_sync.py +++ b/tests/agent/test_memory_async_sync.py @@ -15,6 +15,8 @@ The fix dispatches provider work to a single-worker background executor. for session boundaries and deterministic tests. ``shutdown_all`` drains the executor with a bounded timeout so a wedged provider can't hang teardown. """ +import logging +import threading import time import pytest @@ -136,3 +138,81 @@ def test_writes_are_serialized_in_order(): mgr.sync_all(f"turn-{i}", "resp", session_id="s1") assert mgr.flush_pending(timeout=10) is True assert order == [f"turn-{i}" for i in range(5)] + + +def test_shutdown_drains_queued_writes_and_boundary_in_fifo_order(): + """Shutdown must not cancel durable work merely because it is still queued.""" + started = threading.Event() + release = threading.Event() + calls = [] + + class _BlockingProvider(_SlowProvider): + def sync_turn(self, user_content, assistant_content, *, session_id="", messages=None): + if user_content == "turn-0": + started.set() + assert release.wait(timeout=2) + calls.append(("sync", user_content)) + + def on_session_end(self, messages): + calls.append(("end", messages[0]["content"])) + + def on_session_switch(self, new_session_id, **kwargs): + calls.append(("switch", new_session_id)) + + mgr = MemoryManager() + mgr.add_provider(_BlockingProvider(delay=0)) + mgr.sync_all("turn-0", "response") + assert started.wait(timeout=1) + mgr.sync_all("turn-1", "response") + mgr.commit_session_boundary_async( + [{"role": "user", "content": "old-session"}], + new_session_id="new-session", + ) + + threading.Timer(0.05, release.set).start() + mgr.shutdown_all() + + assert calls == [ + ("sync", "turn-0"), + ("sync", "turn-1"), + ("end", "old-session"), + ("switch", "new-session"), + ] + assert mgr.shutdown_drain_state["status"] == "drained" + assert mgr.shutdown_drain_state["abandoned_writes"] == 0 + + +def test_shutdown_timeout_abandons_queued_write_with_state_and_log(monkeypatch, caplog): + """A wedged active write bounds shutdown and reports queued data loss.""" + import agent.memory_manager as memory_manager_module + + started = threading.Event() + release = threading.Event() + calls = [] + + class _WedgedProvider(_SlowProvider): + def sync_turn(self, user_content, assistant_content, *, session_id="", messages=None): + if user_content == "active": + started.set() + release.wait(timeout=2) + calls.append(user_content) + + monkeypatch.setattr(memory_manager_module, "_SYNC_DRAIN_TIMEOUT_S", 0.1) + mgr = MemoryManager() + mgr.add_provider(_WedgedProvider(delay=0)) + mgr.sync_all("active", "response") + assert started.wait(timeout=1) + mgr.sync_all("queued", "response") + + with caplog.at_level(logging.WARNING, logger="agent.memory_manager"): + t0 = time.monotonic() + mgr.shutdown_all() + elapsed = time.monotonic() - t0 + + state = mgr.shutdown_drain_state + assert elapsed < 0.5 + assert state["status"] == "timed_out" + assert state["abandoned_writes"] == 1 + assert "queued" not in calls + assert "abandoning 1 queued memory write" in caplog.text + release.set()