Merge pull request #309 from Tencent-Hunyuan/fix/memhub-durable-drain-liveness

fix(memhub): keep durable write drain alive

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:

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)