From 842f5b38b5745462108b6eb92597cd0cd2ea0879 Mon Sep 17 00:00:00 2001 From: Stu Hood Date: Mon, 5 Oct 2026 14:20:43 -0700 Subject: [PATCH 1/2] WIP --- src/shm/in_process.rs | 57 +++++++++++ src/shm/runtime.rs | 5 + src/shm/setup.rs | 14 +-- src/shm/transport.rs | 221 +++++++++++++++++++++++++++++++++++++++--- 4 files changed, 278 insertions(+), 19 deletions(-) diff --git a/src/shm/in_process.rs b/src/shm/in_process.rs index 59fcb52d8..d05bad1f5 100644 --- a/src/shm/in_process.rs +++ b/src/shm/in_process.rs @@ -1359,4 +1359,61 @@ mod tests { let _ = worker_handle.join().await; } + + /// When tasks are multiplexed across fewer worker processes than logical tasks, peer workers + /// may complete early, exit, and drop their DSM senders. This detaches the local process's DSM + /// inbox. Local in-process channels must NOT fail when DSM detaches, and must be able to send + /// and receive batches and EOF normally. + #[tokio::test(flavor = "current_thread")] + async fn peer_worker_exit_does_not_fail_multiplexed_local_channels() { + use crate::shm::transport::DrainItem; + + let n_procs = 3; // Proc 0: leader, Proc 1: worker 1 (hosts local channel), Proc 2: peer worker 2 + let boot = bootstrap_mesh(n_procs); + let mut workers = boot.workers; + let (_proc2, worker2_mesh, outbound2) = workers.pop().unwrap(); + let (_proc1, worker1_mesh, _outbound1) = workers.pop().unwrap(); + + // Worker 1 opens a local in-process partition sink (producer). + let stream_key = MppDataStreamKey::new(1, 0, 0); + let mut local_sink = worker1_mesh.open_local_partition_sink(stream_key); + + // Worker 1 also registers the consumer side for that local channel. + let local_consumer_channel = worker1_mesh + .inbound_receiver() + .register_data_channel(worker1_mesh.this_proc, stream_key); + + // Worker 2 (and its outbound senders) exits/drops. + // Dropping worker 2's senders causes worker 1's DSM inbox sender count to drop to zero. + drop(outbound2); + drop(worker2_mesh); + + // Worker 1 executes try_drain_pass, observing DSM detachment. + let _ = worker1_mesh.inbound_receiver().try_drain_pass(); + + // The local in-process channel on worker 1 must NOT be marked failed. + assert!(local_consumer_channel.try_pop().is_none()); + + // The local channel can continue sending batches and cleanly finish with EOF. + let schema = table_schema(); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from(vec![1, 2])), + Arc::new(Int32Array::from(vec![10, 20])), + ], + ) + .unwrap(); + local_sink.send(&batch).await.unwrap(); + local_sink.finish().await.unwrap(); + + match local_consumer_channel.try_pop() { + Some(DrainItem::Batch(b)) => assert_eq!(b.num_rows(), 2), + other => panic!("expected batch, got {other:?}"), + } + assert!(matches!( + local_consumer_channel.try_pop(), + Some(DrainItem::Eof) + )); + } } diff --git a/src/shm/runtime.rs b/src/shm/runtime.rs index 5a2c3cbef..4c9a3e40f 100644 --- a/src/shm/runtime.rs +++ b/src/shm/runtime.rs @@ -115,6 +115,7 @@ impl MppMesh { interrupt: Arc, alive: AliveFlag, ) -> Self { + inbound_receiver.set_this_proc(this_proc); Self { this_proc, n_procs, @@ -236,6 +237,10 @@ impl MppMesh { /// Stream-level so any consumer can cancel its own input: every `(producer_proc, stage, task, /// partition)` channel has a single consumer, so one stream's drop never cuts off a sibling's. pub fn cancel_stream(&self, producer_proc: u32, stream: MppDataStreamKey) { + if producer_proc == self.this_proc { + self.inbound_receiver.cancel_stream(producer_proc, stream); + return; + } let guard = self.cancel_senders.lock().unwrap(); let Some(senders) = guard.as_ref() else { return; diff --git a/src/shm/setup.rs b/src/shm/setup.rs index 31391f52d..780a09337 100644 --- a/src/shm/setup.rs +++ b/src/shm/setup.rs @@ -175,9 +175,10 @@ pub unsafe fn leader_setup( let inbox = DsmInboxReceiver::new(attach.inbound_receiver); inbox.set_receiver(receiver_token); - let inbound = Arc::new(DrainHandle::cooperative(vec![MppReceiver::new(Box::new( - inbox, - ))])); + let inbound = Arc::new(DrainHandle::cooperative_with_proc( + Some(0), + vec![MppReceiver::new(Box::new(inbox))], + )); // The leader hosts no producer fragments, but its senders carry the control plane: // work-unit frames (and later dynamic filters) flow leader -> worker through them. Empty // when the embedder did not opt in: a ring latches `detached` once its sender count hits @@ -295,9 +296,10 @@ pub unsafe fn worker_setup( let inbox = DsmInboxReceiver::new(attach.inbound_receiver); inbox.set_receiver(receiver_token); - let inbound = Arc::new(DrainHandle::cooperative(vec![MppReceiver::new(Box::new( - inbox, - ))])); + let inbound = Arc::new(DrainHandle::cooperative_with_proc( + Some(proc_idx), + vec![MppReceiver::new(Box::new(inbox))], + )); let mesh = Arc::new(MppMesh::new( proc_idx, total_procs, diff --git a/src/shm/transport.rs b/src/shm/transport.rs index 3679bfb69..1d1d9b80e 100644 --- a/src/shm/transport.rs +++ b/src/shm/transport.rs @@ -28,6 +28,7 @@ use async_trait::async_trait; use std::collections::VecDeque; +use std::sync::atomic::{AtomicU32, Ordering}; use std::sync::{Arc, Condvar, Mutex, MutexGuard}; use datafusion::common::{HashMap, HashSet}; @@ -1518,11 +1519,24 @@ impl crate::PartitionSink for MppPartitionSink { /// by another fragment on the same worker process. pub struct LocalDrainPartitionSink { buffer: Arc, + finished: bool, } impl LocalDrainPartitionSink { pub(super) fn new(buffer: Arc) -> Self { - Self { buffer } + Self { + buffer, + finished: false, + } + } +} + +impl Drop for LocalDrainPartitionSink { + fn drop(&mut self) { + if !self.finished { + self.buffer + .fail_pending("local partition sink dropped before EOF"); + } } } @@ -1533,7 +1547,8 @@ impl crate::PartitionSink for LocalDrainPartitionSink { Ok(()) } - async fn finish(self: Box) -> datafusion::common::Result<()> { + async fn finish(mut self: Box) -> datafusion::common::Result<()> { + self.finished = true; self.buffer.notify_source_done(); Ok(()) } @@ -1747,6 +1762,8 @@ pub(super) enum RecvBatchOutcome { /// Each entry is a `DrainBuffer::new(1)`: exactly one sender_proc emits frames for any given /// channel. Per-channel EOF flows via the `Eof` frame demuxed onto the matching buffer; query- /// teardown unblock flows via [`DrainHandle::cancel_channel_buffers`] from the handle's `Drop`. +const UNSET_PROC: u32 = u32::MAX; + #[derive(Default)] struct ChannelBufferRegistry { /// Keyed by `(sender_proc, stage_id, task_id, partition)`. The unified inbox carries frames @@ -1754,10 +1771,12 @@ struct ChannelBufferRegistry { /// buffer. This preserves the implicit "one stream per sender" semantics that /// `WorkerConnection::execute` consumers rely on. map: HashMap>, - /// Whether the receiver detached (or errored) before draining cleanly. Channels + /// Whether the remote inbox detached before draining cleanly. Remote channels /// fail at registration time too, so a consumer that registers after the detach /// does not wait on a channel nothing will ever fill. dead_inbox: bool, + /// Hard failure recorded for the entire scope (e.g. from `fail_scope`). + failed: Option, } /// Per-sender-proc drain: stashes the receivers and polls them inline from the cooperative spin @@ -1771,6 +1790,9 @@ struct ChannelBufferRegistry { /// On drop, the handle cancels every channel buffer so any consumer blocked on `try_pop` unblocks /// with `Eof` — the drain can therefore never outlive its query, even on a panicked teardown. pub struct DrainHandle { + /// Local process index owning this drain handle, if known. Used to distinguish + /// remote DSM channels from local in-process channels. + this_proc: AtomicU32, /// Per-(stage_id, task_id, partition) channel buffer registry. Populated lazily on first frame for a /// channel, or up-front by callers (e.g. `WorkerConnection::execute`) that need a /// buffer to wait on before any frame arrives. @@ -1893,10 +1915,20 @@ impl DrainHandle { /// Construct a cooperative drain handle. Channel buffers are populated lazily by /// [`Self::try_drain_pass`] when a frame arrives, or up-front by [`Self::register_channel`] /// when a consumer needs a buffer to wait on before any frame has come in. + #[cfg(test)] pub(super) fn cooperative(receivers: Vec) -> Self { + Self::cooperative_with_proc(None, receivers) + } + + /// Construct a cooperative drain handle with a known local process index. + pub(super) fn cooperative_with_proc( + this_proc: Option, + receivers: Vec, + ) -> Self { let wrapped = receivers.into_iter().map(Some).collect(); let (task_metrics_tx, task_metrics_rx) = tokio::sync::mpsc::unbounded_channel(); Self { + this_proc: AtomicU32::new(this_proc.unwrap_or(UNSET_PROC)), channel_buffers: Mutex::new(ChannelBufferRegistry::default()), coop_receivers: Mutex::new(wrapped), feed_registry: Mutex::new(FeedRegistry::default()), @@ -1908,6 +1940,17 @@ impl DrainHandle { } } + /// Set the local process index owning this handle. + pub fn set_this_proc(&self, proc: u32) { + self.this_proc.store(proc, Ordering::Release); + } + + /// The local process index owning this handle, if known. + pub fn this_proc(&self) -> Option { + let val = self.this_proc.load(Ordering::Acquire); + if val == UNSET_PROC { None } else { Some(val) } + } + /// Record a `Cancel` frame: the consumer abandoned this data stream, so this proc's producer /// of it stops. Idempotent. fn note_cancel(&self, stream: MppDataStreamKey) { @@ -1920,6 +1963,19 @@ impl DrainHandle { self.cancelled_streams.lock().unwrap().contains(&stream) } + /// Cancel a stream buffer and record the cancellation. For local channels where both producer + /// and consumer reside on this proc, bypassing DSM. + pub(super) fn cancel_stream(&self, sender_proc: u32, stream: MppDataStreamKey) { + self.note_cancel(stream); + let guard = self + .channel_buffers + .lock() + .expect("DrainHandle channel_buffers mutex poisoned"); + if let Some(buf) = guard.map.get(&PhysicalStreamKey::new(sender_proc, stream)) { + buf.cancel(); + } + } + /// Take the receiving end of the `TaskMetrics` frame stream. The embedder (the leader) /// drains it into its metrics store; the first caller gets it, later calls get `None`. pub(super) fn take_task_metrics_receiver( @@ -2070,7 +2126,11 @@ impl DrainHandle { /// senders, which never join a ring's `sender_count`, so "the last data sender dropped" /// says nothing about them: a peer that completed normally has EOF'd everything it served, /// while this proc may still owe the leader work it has not been asked for yet. + /// + /// Local in-process channels (`sender_proc == this_proc`) bypass the DSM inbox entirely + /// and are not failed on DSM detachment; they remain live until the local producer finishes. fn detach_data_scope(&self, reason: &str) { + let this_proc = self.this_proc(); let to_fail = { let mut guard = self .channel_buffers @@ -2080,18 +2140,34 @@ impl DrainHandle { return; } guard.dead_inbox = true; - guard.map.values().cloned().collect::>() + guard + .map + .iter() + .filter(|(key, _)| this_proc != Some(key.sender_proc)) + .map(|(_, buf)| Arc::clone(buf)) + .collect::>() }; for buf in to_fail { buf.fail_pending(reason); } } - /// Hard scope death: [`Self::detach_data_scope`] plus the control-plane registries. For - /// ring corruption and embedder-driven teardown, where no further frame of any kind can be - /// trusted to arrive. + /// Hard scope death: fails all channel buffers (both remote and local) plus the + /// control-plane registries. For ring corruption and embedder-driven teardown, + /// where no further frame of any kind can be trusted to arrive. fn fail_scope(&self, reason: &str) { - self.detach_data_scope(reason); + let to_fail = { + let mut guard = self + .channel_buffers + .lock() + .expect("DrainHandle channel_buffers mutex poisoned"); + guard.dead_inbox = true; + guard.failed = Some(reason.to_string()); + guard.map.values().cloned().collect::>() + }; + for buf in to_fail { + buf.fail_pending(reason); + } let mut registry = self.feed_registry.lock().unwrap(); registry.dead = Some(reason.to_string()); for (_, slot) in registry.map.drain() { @@ -2185,7 +2261,16 @@ impl DrainHandle { .channel_buffers .lock() .expect("DrainHandle channel_buffers mutex poisoned"); - let scope_dead = guard.dead_inbox; + let failure_reason = if let Some(err) = &guard.failed { + Some(err.clone()) + } else if guard.dead_inbox && self.this_proc() != Some(sender_proc) { + Some( + "transport receiver detached before this channel's EOF; the producer went away" + .to_string(), + ) + } else { + None + }; let buf = guard .map .entry(PhysicalStreamKey::new(sender_proc, stream)) @@ -2197,10 +2282,8 @@ impl DrainHandle { }) .clone(); drop(guard); - if scope_dead { - buf.fail_pending( - "transport receiver detached before this channel's EOF; the producer went away", - ); + if let Some(reason) = failure_reason { + buf.fail_pending(&reason); } buf } @@ -3126,6 +3209,118 @@ mod tests { assert!(matches!(late.try_pop(), Some(DrainItem::Failed(_)))); } + /// A peer process departs, detaching the DSM data plane. Remote channels that were waiting on + /// their `Eof` fail, but local in-process channels (`sender_proc == this_proc`) must NOT fail + /// because their producers do not communicate over the DSM inbox. + #[test] + fn peer_detach_does_not_fail_local_channels() { + let (_region, peer, _leader, drain) = detached_test_inbox(); + drain.set_this_proc(1); + + // Remote channel awaiting EOF from peer proc 2. + let remote_dep = drain.register_data_channel(2, MppDataStreamKey::new(2, 1, 0)); + // Local channel on this proc (proc 1). + let local_channel = drain.register_data_channel(1, MppDataStreamKey::new(2, 1, 1)); + + drop(peer); + drain.try_drain_pass().expect("drain over the detach"); + + // The remote channel awaiting EOF from peer failed on detach. + assert!(matches!(remote_dep.try_pop(), Some(DrainItem::Failed(_)))); + // The local channel on this proc was NOT failed. + assert!(local_channel.try_pop().is_none()); + + // The local channel can continue to receive batches and EOF normally. + local_channel.push_batch(sample_batch(3)); + local_channel.notify_source_done(); + match local_channel.try_pop() { + Some(DrainItem::Batch(b)) => assert_eq!(b.num_rows(), 3), + other => panic!("expected batch, got {other:?}"), + } + assert!(matches!(local_channel.try_pop(), Some(DrainItem::Eof))); + + // Late registration for remote proc fails immediately. + let late_remote = drain.register_data_channel(2, MppDataStreamKey::new(2, 1, 2)); + assert!(matches!(late_remote.try_pop(), Some(DrainItem::Failed(_)))); + + // Late registration for local proc succeeds and does not fail. + let late_local = drain.register_data_channel(1, MppDataStreamKey::new(2, 1, 2)); + assert!(late_local.try_pop().is_none()); + late_local.notify_source_done(); + assert!(matches!(late_local.try_pop(), Some(DrainItem::Eof))); + } + + /// Hard scope failure (e.g. from `fail_scope`) must fail both remote and local channels. + #[test] + fn fail_scope_fails_both_remote_and_local_channels() { + let drain = DrainHandle::cooperative(vec![]); + drain.set_this_proc(1); + + let remote = drain.register_data_channel(2, MppDataStreamKey::new(2, 0, 0)); + let local = drain.register_data_channel(1, MppDataStreamKey::new(2, 0, 1)); + + drain.fail_scope("worker 2 panicked"); + + assert!(matches!( + remote.try_pop(), + Some(DrainItem::Failed(msg)) if msg.contains("worker 2 panicked") + )); + assert!(matches!( + local.try_pop(), + Some(DrainItem::Failed(msg)) if msg.contains("worker 2 panicked") + )); + + // Late registrations also fail. + let late_remote = drain.register_data_channel(2, MppDataStreamKey::new(2, 0, 2)); + assert!(matches!( + late_remote.try_pop(), + Some(DrainItem::Failed(msg)) if msg.contains("worker 2 panicked") + )); + + let late_local = drain.register_data_channel(1, MppDataStreamKey::new(2, 0, 3)); + assert!(matches!( + late_local.try_pop(), + Some(DrainItem::Failed(msg)) if msg.contains("worker 2 panicked") + )); + } + + /// Dropping a `LocalDrainPartitionSink` without calling `finish` (e.g. on task panic or error) + /// must fail the pending buffer so the consumer unblocks with an error instead of hanging. + #[test] + fn local_drain_partition_sink_fails_buffer_on_drop_without_finish() { + let buffer = DrainBuffer::new(1); + let sink = LocalDrainPartitionSink::new(Arc::clone(&buffer)); + drop(sink); + assert!(matches!( + buffer.try_pop(), + Some(DrainItem::Failed(msg)) if msg.contains("local partition sink dropped before EOF") + )); + } + + /// Cleanly finishing a `LocalDrainPartitionSink` completes the buffer with EOF on drop. + #[tokio::test] + async fn local_drain_partition_sink_completes_cleanly_on_finish() { + use crate::PartitionSink; + let buffer = DrainBuffer::new(1); + let sink = Box::new(LocalDrainPartitionSink::new(Arc::clone(&buffer))); + sink.finish().await.expect("finish should succeed"); + assert!(matches!(buffer.try_pop(), Some(DrainItem::Eof))); + } + + /// Cancelling a local stream on `DrainHandle` marks the local buffer cancelled so the sink observes it. + #[test] + fn local_stream_cancellation_marks_buffer_cancelled() { + let drain = DrainHandle::cooperative_with_proc(Some(1), vec![]); + let stream = MppDataStreamKey::new(2, 0, 1); + let buffer = drain.register_data_channel(1, stream); + let sink = LocalDrainPartitionSink::new(Arc::clone(&buffer)); + assert!(!crate::PartitionSink::cancelled(&sink)); + + drain.cancel_stream(1, stream); + assert!(crate::PartitionSink::cancelled(&sink)); + assert!(drain.stream_cancelled(stream)); + } + #[test] fn frame_rejects_short_message() { let too_short = vec![0u8; MPP_FRAME_HEADER_SIZE - 1]; From 105b3ca0fca579448ac1e08d333cf2c67a9925d3 Mon Sep 17 00:00:00 2001 From: Stu Hood Date: Tue, 6 Oct 2026 15:17:21 -0700 Subject: [PATCH 2/2] Review feedback. --- src/shm/in_process.rs | 2 +- src/shm/runtime.rs | 13 +-- src/shm/setup.rs | 14 ++- src/shm/transport.rs | 208 ++++++++++++++++++++++-------------------- 4 files changed, 122 insertions(+), 115 deletions(-) diff --git a/src/shm/in_process.rs b/src/shm/in_process.rs index d05bad1f5..f498131fc 100644 --- a/src/shm/in_process.rs +++ b/src/shm/in_process.rs @@ -1381,7 +1381,7 @@ mod tests { // Worker 1 also registers the consumer side for that local channel. let local_consumer_channel = worker1_mesh .inbound_receiver() - .register_data_channel(worker1_mesh.this_proc, stream_key); + .register_local_channel(stream_key); // Worker 2 (and its outbound senders) exits/drops. // Dropping worker 2's senders causes worker 1's DSM inbox sender count to drop to zero. diff --git a/src/shm/runtime.rs b/src/shm/runtime.rs index 4c9a3e40f..ca9435b99 100644 --- a/src/shm/runtime.rs +++ b/src/shm/runtime.rs @@ -115,7 +115,6 @@ impl MppMesh { interrupt: Arc, alive: AliveFlag, ) -> Self { - inbound_receiver.set_this_proc(this_proc); Self { this_proc, n_procs, @@ -238,7 +237,7 @@ impl MppMesh { /// partition)` channel has a single consumer, so one stream's drop never cuts off a sibling's. pub fn cancel_stream(&self, producer_proc: u32, stream: MppDataStreamKey) { if producer_proc == self.this_proc { - self.inbound_receiver.cancel_stream(producer_proc, stream); + self.inbound_receiver.cancel_local_stream(stream); return; } let guard = self.cancel_senders.lock().unwrap(); @@ -284,9 +283,7 @@ impl MppMesh { /// Opens a [`PartitionSink`] for a local in-process data stream (self-loop). pub fn open_local_partition_sink(&self, stream: MppDataStreamKey) -> Box { - let buffer = self - .inbound_receiver - .register_data_channel(self.this_proc, stream); + let buffer = self.inbound_receiver.register_local_channel(stream); Box::new(LocalDrainPartitionSink::new(buffer)) } @@ -593,7 +590,11 @@ fn pull_partition_stream( stream.task_id, stream.partition, ); - let buffer = drain.register_data_channel(producer_proc, stream); + let buffer = if producer_proc == mesh.this_proc { + drain.register_local_channel(stream) + } else { + drain.register_data_channel(producer_proc, stream) + }; let stream = async_stream::stream! { // Any consumer cancels its own input stream when it drops early. The mesh no-ops the send // until the embedder wires this proc's outbound senders. diff --git a/src/shm/setup.rs b/src/shm/setup.rs index 780a09337..31391f52d 100644 --- a/src/shm/setup.rs +++ b/src/shm/setup.rs @@ -175,10 +175,9 @@ pub unsafe fn leader_setup( let inbox = DsmInboxReceiver::new(attach.inbound_receiver); inbox.set_receiver(receiver_token); - let inbound = Arc::new(DrainHandle::cooperative_with_proc( - Some(0), - vec![MppReceiver::new(Box::new(inbox))], - )); + let inbound = Arc::new(DrainHandle::cooperative(vec![MppReceiver::new(Box::new( + inbox, + ))])); // The leader hosts no producer fragments, but its senders carry the control plane: // work-unit frames (and later dynamic filters) flow leader -> worker through them. Empty // when the embedder did not opt in: a ring latches `detached` once its sender count hits @@ -296,10 +295,9 @@ pub unsafe fn worker_setup( let inbox = DsmInboxReceiver::new(attach.inbound_receiver); inbox.set_receiver(receiver_token); - let inbound = Arc::new(DrainHandle::cooperative_with_proc( - Some(proc_idx), - vec![MppReceiver::new(Box::new(inbox))], - )); + let inbound = Arc::new(DrainHandle::cooperative(vec![MppReceiver::new(Box::new( + inbox, + ))])); let mesh = Arc::new(MppMesh::new( proc_idx, total_procs, diff --git a/src/shm/transport.rs b/src/shm/transport.rs index 1d1d9b80e..b03e21535 100644 --- a/src/shm/transport.rs +++ b/src/shm/transport.rs @@ -28,7 +28,6 @@ use async_trait::async_trait; use std::collections::VecDeque; -use std::sync::atomic::{AtomicU32, Ordering}; use std::sync::{Arc, Condvar, Mutex, MutexGuard}; use datafusion::common::{HashMap, HashSet}; @@ -1519,24 +1518,18 @@ impl crate::PartitionSink for MppPartitionSink { /// by another fragment on the same worker process. pub struct LocalDrainPartitionSink { buffer: Arc, - finished: bool, } impl LocalDrainPartitionSink { pub(super) fn new(buffer: Arc) -> Self { - Self { - buffer, - finished: false, - } + Self { buffer } } } impl Drop for LocalDrainPartitionSink { fn drop(&mut self) { - if !self.finished { - self.buffer - .fail_pending("local partition sink dropped before EOF"); - } + self.buffer + .fail_pending("local partition sink dropped before EOF"); } } @@ -1547,8 +1540,7 @@ impl crate::PartitionSink for LocalDrainPartitionSink { Ok(()) } - async fn finish(mut self: Box) -> datafusion::common::Result<()> { - self.finished = true; + async fn finish(self: Box) -> datafusion::common::Result<()> { self.buffer.notify_source_done(); Ok(()) } @@ -1761,22 +1753,25 @@ pub(super) enum RecvBatchOutcome { /// /// Each entry is a `DrainBuffer::new(1)`: exactly one sender_proc emits frames for any given /// channel. Per-channel EOF flows via the `Eof` frame demuxed onto the matching buffer; query- -/// teardown unblock flows via [`DrainHandle::cancel_channel_buffers`] from the handle's `Drop`. -const UNSET_PROC: u32 = u32::MAX; +#[derive(Default, Debug, Clone, PartialEq, Eq)] +enum ChannelBufferState { + #[default] + Live, + /// The remote inbox detached before draining cleanly. Remote channels + /// fail at registration time, but local in-process channels remain live. + Detached(String), + /// Hard failure recorded for the entire scope (e.g. from `fail_scope`). + Failed(String), +} #[derive(Default)] struct ChannelBufferRegistry { - /// Keyed by `(sender_proc, stage_id, task_id, partition)`. The unified inbox carries frames - /// from every peer, so each `(stage, task, partition)` consumer gets its own per-sender - /// buffer. This preserves the implicit "one stream per sender" semantics that - /// `WorkerConnection::execute` consumers rely on. - map: HashMap>, - /// Whether the remote inbox detached before draining cleanly. Remote channels - /// fail at registration time too, so a consumer that registers after the detach - /// does not wait on a channel nothing will ever fill. - dead_inbox: bool, - /// Hard failure recorded for the entire scope (e.g. from `fail_scope`). - failed: Option, + /// Remote inbound channels arriving over DSM from peer processes, keyed by + /// `(sender_proc, stage_id, task_id, partition)`. + remote: HashMap>, + /// Local in-process channels on this process, keyed by `(stage_id, task_id, partition)`. + local: HashMap>, + state: ChannelBufferState, } /// Per-sender-proc drain: stashes the receivers and polls them inline from the cooperative spin @@ -1790,9 +1785,6 @@ struct ChannelBufferRegistry { /// On drop, the handle cancels every channel buffer so any consumer blocked on `try_pop` unblocks /// with `Eof` — the drain can therefore never outlive its query, even on a panicked teardown. pub struct DrainHandle { - /// Local process index owning this drain handle, if known. Used to distinguish - /// remote DSM channels from local in-process channels. - this_proc: AtomicU32, /// Per-(stage_id, task_id, partition) channel buffer registry. Populated lazily on first frame for a /// channel, or up-front by callers (e.g. `WorkerConnection::execute`) that need a /// buffer to wait on before any frame arrives. @@ -1915,20 +1907,10 @@ impl DrainHandle { /// Construct a cooperative drain handle. Channel buffers are populated lazily by /// [`Self::try_drain_pass`] when a frame arrives, or up-front by [`Self::register_channel`] /// when a consumer needs a buffer to wait on before any frame has come in. - #[cfg(test)] pub(super) fn cooperative(receivers: Vec) -> Self { - Self::cooperative_with_proc(None, receivers) - } - - /// Construct a cooperative drain handle with a known local process index. - pub(super) fn cooperative_with_proc( - this_proc: Option, - receivers: Vec, - ) -> Self { let wrapped = receivers.into_iter().map(Some).collect(); let (task_metrics_tx, task_metrics_rx) = tokio::sync::mpsc::unbounded_channel(); Self { - this_proc: AtomicU32::new(this_proc.unwrap_or(UNSET_PROC)), channel_buffers: Mutex::new(ChannelBufferRegistry::default()), coop_receivers: Mutex::new(wrapped), feed_registry: Mutex::new(FeedRegistry::default()), @@ -1940,17 +1922,6 @@ impl DrainHandle { } } - /// Set the local process index owning this handle. - pub fn set_this_proc(&self, proc: u32) { - self.this_proc.store(proc, Ordering::Release); - } - - /// The local process index owning this handle, if known. - pub fn this_proc(&self) -> Option { - let val = self.this_proc.load(Ordering::Acquire); - if val == UNSET_PROC { None } else { Some(val) } - } - /// Record a `Cancel` frame: the consumer abandoned this data stream, so this proc's producer /// of it stops. Idempotent. fn note_cancel(&self, stream: MppDataStreamKey) { @@ -1963,15 +1934,14 @@ impl DrainHandle { self.cancelled_streams.lock().unwrap().contains(&stream) } - /// Cancel a stream buffer and record the cancellation. For local channels where both producer - /// and consumer reside on this proc, bypassing DSM. - pub(super) fn cancel_stream(&self, sender_proc: u32, stream: MppDataStreamKey) { - self.note_cancel(stream); + /// Cancel a stream buffer for a local channel where both producer and consumer reside + /// on this proc, bypassing DSM. + pub(super) fn cancel_local_stream(&self, stream: MppDataStreamKey) { let guard = self .channel_buffers .lock() .expect("DrainHandle channel_buffers mutex poisoned"); - if let Some(buf) = guard.map.get(&PhysicalStreamKey::new(sender_proc, stream)) { + if let Some(buf) = guard.local.get(&stream) { buf.cancel(); } } @@ -2127,25 +2097,21 @@ impl DrainHandle { /// says nothing about them: a peer that completed normally has EOF'd everything it served, /// while this proc may still owe the leader work it has not been asked for yet. /// - /// Local in-process channels (`sender_proc == this_proc`) bypass the DSM inbox entirely - /// and are not failed on DSM detachment; they remain live until the local producer finishes. + /// Local in-process channels bypass the DSM inbox entirely and are not failed on DSM + /// detachment; they remain live until the local producer finishes. fn detach_data_scope(&self, reason: &str) { - let this_proc = self.this_proc(); let to_fail = { let mut guard = self .channel_buffers .lock() .expect("DrainHandle channel_buffers mutex poisoned"); - if guard.dead_inbox { - return; + match &guard.state { + ChannelBufferState::Detached(_) | ChannelBufferState::Failed(_) => return, + ChannelBufferState::Live => { + guard.state = ChannelBufferState::Detached(reason.to_string()); + } } - guard.dead_inbox = true; - guard - .map - .iter() - .filter(|(key, _)| this_proc != Some(key.sender_proc)) - .map(|(_, buf)| Arc::clone(buf)) - .collect::>() + guard.remote.values().cloned().collect::>() }; for buf in to_fail { buf.fail_pending(reason); @@ -2161,15 +2127,22 @@ impl DrainHandle { .channel_buffers .lock() .expect("DrainHandle channel_buffers mutex poisoned"); - guard.dead_inbox = true; - guard.failed = Some(reason.to_string()); - guard.map.values().cloned().collect::>() + if let ChannelBufferState::Failed(_) = &guard.state { + // First error wins; keep the original failure reason for late registrations. + } else { + guard.state = ChannelBufferState::Failed(reason.to_string()); + } + let mut buffers = guard.remote.values().cloned().collect::>(); + buffers.extend(guard.local.values().cloned()); + buffers }; for buf in to_fail { buf.fail_pending(reason); } let mut registry = self.feed_registry.lock().unwrap(); - registry.dead = Some(reason.to_string()); + if registry.dead.is_none() { + registry.dead = Some(reason.to_string()); + } for (_, slot) in registry.map.drain() { if let FeedSlot::Active(senders) = slot { fail_feed_senders(&senders, reason); @@ -2177,14 +2150,18 @@ impl DrainHandle { } drop(registry); let mut plans = self.set_plan_registry.lock().unwrap(); - plans.dead = Some(reason.to_string()); + if plans.dead.is_none() { + plans.dead = Some(reason.to_string()); + } for (_, slot) in plans.map.drain() { if let SetPlanSlot::Waiting(tx) = slot { let _ = tx.send(Err(DataFusionError::Execution(reason.to_string()))); } } let mut execs = self.execute_task_registry.lock().unwrap(); - execs.dead = Some(reason.to_string()); + if execs.dead.is_none() { + execs.dead = Some(reason.to_string()); + } for (_, slot) in execs.map.drain() { if let ExecuteTaskSlot::Active(tx) = slot { let _ = tx.send(Err(DataFusionError::Execution(reason.to_string()))); @@ -2248,8 +2225,8 @@ impl DrainHandle { Ok(rx) } - /// Register (or look up) the channel buffer for one producer's data stream. - /// The returned `Arc` is the canonical destination for frames matching + /// Register (or look up) the channel buffer for one remote producer's data stream arriving + /// over DSM. The returned `Arc` is the canonical destination for frames matching /// that key: `try_drain_pass` pushes into the same entry on every `Batch { header, .. }` /// whose `header.sender_proc()` / `stage_id` / `partition` matches. pub(super) fn register_data_channel( @@ -2261,18 +2238,14 @@ impl DrainHandle { .channel_buffers .lock() .expect("DrainHandle channel_buffers mutex poisoned"); - let failure_reason = if let Some(err) = &guard.failed { - Some(err.clone()) - } else if guard.dead_inbox && self.this_proc() != Some(sender_proc) { - Some( - "transport receiver detached before this channel's EOF; the producer went away" - .to_string(), - ) - } else { - None + let failure_reason = match &guard.state { + ChannelBufferState::Live => None, + ChannelBufferState::Detached(reason) | ChannelBufferState::Failed(reason) => { + Some(reason.clone()) + } }; let buf = guard - .map + .remote .entry(PhysicalStreamKey::new(sender_proc, stream)) .or_insert_with(|| { // num_sources stays 1: each (sender_proc, stage, task, partition) tuple has @@ -2288,6 +2261,30 @@ impl DrainHandle { buf } + /// Register (or look up) the channel buffer for a local in-process data stream. + /// Local channels bypass DSM and survive remote inbox detachment, failing only on + /// hard scope failure. + pub(super) fn register_local_channel(&self, stream: MppDataStreamKey) -> Arc { + let mut guard = self + .channel_buffers + .lock() + .expect("DrainHandle channel_buffers mutex poisoned"); + let failure_reason = match &guard.state { + ChannelBufferState::Failed(reason) => Some(reason.clone()), + ChannelBufferState::Live | ChannelBufferState::Detached(_) => None, + }; + let buf = guard + .local + .entry(stream) + .or_insert_with(|| DrainBuffer::new(1)) + .clone(); + drop(guard); + if let Some(reason) = failure_reason { + buf.fail_pending(&reason); + } + buf + } + /// Cancel every registered channel buffer. Called from `Drop` to unblock any consumer waiting on /// a channel buffer when the handle goes away mid-query. /// @@ -2302,7 +2299,9 @@ impl DrainHandle { .channel_buffers .lock() .expect("DrainHandle channel_buffers mutex poisoned"); - guard.map.values().cloned().collect::>() + let mut buffers = guard.remote.values().cloned().collect::>(); + buffers.extend(guard.local.values().cloned()); + buffers }; for buf in to_cancel { buf.cancel(); @@ -3210,17 +3209,17 @@ mod tests { } /// A peer process departs, detaching the DSM data plane. Remote channels that were waiting on - /// their `Eof` fail, but local in-process channels (`sender_proc == this_proc`) must NOT fail - /// because their producers do not communicate over the DSM inbox. + /// A peer process departs, detaching the DSM data plane. Remote channels that were waiting on + /// their `Eof` fail, but local in-process channels must NOT fail because their producers do not + /// communicate over the DSM inbox. #[test] fn peer_detach_does_not_fail_local_channels() { let (_region, peer, _leader, drain) = detached_test_inbox(); - drain.set_this_proc(1); // Remote channel awaiting EOF from peer proc 2. let remote_dep = drain.register_data_channel(2, MppDataStreamKey::new(2, 1, 0)); - // Local channel on this proc (proc 1). - let local_channel = drain.register_data_channel(1, MppDataStreamKey::new(2, 1, 1)); + // Local channel on this proc. + let local_channel = drain.register_local_channel(MppDataStreamKey::new(2, 1, 1)); drop(peer); drain.try_drain_pass().expect("drain over the detach"); @@ -3244,20 +3243,20 @@ mod tests { assert!(matches!(late_remote.try_pop(), Some(DrainItem::Failed(_)))); // Late registration for local proc succeeds and does not fail. - let late_local = drain.register_data_channel(1, MppDataStreamKey::new(2, 1, 2)); + let late_local = drain.register_local_channel(MppDataStreamKey::new(2, 1, 2)); assert!(late_local.try_pop().is_none()); late_local.notify_source_done(); assert!(matches!(late_local.try_pop(), Some(DrainItem::Eof))); } - /// Hard scope failure (e.g. from `fail_scope`) must fail both remote and local channels. + /// Hard scope failure (e.g. from `fail_scope`) must fail both remote and local channels, + /// and preserve the first failure reason. #[test] fn fail_scope_fails_both_remote_and_local_channels() { let drain = DrainHandle::cooperative(vec![]); - drain.set_this_proc(1); let remote = drain.register_data_channel(2, MppDataStreamKey::new(2, 0, 0)); - let local = drain.register_data_channel(1, MppDataStreamKey::new(2, 0, 1)); + let local = drain.register_local_channel(MppDataStreamKey::new(2, 0, 1)); drain.fail_scope("worker 2 panicked"); @@ -3270,18 +3269,26 @@ mod tests { Some(DrainItem::Failed(msg)) if msg.contains("worker 2 panicked") )); - // Late registrations also fail. + // Late registrations also fail with the first failure reason. let late_remote = drain.register_data_channel(2, MppDataStreamKey::new(2, 0, 2)); assert!(matches!( late_remote.try_pop(), Some(DrainItem::Failed(msg)) if msg.contains("worker 2 panicked") )); - let late_local = drain.register_data_channel(1, MppDataStreamKey::new(2, 0, 3)); + let late_local = drain.register_local_channel(MppDataStreamKey::new(2, 0, 3)); assert!(matches!( late_local.try_pop(), Some(DrainItem::Failed(msg)) if msg.contains("worker 2 panicked") )); + + // A second fail_scope does not overwrite the first reason (first error wins). + drain.fail_scope("worker 3 panicked"); + let late_local_2 = drain.register_local_channel(MppDataStreamKey::new(2, 0, 4)); + assert!(matches!( + late_local_2.try_pop(), + Some(DrainItem::Failed(msg)) if msg.contains("worker 2 panicked") + )); } /// Dropping a `LocalDrainPartitionSink` without calling `finish` (e.g. on task panic or error) @@ -3310,15 +3317,14 @@ mod tests { /// Cancelling a local stream on `DrainHandle` marks the local buffer cancelled so the sink observes it. #[test] fn local_stream_cancellation_marks_buffer_cancelled() { - let drain = DrainHandle::cooperative_with_proc(Some(1), vec![]); + let drain = DrainHandle::cooperative(vec![]); let stream = MppDataStreamKey::new(2, 0, 1); - let buffer = drain.register_data_channel(1, stream); + let buffer = drain.register_local_channel(stream); let sink = LocalDrainPartitionSink::new(Arc::clone(&buffer)); assert!(!crate::PartitionSink::cancelled(&sink)); - drain.cancel_stream(1, stream); + drain.cancel_local_stream(stream); assert!(crate::PartitionSink::cancelled(&sink)); - assert!(drain.stream_cancelled(stream)); } #[test] @@ -4381,10 +4387,12 @@ mod tests { let buf_a = handle.register_data_channel(0, MppDataStreamKey::new(0, 0, 0)); let buf_b = handle.register_data_channel(0, MppDataStreamKey::new(7, 0, 3)); + let buf_c = handle.register_local_channel(MppDataStreamKey::new(1, 0, 1)); // No data ever flows; the handle is just dropped. drop(handle); assert!(matches!(buf_a.try_pop(), Some(DrainItem::Eof))); assert!(matches!(buf_b.try_pop(), Some(DrainItem::Eof))); + assert!(matches!(buf_c.try_pop(), Some(DrainItem::Eof))); } }