fix(quota): arbitrate the reserved file-slot release instead of leaking it - #700
fix(quota): arbitrate the reserved file-slot release instead of leaking it#700Ahmath-Gadji wants to merge 1 commit into
Conversation
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (10)
🚧 Files skipped from review as they are similar to previous changes (7)
📝 WalkthroughWalkthroughThe change adds persistent quota reservation and release-claim state, coordinates worker and cancellation cleanup through one-shot arbitration, wires ChangesQuota release arbitration
Estimated code review effort: 3 (Moderate) | ~25 minutes Sequence Diagram(s)sequenceDiagram
participant IndexerWorker
participant TaskStateManager
participant WorkerDispatcher
participant UserRepository
IndexerWorker->>TaskStateManager: claim_quota_release(task_id)
TaskStateManager-->>IndexerWorker: claim result
WorkerDispatcher->>TaskStateManager: claim_quota_release(task_id) after cancellation
TaskStateManager-->>WorkerDispatcher: claim result
WorkerDispatcher->>UserRepository: release_file_slot(user_id) when claimed
Possibly related PRs
Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
⚔️ Resolve merge conflicts
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
openrag/services/workers/dispatcher.py (1)
368-445: 🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick winCancellation path leaks the slot instead of releasing it when the arbiter is unreachable.
_release_cancelled_slotwraps theclaim_quota_releaseRPC in a baretry/exceptthat only logs on failure — it never falls back to releasing. This contradicts both the stated PR intent ("favor recoverable undercounts over permanent leaks") and the parallel worker-side helperclaim_quota_releaseinindexer_actor.py, which explicitly treats an unreachable/erroring actor as "release anyway." As written, if the state actor is unreachable exactly during a cancel-before-start (the scenario#664targets), the slot leaks silently until manual reconciliation — the opposite of the documented tradeoff. This path also lacks test coverage for the RPC-failure case (existing tests only coverrelease_file_slotraising, andclaim_quota_releasereturningFalse).🛡️ Proposed fix: fall back to releasing on arbitration failure
async def _release_cancelled_slot(self, task_id: str) -> None: if self._user_repo is None: return try: details = await self._call( self._tsm.get_details.remote(task_id), task_description=f"get_details({task_id})", ) user_id = (details or {}).get("user_id") if user_id is None: return - claimed = await self._call( - self._tsm.claim_quota_release.remote(task_id), - task_description=f"claim_quota_release({task_id})", - ) + try: + claimed = await self._call( + self._tsm.claim_quota_release.remote(task_id), + task_description=f"claim_quota_release({task_id})", + ) + except Exception: # noqa: BLE001 - prefer a recoverable undercount to a leak + claimed = True if not claimed: return await self._user_repo.release_file_slot(user_id) except Exception as exc: # noqa: BLE001 - the cancellation itself already succeeded ...Want me to add a regression test (
claim_quota_release.remoteraising duringcancel_task) alongside this fix, or open a follow-up issue?🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@openrag/services/workers/dispatcher.py` around lines 368 - 445, Update _release_cancelled_slot so failures reaching get_details or claim_quota_release, including RPC exceptions, fall back to releasing the user’s file slot rather than returning without release; preserve the existing no-op when the user repository or user_id is unavailable, and keep claim_quota_release returning false as the non-owner path. Add regression coverage through cancel_task for a claim_quota_release RPC failure, verifying release_file_slot is still attempted.Source: Coding guidelines
🧹 Nitpick comments (1)
openrag/services/workers/indexer_actor.py (1)
348-367: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winConsider a timeout on the arbitration RPC.
claim_quota_releaseawaitstsm.claim_quota_release.remote(task_id)directly with no timeout. This call runs unconditionally inIndexerWorker.process_file'sfinally(line 170) and inIndexerWorkerActor.process_file's setup-failure handler (indexer_pool.py, line 244) — both cleanup paths. If theTaskStateManageractor is merely slow (not "unreachable" in the sense that raises), thisawaitcan hang indefinitely, blocking the worker's cleanup and tying up its concurrency slot.dispatcher.pyalready wraps equivalent actor calls withcall_ray_actor_with_timeout; this helper doesn't.♻️ Proposed fix: wrap the RPC with a timeout
async def claim_quota_release(tsm: Any, task_id: str) -> bool: try: - return bool(await tsm.claim_quota_release.remote(task_id)) + from services.workers.ray_utils import call_ray_actor_with_timeout + + return bool( + await call_ray_actor_with_timeout( + future=tsm.claim_quota_release.remote(task_id), + timeout=DEFAULT_TIMEOUT, + task_description=f"claim_quota_release({task_id})", + ) + ) except Exception: # noqa: BLE001 - cleanup must never mask the indexing errorPlease confirm whether
call_ray_actor_with_timeoutis safe to import here without introducing a circular import betweenindexer_actor.pyandray_utils.py.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@openrag/services/workers/indexer_actor.py` around lines 348 - 367, Update claim_quota_release to wrap the tsm.claim_quota_release.remote(task_id) arbitration call with call_ray_actor_with_timeout, using the existing timeout/error-handling conventions from dispatcher.py. Verify the import from ray_utils.py does not create a circular dependency; if it does, use an equivalent local timeout mechanism. Preserve the current fallback that logs and returns True for timeout or other arbitration failures.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Outside diff comments:
In `@openrag/services/workers/dispatcher.py`:
- Around line 368-445: Update _release_cancelled_slot so failures reaching
get_details or claim_quota_release, including RPC exceptions, fall back to
releasing the user’s file slot rather than returning without release; preserve
the existing no-op when the user repository or user_id is unavailable, and keep
claim_quota_release returning false as the non-owner path. Add regression
coverage through cancel_task for a claim_quota_release RPC failure, verifying
release_file_slot is still attempted.
---
Nitpick comments:
In `@openrag/services/workers/indexer_actor.py`:
- Around line 348-367: Update claim_quota_release to wrap the
tsm.claim_quota_release.remote(task_id) arbitration call with
call_ray_actor_with_timeout, using the existing timeout/error-handling
conventions from dispatcher.py. Verify the import from ray_utils.py does not
create a circular dependency; if it does, use an equivalent local timeout
mechanism. Preserve the current fallback that logs and returns True for timeout
or other arbitration failures.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: 40ae186d-7268-4877-a2a6-7b234593d82d
📒 Files selected for processing (9)
CLAUDE.mdopenrag/di/container.pyopenrag/services/workers/dispatcher.pyopenrag/services/workers/indexer_actor.pyopenrag/services/workers/indexer_pool.pyopenrag/services/workers/task_state.pytests/unit/services/workers/test_dispatcher.pytests/unit/services/workers/test_indexer_worker_quota_release.pytests/unit/services/workers/test_task_state.py
hedhoud
left a comment
hedhoud
left a comment
There was a problem hiding this comment.
The arbitration direction makes sense, but two correctness blockers and one degraded-path leak remain. I’ve left the details inline.
| user_id = (details or {}).get("user_id") | ||
| if user_id is None: | ||
| return | ||
| claimed = await self._call( |
There was a problem hiding this comment.
This claim does not distinguish a task that never started from one that has already created its catalog row. If cancellation lands after add_file_to_partition succeeds but before COMPLETED is reported, the dispatcher can win this claim and decrement file_count even though the file remains stored. Could the worker take ownership when processing starts, or otherwise mark the reservation consumed atomically, so only a true cancel-before-start releases here? This also needs a regression test for that timing window.
| user_id=user.get("id") if user else None, | ||
| # Tells the actor a slot is outstanding for this task, so | ||
| # ``claim_quota_release`` can arbitrate who gives it back. | ||
| quota_reserved=quota_reserved, |
There was a problem hiding this comment.
TaskStateManager is detached and reused under the same actor name. During a rolling upgrade, the existing actor will not accept the new quota_reserved argument and also will not expose claim_quota_release, so new indexing requests can fail until Ray is manually restarted. Could we version the actor name or add an explicit compatibility/recreation guard before changing this contract?
| if not claimed: | ||
| return | ||
| await self._user_repo.release_file_slot(user_id) | ||
| except Exception as exc: # noqa: BLE001 - the cancellation itself already succeeded |
There was a problem hiding this comment.
If get_details succeeds but this arbitration RPC raises or times out, the outer handler only logs and leaves the slot reserved. That reintroduces the permanent leak this PR is meant to close and differs from the worker-side fallback. Could we release with the already-resolved user_id when the claim cannot be made, and cover that path with a regression test?
a0286dd to
4402d7e
Compare
Cancelling a task that had not started yet leaked its reserved file slot outright. After dispatch the worker owns the slot and only releases it in `IndexerWorker.process_file`'s `finally`; when `ray.cancel` retires the task before that body runs, the `finally` never executes and nothing else gives the slot back. The CANCELLED row is terminal, so the loss is indistinguishable from a clean cancel, and the uploader's quota is permanently one narrower. `cancel_task` could not simply release, because a task cancelled *mid-flight* does run its `finally` — releasing in both places would drive `file_count` below reality and hand out free quota. The two outcomes were only distinguishable by how far the worker happened to get. `TaskStateManager` settles it. It is a single-threaded Ray actor and already arbitrates this way (`set_cancelled_if_active`, `set_failed_if_not_cancelled`), so a one-shot `claim_quota_release` compare-and-set is the idiomatic primitive here: the worker and the canceller both ask, exactly one is told to release. The dispatcher records `quota_reserved` on the task at admission so the actor knows whether a slot is outstanding at all — `put_file` reuses an existing row and reserves nothing. The setup-failure release in `indexer_pool` goes through the same claim: a cancellation *during* setup reaches both it and `cancel_task`. This also closes the dispatch-partial-success double-release noted in the PR description — the worker and the request teardown can no longer both consume the same reservation. Failure handling stays biased the way the rest of the branch is: an unreachable arbiter releases anyway. An undercount is recoverable and self-heals under the is best-effort throughout, because a cancellation the user already asked for and already saw must not fail on bookkeeping. Worth doing now rather than deferring to #676: #685 shipped a cancel-all UI that fires one cancel per selected task, so the leak had become reachable at scale from a single click.
a42fbb1 to
690f0ae
Compare
3 participants
Closes the cancel-before-start quota-slot leak that #677 documents as a known gap.
The leak
After dispatch the worker owns the uploader's reserved file slot and hands it back in
IndexerWorker.process_file'sfinally. Whenray.cancelretires a task before that body runs, thefinallynever executes and nothing else gives the slot back. TheCANCELLEDrow written bycancel_taskis terminal, so the loss is indistinguishable from a clean cancel — the uploader's quota is permanently one narrower, silently.cancel_taskcould not simply release. A task cancelled mid-flight does run itsfinally, so releasing in both places would drivefile_countbelow reality and hand the user free quota. The two cases were distinguishable only by how far the worker happened to get — which is exactly the kind of thing the caller cannot observe.The fix
TaskStateManageris a single-threaded Ray actor that already arbitrates precisely this shape of race (set_cancelled_if_active,set_failed_if_not_cancelled). A one-shotclaim_quota_releasecompare-and-set is the idiomatic primitive here:process_file'sfinallyasks before releasing.cancel_taskasks after claiming the cancellation.indexer_poolsetup-failure release asks too — a cancellation during setup reaches both it andcancel_task.Exactly one caller is told to release; the outcome no longer depends on worker progress. The dispatcher records
quota_reservedon the task at admission so the actor knows whether a slot is outstanding at all (put_filereuses an existing row and reserves nothing).This also closes the dispatch-partial-success double-release that #677 accepts as a deliberate trade-off: the worker and the request teardown can no longer both consume the same reservation.
Failure behaviour
Biased the same way the rest of the reserve/release design is — an unreachable arbiter releases anyway. An undercount is recoverable and self-heals under the #676 recount; a leak is permanent and silent.
_release_cancelled_slotis best-effort throughout, because a cancellation the user already asked for and already saw must not fail on quota bookkeeping.The one case that stays open is genuinely un-arbitrable and belongs to #676: outright process death between reserve and dispatch, and an unreachable Postgres inside the release itself.
Why now rather than in #676
#685 shipped a cancel-all UI that fires one
cancelTaskper selected task (Promise.allSettled(activeTaskIds.map(...))). The leak was previously one slot per manual cancel; it is now N slots from a single click, which moves it out of "acceptable known gap" territory.Verification
ruff check/ruff format --checkclean;scripts/check_layer_imports.py→layer import guard: OK.test_create_is_idempotent_per_name, fixed separately in test(repos): assert create_partition rejects duplicates, as it does #691.Summary by CodeRabbit