fix(lance-io): drop reservation when future is cancelled#7638
Merged
wjones127 merged 3 commits intoJul 9, 2026
Merged
Conversation
github-actions
Bot
added
A-encoding
valkum
force-pushed
the
fix/scheduler-caller-drop-reservation-leak
branch
from
3e4d6d8 to
929ec28
Compare
Contributor
Author
|
#7588 touched some common parts. I rebased this PR correctly now. |
valkum
force-pushed
the
fix/scheduler-caller-drop-reservation-leak
branch
from
929ec28 to
825247d
Compare
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
wjones127
approved these changes
Jul 9, 2026
wjones127
left a comment
wjones127
left a comment
Contributor
There was a problem hiding this comment.
Looks good. Needs another rebase but then I'm happy to merge.
Comment on lines
+2084
to
+2104
| #[tokio::test(flavor = "multi_thread")] | ||
| async fn standard_scheduler_refunds_reservation_on_caller_drop() { | ||
| let (completed, elapsed) = run_caller_drop_scenario(false).await; | ||
| assert!( | ||
| completed, | ||
| "standard scheduler deadlocked the follow-up request (elapsed {elapsed:?}); \ | ||
| the dropped request's reservation was not refunded" | ||
| ); | ||
| } | ||
|
|
||
| /// Same guarantee for the lite scheduler: dropping a request future mid-read | ||
| /// releases its reservation via the `TaskHandle` drop path. | ||
| #[tokio::test(flavor = "multi_thread")] | ||
| async fn lite_scheduler_refunds_reservation_on_caller_drop() { | ||
| let (completed, elapsed) = run_caller_drop_scenario(true).await; | ||
| assert!( | ||
| completed, | ||
| "lite scheduler deadlocked the follow-up request (elapsed {elapsed:?}); \ | ||
| the dropped request's reservation was not refunded" | ||
| ); | ||
| } |
Contributor
There was a problem hiding this comment.
praise: I appreciate the regression tests here.
Comment on lines
+595
to
+609
| let reservation = match task.state { | ||
| TaskState::Reserved { | ||
| backpressure_reservation, | ||
| .. | ||
| } | ||
| | TaskState::Running { | ||
| backpressure_reservation, | ||
| .. | ||
| } | ||
| | TaskState::Finished { | ||
| backpressure_reservation, | ||
| .. | ||
| } => Some(backpressure_reservation), | ||
| TaskState::Initial { .. } | TaskState::Broken => None, | ||
| }; |
Contributor
There was a problem hiding this comment.
suggestion: if this pattern shows up elsewhere, it might be nice to create a method .backpressure_reservation(&self) -> Option<Reservation>.
Contributor
Author
There was a problem hiding this comment.
Pulled this out to a TaskState impl, so it can be reused in BrokenTaskError::new as well.
Both the standard and lite ScanScheduler leak their backpressure reservation when a caller drops the future returned by submit_request before the underlying read completes. The refund is tied to the caller polling the future to completion, so an aborted request never returns its bytes to the budget and a later request that would exceed the leaked-down ceiling deadlocks forever. Add two tests that reproduce the deadlock against a small io_buffer budget using a semaphore-gated BlockingReader to control exactly when the in-flight read completes. Both assert the dropped request's reservation is refunded, so they fail until the accompanying fix lands. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…quest Both ScanScheduler variants tied the backpressure refund to the caller polling the request future to completion, so dropping the future while a read was still in flight leaked the reservation. Because the scheduler is cached per-dataset, the leaked-down budget persisted and a later request that would exceed it deadlocked forever. Standard scheduler: move the on_bytes_consumed call into a Drop impl on Response (with data wrapped in Option so the caller can still take it). The refund now fires exactly once when the response is delivered or, if the caller has gone away, when the undelivered response is discarded. Lite scheduler: add a Drop impl to TaskHandle that abandons the task, releasing any reservation it still holds and re-checking the queue so newly-affordable tasks can start. Harden on_task_complete to discard stale pending entries left behind when a queued task's handle is dropped, which also avoids a latent spin on a missing task. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
valkum
force-pushed
the
fix/scheduler-caller-drop-reservation-leak
branch
from
825247d to
d97aa3b
Compare
📝 WalkthroughWalkthroughBoth schedulers now refund backpressure reservations when caller-side futures or task handles are dropped early. The standard path moves refunding to ChangesCancellation-safe backpressure refunds
Estimated code review effort: 4 (Complex) | ~45 minutes 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 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.
Inline comments:
In `@rust/lance-io/src/scheduler.rs`:
- Around line 1027-1032: The hot-path `rx.map` logic in `scheduler.rs` uses bare
`.unwrap()` twice in library code; replace both with `.expect(...)` so the
documented invariants are enforced at the panic site. Update the
`wrapped_rsp.unwrap()` and `rsp.data.take().unwrap()` calls in this closure with
clear reason messages that match the surrounding comments about the sender
always sending before dropping and the reservation being refunded on `Response`
drop.
In `@rust/lance-io/src/scheduler/lite.rs`:
- Around line 700-702: The cleanup path in `lite.rs` is swallowing a failure
from `on_task_complete` by assigning it to `_`, which hides observable errors
during `abandon`/`Drop`. Update the `on_task_complete` call site in `abandon` to
capture the result and emit a `warn!` when it returns an error, using the
existing `state`/task context if available, so best-effort cleanup failures are
logged instead of silently discarded.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: ASSERTIVE
Plan: Pro Plus
Run ID: 811ba4a8-4f41-489f-9831-83a99e2dea9a
📒 Files selected for processing (2)
rust/lance-io/src/scheduler.rsrust/lance-io/src/scheduler/lite.rs
Comment on lines
+1027
to
1032
| rx.map(|wrapped_rsp| { | ||
| // A cancel error can't occur: the sender always sends before dropping. | ||
| // The reservation is refunded on `Response` drop, so just take the data. | ||
| let mut rsp = wrapped_rsp.unwrap(); | ||
| rsp.data.take().unwrap() | ||
| }) |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win
Prefer .expect("reason") over bare .unwrap() in library code.
Both unwraps here are on the caller/hot path in non-test code. Encoding the invariants documented in the comments as .expect() messages preserves them at the panic site.
♻️ Proposed change
- let mut rsp = wrapped_rsp.unwrap();
- rsp.data.take().unwrap()
+ let mut rsp = wrapped_rsp.expect("sender always sends before dropping");
+ rsp.data.take().expect("response data is taken exactly once")As per coding guidelines: "Avoid bare .unwrap(); ... If unavoidable, use .expect(\"reason\")."
📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| rx.map(|wrapped_rsp| { | |
| // A cancel error can't occur: the sender always sends before dropping. | |
| // The reservation is refunded on `Response` drop, so just take the data. | |
| let mut rsp = wrapped_rsp.unwrap(); | |
| rsp.data.take().unwrap() | |
| }) | |
| rx.map(|wrapped_rsp| { | |
| // A cancel error can't occur: the sender always sends before dropping. | |
| // The reservation is refunded on `Response` drop, so just take the data. | |
| let mut rsp = wrapped_rsp.expect("sender always sends before dropping"); | |
| rsp.data.take().expect("response data is taken exactly once") | |
| }) |
🤖 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 `@rust/lance-io/src/scheduler.rs` around lines 1027 - 1032, The hot-path
`rx.map` logic in `scheduler.rs` uses bare `.unwrap()` twice in library code;
replace both with `.expect(...)` so the documented invariants are enforced at
the panic site. Update the `wrapped_rsp.unwrap()` and `rsp.data.take().unwrap()`
calls in this closure with clear reason messages that match the surrounding
comments about the sender always sending before dropping and the reservation
being refunded on `Response` drop.
Source: Coding guidelines
Comment on lines
+700
to
+702
| // Freed budget may make queued tasks runnable; there is no caller to surface | ||
| // an error to here. | ||
| let _ = self.on_task_complete(state); |
There was a problem hiding this comment.
🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick win
Log the swallowed error instead of discarding it silently.
on_task_complete can fail here (e.g. a BrokenTaskError while starting newly-runnable queued tasks). Since abandon runs from Drop and cannot propagate, discarding via let _ hides a real failure. Emit a warn! so the condition is observable.
♻️ Proposed change
- // Freed budget may make queued tasks runnable; there is no caller to surface
- // an error to here.
- let _ = self.on_task_complete(state);
+ // Freed budget may make queued tasks runnable; there is no caller to surface
+ // an error to here, so log on failure instead of propagating.
+ if let Err(e) = self.on_task_complete(state) {
+ log::warn!("Error re-checking queue after abandoning task {task_id}: {e}");
+ }As per coding guidelines: "Log warnings on best-effort or cleanup failures instead of silently swallowing or propagating errors."
📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| // Freed budget may make queued tasks runnable; there is no caller to surface | |
| // an error to here. | |
| let _ = self.on_task_complete(state); | |
| // Freed budget may make queued tasks runnable; there is no caller to surface | |
| // an error to here, so log on failure instead of propagating. | |
| if let Err(e) = self.on_task_complete(state) { | |
| log::warn!("Error re-checking queue after abandoning task {task_id}: {e}"); | |
| } |
🤖 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 `@rust/lance-io/src/scheduler/lite.rs` around lines 700 - 702, The cleanup path
in `lite.rs` is swallowing a failure from `on_task_complete` by assigning it to
`_`, which hides observable errors during `abandon`/`Drop`. Update the
`on_task_complete` call site in `abandon` to capture the result and emit a
`warn!` when it returns an error, using the existing `state`/task context if
available, so best-effort cleanup failures are logged instead of silently
discarded.
Source: Coding guidelines
Cleans up callsites a bit.
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 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.
Inline comments:
In `@rust/lance-io/src/scheduler/lite.rs`:
- Around line 110-121: Collapse the redundant match in BackpressureError::new
into a single struct literal by assigning task_state.backpressure_reservation()
directly to backpressure_reservation; the None/Some(reservation) branching is
unnecessary because the method already returns the correct Option value. Keep
the message field unchanged and preserve the existing constructor behavior in
new.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: ASSERTIVE
Plan: Pro Plus
Run ID: a3521bcb-9274-4447-88cd-2a8324b3dc70
📒 Files selected for processing (1)
rust/lance-io/src/scheduler/lite.rs
Comment on lines
110
to
121
| fn new(task_state: TaskState, message: String) -> Self { | ||
| match task_state { | ||
| TaskState::Reserved { | ||
| backpressure_reservation, | ||
| .. | ||
| } | ||
| | TaskState::Running { | ||
| backpressure_reservation, | ||
| .. | ||
| } | ||
| | TaskState::Finished { | ||
| backpressure_reservation, | ||
| .. | ||
| } => Self { | ||
| match task_state.backpressure_reservation() { | ||
| None => Self { | ||
| message, | ||
| backpressure_reservation: Some(backpressure_reservation), | ||
| backpressure_reservation: None, | ||
| }, | ||
| TaskState::Broken | TaskState::Initial { .. } => Self { | ||
| Some(reservation) => Self { | ||
| message, | ||
| backpressure_reservation: None, | ||
| backpressure_reservation: Some(reservation), | ||
| }, | ||
| } | ||
| } |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value
Collapse the redundant match into a single struct literal.
Both arms build the same Self, differing only in None/Some(reservation), which is exactly what backpressure_reservation() already returns.
♻️ Proposed simplification
fn new(task_state: TaskState, message: String) -> Self {
- match task_state.backpressure_reservation() {
- None => Self {
- message,
- backpressure_reservation: None,
- },
- Some(reservation) => Self {
- message,
- backpressure_reservation: Some(reservation),
- },
- }
+ Self {
+ backpressure_reservation: task_state.backpressure_reservation(),
+ message,
+ }
}📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| fn new(task_state: TaskState, message: String) -> Self { | |
| match task_state { | |
| TaskState::Reserved { | |
| backpressure_reservation, | |
| .. | |
| } | |
| | TaskState::Running { | |
| backpressure_reservation, | |
| .. | |
| } | |
| | TaskState::Finished { | |
| backpressure_reservation, | |
| .. | |
| } => Self { | |
| match task_state.backpressure_reservation() { | |
| None => Self { | |
| message, | |
| backpressure_reservation: Some(backpressure_reservation), | |
| backpressure_reservation: None, | |
| }, | |
| TaskState::Broken | TaskState::Initial { .. } => Self { | |
| Some(reservation) => Self { | |
| message, | |
| backpressure_reservation: None, | |
| backpressure_reservation: Some(reservation), | |
| }, | |
| } | |
| } | |
| fn new(task_state: TaskState, message: String) -> Self { | |
| Self { | |
| backpressure_reservation: task_state.backpressure_reservation(), | |
| message, | |
| } | |
| } |
🤖 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 `@rust/lance-io/src/scheduler/lite.rs` around lines 110 - 121, Collapse the
redundant match in BackpressureError::new into a single struct literal by
assigning task_state.backpressure_reservation() directly to
backpressure_reservation; the None/Some(reservation) branching is unnecessary
because the method already returns the correct Option value. Keep the message
field unchanged and preserve the existing constructor behavior in new.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What
This fixes #7619.
lance-ioimplements a reservation system to create backpressure. A future needed to be driven to completion for a reservation to be returned, leading to a deadlock when a lot of futures got cancelled. Both schedulers are affected.Real-world trigger: an in-flight read cancelled mid-flight, e.g. an async request handler dropped on client disconnect, a
timeout/select!cancelling a query, or an early stream drop.Root cause
scheduler.rs):on_bytes_consumedis called only from the caller-siderx.mapclosure insubmit_request_standard. If the caller drops the future, the closure never runs; the spawned server task still completes and itstx.sendfails silently, so therefund is skipped.
scheduler/lite.rs): the in-flightIoTask(owning itsBackpressureReservation) lives inIoQueueState.tasksand is released only inIoQueue::pollon theFinishedtransition.TaskHandlehad noDrop, so aborting the caller stranded the taskRunningwith its reservation held.Fix
Return reservations on drop.
Note
I used Claude Code Opus 4.8 to create this fix.
Summary by CodeRabbit