From 0953c5af357f60c2e9f503d69b177924a952ae3e Mon Sep 17 00:00:00 2001 From: kristinhe Date: Wed, 4 Nov 2026 15:05:16 +0800 Subject: [PATCH] fix(memhub): keep durable write drain alive Co-authored-by: Cursor --- .../one_memory/memory_server/durable_queue.py | 66 ++++++++++++-- .../tests/test_memory_server_durable_queue.py | 86 +++++++++++++++++++ 2 files changed, 144 insertions(+), 8 deletions(-) diff --git a/memhub/apps/core/src/one_memory/memory_server/durable_queue.py b/memhub/apps/core/src/one_memory/memory_server/durable_queue.py index 62c503a..20ffc78 100644 --- a/memhub/apps/core/src/one_memory/memory_server/durable_queue.py +++ b/memhub/apps/core/src/one_memory/memory_server/durable_queue.py @@ -3261,11 +3261,9 @@ class DurableWriteDrain: "durable_write_drain_job_ceiling_exceeded job_id=%s", job_id, ) - await in_the_background( - self.queue.settle_worker_error, + await self._settle_worker_error( job_id, error_code="durable_work_ceiling_exceeded", - owner=self.owner, ) except Exception as error: logger.error( @@ -3274,11 +3272,9 @@ class DurableWriteDrain: type(error).__name__, ) self._record_database_error(error) - await in_the_background( - self.queue.settle_worker_error, + await self._settle_worker_error( job_id, error_code="durable_executor_error", - owner=self.owner, ) finally: execution.cancel() @@ -3287,6 +3283,36 @@ class DurableWriteDrain: with contextlib.suppress(asyncio.CancelledError, Exception): await task + async def _settle_worker_error(self, job_id: str, *, error_code: str) -> None: + """Best-effort settlement must not kill the long-lived drain. + + The row still has a durable lease when this fallback runs. If the + single SQLite writer is busy, waiting and retrying here only extends + the same contention that made the executor fail. Leave the row where + it is instead: lease recovery later moves a never-invoked job back to + pending and a possibly-invoked job to unknown, preserving the + no-blind-retry boundary. + + Only transient lock/busy refusals are deferred. A malformed query or + other programming error still escapes to the outer per-job fence. + """ + try: + await in_the_background( + self.queue.settle_worker_error, + job_id, + error_code=error_code, + owner=self.owner, + ) + except sqlite3.OperationalError as error: + if "locked" not in str(error).lower() and "busy" not in str(error).lower(): + raise + self._sample_database_failure("settle_worker_error", error) + logger.warning( + "durable_write_worker_settlement_deferred job_id=%s error_type=%s", + job_id, + type(error).__name__, + ) + async def _execute_with_ceiling(self, job_id: str) -> None: ceiling = ( await in_the_background(self.work_ceiling) @@ -3324,7 +3350,26 @@ class DurableWriteDrain: finished = {task for task in active if task.done()} active.difference_update(finished) for task in finished: - await task + try: + await task + except asyncio.CancelledError: + raise + except _LeaseRenewalFatalError: + # Renewal is the ownership fence. If it is broken, + # continuing could let another worker reclaim a job + # while this process is still executing it. + raise + except Exception as error: + # A job task is a fault-containment boundary. Its row + # retains the durable lease and is recovered according + # to hy_invoked; killing this only consumer would strand + # every unrelated write behind it. + self._record_database_error(error) + logger.error( + "durable_write_drain_job_task_failed task_name=%s error_type=%s", + task.get_name(), + type(error).__name__, + ) if self.budget is not None and work_deadline is not None: now = _time.monotonic() # Always admit the first unit of a cycle even when the @@ -3349,7 +3394,12 @@ class DurableWriteDrain: else () ) for job_id in ids: - active.add(asyncio.create_task(self._run_claimed_job(job_id))) + active.add( + asyncio.create_task( + self._run_claimed_job(job_id), + name=f"durable-write:{job_id}", + ) + ) executed = len(ids) self._database_failure_streaks.pop("claim", None) except asyncio.CancelledError: diff --git a/memhub/apps/core/tests/test_memory_server_durable_queue.py b/memhub/apps/core/tests/test_memory_server_durable_queue.py index 9916a8a..c4991b1 100644 --- a/memhub/apps/core/tests/test_memory_server_durable_queue.py +++ b/memhub/apps/core/tests/test_memory_server_durable_queue.py @@ -4138,6 +4138,71 @@ def test_a_failing_renewal_does_not_kill_the_drain(tmp_path) -> None: assert queue.get("idem-1").state == COMPLETED +def test_a_locked_fallback_settlement_does_not_kill_the_drain(tmp_path, monkeypatch) -> None: + """The incident shape: executor and its fallback both lose the writer. + + The first row keeps its lease for normal reclaim; an unrelated row must + still run instead of being stranded behind a dead drain task. + """ + queue = DurableWriteQueue(tmp_path / "one-memory.sqlite3") + first = queue.submit(_request(key="first", source="cursor")) + second = queue.submit(_request(key="second", source="codex")) + executed: list[str] = [] + + async def execute(job_id: str) -> None: + executed.append(job_id) + if job_id == first.job_id: + raise RuntimeError("executor failed") + queue.settle(job_id, state=COMPLETED) + drain.stop() + + drain = DurableWriteDrain(queue, execute=execute, max_concurrent_sources=1) + + def locked_settlement(*_args, **_kwargs): + raise sqlite3.OperationalError("database is locked") + + monkeypatch.setattr(queue, "settle_worker_error", locked_settlement) + + asyncio.run(asyncio.wait_for(drain.run_forever(idle_pause=0.005), 1)) + + assert executed == [first.job_id, second.job_id] + assert queue.get("first").state == RUNNING + assert queue.get("second").state == COMPLETED + + +def test_an_unexpected_job_task_failure_is_contained_by_the_drain( + tmp_path, monkeypatch, caplog +) -> None: + """A programming error stays visible without taking down the consumer.""" + queue = DurableWriteQueue(tmp_path / "one-memory.sqlite3") + first = queue.submit(_request(key="first", source="cursor")) + second = queue.submit(_request(key="second", source="codex")) + executed: list[str] = [] + + async def run_job(job_id: str) -> None: + executed.append(job_id) + if job_id == first.job_id: + raise RuntimeError("unexpected job task failure") + drain.stop() + + drain = DurableWriteDrain( + queue, + execute=_never_execute, + max_concurrent_sources=1, + ) + monkeypatch.setattr(drain, "_run_claimed_job", run_job) + + with caplog.at_level(logging.ERROR, logger="one_memory.memory_server.durable_queue"): + asyncio.run(asyncio.wait_for(drain.run_forever(idle_pause=0.005), 1)) + + assert executed == [first.job_id, second.job_id] + assert any( + "durable_write_drain_job_task_failed" in message.getMessage() + and "error_type=RuntimeError" in message.getMessage() + for message in caplog.records + ) + + @pytest.mark.parametrize("terminal", [COMPLETED, FAILED, DEAD]) def test_a_terminal_row_is_sticky_against_a_late_settle(tmp_path, terminal: str) -> None: """A late settle must not overwrite an outcome that already happened. @@ -4980,6 +5045,27 @@ def test_permanent_renewal_error_stops_the_worker(tmp_path, monkeypatch) -> None asyncio.run(drain.run_one()) +def test_permanent_renewal_error_stops_the_long_lived_drain(tmp_path, monkeypatch) -> None: + queue = DurableWriteQueue(tmp_path / "one-memory.sqlite3") + queue.submit(_request()) + drain = DurableWriteDrain( + queue, + execute=lambda _job_id: asyncio.sleep(0.03), + lease=timedelta(seconds=1), + renew_every=timedelta(milliseconds=10), + ) + monkeypatch.setattr( + drain, + "renew", + lambda *_args, **_kwargs: (_ for _ in ()).throw( + sqlite3.OperationalError("database schema is corrupt") + ), + ) + + with pytest.raises(RuntimeError): + asyncio.run(asyncio.wait_for(drain.run_forever(idle_pause=0.005), 1)) + + def test_lost_renewal_waits_for_execution_before_releasing_source(tmp_path, monkeypatch) -> None: path = tmp_path / "one-memory.sqlite3" queue = DurableWriteQueue(path)