diff --git a/src/shm/in_process.rs b/src/shm/in_process.rs index 59fcb52d8..f498131fc 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_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. + 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..ca9435b99 100644 --- a/src/shm/runtime.rs +++ b/src/shm/runtime.rs @@ -236,6 +236,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_local_stream(stream); + return; + } let guard = self.cancel_senders.lock().unwrap(); let Some(senders) = guard.as_ref() else { return; @@ -279,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)) } @@ -588,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/transport.rs b/src/shm/transport.rs index 3679bfb69..b03e21535 100644 --- a/src/shm/transport.rs +++ b/src/shm/transport.rs @@ -1526,6 +1526,13 @@ impl LocalDrainPartitionSink { } } +impl Drop for LocalDrainPartitionSink { + fn drop(&mut self) { + self.buffer + .fail_pending("local partition sink dropped before EOF"); + } +} + #[async_trait] impl crate::PartitionSink for LocalDrainPartitionSink { async fn send(&mut self, batch: &RecordBatch) -> datafusion::common::Result<()> { @@ -1746,18 +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`. +#[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 receiver detached (or errored) before draining cleanly. 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, + /// 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 @@ -1920,6 +1934,18 @@ impl DrainHandle { self.cancelled_streams.lock().unwrap().contains(&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.local.get(&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,30 +2096,53 @@ 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 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 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.values().cloned().collect::>() + guard.remote.values().cloned().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"); + 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); @@ -2101,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()))); @@ -2172,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( @@ -2185,9 +2238,14 @@ impl DrainHandle { .channel_buffers .lock() .expect("DrainHandle channel_buffers mutex poisoned"); - let scope_dead = guard.dead_inbox; + 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 @@ -2197,10 +2255,32 @@ 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 + } + + /// 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 } @@ -2219,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(); @@ -3126,6 +3208,125 @@ 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 + /// 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(); + + // 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. + let local_channel = drain.register_local_channel(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_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, + /// and preserve the first failure reason. + #[test] + fn fail_scope_fails_both_remote_and_local_channels() { + let drain = DrainHandle::cooperative(vec![]); + + let remote = drain.register_data_channel(2, MppDataStreamKey::new(2, 0, 0)); + let local = drain.register_local_channel(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 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_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) + /// 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(vec![]); + let stream = MppDataStreamKey::new(2, 0, 1); + let buffer = drain.register_local_channel(stream); + let sink = LocalDrainPartitionSink::new(Arc::clone(&buffer)); + assert!(!crate::PartitionSink::cancelled(&sink)); + + drain.cancel_local_stream(stream); + assert!(crate::PartitionSink::cancelled(&sink)); + } + #[test] fn frame_rejects_short_message() { let too_short = vec![0u8; MPP_FRAME_HEADER_SIZE - 1]; @@ -4186,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))); } }