[andygrove]

Which issue does this PR close?

Part of #2652. This is the second of a series of smaller PRs split out of #2657 (pipelined execution). It is independent of #2669.

Rationale for this change

In pipelined execution, every task of a job runs at once and stages stream record batches to each other instead of exchanging shuffle files. This PR adds the exchange itself: the operators at either end of a pipelined stage boundary, the per-executor registry that connects them, and the connection handling a wide job needs. Nothing plans these operators yet, so this PR can't change how existing jobs run.

#2657 describes the whole feature, its design notes and SF1000 results.

What changes are included in this PR?

• execution_plans/pipelined/
  • PipelinedRegistry holds a bounded queue for each output partition of each running producer task, keyed by job, stage and task.
    • Broadcast outputs, which every consumer task reads in full, are kept in replay buffers instead.
    • It remembers recently removed jobs and stages, so a leftover task fails at once instead of waiting out its timeout.
  • PipelinedShuffleWriterExec is a stage root. It drives every input partition and routes batches into the queues: hash-partitioned, passthrough or broadcast.
    • Each queue ends with an explicit end marker, so a consumer can tell a finished partition from a failed producer.
    • The writer stops reading its input once every consumer has gone away.
    • It fails if no consumer opens a full queue within ballista.pipelined.registration_timeout_seconds.
  • PipelinedShuffleReaderExec is a stage leaf. It drains all of its producers concurrently: directly from the registry when the producer is on the same executor, and over Flight otherwise. Reading concurrently is what keeps bounded queues free of deadlock for tree-shaped plans.
    • Remote streams go over one connection per executor pair. The executor's client pool keeps it (BallistaClientPool::shared), and the streams of every task on the executor multiplex over it, so an executor holds one connection to each other executor however wide the job. The first version opened a connection per reader partition and producer executor, which at TPC-H SF1000 on 32 executors asked each executor to accept thousands of connections within a second and overflowed its listen queue.
    • Sharing a connection is safe because of the flow-control windows. A stream nobody is reading, such as a join's probe side while the build side loads, holds up to its stream window of data on the connection, and stalled streams that used up the connection window would starve the streams something is waiting on. The shared connection advertises the largest connection window HTTP/2 allows and a small stream window, ballista.pipelined.stream_window_bytes (512 KiB by default), so stalled streams hold at most one window each and a connection carries over 4000 streams before any could starve the rest.
    • A stream open whose connection fails before the producer's executor answers (a reset, a refused connection, a connection closed under a queued request) is retried on a new connection, ballista.client.io_retries_times times with ballista.client.io_retry_wait_time_ms between attempts. BallistaClient::fetch_pipelined_partition reports those as GrpcConnectionError, and anything the server answered as GrpcActionError, which is never retried: the producer may already have handed the partition to the failed stream.
  • tests.rs runs writers against readers in-process.
• Executor servers. create_grpc_server_incoming listens with GrpcServerConfig::listen_backlog, 4096 by default, instead of the 128 the standard library gives tonic's serve, and the Flight server now binds through it like the gRPC server already did. The kernel caps the value at net.core.somaxconn. create_grpc_server explicitly sets no limit on streams per connection, since a stream the client had to queue behind a limit would leave its producer blocked on a full queue.
• Protobuf and serde: PipelinedShuffleWriterExecNode, PipelinedShuffleReaderExecNode, PipelinedPartitionLocation, and the FetchPipelinedPartition action, plus BallistaClient::fetch_pipelined_partition. The Flight proxy in serving.rs forwards the new action.
• Settings: ballista.execution.mode (blocking or pipelined), ballista.pipelined.queue_capacity, ballista.pipelined.registration_timeout_seconds and ballista.pipelined.stream_window_bytes. They are documented in a new section of the config reference.
• Guards until the rest lands:
  • The scheduler rejects a job with ballista.execution.mode = pipelined ("not supported yet").
  • The executor's Flight service answers FetchPipelinedPartition with unimplemented.
  • Later PRs in the series replace both guards.

How are these changes tested?

• Unit tests for the registry, covering open and claim, timeouts, replay, cancellation and removed jobs.
• Unit tests for the writer and reader, run against each other in-process, covering hash, passthrough and broadcast outputs, end markers, producer failure, and early consumer drop.
• Reader tests for the connection handling: streams to one executor share one connection; a stream open whose first connection is reset (through a proxy that resets it) is retried on a second connection and delivers its rows; a failure the server answered is not retried.
• A test for the executor pool's shared connection: one per key, another per configuration, replaced after a discard.
• Serde round trips and config parsing.
• A scheduler test checks that a pipelined job is rejected and no graph is created.

Are there any user-facing changes?

There are four new session settings, documented in the config reference. Setting ballista.execution.mode = pipelined currently makes the job fail with a "not supported yet" error. The default behavior is unchanged.

Two public types change: GrpcServerConfig gains a listen_backlog field, which affects code that builds it by struct literal rather than Default, and BallistaClientPool gains a shared method with a default implementation, so existing implementations keep compiling. The executor's gRPC and Flight listeners get a 4096 backlog in every mode.

[andygrove]

cc @phillipleblanc

[comphead]

Thanks @andygrove, splitting this out of #2657 makes it much easier to review. Notes below, most important first. Items marked (unverified) are expectations I did not test.

Correctness and design

1. remove_job does not release a producer blocked on an unclaimed queue in the real writer. registry.rs documents that remove_job closes unclaimed queues, and removing_the_job_releases_a_producer_blocked_on_an_unclaimed_queue passes because the test drops its Arc<ProducerOutputs>. In execute_pipelined (writer.rs:232-272) the writer keeps outputs alive until the task ends, so the receivers stay open and the producer waits out the claim timeout instead. Closing the receivers explicitly in remove_job and remove_stage (rx.close()) would fix it. The test should hold the Arc the way the writer does.
2. fetch_pipelined_partition may retry after the server has answered. The stream.message().await error arm goes through is_transport_failure, so a reset after the response headers arrive becomes GrpcConnectionError and is retried. If the executor claims the queue before it answers (the docs imply it does, but the server side is not in this PR), the retry finds a claimed or released queue. It then waits the full registration timeout and reports "timed out waiting for pipelined producer to start", and the batches the first stream took are lost. Consider mapping only errors from do_get().await to GrpcConnectionError. A test that resets the connection after the headers would show which case applies.
3. is_transport_failure matches status.message() == "transport error". String matching is fragile. Unavailable is also retried even when the server sent it, for example from a draining executor. Source-chain detection alone may be enough.
4. Memory is outside the DataFusion MemoryPool. Queues (queue_capacity batches per output partition per task), the per-output BatchCoalescers and the broadcast ReplayBuffer (kept until remove_job) are untracked. Is accounting planned for a later PR? Also, where is the broadcast size bound enforced? The comment points at the planner threshold, which is not in this PR.
5. The claim timeout is an absolute deadline from ProducerOutputs::new, not an idle timeout. A consumer plan that calls execute on a child lazily, after a long first phase, would hit "no consumer opened output partition" while the producer's queue is full. The hash join probe side is opened eagerly, so it is fine. Other plan shapes may not be (unverified). A doc note on the plan shapes this relies on would help.
6. A second open of a released queue gives a misleading error. Once all queues are claimed, release drops the slot. A repeated open then creates a new Waiting slot and fails after the timeout with "timed out ... to start". An explicit "already read" error would be clearer, and a test for it is missing.
7. create_grpc_server: .max_concurrent_streams(None) appears to be a no-op. tonic 0.14.6 defaults it to None (transport/server/mod.rs:128) and passes it through to hyper unchanged. It is fine as documentation, but the comment reads as if it changes behavior for the scheduler and executor servers.
8. SharedMap in executor/src/client_pool.rs is unbounded. Slots are never removed, and evict_idle ignores them. With executor churn, slots for dead (host, port, config) keys accumulate for the life of the process. Removing the slot on discard, and evicting slots with no outside Arc after the idle timeout, would bound it.
9. No task or executor loss recovery. Queues are one-shot, so any loss fails the job. This is documented in the config text and looks like a deliberate first step. It would be worth stating in Design: optional pipelined execution mode (all tasks placed up front, streaming exchange) #2652 as a non-goal. For comparison, Flink restarts the pipelined region and Spark retries a stage from shuffle files.

Reuse and simplification

• open_remote copies with_retry and is_retriable_fetch_error (shuffle_reader.rs:943-980). Making those pub(crate) removes a second retry policy.
• fetch_pipelined_partition repeats the do_get and schema-read sequence in execute_do_get. A shared helper that takes the error mapper would do.
• ProducerKey duplicates TaskKey (serde/scheduler/mod.rs). Using TaskKey avoids a String allocation per lookup.
• output_partition_ids (writer and Routing) re-derives what compute_global_output_partition_ids computes, and nothing outside writer.rs calls it. Keep one.
• The hash-partitioning proto encode is copied from SortShuffleWriterExec (serde/mod.rs ~1222 vs ~1176). A small helper would serve both.
• with_partitions repeats the slicing in task_builder.rs, and filter_map silently drops out-of-range indexes where slice_by_global_partition errors. Only tests call it today, so it could move to the scheduler PR.
• Writer plumbing: Arc<HashMap<usize, Output>> and Arc<AtomicUsize> are threaded through five signatures, so one Outputs struct would do. Option<Mutex<BatchCoalescer>> plus .expect("hash outputs always coalesce") carries an invariant at runtime. finish_outputs does not use &self. Routing::Hash rebuilds Partitioning::Hash for every input.
• Small items:
  • try_new_broadcast returns Result but cannot fail.
  • ProducerOutputs::schema() is pub and used once.
  • producer_count is only used by tests.
  • The producer failure is stored twice, in ProducerOutputs.failure and in ReplayBuffer.outcome.
  • The io-util tokio feature is only needed by copy_bidirectional in tests.rs, so it can go in dev-dependencies.
• Metric names differ from the siblings (ShuffleWriteMetrics, local_partitions and remote_partitions).
• Question on design: DataFusion's RepartitionExec already hash-partitions and coalesces. Was reimplementing it in the writer (about 120 lines plus the proto field and serde) a deliberate choice to avoid a channel hop?

Efficiency

• Each output's Mutex<BatchCoalescer> is shared by all input tasks and held during the copy. Rows are also copied twice, once by partition_iter and once by the coalescer. Per-input coalescers would avoid the lock at a memory cost (unverified, needs a benchmark).
• drive_hashed awaits sends in order, so one full queue (default capacity 2) stalls the other partitions of that input.
• shared() allocates the key (host String and config clone) and takes the DashMap shard lock and the connecting mutex on every open, even when a connection exists. Checking the existing connection first would keep the hot path read-only.
• execute() builds GrpcClientConfig and shared_connection_config even when every producer is local.
• RegistryState::closed clones job_id for each lookup under the registry mutex.
• Each PipelinedPartitionLocation carries a full ExecutorMetadata, and execute deep-clones the Vec per partition. This follows PartitionLocation, but plan size grows with P x T on wide jobs.
• Small items: a HashMap lookup per piece where partition ids are dense, and get_array_memory_size() per batch for the bytes metric.
• Per-stream throughput is roughly stream window / RTT, so about 400 Mbit/s at 10 ms RTT with the 512 KiB default (my arithmetic, not measured). A note on how the default was chosen would help.

Tests

• Overlap: tests::producer_fails_when_no_consumer_opens_its_full_queue vs registry::producer_gives_up_on_a_full_queue_nobody_opens, and tests::broadcast_outputs_replay_to_every_consumer_task vs registry::broadcast_output_replays_to_every_consumer.
• queue_can_be_opened_once_and_is_released_once_claimed never opens a second time, so the "once" half is untested.
• pipelined_queue_capacity_is_never_zero only checks .max(1).
• unused_port() binds, drops and reuses the port, which is racy. Keeping the listener alive avoids that.
• ProducerOutputs::new(schema(), [0], false, 1, None, WAIT) is repeated about ten times. A helper would shorten the tests.

[andygrove]

Addressed in 0d53dfc. Fixed: unclaimed queues are closed on cleanup (the writer keeps the outputs alive, as you said), a second open of a claimed queue fails at once, only a failure before the server answers is retried, and is_transport_failure uses the source chain alone (a reset shows as Cancelled and a refused connection as Unavailable, both with the transport error as source). The simplifications are in too: TaskKey, one Outputs struct, with_retry and a shared do_get helper, the hash partitioning encoder, single failure store, metric names, dev-dependency, pool fast path, test helpers and listeners kept open.

Filed as follow-ups rather than grow this PR: memory accounting #2672, shared slot bounding #2673, writer send path and the RepartitionExec question #2674, location size #2675. Task loss as a non-goal is already in #2652's comparison table. with_partitions matches how task_builder.rs restricts ShuffleReaderExec today and gets its caller in the scheduler PR. The two overlapping test pairs cover different layers, so I kept both.

[comphead]

Thanks for the quick turnaround, the new head reads much tighter. A second pass on 0d53dfca, focused on simplifications. I prototyped the first group as a patch against this head. It is not compiled or tested, so please treat it as a sketch. I can share it if useful.

Prototyped (untested)

1. SharedClient as a flag instead of a closure. discard() only sets discarded, and the pool's shared() skips a discarded entry and reconnects. That removes the Box<dyn Fn>, the Weak<SharedSlot>, the generations counter and the two-mutex SharedSlot. The slot becomes one tokio::sync::Mutex<Option<Arc<SharedClient>>> per key. CountingPool in tests.rs shrinks the same way, and open_remote_once no longer needs (client, Option<shared>). The no-pool branch can wrap a fresh client in SharedClient. One cost is that the hot path now takes an async lock per open instead of a sync one, which probably does not matter but I have not measured it.
2. Registry slot on tokio::sync::watch. ProducerSlot::{Waiting(Vec<oneshot::Sender>), Ready} becomes watch::Sender<Option<Arc<ProducerOutputs>>>. That drops the waiter list, the is_closed() retain and the double match in register and wait_for.
3. Broadcast replay on watch. Notify plus Mutex<ReplayState> becomes watch::Sender<ReplayState>. This removes the "register for a wakeup before looking" protocol, the extra Arc handle, the done flag and .boxed(), and replay_stream roughly halves.
4. failure as OnceLock<String> instead of Mutex<Option<String>>.
5. Writer: coalescer inside Output. Output::write and Output::finish replace the parallel coalescers vector, the paired .get() with its Internal error, drain and Outputs::send(output, ..). The open counter is derivable, so all_closed() can be slots.iter().all(..). That is O(P) per input batch instead of O(1), and the write_time timer would now also cover coalescing. Both look negligible next to the hash and take, but I have not measured.
6. shared_connection_config takes ballista_config twice and can take it once.

Net for the prototype is about 190 added and 250 removed lines across 6 files.

Proposed, not prototyped

• Drop registry: Option<Arc<PipelinedRegistry>> and with_registry from the writer and reader, and use the session extension only. That removes two fields, two setters and the re-apply branches in with_new_children and with_partitions, at the cost of attaching the registry through SessionConfig in about ten test call sites. with_client_pool can stay, since ShuffleReaderExec already uses that pattern for a dyn pool.
• Merge closed_jobs and closed_stages into one bounded map keyed by JobId, each entry holding an all flag and a stage set. That gives one capacity constant, and closed() no longer clones a JobId into a tuple key on every register and open.
• With (1), evicting discarded or never-connected shared slots in evict_idle is a single retain, which would cover most of Bound the executor client pool's shared connection slots #2673 cheaply.

[andygrove]

Addressed in 353e907, thanks for the prototype notes. All six prototyped items are in: SharedClient is a flag now, the registry slot and broadcast replay use watch, failure is a OnceLock, the coalescer lives in Output with all_closed() derived, and shared_connection_config takes the config once. open_remote_once wraps the no-pool client in SharedClient too. One difference: I kept a sync fast path in shared() (plain mutex for the connection, tokio mutex only while connecting) rather than one async lock, since the first round asked for a hot path without it. The proposals are in as well: the registry comes from the session config only, closed_jobs and closed_stages are one bounded map keyed by job, and evict_idle drops discarded shared slots. That covers dead executors from #2673, and the part left there is idle slots that are still connected.

[comphead]

Thanks @andygrove lets give it a try

[andygrove]

Thanks @comphead !