andygrove commented on code in PR #2780:
URL:
https://github.com/apache/datafusion-ballista/pull/2780#discussion_r4494785986
##########
ballista/executor/src/execution_engine.rs:
##########
@@ -538,7 +541,7 @@ async fn drive_shuffle_writer_stage(
for n in 0..k {
let plan = plan.clone();
let ctx = context.clone();
- stream_futures.push(tokio::spawn(async move {
+ stream_futures.push(SpawnedTask::spawn(async move {
Review Comment:
Question about the partial file after a cancel.
`utils::write_stream_to_disk` does its file I/O in `spawn_blocking`, which
`abort` can't stop. When the future is dropped the sender closes, so the
blocking thread leaves its loop and calls `finish()`. That leaves a truncated
but valid IPC file at the path, where before it would have been complete.
I think that's harmless as long as the task reports failure and job cleanup
removes the file. Could you confirm that, and maybe add a short comment here
saying so?
##########
ballista/executor/src/execution_engine.rs:
##########
@@ -522,6 +523,8 @@ impl QueryStageExecutor for PipelinedQueryStageExec {
/// collect metadata batches from each, and turn them back into
/// `Vec<ShuffleWritePartition>`. All K streams must be driven concurrently
/// so the writer's internal coordinator sees every oneshot receiver taken.
+/// A cancelled task drops the returned future. That aborts the K tasks,
+/// drops their streams and stops the writer's coordinator.
Review Comment:
Small behavior change worth mentioning in the PR description. If an early
partition's task returns an error, the `?` here now drops the remaining
`SpawnedTask`s and aborts them. Before, they kept running detached. That seems
strictly better to me, but it goes beyond the cancel case the description talks
about.
##########
ballista/core/src/execution_plans/shuffle_writer.rs:
##########
@@ -1092,45 +1184,12 @@ pub(crate) async fn run_coordinator<F>(
}
}
-/// Takes the handoff receiver for `partition`. The first call creates one
-/// oneshot per output partition and passes the senders to
-/// `spawn_coordinator`. `name` is the calling exec, used in the error for a
-/// repeated or out-of-range partition.
-pub(crate) fn take_handoff(
- name: &str,
- initialized: &mut bool,
- handoffs: &mut
[Option<oneshot::Receiver<Result<Vec<ShuffleWritePartition>>>>],
- partition: usize,
- spawn_coordinator: impl
FnOnce(Vec<oneshot::Sender<Result<Vec<ShuffleWritePartition>>>>),
-) -> Result<oneshot::Receiver<Result<Vec<ShuffleWritePartition>>>> {
- if !*initialized {
- *initialized = true;
- let mut senders:
Vec<oneshot::Sender<Result<Vec<ShuffleWritePartition>>>> =
- Vec::with_capacity(handoffs.len());
- for slot in handoffs.iter_mut() {
- let (tx, rx) = oneshot::channel();
- senders.push(tx);
- *slot = Some(rx);
- }
- spawn_coordinator(senders);
- }
-
- let k = handoffs.len();
- handoffs
- .get_mut(partition)
- .and_then(Option::take)
- .ok_or_else(|| {
- DataFusionError::Internal(format!(
- "{name}: execute({partition}) called twice or out of range
(have {k} partitions)"
- ))
- })
-}
-
-/// Wraps a receiver from [`take_handoff`] in the metadata stream returned by
-/// `execute(N)`.
+/// Wraps a handoff from [`WriterState::take_handoff`] in the metadata stream
returned by
Review Comment:
Nit: this first doc line runs past the line width. rustfmt doesn't wrap
comments, so it might be worth rewrapping by hand.
##########
ballista/executor/src/executor.rs:
##########
@@ -599,6 +611,22 @@ mod test {
assert_eq!(executor.active_task_count(), expected);
}
+ /// Waits until `refs` has `expected` strong references, and fails after
+ /// 5 seconds.
+ async fn assert_strong_count_reaches<T>(refs: &Weak<T>, expected: usize) {
Review Comment:
Nit: this helper is a copy of `assert_strong_count_reaches` in the
`ballista-core` tests, with a 5s timeout here and 10s there. They live in
separate crates so the duplication is fine, but it would be nice to use the
same timeout in both.
##########
ballista/core/src/execution_plans/shuffle_writer.rs:
##########
@@ -234,20 +235,115 @@ pub const DEFAULT_SHUFFLE_CHANNEL_CAPACITY: usize = 8;
/// Shared state for the coordinator+handoff pattern (see range_repartition.rs
/// upstream for the reference design). The first `execute(N)` call spawns a
/// single coordinator task that owns all K writes; every `execute(N)` call
-/// takes its own `oneshot::Receiver` and awaits the summaries for partition N.
+/// takes its own [`Handoff`] and awaits the summaries for partition N.
///
/// The receiver payload is a `Vec` so writers that produce multiple files per
/// output partition (e.g. `SortShuffleWriterExec` under a multi-partition
/// task, where each slice member yields one file each containing K logical
/// partitions) can hand the full set to a single `execute(N)` stream.
/// Passthrough / hash-repart writers produce at most one summary per slot.
pub(crate) struct WriterState {
- pub(crate) initialized: bool,
+ /// The coordinator spawned by the first `execute(N)` call, or `None`
+ /// before that call. The state holds it weakly, so only the handoffs
+ /// keep it running (see [`Handoff`]).
+ coordinator: Option<Weak<SpawnedTask<()>>>,
/// One receiver per output partition. `execute(N)` takes `handoffs[N]`;
/// the coordinator holds the matching sender and pushes the summaries
/// once partition N's files are closed.
- pub(crate) handoffs:
- Vec<Option<oneshot::Receiver<Result<Vec<ShuffleWritePartition>>>>>,
+ handoffs:
Vec<Option<oneshot::Receiver<Result<Vec<ShuffleWritePartition>>>>>,
+}
+
+impl WriterState {
+ /// Creates the state of a writer with `output_partition_count` output
+ /// partitions, before its first `execute(N)` call.
+ pub(crate) fn new(output_partition_count: usize) -> Self {
+ Self {
+ coordinator: None,
+ handoffs: (0..output_partition_count).map(|_| None).collect(),
+ }
+ }
+
+ /// Takes the handoff for `partition`, spawning the coordinator on the
+ /// first call.
+ ///
+ /// The first call creates one oneshot per output partition and passes the
+ /// senders to `spawn_coordinator`, which spawns the coordinator and
returns
+ /// its task. Fails when `partition` is out of range or the handoff of
+ /// `partition` was already taken. The range check runs before the
+ /// coordinator is spawned. `name`, the calling exec's name, starts the
+ /// error message.
+ pub(crate) fn take_handoff(
+ &mut self,
+ name: &str,
+ partition: usize,
+ spawn_coordinator: impl FnOnce(
+ Vec<oneshot::Sender<Result<Vec<ShuffleWritePartition>>>>,
+ ) -> SpawnedTask<()>,
+ ) -> Result<Handoff> {
+ let k = self.handoffs.len();
+ let taken_or_out_of_range = || {
+ DataFusionError::Internal(format!(
+ "{name}: execute({partition}) called twice or out of range
(have {k} partitions)"
+ ))
+ };
+ if partition >= k {
+ return Err(taken_or_out_of_range());
+ }
+
+ let coordinator = match &self.coordinator {
+ Some(coordinator) => coordinator.upgrade(),
Review Comment:
The `upgrade()` returning `None` branch is subtle but looks right to me.
When the coordinator has already finished, the oneshot already holds the
summaries, so a later `execute(N)` still succeeds. It also relies on the
coordinator sending every summary in one synchronous loop, so an abort can't
land between sends.
Would you add a small test that fully consumes `execute(0)` and then calls
`execute(1)`? That would guard the "coordinator already finished" case against
future refactors.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
