1537: feat: move shuffle writer disk I/O off tokio worker threads - #10
1537: feat: move shuffle writer disk I/O off tokio worker threads#10martin-augment wants to merge 3 commits into
Conversation
Synchronous file I/O inside async shuffle writer functions blocks tokio workers, hurting throughput and responsiveness under concurrent loads. This wraps all disk writes in spawn_blocking tasks fed through bounded mpsc channels, so workers stay free while Arrow StreamWriter and BatchPartitioner run on the blocking thread pool. Both the no-repartition and hash repartition paths use the same pattern: the async side reads the input stream and sends batches through a channel with capacity 2, while the blocking side writes them to disk. The sender is always dropped and the blocking handle is always awaited before returning, even on error. Also fixes write_time double-counting on the no-repartition path and adds write-failure regression tests for both code paths.
Adds a --concurrency flag that runs multiple input partitions in parallel using tokio::spawn with a Semaphore for backpressure. Defaults to 1, preserving the original sequential behavior. This makes it possible to measure how shuffle writes perform under realistic concurrent workloads, where the async I/O changes have the most impact.
WalkthroughThis pull request adds AI agent guidance rules prohibiting links to GitHub issues/pull requests in code reviews and review of agent configuration files. It refactors ballista shuffle and file I/O operations from synchronous to async patterns using tokio's ✨ Finishing Touches📝 Generate docstrings
🧪 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 |
|
Code Review Summary This PR correctly addresses blocking I/O on tokio worker threads by offloading file writes to spawn_blocking tasks and using bounded mpsc channels to bridge the async and blocking worlds. The core approach is sound. Below are specific observations. Issues
The outer write_time timer that wrapped both create_dir_all and the write_stream_to_disk call was removed. Now write_time is only incremented inside write_stream_to_disk per-batch. The original code double-counted write time (the outer timer ran concurrently with the inner per-batch timers). The fix silently resolves that double-count without a comment - worth noting in a code comment that the inner metric is now the canonical measure.
In both write_stream_to_disk and the hash partition path in shuffle_writer.rs, when stream_err.is_some() the write task result is discarded. If the write task has an independent failure (e.g. disk full detected during writer.finish() after the stream errored), that error is silently swallowed. Consider at minimum logging it with warn!().
A buffer of 2 means the async reader blocks after 2 un-consumed batches. For small batches with slow I/O this forces frequent context switches. A value of 4-8 would reduce round-trips without meaningfully increasing memory pressure. Alternatively, expose it as a configurable constant.
The semaphore-based concurrent execution block is copy-pasted verbatim between benchmark_hash_shuffle and benchmark_sort_shuffle. Extracting a generic helper would make future changes easier to apply consistently.
num_bytes += batch.get_array_memory_size() measures uncompressed in-memory size. After LZ4 compression the on-disk size is smaller. This was a pre-existing issue, but this refactor is a good opportunity to fix it - use fs::metadata(&path)?.len() after writing (as already done in the hash-partition finalizer) or track the count inside the blocking task. Minor Observations
Test Coverage The two new tests (test_no_repart_write_failure_propagates and test_hash_repart_write_failure_propagates) cover the primary error paths introduced by this refactor. Good addition. It would also be valuable to add a test that verifies metrics (input_rows, output_rows) are correctly recorded after the async/blocking split, since those counters now cross a thread boundary via cloned handles. Overall The design is correct and the change achieves its goal. Addressing the silent error discard (issue 2) and the in-memory vs on-disk byte count (issue 5) before merging would be worthwhile. The rest are low-priority polish items. |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request refactors the shuffle writer and disk writing utilities to offload synchronous file I/O to blocking tasks using tokio::task::spawn_blocking and mpsc channels, preventing the async executor from being blocked. It also adds a concurrency parameter to the shuffle benchmarks and introduces unit tests for error propagation during write failures. Feedback suggests increasing the mpsc channel capacity from 2 to 16 in both shuffle_writer.rs and utils.rs to improve throughput and pipelining by reducing context switching.
| path, | ||
| }); | ||
| let schema = stream.schema(); | ||
| let (tx, mut rx) = tokio::sync::mpsc::channel::<RecordBatch>(2); |
There was a problem hiding this comment.
The channel capacity of 2 is quite low for a shuffle operation. While it provides backpressure, such a small buffer might lead to excessive context switching between the async stream reader and the blocking writer task, potentially reducing throughput. Consider increasing this to a larger value like 16 or 32 to allow for better pipelining of record batches.
| let (tx, mut rx) = tokio::sync::mpsc::channel::<RecordBatch>(2); | |
| let (tx, mut rx) = tokio::sync::mpsc::channel::<RecordBatch>(16); |
There was a problem hiding this comment.
value:good-to-have; category:bug; feedback: The Gemini AI reviewer is correct! The channel used for passing back the data from the blocking task may contain at most two items. If the receiver is slow for any reason then the sender will block and the thread context will switch. A channel with more slots will allow to reduce the thread context switching by using more memory.
| let path_owned = path.to_owned(); | ||
| let write_metric = disk_write_metric.clone(); | ||
|
|
||
| let (tx, mut rx) = tokio::sync::mpsc::channel::<RecordBatch>(2); |
There was a problem hiding this comment.
The channel capacity of 2 is quite low for high-throughput disk I/O. A larger buffer (e.g., 16 or 32) would allow the async stream to continue reading and processing the next batches while the blocking task is busy with synchronous file writes, improving overall performance through better pipelining.
| let (tx, mut rx) = tokio::sync::mpsc::channel::<RecordBatch>(2); | |
| let (tx, mut rx) = tokio::sync::mpsc::channel::<RecordBatch>(16); |
There was a problem hiding this comment.
value:good-to-have; category:bug; feedback: The Gemini AI reviewer is correct! The channel used for passing back the data from the blocking task may contain at most two items. If the receiver is slow for any reason then the sender will block and the thread context will switch. A channel with more slots will allow to reduce the thread context switching by using more memory.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (1)
benchmarks/src/bin/shuffle_bench.rs (1)
293-326: Code duplication with hash shuffle benchmark.The concurrent execution logic (lines 293-326) is nearly identical to
benchmark_hash_shuffle(lines 199-232). For a benchmark utility this is acceptable, but consider extracting a helper if this pattern is reused further.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@benchmarks/src/bin/shuffle_bench.rs` around lines 293 - 326, Extract the repeated concurrent-partition execution block used in both the current shuffle benchmark and benchmark_hash_shuffle into a shared helper function (e.g., run_concurrent_partitions) that takes the shuffle writer (shuffle_writer), task context (task_ctx), input_partition_count, and concurrency and returns Result<total_files, Box<dyn std::error::Error>>; inside the helper reuse the same logic: create Arc<tokio::sync::Semaphore>, spawn tasks that acquire_owned a permit, call writer.execute(partition, ctx) and collect via utils::collect_stream, count rows via batches.first().map_or(0, |b| b.num_rows()), drop the permit, gather results from handles and return the first error if any, otherwise the total_files; replace the duplicated blocks in both functions with calls to this helper to remove duplication while keeping existing symbols (shuffle_writer, task_ctx, utils::collect_stream, execute, acquire_owned).
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Nitpick comments:
In `@benchmarks/src/bin/shuffle_bench.rs`:
- Around line 293-326: Extract the repeated concurrent-partition execution block
used in both the current shuffle benchmark and benchmark_hash_shuffle into a
shared helper function (e.g., run_concurrent_partitions) that takes the shuffle
writer (shuffle_writer), task context (task_ctx), input_partition_count, and
concurrency and returns Result<total_files, Box<dyn std::error::Error>>; inside
the helper reuse the same logic: create Arc<tokio::sync::Semaphore>, spawn tasks
that acquire_owned a permit, call writer.execute(partition, ctx) and collect via
utils::collect_stream, count rows via batches.first().map_or(0, |b|
b.num_rows()), drop the permit, gather results from handles and return the first
error if any, otherwise the total_files; replace the duplicated blocks in both
functions with calls to this helper to remove duplication while keeping existing
symbols (shuffle_writer, task_ctx, utils::collect_stream, execute,
acquire_owned).
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: b5806916-4246-4aff-9ab6-eb01b430bdea
📒 Files selected for processing (8)
.cursor/rules.mdAGENTS.mdCLAUDE.mdballista/core/Cargo.tomlballista/core/src/execution_plans/shuffle_writer.rsballista/core/src/utils.rsbenchmarks/Cargo.tomlbenchmarks/src/bin/shuffle_bench.rs
value:good-to-have; category:bug; feedback: The CodeRabbit AI reviewer is correct! The logic for setup is duplicated for the two benchmark tests. It would be good to extract it to a helper method and reuse it. |
value:good-to-have; category:bug; feedback: The Claude AI reviewer is correct! The logic for setup is duplicated for the two benchmark tests. It would be good to extract it to a helper method and reuse it. |
value:useful; category:bug; feedback: The Claude AI reviewer is correct! Two unrelated operations are executed and both could produce their own errors. One of them is threaten with priority and the other one is just silently dropped. It would be better to at least log it for full context. |
value:good-to-have; category:bug; feedback: The Claude AI reviewer is correct! The channel used for passing back the data from the blocking task may contain at most two items. If the receiver is slow for any reason then the sender will block and the thread context will switch. A channel with more slots will allow to reduce the thread context switching by using more memory. |
🤖 Augment PR SummarySummary: This PR moves shuffle-writer disk I/O off tokio worker threads by routing record batches through bounded channels to blocking writer tasks. Changes:
Technical Notes: Both shuffle paths now use small bounded channels (size 2) to provide backpressure while keeping blocking filesystem and IPC writes off the async runtime threads. 🤖 Was this summary useful? React with 👍 or 👎 |
![augmentcode[bot]](https://avatars.githubusercontent.com/in/1027498?s=60&v=4){kind=small}
| let timer = disk_write_metric.timer(); | ||
| writer.write(&batch)?; | ||
| timer.done(); | ||
| if let Some(e) = stream_err { |
There was a problem hiding this comment.
ballista/core/src/utils.rs:225: When stream_err is Some(_), this returns the stream error without checking write_result, which can mask a concurrent disk-writer failure from the spawn_blocking task. If both can happen, it may be worth deciding which error should win (or logging/combining them) to avoid losing the underlying I/O failure.
Severity: medium
Other Locations
ballista/core/src/execution_plans/shuffle_writer.rs:359
🤖 Was this useful? React with 👍 or 👎, or 🚀 if it prevented an incident/outage.
There was a problem hiding this comment.
value:useful; category:bug; feedback: The Augment AI reviewer is correct! Two unrelated operations are executed and both could produce their own errors. One of them is threaten with priority and the other one is just silently dropped. It would be better to at least log it for full context.
3 participants
1537: To review by AI