diff --git a/app/chain/transfer/queue.py b/app/chain/transfer/queue.py index 3214b602f..fcd8f8a6a 100644 --- a/app/chain/transfer/queue.py +++ b/app/chain/transfer/queue.py @@ -496,6 +496,24 @@ class TransferQueueOwner(_TransferOwnerBase): if not self._closing: self._TransferChain__start_lease_heartbeat_owner_locked() + def _TransferChain__claim_recovery_batch(self) -> List[TransferAdmission]: + """按 worker 并发度 claim 当前可安全占用的恢复租约槽位。""" + runtime_config = getattr(self, "runtime_config", None) + worker_capacity = max( + 1, + int(getattr(runtime_config, "transfer_threads", 1)), + ) + with self._worker_state_lock: + owned_count = len(self._owned_leases) + recovery_capacity = max(0, worker_capacity - owned_count) + if recovery_capacity <= 0: + return [] + return self._transfer_admissions.claim_recoverable( + owner_id=self._worker_owner_id, + limit=min(self._RECOVERY_CLAIM_LIMIT, recovery_capacity), + lease_seconds=self._WORKER_LEASE_SECONDS, + ) + def _TransferChain__run_lease_heartbeat(self, stop_event: threading.Event) -> None: """按固定周期续期本进程已 claim 且尚未结算的任务。""" try: @@ -745,11 +763,7 @@ class TransferQueueOwner(_TransferOwnerBase): return self._TransferChain__ensure_lease_runtime_state() try: - pendings = self._transfer_admissions.claim_recoverable( - owner_id=self._worker_owner_id, - limit=self._RECOVERY_CLAIM_LIMIT, - lease_seconds=self._WORKER_LEASE_SECONDS, - ) + pendings = self._TransferChain__claim_recovery_batch() except Exception as err: logger.error(f"读取待整理文件登记失败:{err}") return diff --git a/tests/test_transfer_pending_replay.py b/tests/test_transfer_pending_replay.py index 9991771aa..4d102b42c 100644 --- a/tests/test_transfer_pending_replay.py +++ b/tests/test_transfer_pending_replay.py @@ -41,6 +41,7 @@ def _build_chain(admissions) -> TransferChain: chain._transfer_executions.get_snapshot.side_effect = ( lambda *, task_id: _execution_snapshot(task_id=task_id) ) + chain.runtime_config = SimpleNamespace(transfer_threads=1) chain._worker_owner_id = "test-owner" chain._owned_leases = {} chain._queued_lease_tokens = set() @@ -385,6 +386,40 @@ def test_replay_is_noop_without_registrations(): chain._execute_transfer.assert_not_called() +def test_replay_waits_for_a_free_worker_capacity_slot(): + """已有 worker 槽位被租约占用时,恢复器不得继续批量 claim。""" + admissions = MagicMock() + chain = _build_chain(admissions) + chain._owned_leases["running-task"] = ( + "running-token", + time.monotonic() + 60, + ) + + chain._TransferChain__replay_pending() + + admissions.claim_recoverable.assert_not_called() + + +def test_replay_claims_only_free_worker_capacity_slots(): + """恢复 claim 数量应受整理并发度与当前持有租约共同限制。""" + admissions = MagicMock() + admissions.claim_recoverable.return_value = [] + chain = _build_chain(admissions) + chain.runtime_config.transfer_threads = 2 + chain._owned_leases["running-task"] = ( + "running-token", + time.monotonic() + 60, + ) + + chain._TransferChain__replay_pending() + + admissions.claim_recoverable.assert_called_once_with( + owner_id="test-owner", + limit=1, + lease_seconds=120, + ) + + def test_replay_survives_db_failure(): """ 读取登记失败不能让启动流程报错。