0xdeafbeef commented on code in PR #2780:
URL:
https://github.com/apache/datafusion-ballista/pull/2780#discussion_r4512544333
##########
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:
Added `later_calls_succeed_after_the_first_stream_completed`: it drains
`execute(0)` inside a block so the stream is dropped, checks that the
coordinator's strong count is 0, then drains `execute(1)`.
##########
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:
Wrapped at 80.
##########
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:
Confirmed. `execute_query_stage` maps the abort to
`BallistaError::Cancelled`, so no `PartitionLocation` is registered and nothing
reads the file. When the job is cancelled or fails, the scheduler tells every
executor to remove the job's data at once; when it succeeds, after
`finished_job_data_clean_up_interval_seconds`. The client can also send
`CleanJobData`, and the executor's TTL cleaner covers the rest when it is
enabled; all of them remove the whole job directory.
The range writer behaves the same way and can also leave its data file
without an index, because `write_index_file` runs after the await. The sort
writer writes its data and index files synchronously at the end of the
coordinator, so it leaves either nothing or a complete file.
Added a comment at the spawn.
##########
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:
Added to the description: when one `execute(N)` task fails,
`drive_shuffle_writer_stage` now aborts the others instead of leaving them
running.
##########
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:
Both use 10 seconds now.
--
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]
