fix(mem-wal): propagate task shutdown failures - #7915
Conversation
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthrough
ChangesTask shutdown error propagation
Estimated code review effort: 3 (Moderate) | ~20 minutes Possibly related PRs
Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
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 |
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/src/dataset/mem_wal/write.rs`:
- Around line 544-565: Update the public rustdoc for TaskExecutor::shutdown_all
to document that it cancels tasks, joins every task, and returns the first
cleanup or panic failure after all joins complete. Add a representative Rust
example and links to TaskExecutor::add_handler and the handler lifecycle
documentation, following the project’s public API documentation conventions.
- Around line 4726-4825: The shutdown tests around
test_task_executor_shutdown_propagates_cleanup_error_and_joins_all_tasks and
test_task_executor_shutdown_propagates_task_panic need coverage for first-error
precedence. Add two failing handlers ordered so the first produces a cleanup
error and the later panics, then assert shutdown_all returns the first cleanup
error while also verifying the later task was joined and cleanup executed.
🪄 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: QUIET
Plan: Pro Plus
Run ID: d8f63b66-e679-4ceb-b150-15d194813406
📒 Files selected for processing (1)
rust/lance/src/dataset/mem_wal/write.rs
1a61061 to
d78b693
Compare
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 (2)
rust/lance/src/dataset/mem_wal/write.rs (2)
589-616: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winUnbounded
handle.awaitcan hang shutdown indefinitely.The join loop has no timeout on
handle.await; if any registered handler'scleanup()never returns (blocks, deadlocks, or loops),shutdown_all()— and thereforeShardWriter::close, per the PR objectives — will hang forever with no diagnostic. This turns a single misbehaving handler into a full shutdown/close stall.Consider wrapping each join in
tokio::time::timeout(..), logging/recording a timeout as an error (and optionally callinghandle.abort()), so shutdown always makes forward progress.🛡️ Proposed fix sketch
for (name, handle) in tasks { - match handle.await { + match tokio::time::timeout(SHUTDOWN_JOIN_TIMEOUT, handle).await { + Err(_) => { + error!("Task '{}' did not shut down within timeout", name); + if first_error.is_none() { + first_error = Some(Error::internal(format!( + "Task '{name}' did not shut down within {SHUTDOWN_JOIN_TIMEOUT:?}" + ))); + } + } + Ok(join_result) => match join_result { Ok(Ok(())) => debug!("Task '{}' completed successfully", name), Ok(Err(e)) => { warn!("Task '{}' completed with error: {}", name, e); if first_error.is_none() { first_error = Some(e); } } Err(e) => { error!("Task '{}' panicked: {}", name, e); if first_error.is_none() { first_error = Some(Error::internal(format!( "Task '{name}' panicked during shutdown: {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/src/dataset/mem_wal/write.rs` around lines 589 - 616, Update shutdown_all to await each task handle with a bounded tokio::time::timeout instead of unbounded handle.await. Treat an elapsed timeout as the first shutdown error, log which task exceeded the limit, and abort the timed-out handle if appropriate, while preserving the existing success, task-error, and panic handling.
4776-4875: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick winAdd the two-failing-handlers shutdown test. The current cases cover one cleanup failure plus success, and a panic in isolation, but not the precedence case where both handlers fail. Add a test with a cleanup error first and a panic second, asserting
shutdown_all()returns the first registered error while still joining the later task.🤖 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/src/dataset/mem_wal/write.rs` around lines 4776 - 4875, Add a test alongside test_task_executor_shutdown_propagates_cleanup_error_and_joins_all_tasks and test_task_executor_shutdown_propagates_task_panic with two handlers whose cleanup operations fail: register an error-returning handler first and a panicking handler second. Await shutdown_all(), assert it returns the first handler’s cleanup error rather than the panic, and verify the later panicking task was still joined and all tasks were removed.
🤖 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 `@rust/lance/src/dataset/mem_wal/write.rs`:
- Around line 589-616: Update shutdown_all to await each task handle with a
bounded tokio::time::timeout instead of unbounded handle.await. Treat an elapsed
timeout as the first shutdown error, log which task exceeded the limit, and
abort the timed-out handle if appropriate, while preserving the existing
success, task-error, and panic handling.
- Around line 4776-4875: Add a test alongside
test_task_executor_shutdown_propagates_cleanup_error_and_joins_all_tasks and
test_task_executor_shutdown_propagates_task_panic with two handlers whose
cleanup operations fail: register an error-returning handler first and a
panicking handler second. Await shutdown_all(), assert it returns the first
handler’s cleanup error rather than the panic, and verify the later panicking
task was still joined and all tasks were removed.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: QUIET
Plan: Pro Plus
Run ID: fa1910c3-250f-4610-b4a1-579a9d86a89c
📒 Files selected for processing (1)
rust/lance/src/dataset/mem_wal/write.rs
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
hamersaw
left a comment
There was a problem hiding this comment.
Lets just remove the doc example, otherwise looks great. Thanks!
c3030d0 to
dd2c506
Compare
thanks! already done. |
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/src/dataset/mem_wal/write.rs`:
- Around line 547-549: Update the lifecycle coordination between shutdown_all
and add_handler so shutdown atomically marks the executor as shutting down while
taking the task list, and add_handler rejects registrations after that state is
set with a descriptive error. Ensure rejected tasks are not spawned or lost, and
add a regression test covering registration racing with shutdown.
🪄 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: QUIET
Plan: Pro Plus
Run ID: ee8ef9cc-a492-4323-b0be-9d3953eda24d
📒 Files selected for processing (1)
rust/lance/src/dataset/mem_wal/write.rs
| self.cancellation_token.cancel(); | ||
|
|
||
| let tasks = std::mem::take(&mut *self.tasks.write().unwrap()); |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | 🏗️ Heavy lift
Prevent handlers from being registered after shutdown begins.
At Line 549, a concurrent add_handler can push a newly spawned task after shutdown_all takes tasks. That task is neither joined nor included in the returned error, despite already receiving the cancelled token. Serialize registration and shutdown with a lifecycle state: atomically mark the executor as shutting down while taking tasks, and reject later registrations with a descriptive error. Add a regression test for this interleaving.
🤖 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/src/dataset/mem_wal/write.rs` around lines 547 - 549, Update the
lifecycle coordination between shutdown_all and add_handler so shutdown
atomically marks the executor as shutting down while taking the task list, and
add_handler rejects registrations after that state is set with a descriptive
error. Ensure rejected tasks are not spawned or lost, and add a regression test
covering registration racing with shutdown.
hamersaw
left a comment
There was a problem hiding this comment.
Fantastic, thanks! Will merge when the CI passes.
Summary
TaskExecutor::shutdown_allCloses #7914
Testing
cargo test -p lance task_executor_shutdown --libcargo test -p lance test_task_dispatcher_survives_handle_error --libcargo fmt --all -- --checkcargo clippy --all --tests --benches -- -D warnings