fix(db): release worker capacity before publishing results

app/db/worker.py

@@ -3,19 +3,19 @@
 from __future__ import annotations
 
 import asyncio
+import concurrent.futures
 import threading
 import time
 from concurrent.futures import Future, ThreadPoolExecutor
 from contextvars import copy_context
 from dataclasses import dataclass
-from typing import Callable, TypeVar
+from typing import Callable, TypeVar, cast
 
+from app.runtime.observability import record_metric
 from app.schemas.exception import (
     DatabaseWorkerClosedError,
     DatabaseWorkerOverloadedError,
 )
-from app.runtime.observability import record_metric
-
 
 T = TypeVar("T")
 
@@ -99,7 +99,7 @@ class DatabaseWorker:
             )
 
     async def run(self, operation: Callable[[], T]) -> T:
-        """执行短事务，取消时仍等待已开始的事务取得最终结果。"""
+        """执行短事务；结果发布和取消返回均晚于容量回收，取消仍等待已开始事务收口。"""
         loop = asyncio.get_running_loop()
         executor = self._executor
         if executor is None or self._loop is not loop or self._closing:
@@ -124,7 +124,8 @@ class DatabaseWorker:
             self._record_depth()
             raise
 
-        wrapped = asyncio.wrap_future(future, loop=loop)
+        # 不能直接 wrap native Future：其结果转发可能先于完成记账，晚注册回调还会立即执行。
+        wrapped: asyncio.Future[object] = loop.create_future()
         with self._state_lock:
             self._futures[future] = (wrapped, item)
         future.add_done_callback(
@@ -133,10 +134,10 @@ class DatabaseWorker:
         self._record_depth()
 
         try:
-            return await asyncio.shield(wrapped)
+            return cast(T, await asyncio.shield(wrapped))
         except asyncio.CancelledError:
-            if not future.cancel():
-                await self._wait_until_done(wrapped)
+            future.cancel()
+            await self._wait_until_done(wrapped)
             if wrapped.done() and not wrapped.cancelled():
                 wrapped.exception()
             raise
@@ -181,19 +182,14 @@ class DatabaseWorker:
             pass
 
     def _complete(self, future: Future[object], item: _WorkItem) -> None:
-        """释放 admission，并记录任务的最终结果。"""
+        """先释放 admission，再发布最终结果；同一事件循环回调保证调用方观察顺序。"""
         with self._state_lock:
-            self._futures.pop(future, None)
+            wrapped, _ = self._futures.pop(future)
             if future.running() or future.done() and not future.cancelled():
                 self._running -= 1
             else:
                 self._queued -= 1
-        outcome = "cancelled" if future.cancelled() else "success"
-        if not future.cancelled():
-            try:
-                future.result()
-            except BaseException:
-                outcome = "error"
+        outcome = self._publish_result(future, wrapped)
         started_at = item.started_at or item.submitted_at
         record_metric(
             "db.worker.duration",
@@ -202,6 +198,27 @@ class DatabaseWorker:
         )
         self._record_depth()
 
+    @staticmethod
+    def _publish_result(future: Future[object], wrapped: asyncio.Future[object]) -> str:
+        """保留原 wrap_future 的取消和异常翻译合同，不调用 asyncio 私有实现。"""
+        if future.cancelled():
+            wrapped.cancel()
+            return "cancelled"
+        try:
+            result = future.result()
+        except BaseException as error:
+            exception_types: dict[type[BaseException], type[BaseException]] = {
+                concurrent.futures.CancelledError: asyncio.CancelledError,
+                concurrent.futures.InvalidStateError: asyncio.InvalidStateError,
+            }
+            exception_type = exception_types.get(type(error))
+            if exception_type is not None:
+                error = exception_type(*error.args).with_traceback(error.__traceback__)
+            wrapped.set_exception(error)
+            return "error"
+        wrapped.set_result(result)
+        return "success"
+
     async def _wait_until_done(
         self,
         future: asyncio.Future[object],

docs/architecture/optimization-checklist.md

@@ -103,7 +103,7 @@ ARCH-201 至 ARCH-204 均达到实现、验证、提交、推送和远端门禁
 | Python 源码量 | 305,884 行 | 排除 `app/plugins/**`；61 个文件超过 1,000 行，11 个超过 2,000 行 |
 | 长方法 | 290 个超过 80 行 | AST 统计排除 `app/plugins/**`；65 个超过 150 行，21 个超过 250 行 |
 | 全量 mypy 历史债务 | 9,408 / 510 文件 | Agent API 重构后的现状基线；canonical Facade 与 endpoint 类型边界已补齐，低水位只允许继续下降 |
-| Ruff 历史诊断 | 533 | 低水位门禁通过，但规则集只覆盖 `E4/E7/E9/F/I` |
+| Ruff 历史诊断 | 532 | 低水位门禁通过，但规则集只覆盖 `E4/E7/E9/F/I` |
 | 覆盖率固定基线 | Application 80.00%，Domain 80.00% | Chain、Runtime、Agent、Adapter、Startup 未进入包级覆盖率门禁 |
 
 ### 3.3 热点文件

tests/fixtures/architecture/ruff-baseline.json

@@ -119,9 +119,6 @@
   "app/db/uow.py": {
     "I001": 1
   },
-  "app/db/worker.py": {
-    "I001": 1
-  },
   "app/doctor/__init__.py": {
     "I001": 1
   },

tests/test_database_worker.py

@@ -1,7 +1,10 @@
 """数据库短事务 worker 的容量、取消与关闭合同测试。"""
 
 import asyncio
+import concurrent.futures
 import threading
+from concurrent.futures import Future
+from typing import Any
 from unittest.mock import patch
 
 import pytest
@@ -31,6 +34,7 @@ async def test_worker_rejects_work_beyond_running_and_queue_capacity() -> None:
     release = threading.Event()
 
     def block() -> None:
+        """占住唯一线程，以便验证排队额度与拒绝边界。"""
         started.set()
         release.wait(1)
 
@@ -61,6 +65,7 @@ async def test_cancelling_queued_work_prevents_execution() -> None:
     queued_executed = threading.Event()
 
     def block() -> None:
+        """保留一个运行中事务，使下一次调用稳定处于排队状态。"""
         started.set()
         release.wait(1)
 
@@ -71,6 +76,8 @@ async def test_cancelling_queued_work_prevents_execution() -> None:
     queued.cancel()
     with pytest.raises(asyncio.CancelledError):
         await queued
+    assert worker.snapshot().queued == 0
+    assert worker.snapshot().running == 1
 
     release.set()
     await running
@@ -91,6 +98,7 @@ async def test_cancelling_running_work_waits_for_transaction_terminal_state() ->
     completed = threading.Event()
 
     def operation() -> None:
+        """取消期间继续持有真实事务，直到调用方明确允许收口。"""
         started.set()
         release.wait(1)
         completed.set()
@@ -108,6 +116,8 @@ async def test_cancelling_running_work_waits_for_transaction_terminal_state() ->
         await task
 
     assert completed.is_set() is True
+    assert worker.snapshot().queued == 0
+    assert worker.snapshot().running == 0
     await worker.shutdown()
 
 
@@ -120,6 +130,7 @@ async def test_shutdown_rejects_new_work_and_waits_for_running_work() -> None:
     release = threading.Event()
 
     def operation() -> None:
+        """保持已开始任务运行，使测试能观察关闭期间的拒绝策略。"""
         started.set()
         release.wait(1)
 
@@ -150,6 +161,7 @@ async def test_shutdown_timeout_keeps_running_owner_until_transaction_finishes()
     release = threading.Event()
 
     def operation() -> None:
+        """让事务跨过关闭期限，验证执行器仍保有它的 owner。"""
         started.set()
         release.wait(1)
 
@@ -177,6 +189,7 @@ async def test_worker_depth_metrics_emit_deltas_and_return_to_zero() -> None:
     release = threading.Event()
 
     def operation() -> None:
+        """控制开始和结束两个观测点，验证指标最终归零。"""
         started.set()
         release.wait(1)
 
@@ -200,3 +213,226 @@ async def test_worker_depth_metrics_emit_deltas_and_return_to_zero() -> None:
     ]
     assert queue_values == [1.0, -1.0]
     assert active_values == [1.0, -1.0]
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("fails", [False, True])
+async def test_completion_releases_capacity_before_publishing_result(monkeypatch, fails: bool) -> None:
+    """线程完成与事件循环发布之间的握手确保结果或异常不能早于容量回收。"""
+    worker = DatabaseWorker(max_workers=1, capacity=1)
+    await worker.start()
+    loop = asyncio.get_running_loop()
+    operation_started = threading.Event()
+    release_operation = threading.Event()
+    release_completion = threading.Event()
+    observed = loop.create_future()
+    original_schedule = worker._schedule_completion
+
+    def pause_completion(completed: Future, item: Any) -> None:
+        """模拟线程在完成 callback 中暂时失去调度，结果转发不能越过这道边界。"""
+        def observe() -> None:
+            """在所属事件循环核验 Future 可见性，而非等待一个猜测性的延迟。"""
+            if not observed.done():
+                wrapped, _ = worker._futures[completed]
+                observed.set_result((wrapped.done(), worker.snapshot()))
+
+        loop.call_soon_threadsafe(observe)
+        try:
+            assert release_completion.wait(5)
+        finally:
+            original_schedule(completed, item)
+
+    def operation() -> str:
+        """先等完成回调注册完毕，再交付成功值或原始业务异常。"""
+        operation_started.set()
+        assert release_operation.wait(5)
+        if fails:
+            raise ValueError("operation failed")
+        return "committed-result"
+
+    monkeypatch.setattr(worker, "_schedule_completion", pause_completion)
+    task = asyncio.create_task(worker.run(operation))
+    try:
+        assert await asyncio.to_thread(operation_started.wait, 5)
+        release_operation.set()
+        published, pending = await asyncio.wait_for(observed, timeout=5)
+        assert pending.running == 1 and pending.queued == 0
+        assert published is False
+        release_completion.set()
+        if fails:
+            with pytest.raises(ValueError, match="operation failed"):
+                await task
+        else:
+            assert await task == "committed-result"
+        assert worker.snapshot().running == 0
+        assert worker.snapshot().queued == 0
+        assert worker._futures == {}
+        assert await worker.run(lambda: "capacity available") == "capacity available"
+    finally:
+        release_operation.set()
+        release_completion.set()
+        await asyncio.gather(task, return_exceptions=True)
+        await worker.shutdown()
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("outcome", ["success", "error", "cancelled"])
+async def test_completion_callback_handles_future_finished_before_registration(monkeypatch, outcome: str) -> None:
+    """submit 已返回终态 Future 时 callback 会同步触发，仍需先登记 owner 再有序释放。"""
+    worker = DatabaseWorker(max_workers=1, capacity=1)
+    await worker.start()
+    submitted = []
+
+    def submit(function, *args: Any) -> Future:
+        """模拟极快线程在 submit 返回前结束，保留真实 _execute 的开始记账路径。"""
+        future: Future = Future()
+        if outcome == "cancelled":
+            future.cancel()
+        else:
+            future.set_running_or_notify_cancel()
+            try:
+                result = function(*args)
+            except BaseException as error:
+                future.set_exception(error)
+            else:
+                future.set_result(result)
+        assert future.done()
+        submitted.append(future)
+        return future
+
+    def operation() -> str:
+        """测试业务异常必须原样传递，完成顺序不能吞掉失败。"""
+        if outcome == "error":
+            raise ValueError("immediate failure")
+        return "immediate result"
+
+    monkeypatch.setattr(worker._executor, "submit", submit)
+    try:
+        if outcome == "error":
+            with pytest.raises(ValueError, match="immediate failure"):
+                await worker.run(operation)
+        elif outcome == "cancelled":
+            with pytest.raises(asyncio.CancelledError):
+                await worker.run(operation)
+        else:
+            assert await worker.run(operation) == "immediate result"
+        assert len(submitted) == 1
+        assert worker.snapshot().queued == 0 and worker.snapshot().running == 0
+        assert worker._futures == {}
+    finally:
+        await worker.shutdown()
+
+
+@pytest.mark.asyncio
+async def test_finished_future_cannot_publish_while_first_callback_is_pending(monkeypatch) -> None:
+    """Future 先变 FINISHED 再执行 callback；晚注册的转发不能绕过仍未排队的完成记账。"""
+    worker = DatabaseWorker(max_workers=1, capacity=1)
+    await worker.start()
+    loop = asyncio.get_running_loop()
+    release_operation = threading.Event()
+    first_callback_entered = threading.Event()
+    release_callback = threading.Event()
+    original_submit = worker._executor.submit
+    native_futures = []
+
+    def submit(function, *args: Any) -> Future:
+        """保留真实线程/Future，只控制首次回调注册与后续注册之间的合法时序。"""
+        future = original_submit(function, *args)
+        original_add = future.add_done_callback
+        first = True
+
+        def add(callback) -> None:
+            """首个回调已阻塞且 Future 已结束后，才允许主线程继续注册其余回调。"""
+            nonlocal first
+            if not first:
+                original_add(callback)
+                return
+            first = False
+
+            def paused(completed: Future) -> None:
+                """线程端在 FINISHED 状态暂停，后注册回调仍可能在调用线程立即执行。"""
+                first_callback_entered.set()
+                assert release_callback.wait(5)
+                callback(completed)
+
+            original_add(paused)
+            release_operation.set()
+            assert first_callback_entered.wait(5)
+            assert future.done()
+
+        monkeypatch.setattr(future, "add_done_callback", add)
+        native_futures.append(future)
+        return future
+
+    def operation() -> str:
+        """用输入握手保证首次回调先注册，任务随后才结束。"""
+        assert release_operation.wait(5)
+        return "finished value"
+
+    monkeypatch.setattr(worker._executor, "submit", submit)
+    task = asyncio.create_task(worker.run(operation))
+    try:
+        assert await asyncio.to_thread(first_callback_entered.wait, 5)
+        barrier = loop.create_future()
+        loop.call_soon(barrier.set_result, None)
+        await barrier
+        state = worker._futures.get(native_futures[0])
+        if state is not None:
+            assert state[0].done() is False
+        assert task.done() is False
+        release_callback.set()
+        assert await task == "finished value"
+        assert worker.snapshot().queued == 0 and worker.snapshot().running == 0
+    finally:
+        release_operation.set()
+        release_callback.set()
+        await asyncio.gather(task, return_exceptions=True)
+        await worker.shutdown()
+
+
+class _DerivedCancellation(concurrent.futures.CancelledError):
+    """并发 Future 异常子类不属于 asyncio 的精确类型翻译范围。"""
+
+
+class _DerivedInvalidState(concurrent.futures.InvalidStateError):
+    """业务声明的状态错误子类必须保持原对象，而非统一换成 asyncio 异常。"""
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(("source_type", "expected_type"), [
+    (concurrent.futures.CancelledError, asyncio.CancelledError),
+    (concurrent.futures.InvalidStateError, asyncio.InvalidStateError),
+    (ValueError, ValueError),
+    (_DerivedCancellation, _DerivedCancellation),
+    (_DerivedInvalidState, _DerivedInvalidState),
+])
+async def test_result_publication_preserves_future_exception_contract(source_type, expected_type) -> None:
+    """精确翻译原有 Future 异常，保留参数、业务 traceback 和其他异常的对象身份。"""
+    worker = DatabaseWorker(max_workers=1, capacity=1)
+    await worker.start()
+    source = source_type("native future error", 42)
+
+    def operation() -> None:
+        """在真实 worker 线程抛出异常，验证跨线程发布后的 traceback。"""
+        raise source
+
+    try:
+        with pytest.raises(expected_type) as raised:
+            await worker.run(operation)
+        assert type(raised.value) is expected_type
+        assert raised.value.args == source.args
+        if source_type is expected_type:
+            assert raised.value is source
+        else:
+            assert raised.value is not source
+        traceback = raised.value.__traceback__
+        codes = []
+        while traceback is not None:
+            codes.append(traceback.tb_frame.f_code)
+            traceback = traceback.tb_next
+        assert operation.__code__ in codes
+        assert worker.snapshot().queued == 0 and worker.snapshot().running == 0
+        assert worker._futures == {}
+        assert await worker.run(lambda: "still usable") == "still usable"
+    finally:
+        await worker.shutdown()