You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by su...@apache.org on 2026/08/28 16:57:54 UTC
(datafusion-comet) branch main updated: feat: complete Celeborn map-side shuffle push lifecycle (6/n) (#5513)
This is an automated email from the ASF dual-hosted git repository.
sunchao pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
The following commit(s) were added to refs/heads/main by this push:
new c067e4e0db feat: complete Celeborn map-side shuffle push lifecycle (6/n) (#5513)
c067e4e0db is described below
commit c067e4e0db4e2fcae0cf99e60c54a673f59c75eb
Author: Ping Zhang <pi...@gmail.com>
AuthorDate: Fri Aug 28 09:57:49 2026 -0700
feat: complete Celeborn map-side shuffle push lifecycle (6/n) (#5513)
---
.github/workflows/pr_build_linux.yml | 1 +
.github/workflows/pr_build_macos.yml | 1 +
native/Cargo.lock | 1 +
native/Cargo.toml | 1 +
native/core/src/execution/planner.rs | 40 +-
native/jni-bridge/src/shuffle_partition_pusher.rs | 156 +++
native/shuffle/Cargo.toml | 3 +-
native/shuffle/src/shuffle_writer.rs | 4 +-
native/shuffle/src/writers/rss/mod.rs | 891 +++++++++++++-
.../src/writers/rss/rss_partition_writer.rs | 1027 +++++++++++++++-
native/shuffle/src/writers/shuffle_block_writer.rs | 127 +-
.../shuffle/CelebornShufflePartitionPusher.java | 939 +++++++++++++-
.../shuffle/CelebornTransportCallbackTracker.java | 1000 +++++++++++++++
.../shuffle/ExecutorShufflePushAdmission.java | 103 ++
.../comet/shuffle/ShufflePartitionPusher.java | 20 +
.../main/scala/org/apache/comet/CometConf.scala | 28 +
.../shuffle/CelebornShufflePusherFactory.scala | 380 ++++++
.../shuffle/CometCelebornShuffleManager.scala | 655 +++++++++-
.../shuffle/CometNativeShuffleWriter.scala | 314 ++++-
.../CelebornTransportOwnershipTestHelper.java | 382 ++++++
.../scala/org/apache/comet/CometConfSuite.scala | 35 +
.../comet/exec/CometNativeShuffleSuite.scala | 70 +-
.../CelebornShufflePartitionPusherSuite.scala | 1293 +++++++++++++++++++-
.../CometCelebornNativeShuffleWriterSuite.scala | 481 ++++++++
.../shuffle/CometCelebornShuffleManagerSuite.scala | 249 +++-
25 files changed, 8007 insertions(+), 194 deletions(-)
diff --git a/.github/workflows/pr_build_linux.yml b/.github/workflows/pr_build_linux.yml
index 29159dde31..38e0c7dd8a 100644
--- a/.github/workflows/pr_build_linux.yml
+++ b/.github/workflows/pr_build_linux.yml
@@ -329,6 +329,7 @@ jobs:
org.apache.comet.exec.CometNativeShuffleSuite
org.apache.comet.shuffle.CelebornShufflePartitionPusherSuite
org.apache.spark.sql.comet.execution.shuffle.CometCelebornShuffleManagerSuite
+ org.apache.spark.sql.comet.execution.shuffle.CometCelebornNativeShuffleWriterSuite
org.apache.spark.sql.comet.execution.shuffle.CometNativeShuffleInputRDDSuite
org.apache.comet.exec.CometShuffleEncryptionSuite
org.apache.comet.exec.CometShuffleManagerSuite
diff --git a/.github/workflows/pr_build_macos.yml b/.github/workflows/pr_build_macos.yml
index 423953130d..69cded40d0 100644
--- a/.github/workflows/pr_build_macos.yml
+++ b/.github/workflows/pr_build_macos.yml
@@ -145,6 +145,7 @@ jobs:
org.apache.comet.exec.CometNativeShuffleSuite
org.apache.comet.shuffle.CelebornShufflePartitionPusherSuite
org.apache.spark.sql.comet.execution.shuffle.CometCelebornShuffleManagerSuite
+ org.apache.spark.sql.comet.execution.shuffle.CometCelebornNativeShuffleWriterSuite
org.apache.spark.sql.comet.execution.shuffle.CometNativeShuffleInputRDDSuite
org.apache.comet.exec.CometShuffleEncryptionSuite
org.apache.comet.exec.CometShuffleManagerSuite
diff --git a/native/Cargo.lock b/native/Cargo.lock
index eadd9d995c..e835cdced3 100644
--- a/native/Cargo.lock
+++ b/native/Cargo.lock
@@ -2040,6 +2040,7 @@ name = "datafusion-comet-shuffle"
version = "1.1.0"
dependencies = [
"arrow",
+ "arrow-select",
"async-trait",
"bytes",
"clap",
diff --git a/native/Cargo.toml b/native/Cargo.toml
index 471060bc82..fba8e3c485 100644
--- a/native/Cargo.toml
+++ b/native/Cargo.toml
@@ -39,6 +39,7 @@ rust-version = "1.88"
[workspace.dependencies]
arrow = { version = "58.4.0", features = ["prettyprint", "ffi", "chrono-tz"] }
+arrow-select = { version = "58.4.0" }
async-trait = { version = "0.1" }
bytes = { version = "1.11.1" }
parquet = { version = "58.4.0", default-features = false, features = ["experimental"] }
diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs
index 9d2a03f29a..8679b16928 100644
--- a/native/core/src/execution/planner.rs
+++ b/native/core/src/execution/planner.rs
@@ -47,7 +47,7 @@ use crate::execution::{
serde::to_arrow_datatype,
shuffle::{SchemaAlignExec, ShuffleWriterDestination, ShuffleWriterExec},
};
-use crate::jvm_bridge::{jni_call, JVMClasses, JavaShufflePartitionPusher, ShufflePartitionPusher};
+use crate::jvm_bridge::{jni_call, JVMClasses, ShufflePartitionPusher};
use arrow::compute::CastOptions;
use arrow::datatypes::{DataType, Field, FieldRef, Schema, TimeUnit, DECIMAL128_MAX_PRECISION};
use arrow::ffi_stream::FFI_ArrowArrayStream;
@@ -3994,7 +3994,7 @@ fn shuffle_writer_destination(
Ok(ShuffleWriterDestination::Rss {
pusher: Arc::clone(pusher),
- max_frame_size: JavaShufflePartitionPusher::MAX_PAYLOAD_SIZE,
+ max_frame_size: pusher.max_frame_size(),
})
}
None => Err(GeneralError(
@@ -4860,6 +4860,24 @@ mod tests {
}
}
+ struct BoundedShufflePartitionPusher {
+ max_frame_size: usize,
+ }
+
+ impl ShufflePartitionPusher for BoundedShufflePartitionPusher {
+ fn push_partition_data(
+ &self,
+ _partition_id: i32,
+ _data: &[u8],
+ ) -> datafusion::common::Result<()> {
+ Ok(())
+ }
+
+ fn max_frame_size(&self) -> usize {
+ self.max_frame_size
+ }
+ }
+
fn local_shuffle_partition_writer(
output_data_file: &str,
output_index_file: &str,
@@ -5055,6 +5073,24 @@ mod tests {
}
}
+ #[test]
+ fn shuffle_partition_writer_uses_its_task_callback_frame_limit() {
+ let writer = spark_operator::ShuffleWriter {
+ partition_writer: Some(rss_shuffle_partition_writer()),
+ ..Default::default()
+ };
+ let callback: Arc<dyn ShufflePartitionPusher> = Arc::new(BoundedShufflePartitionPusher {
+ max_frame_size: 4096,
+ });
+
+ match super::shuffle_writer_destination(&writer, Some(&callback)).unwrap() {
+ ShuffleWriterDestination::Rss { max_frame_size, .. } => {
+ assert_eq!(max_frame_size, 4096);
+ }
+ destination => panic!("expected an RSS shuffle destination, got {destination:?}"),
+ }
+ }
+
#[test]
fn shuffle_partition_writer_rejects_rss_with_legacy_data_path() {
let writer = spark_operator::ShuffleWriter {
diff --git a/native/jni-bridge/src/shuffle_partition_pusher.rs b/native/jni-bridge/src/shuffle_partition_pusher.rs
index e5269b002a..a2cf2e8ee0 100644
--- a/native/jni-bridge/src/shuffle_partition_pusher.rs
+++ b/native/jni-bridge/src/shuffle_partition_pusher.rs
@@ -29,6 +29,28 @@ const MAX_JVM_ARRAY_LENGTH: i32 = i32::MAX - 8;
/// Implementations must remain safe when invoked from native execution
/// threads that do not inherit Spark's task-local JVM state.
pub trait ShufflePartitionPusher: Send + Sync {
+ /// Reserves encoding scratch before a complete shuffle frame is materialized.
+ fn reserve_partition_data(&self, _reservation_bytes: usize) -> Result<()> {
+ Ok(())
+ }
+
+ /// Acknowledges that encoding buffers and JNI references have been released, even on success.
+ /// An asynchronous implementation must also wait for its transport buffers before releasing
+ /// the reservation. Reserve, push, and release form one synchronous invocation on a worker.
+ fn release_partition_data_reservation(&self) -> Result<()> {
+ Ok(())
+ }
+
+ /// Returns the largest complete frame accepted by this callback.
+ fn max_frame_size(&self) -> usize {
+ JavaShufflePartitionPusher::MAX_PAYLOAD_SIZE
+ }
+
+ /// Returns the largest encoding reservation this callback can admit before allocation.
+ fn max_reservation_size(&self) -> usize {
+ usize::MAX
+ }
+
/// Sends one complete, length-prefixed Arrow IPC shuffle block.
fn push_partition_data(&self, partition_id: i32, data: &[u8]) -> Result<()>;
}
@@ -41,6 +63,10 @@ pub trait ShufflePartitionPusher: Send + Sync {
pub struct JavaShufflePartitionPusher {
callback: Global<JObject<'static>>,
push_method: JMethodID,
+ reserve_method: JMethodID,
+ release_method: JMethodID,
+ max_frame_size: usize,
+ max_reservation_size: usize,
}
impl JavaShufflePartitionPusher {
@@ -66,11 +92,85 @@ impl JavaShufflePartitionPusher {
jni::jni_sig!("(I[BI)V"),
)
.map_err(CometError::from)?;
+ let reserve_method = env
+ .get_method_id(
+ &callback_class,
+ jni::jni_str!("reservePartitionData"),
+ jni::jni_sig!("(I)V"),
+ )
+ .map_err(CometError::from)?;
+ let release_method = env
+ .get_method_id(
+ &callback_class,
+ jni::jni_str!("releasePartitionDataReservation"),
+ jni::jni_sig!("()V"),
+ )
+ .map_err(CometError::from)?;
+ let max_frame_method = env
+ .get_method_id(
+ &callback_class,
+ jni::jni_str!("maxFrameBytes"),
+ jni::jni_sig!("()I"),
+ )
+ .map_err(CometError::from)?;
+ // SAFETY: the cached ID was resolved on this callback class with a no-argument int ABI.
+ let configured_maximum = unsafe {
+ env.call_method_unchecked(
+ callback,
+ max_frame_method,
+ ReturnType::Primitive(Primitive::Int),
+ &[],
+ )
+ };
+ if let Some(exception) = check_exception(env)? {
+ return Err(exception.into());
+ }
+ let configured_maximum = configured_maximum
+ .map_err(CometError::from)?
+ .i()
+ .map_err(CometError::from)?;
+ if configured_maximum <= 0 || configured_maximum > MAX_JVM_ARRAY_LENGTH {
+ return Err(DataFusionError::Execution(format!(
+ "Remote shuffle maximum frame size {configured_maximum} is outside the JVM array limit of {MAX_JVM_ARRAY_LENGTH} bytes"
+ )));
+ }
+ let max_reservation_method = env
+ .get_method_id(
+ &callback_class,
+ jni::jni_str!("maxReservationBytes"),
+ jni::jni_sig!("()I"),
+ )
+ .map_err(CometError::from)?;
+ // SAFETY: the ID was resolved on this callback class with a no-argument int ABI.
+ let reservation_maximum = unsafe {
+ env.call_method_unchecked(
+ callback,
+ max_reservation_method,
+ ReturnType::Primitive(Primitive::Int),
+ &[],
+ )
+ };
+ if let Some(exception) = check_exception(env)? {
+ return Err(exception.into());
+ }
+ let reservation_maximum = reservation_maximum
+ .map_err(CometError::from)?
+ .i()
+ .map_err(CometError::from)?;
+ if reservation_maximum <= 0 {
+ return Err(DataFusionError::Execution(
+ "Remote shuffle maximum encoding reservation must be positive".to_string(),
+ ));
+ }
let callback = env.new_global_ref(callback).map_err(CometError::from)?;
Ok(Self {
callback,
push_method,
+ reserve_method,
+ release_method,
+ max_frame_size: configured_maximum as usize,
+ max_reservation_size: reservation_maximum as usize,
})
}
@@ -92,6 +192,62 @@ impl JavaShufflePartitionPusher {
}
impl ShufflePartitionPusher for JavaShufflePartitionPusher {
+ fn reserve_partition_data(&self, reservation_bytes: usize) -> Result<()> {
+ let bytes = i32::try_from(reservation_bytes).map_err(|_| {
+ DataFusionError::Execution(format!(
+ "Remote shuffle encoding reservation {reservation_bytes} exceeds the JVM integer limit"
+ ))
+ })?;
+ if bytes <= 0 {
+ return Err(DataFusionError::Execution(
+ "Remote shuffle encoding reservation must be positive".to_string(),
+ ));
+ }
+ JVMClasses::with_env(|env| {
+ // SAFETY: this default callback method has the cached `(I)V` signature.
+ let result = unsafe {
+ env.call_method_unchecked(
+ self.callback.as_obj(),
+ self.reserve_method,
+ ReturnType::Primitive(Primitive::Void),
+ &[JValue::Int(bytes).as_jni()],
+ )
+ };
+ if let Some(exception) = check_exception(env)? {
+ return Err(exception.into());
+ }
+ result.map_err(CometError::from)?;
+ Ok(())
+ })
+ }
+
+ fn release_partition_data_reservation(&self) -> Result<()> {
+ JVMClasses::with_env(|env| {
+ // SAFETY: this default callback method has the cached `()V` signature.
+ let result = unsafe {
+ env.call_method_unchecked(
+ self.callback.as_obj(),
+ self.release_method,
+ ReturnType::Primitive(Primitive::Void),
+ &[],
+ )
+ };
+ if let Some(exception) = check_exception(env)? {
+ return Err(exception.into());
+ }
+ result.map_err(CometError::from)?;
+ Ok(())
+ })
+ }
+
+ fn max_frame_size(&self) -> usize {
+ self.max_frame_size
+ }
+
+ fn max_reservation_size(&self) -> usize {
+ self.max_reservation_size
+ }
+
fn push_partition_data(&self, partition_id: i32, data: &[u8]) -> Result<()> {
let payload_length = Self::checked_payload_length(partition_id, data.len())?;
diff --git a/native/shuffle/Cargo.toml b/native/shuffle/Cargo.toml
index 58d82d2480..71be932422 100644
--- a/native/shuffle/Cargo.toml
+++ b/native/shuffle/Cargo.toml
@@ -30,6 +30,7 @@ publish = false
[dependencies]
arrow = { workspace = true }
+arrow-select = { workspace = true }
async-trait = { workspace = true }
bytes = { workspace = true }
clap = { version = "4", features = ["derive"], optional = true }
@@ -49,7 +50,7 @@ parquet = { workspace = true, optional = true }
simd-adler32 = "0.3.9"
snap = "1.1"
tokio = { version = "1", features = ["rt-multi-thread"] }
-zstd = "0.13.3"
+zstd = { version = "0.13.3", features = ["experimental"] }
[dev-dependencies]
criterion = { version = "0.7", features = ["async", "async_tokio", "async_std"] }
diff --git a/native/shuffle/src/shuffle_writer.rs b/native/shuffle/src/shuffle_writer.rs
index 810dbd626f..21d71dc122 100644
--- a/native/shuffle/src/shuffle_writer.rs
+++ b/native/shuffle/src/shuffle_writer.rs
@@ -267,12 +267,12 @@ async fn external_shuffle(
) -> Result<SendableRecordBatchStream> {
let schema = input.schema();
- let shuffle_block_writer = ShuffleBlockWriter::try_new(schema.as_ref(), codec.clone())?;
let mut repartitioner = match destination {
ShuffleWriterDestination::Local {
output_data_file,
output_index_file,
} => {
+ let shuffle_block_writer = ShuffleBlockWriter::try_new(schema.as_ref(), codec.clone())?;
let writer = LocalPartitionWriter::try_new(
output_data_file,
output_index_file,
@@ -298,6 +298,8 @@ async fn external_shuffle(
pusher,
max_frame_size,
} => {
+ let shuffle_block_writer =
+ ShuffleBlockWriter::try_new_rss(Arc::clone(&schema), codec.clone())?;
let writer = RssPartitionWriter::try_new(
shuffle_block_writer,
pusher,
diff --git a/native/shuffle/src/writers/rss/mod.rs b/native/shuffle/src/writers/rss/mod.rs
index c23d742e7d..a151b6831f 100644
--- a/native/shuffle/src/writers/rss/mod.rs
+++ b/native/shuffle/src/writers/rss/mod.rs
@@ -23,15 +23,143 @@ mod tests {
use crate::metrics::ShufflePartitionerMetrics;
use crate::writers::PartitionWriter;
use crate::{read_ipc_compressed, CompressionCodec, ShuffleBlockWriter};
- use arrow::array::{Array, DictionaryArray, Int32Array};
+ use arrow::array::{
+ Array, ArrayRef, DictionaryArray, Int32Array, ListArray, MapArray, StringArray, StructArray,
+ };
+ use arrow::buffer::OffsetBuffer;
+ use arrow::compute::cast;
use arrow::datatypes::{DataType, Field, Int32Type, Schema};
use arrow::ipc::writer::CompressionContext;
use arrow::record_batch::RecordBatch;
use datafusion::common::{DataFusionError, Result};
use datafusion::physical_plan::metrics::{ExecutionPlanMetricsSet, Time};
use datafusion_comet_jni_bridge::ShufflePartitionPusher;
- use std::io::Cursor;
- use std::sync::{Arc, Mutex};
+ use std::collections::HashMap;
+ use std::io::{self, Cursor};
+ use std::sync::atomic::{AtomicUsize, Ordering};
+ use std::sync::{Arc, Condvar, Mutex};
+ use std::thread::{self, ThreadId};
+ use std::time::Duration;
+
+ /// Test-only allocation observation on a synchronous encoder thread. Production execution
+ /// does not use thread-local state. Zstd's C allocations are covered separately by its public
+ /// streaming-workspace estimate; this observes Rust buffers and their realloc overlap.
+ mod allocations {
+ use std::alloc::{GlobalAlloc, Layout, System};
+ use std::cell::Cell;
+
+ #[derive(Clone, Copy, Default)]
+ struct Counters {
+ live: isize,
+ peak: usize,
+ admission: Option<(isize, usize)>,
+ excess: usize,
+ }
+
+ thread_local! {
+ static COUNTERS: Cell<Option<Counters>> = const { Cell::new(None) };
+ }
+
+ struct ObservedAllocator;
+
+ #[global_allocator]
+ static ALLOCATOR: ObservedAllocator = ObservedAllocator;
+
+ fn record(added: usize, removed: usize) {
+ let _ = COUNTERS.try_with(|counter| {
+ if let Some(mut value) = counter.get() {
+ value.live = value
+ .live
+ .saturating_add(added as isize)
+ .saturating_sub(removed as isize);
+ value.peak = value.peak.max(value.live.max(0) as usize);
+ if let Some((baseline, admitted)) = value.admission {
+ let used = value.live.saturating_sub(baseline).max(0) as usize;
+ value.excess = value.excess.max(used.saturating_sub(admitted));
+ }
+ counter.set(Some(value));
+ }
+ });
+ }
+
+ // SAFETY: every operation is forwarded to the system allocator with the unchanged
+ // pointer/layout. The observation uses only allocation-free thread-local Cells.
+ unsafe impl GlobalAlloc for ObservedAllocator {
+ unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
+ let pointer = unsafe { System.alloc(layout) };
+ if !pointer.is_null() {
+ record(layout.size(), 0);
+ }
+ pointer
+ }
+
+ unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
+ let pointer = unsafe { System.alloc_zeroed(layout) };
+ if !pointer.is_null() {
+ record(layout.size(), 0);
+ }
+ pointer
+ }
+
+ unsafe fn dealloc(&self, pointer: *mut u8, layout: Layout) {
+ unsafe { System.dealloc(pointer, layout) };
+ record(0, layout.size());
+ }
+
+ unsafe fn realloc(&self, pointer: *mut u8, layout: Layout, size: usize) -> *mut u8 {
+ let result = unsafe { System.realloc(pointer, layout, size) };
+ if !result.is_null() {
+ // A moving realloc may hold old and new buffers at once. Count that peak
+ // even when the system allocator happened to resize in place in this run.
+ record(size, 0);
+ record(0, layout.size());
+ }
+ result
+ }
+ }
+
+ pub(super) fn live() -> isize {
+ COUNTERS.with(|counter| counter.get().unwrap().live)
+ }
+
+ pub(super) fn admit(bytes: usize) {
+ COUNTERS.with(|counter| {
+ if let Some(mut value) = counter.get() {
+ assert!(value.admission.is_none());
+ value.admission = Some((value.live, bytes));
+ counter.set(Some(value));
+ }
+ });
+ }
+
+ pub(super) fn retire() {
+ COUNTERS.with(|counter| {
+ if let Some(mut value) = counter.get() {
+ value.admission = None;
+ counter.set(Some(value));
+ }
+ });
+ }
+
+ pub(super) fn measure<T>(run: impl FnOnce() -> T) -> (T, usize) {
+ struct Reset;
+ impl Drop for Reset {
+ fn drop(&mut self) {
+ COUNTERS.with(|counter| counter.set(None));
+ }
+ }
+ COUNTERS.with(|counter| assert!(counter.replace(Some(Counters::default())).is_none()));
+ let reset = Reset;
+ let result = run();
+ let observed = COUNTERS.with(|counter| counter.get().unwrap());
+ drop(reset);
+ assert_eq!(
+ observed.excess, 0,
+ "a planning or encoding phase exceeded its live reservation"
+ );
+ (result, observed.peak)
+ }
+ }
type RecordedFrame = (i32, Vec<u8>);
@@ -66,6 +194,152 @@ mod tests {
}
}
+ #[derive(Default)]
+ struct ReservationRecordingPusher {
+ frames: Mutex<Vec<RecordedFrame>>,
+ reservations: Mutex<Vec<usize>>,
+ outstanding: Mutex<Option<usize>>,
+ releases: AtomicUsize,
+ capacity: Option<usize>,
+ max_reservation: Option<usize>,
+ }
+
+ impl ShufflePartitionPusher for ReservationRecordingPusher {
+ fn max_reservation_size(&self) -> usize {
+ self.max_reservation.unwrap_or(usize::MAX)
+ }
+
+ fn reserve_partition_data(&self, bytes: usize) -> Result<()> {
+ if self.capacity.is_some_and(|capacity| bytes > capacity) {
+ return Err(DataFusionError::External(Box::new(io::Error::other(
+ "native encoding admission denied",
+ ))));
+ }
+ let mut outstanding = self.outstanding.lock().unwrap();
+ assert!(outstanding.replace(bytes).is_none());
+ allocations::admit(bytes);
+ self.reservations.lock().unwrap().push(bytes);
+ Ok(())
+ }
+
+ fn release_partition_data_reservation(&self) -> Result<()> {
+ if self.outstanding.lock().unwrap().take().is_some() {
+ self.releases.fetch_add(1, Ordering::Relaxed);
+ allocations::retire();
+ }
+ Ok(())
+ }
+
+ fn push_partition_data(&self, partition_id: i32, bytes: &[u8]) -> Result<()> {
+ assert!(
+ self.outstanding.lock().unwrap().is_some(),
+ "encoding must be admitted before a frame reaches its transport"
+ );
+ self.frames
+ .lock()
+ .unwrap()
+ .push((partition_id, bytes.to_vec()));
+ Ok(())
+ }
+ }
+
+ #[derive(Default)]
+ struct ConcurrentAdmissionState {
+ reservations: HashMap<ThreadId, usize>,
+ rounds: HashMap<ThreadId, usize>,
+ total: usize,
+ admitted: usize,
+ failed: bool,
+ }
+
+ struct ConcurrentAdmissionPusher {
+ capacity: usize,
+ state: Mutex<ConcurrentAdmissionState>,
+ ready: Condvar,
+ peak: AtomicUsize,
+ frames: AtomicUsize,
+ }
+
+ impl ConcurrentAdmissionPusher {
+ fn new(capacity: usize) -> Self {
+ Self {
+ capacity,
+ state: Mutex::new(ConcurrentAdmissionState::default()),
+ ready: Condvar::new(),
+ peak: AtomicUsize::new(0),
+ frames: AtomicUsize::new(0),
+ }
+ }
+ }
+
+ impl ShufflePartitionPusher for ConcurrentAdmissionPusher {
+ fn reserve_partition_data(&self, bytes: usize) -> Result<()> {
+ let current = thread::current().id();
+ let mut state = self.state.lock().unwrap();
+ if state.total.saturating_add(bytes) > self.capacity {
+ state.failed = true;
+ self.ready.notify_all();
+ return Err(DataFusionError::Execution(
+ "concurrent native encoding exhausted executor admission".to_string(),
+ ));
+ }
+
+ assert!(state.reservations.insert(current, bytes).is_none());
+ state.total += bytes;
+ self.peak.fetch_max(state.total, Ordering::Relaxed);
+ let round = state.rounds.entry(current).or_default();
+ *round += 1;
+ if *round == 1 {
+ // The first lease only protects descriptor-based size estimation. Synchronize
+ // the following full encoding reservations, which must really coexist.
+ return Ok(());
+ }
+ state.admitted += 1;
+ self.ready.notify_all();
+
+ while state.admitted < 2 && !state.failed {
+ let (updated, timeout) = self
+ .ready
+ .wait_timeout(state, Duration::from_secs(10))
+ .unwrap();
+ state = updated;
+ if timeout.timed_out() && state.admitted < 2 {
+ state.failed = true;
+ self.ready.notify_all();
+ }
+ }
+ if state.failed {
+ state.total -= state.reservations.remove(¤t).unwrap();
+ return Err(DataFusionError::Execution(
+ "concurrent native encoders could not share executor admission".to_string(),
+ ));
+ }
+ Ok(())
+ }
+
+ fn release_partition_data_reservation(&self) -> Result<()> {
+ let mut state = self.state.lock().unwrap();
+ if let Some(bytes) = state.reservations.remove(&thread::current().id()) {
+ state.total -= bytes;
+ }
+ Ok(())
+ }
+
+ fn push_partition_data(&self, _partition_id: i32, data: &[u8]) -> Result<()> {
+ let state = self.state.lock().unwrap();
+ let bytes = state
+ .reservations
+ .get(&thread::current().id())
+ .expect("shuffle frame must keep its encoding reservation");
+ assert!(
+ *bytes >= data.len() * 3,
+ "native IPC, JNI, and transport copies must all be admitted"
+ );
+ self.frames.fetch_add(1, Ordering::Relaxed);
+ Ok(())
+ }
+ }
+
fn sample_batch(start: i32, rows: i32) -> RecordBatch {
let values = Int32Array::from_iter_values(start..start + rows);
let schema = Arc::new(Schema::new(vec![Field::new(
@@ -87,7 +361,7 @@ mod tests {
partitions: usize,
max_frame_size: usize,
) -> RssPartitionWriter {
- let block_writer = ShuffleBlockWriter::try_new(batch.schema().as_ref(), codec).unwrap();
+ let block_writer = ShuffleBlockWriter::try_new_rss(batch.schema(), codec).unwrap();
RssPartitionWriter::try_new(block_writer, pusher, partitions, max_frame_size).unwrap()
}
@@ -347,8 +621,8 @@ mod tests {
}
#[test]
- fn rejects_oversized_frames_without_splitting_or_pushing() {
- let batch = sample_batch(0, 16);
+ fn rejects_irreducible_single_row_frames_without_fragmenting_or_pushing() {
+ let batch = sample_batch(0, 1);
let frame_size = encoded_frame_size(&batch);
let pusher = Arc::new(RecordingPusher::default());
let mut writer = writer(
@@ -362,7 +636,7 @@ mod tests {
assert!(write_batches(&mut writer, 0, vec![batch], &metrics()).is_err());
assert!(
pusher.frames().is_empty(),
- "an oversized Arrow IPC frame must never be fragmented"
+ "a single oversized Arrow IPC row must never be fragmented"
);
}
@@ -469,4 +743,607 @@ mod tests {
assert!(writer.finish_all(&metrics).is_err());
assert!(pusher.frames().is_empty());
}
+
+ #[test]
+ fn splits_oversized_batches_into_complete_row_aligned_frames() {
+ let batch = sample_batch(0, 512);
+ let limit = encoded_frame_size(&sample_batch(0, 32));
+ let pusher = Arc::new(ReservationRecordingPusher::default());
+ let mut writer = writer(&batch, CompressionCodec::None, pusher.clone(), 1, limit);
+ let metrics = metrics();
+
+ finish_partition(&mut writer, 0, vec![batch], &metrics).unwrap();
+ writer.finish_all(&metrics).unwrap();
+
+ let frames = pusher.frames.lock().unwrap();
+ assert!(frames.len() > 1);
+ let mut values = Vec::new();
+ for (partition, frame) in frames.iter() {
+ assert_eq!(*partition, 0);
+ assert!(frame.len() <= limit);
+ let decoded = decode_frame(frame);
+ values.extend(
+ decoded
+ .column(0)
+ .as_any()
+ .downcast_ref::<Int32Array>()
+ .unwrap()
+ .values()
+ .iter()
+ .copied(),
+ );
+ }
+ assert_eq!(values, (0..512).collect::<Vec<_>>());
+ assert!(pusher
+ .reservations
+ .lock()
+ .unwrap()
+ .iter()
+ .all(|reservation| *reservation >= 3 * limit));
+ assert!(pusher.outstanding.lock().unwrap().is_none());
+ }
+
+ #[test]
+ fn small_batches_encode_concurrently_with_default_frame_and_admission_limits() {
+ let frame_limit = 64 * 1024 * 1024;
+ let admission_limit = 256 * 1024 * 1024;
+ let pusher = Arc::new(ConcurrentAdmissionPusher::new(admission_limit));
+
+ thread::scope(|scope| {
+ let tasks = (0..2)
+ .map(|partition| {
+ let pusher = Arc::clone(&pusher);
+ scope.spawn(move || {
+ let batch = sample_batch(partition * 8, 8);
+ let mut writer =
+ writer(&batch, CompressionCodec::None, pusher, 1, frame_limit);
+ write_batches(&mut writer, 0, vec![batch], &metrics())
+ })
+ })
+ .collect::<Vec<_>>();
+
+ for task in tasks {
+ task.join().unwrap().unwrap();
+ }
+ });
+
+ assert_eq!(pusher.frames.load(Ordering::Relaxed), 2);
+ assert!(
+ pusher.peak.load(Ordering::Relaxed) < admission_limit / 2,
+ "tiny frames must not each reserve three copies of the maximum frame"
+ );
+ assert_eq!(pusher.state.lock().unwrap().total, 0);
+ }
+
+ #[test]
+ fn codec_workspace_is_rejected_before_encoding_under_a_tiny_budget() {
+ let batch = sample_batch(0, 1);
+ for codec in [
+ CompressionCodec::Lz4Frame,
+ CompressionCodec::Snappy,
+ CompressionCodec::Zstd(1),
+ ] {
+ let pusher = Arc::new(ReservationRecordingPusher {
+ max_reservation: Some(16_384),
+ capacity: Some(16_384),
+ ..ReservationRecordingPusher::default()
+ });
+ let mut writer = writer(&batch, codec, pusher.clone(), 1, 5_456);
+ let error = write_batches(&mut writer, 0, vec![batch.clone()], &metrics()).unwrap_err();
+ assert!(error.to_string().contains("encoding workspace"));
+ assert!(pusher.reservations.lock().unwrap().is_empty());
+ assert!(pusher.frames.lock().unwrap().is_empty());
+ }
+ }
+
+ #[test]
+ fn rss_construction_retains_schema_without_allocating_schema_sized_buffers() {
+ let schema = Arc::new(Schema::new(vec![Field::new(
+ "large-field-name".repeat(65_536),
+ DataType::Int32,
+ false,
+ )]));
+ let (writer, peak) = allocations::measure(|| {
+ ShuffleBlockWriter::try_new_rss(Arc::clone(&schema), CompressionCodec::Lz4Frame)
+ });
+ writer.unwrap();
+ assert!(
+ peak < 1_024,
+ "RSS constructor pre-encoded or cloned its schema: {peak} bytes"
+ );
+ }
+
+ #[test]
+ fn splits_when_full_encoding_workspace_cannot_fit_one_reservation() {
+ let batch = sample_batch(0, 16_384);
+ let pusher = Arc::new(ReservationRecordingPusher {
+ max_reservation: Some(262_144),
+ capacity: Some(262_144),
+ ..ReservationRecordingPusher::default()
+ });
+ let mut writer = writer(&batch, CompressionCodec::None, pusher.clone(), 1, 67_996);
+ write_batches(&mut writer, 0, vec![batch.clone()], &metrics()).unwrap();
+ let frames = pusher.frames.lock().unwrap();
+ assert!(frames.len() > 1);
+ let decoded = frames
+ .iter()
+ .map(|(_, frame)| decode_frame(frame))
+ .collect::<Vec<_>>();
+ assert_eq!(
+ arrow::compute::concat_batches(&batch.schema(), &decoded).unwrap(),
+ batch
+ );
+ assert!(pusher
+ .reservations
+ .lock()
+ .unwrap()
+ .iter()
+ .all(|bytes| *bytes <= 262_144));
+ assert!(pusher.outstanding.lock().unwrap().is_none());
+ }
+
+ #[test]
+ fn schema_scratch_is_charged_before_a_compressible_dictionary_schema_is_encoded() {
+ let dictionary: DictionaryArray<Int32Type> = ["value"].into_iter().collect();
+ let field = Field::new(
+ "long-field".repeat(16_384),
+ dictionary.data_type().clone(),
+ false,
+ )
+ .with_metadata(HashMap::from([(
+ "metadata".to_string(),
+ "x".repeat(65_536),
+ )]));
+ let batch = RecordBatch::try_new(
+ Arc::new(Schema::new(vec![field])),
+ vec![Arc::new(dictionary)],
+ )
+ .unwrap();
+ let pusher = Arc::new(ReservationRecordingPusher {
+ max_reservation: Some(512 * 1024),
+ capacity: Some(512 * 1024),
+ ..ReservationRecordingPusher::default()
+ });
+ let mut writer = writer(
+ &batch,
+ CompressionCodec::Lz4Frame,
+ pusher.clone(),
+ 1,
+ 64 * 1024,
+ );
+ let error = write_batches(&mut writer, 0, vec![batch], &metrics()).unwrap_err();
+ assert!(error.to_string().contains("encoding workspace"));
+ assert!(pusher.reservations.lock().unwrap().is_empty());
+ assert!(pusher.frames.lock().unwrap().is_empty());
+ }
+
+ #[test]
+ fn measured_encoding_allocations_fit_admitted_workspace_for_all_codecs() {
+ let dictionary: DictionaryArray<Int32Type> = ["value"].into_iter().collect();
+ let dictionary_batch = RecordBatch::try_new(
+ Arc::new(Schema::new(vec![Field::new(
+ "wide-name".repeat(8_192),
+ dictionary.data_type().clone(),
+ false,
+ )
+ .with_metadata(HashMap::from([(
+ "key".to_string(),
+ "metadata".repeat(4_096),
+ )]))])),
+ vec![Arc::new(dictionary)],
+ )
+ .unwrap();
+ let mut nested: ArrayRef = Arc::new(Int32Array::from_iter_values(0..128));
+ for _ in 0..16 {
+ let field = Arc::new(Field::new("item", nested.data_type().clone(), false));
+ let offsets = OffsetBuffer::new(vec![0, nested.len() as i32].into());
+ nested = Arc::new(ListArray::try_new(field, offsets, nested, None).unwrap());
+ }
+ let nested_batch = RecordBatch::try_new(
+ Arc::new(Schema::new(vec![Field::new(
+ "nested",
+ nested.data_type().clone(),
+ false,
+ )])),
+ vec![nested],
+ )
+ .unwrap();
+ let wide_batch = RecordBatch::try_new(
+ Arc::new(Schema::new(
+ (0..256)
+ .map(|index| Field::new(format!("field{index}"), DataType::Int32, false))
+ .collect::<Vec<_>>(),
+ )),
+ (0..256)
+ .map(|_| Arc::new(Int32Array::from(vec![1])) as ArrayRef)
+ .collect(),
+ )
+ .unwrap();
+ for batch in [
+ sample_batch(0, 1),
+ sample_batch(0, 16_384),
+ dictionary_batch,
+ nested_batch,
+ wide_batch,
+ ] {
+ for codec in [
+ CompressionCodec::None,
+ CompressionCodec::Lz4Frame,
+ CompressionCodec::Snappy,
+ CompressionCodec::Zstd(1),
+ ] {
+ let pusher = Arc::new(ReservationRecordingPusher {
+ max_reservation: Some(16 * 1024 * 1024),
+ capacity: Some(16 * 1024 * 1024),
+ ..ReservationRecordingPusher::default()
+ });
+ let mut writer = writer(&batch, codec.clone(), pusher.clone(), 1, 1024 * 1024);
+ let metrics = metrics();
+ let (result, peak) = allocations::measure(|| {
+ write_batches(&mut writer, 0, vec![batch.clone()], &metrics)
+ });
+ result.unwrap();
+ let reserved = *pusher.reservations.lock().unwrap().iter().max().unwrap();
+ assert!(
+ peak <= reserved,
+ "{codec:?}: peak allocation {peak} exceeds reservation {reserved}"
+ );
+ assert_eq!(decode_frame(&pusher.frames.lock().unwrap()[0].1), batch);
+ }
+ }
+ }
+
+ #[test]
+ fn native_capacity_is_freed_before_the_successful_callback_is_acknowledged() {
+ use std::sync::atomic::{AtomicBool, AtomicIsize};
+
+ #[derive(Default)]
+ struct Pusher {
+ at_push: AtomicIsize,
+ frame_length: AtomicUsize,
+ acknowledged: AtomicBool,
+ }
+ impl ShufflePartitionPusher for Pusher {
+ fn push_partition_data(&self, _partition: i32, bytes: &[u8]) -> Result<()> {
+ self.at_push.store(allocations::live(), Ordering::Relaxed);
+ self.frame_length.store(bytes.len(), Ordering::Relaxed);
+ Ok(())
+ }
+ fn release_partition_data_reservation(&self) -> Result<()> {
+ if self.frame_length.load(Ordering::Relaxed) == 0 {
+ // Planning descriptors are released before the actual frame is submitted.
+ return Ok(());
+ }
+ assert!(
+ allocations::live()
+ <= self.at_push.load(Ordering::Relaxed)
+ - self.frame_length.load(Ordering::Relaxed) as isize,
+ "the output allocation must die before native ownership is acknowledged"
+ );
+ self.acknowledged.store(true, Ordering::Relaxed);
+ Ok(())
+ }
+ }
+ let batch = sample_batch(0, 2_048);
+ let pusher = Arc::new(Pusher::default());
+ let mut writer = writer(&batch, CompressionCodec::None, pusher.clone(), 1, 64 * 1024);
+ let metrics = metrics();
+ allocations::measure(|| write_batches(&mut writer, 0, vec![batch], &metrics))
+ .0
+ .unwrap();
+ assert!(pusher.acknowledged.load(Ordering::Relaxed));
+ }
+
+ #[test]
+ fn oversized_schema_metadata_retries_with_a_fresh_larger_atomic_reservation() {
+ // Keep the schema difficult to compress so every supported codec must retry its
+ // underestimated metadata rather than fitting it into the initial 4 KiB allowance.
+ let mut random = 0x1234_5678_9abc_def0u64;
+ let field_name = (0..16_384)
+ .map(|_| {
+ random = random
+ .wrapping_mul(6_364_136_223_846_793_005)
+ .wrapping_add(1);
+ char::from(b'a' + ((random >> 32) % 26) as u8)
+ })
+ .collect::<String>();
+ let schema = Arc::new(Schema::new(vec![Field::new(
+ field_name.clone(),
+ DataType::Int32,
+ false,
+ )]));
+ let plain_batch =
+ RecordBatch::try_new(schema, vec![Arc::new(Int32Array::from(vec![7]))]).unwrap();
+ let dictionary: DictionaryArray<Int32Type> = ["dictionary-value"].into_iter().collect();
+ let dictionary_schema = Arc::new(Schema::new(vec![Field::new(
+ field_name,
+ dictionary.data_type().clone(),
+ false,
+ )]));
+ let dictionary_batch =
+ RecordBatch::try_new(dictionary_schema, vec![Arc::new(dictionary)]).unwrap();
+
+ for batch in [plain_batch, dictionary_batch] {
+ for codec in [
+ CompressionCodec::None,
+ CompressionCodec::Lz4Frame,
+ CompressionCodec::Snappy,
+ CompressionCodec::Zstd(1),
+ ] {
+ let pusher = Arc::new(ReservationRecordingPusher::default());
+ let mut writer = writer(&batch, codec, pusher.clone(), 1, 64 * 1024);
+
+ write_batches(&mut writer, 0, vec![batch.clone()], &metrics()).unwrap();
+
+ let frames = pusher.frames.lock().unwrap();
+ assert_eq!(frames.len(), 1);
+ assert_eq!(decode_frame(&frames[0].1), batch);
+ let reservations = pusher.reservations.lock().unwrap();
+ assert!(reservations.len() > 1);
+ assert!(reservations.windows(2).all(|pair| pair[0] < pair[1]));
+ assert!(reservations.last().unwrap() >= &(frames[0].1.len() * 3));
+ assert_eq!(
+ pusher.releases.load(Ordering::Relaxed),
+ reservations.len(),
+ "every retry and successful frame must acknowledge native ownership release"
+ );
+ assert!(pusher.outstanding.lock().unwrap().is_none());
+ }
+ }
+ }
+
+ #[test]
+ fn denied_pre_encoding_reservation_preserves_original_error_without_pushing() {
+ let batch = sample_batch(0, 8);
+ let limit = 1_024;
+ let pusher = Arc::new(ReservationRecordingPusher {
+ capacity: Some(3 * limit - 1),
+ ..ReservationRecordingPusher::default()
+ });
+ let mut writer = writer(&batch, CompressionCodec::None, pusher.clone(), 1, limit);
+
+ let error = write_batches(&mut writer, 0, vec![batch], &metrics()).unwrap_err();
+ let DataFusionError::External(cause) = error else {
+ panic!("encoding admission failure was wrapped or stringified");
+ };
+ assert_eq!(cause.to_string(), "native encoding admission denied");
+ assert!(pusher.frames.lock().unwrap().is_empty());
+ assert_eq!(pusher.releases.load(Ordering::Relaxed), 0);
+ }
+
+ #[test]
+ fn oversized_single_rows_release_encoding_admission_without_fragmentation() {
+ let schema = Arc::new(Schema::new(vec![Field::new(
+ "value",
+ DataType::Utf8,
+ false,
+ )]));
+ let batch = RecordBatch::try_new(
+ schema,
+ vec![Arc::new(StringArray::from(vec!["x".repeat(16_384)]))],
+ )
+ .unwrap();
+ let pusher = Arc::new(ReservationRecordingPusher::default());
+ let mut writer = writer(&batch, CompressionCodec::None, pusher.clone(), 1, 512);
+
+ let error = write_batches(&mut writer, 0, vec![batch], &metrics()).unwrap_err();
+ assert!(error.to_string().contains("single row"));
+ assert!(pusher.frames.lock().unwrap().is_empty());
+ assert_eq!(pusher.releases.load(Ordering::Relaxed), 2);
+ assert!(pusher.outstanding.lock().unwrap().is_none());
+ }
+
+ #[test]
+ fn compacts_dictionaries_nested_inside_sliced_lists_and_structs() {
+ let values = (0..128)
+ .map(|index| format!("value-{index:03}-{}", "x".repeat(128)))
+ .collect::<Vec<_>>();
+ let dictionary_type =
+ DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8));
+ let dictionary = DictionaryArray::<Int32Type>::try_new(
+ Int32Array::from((0..32).collect::<Vec<_>>()),
+ Arc::new(StringArray::from(values.clone())),
+ )
+ .unwrap();
+ let item = Arc::new(Field::new("item", dictionary_type, false));
+ let list = ListArray::try_new(
+ Arc::clone(&item),
+ OffsetBuffer::new((0..=32).collect::<Vec<i32>>().into()),
+ Arc::new(dictionary),
+ None,
+ )
+ .unwrap();
+ let nested = Arc::new(Field::new("items", DataType::List(item), false));
+ let structure =
+ StructArray::try_new(vec![Arc::clone(&nested)].into(), vec![Arc::new(list)], None)
+ .unwrap();
+ let schema = Arc::new(Schema::new(vec![Field::new(
+ "nested",
+ DataType::Struct(vec![nested].into()),
+ false,
+ )]));
+ let batch = RecordBatch::try_new(schema, vec![Arc::new(structure)])
+ .unwrap()
+ .slice(7, 16);
+ let pusher = Arc::new(ReservationRecordingPusher::default());
+ let mut writer = writer(&batch, CompressionCodec::None, pusher.clone(), 1, 2_048);
+
+ finish_partition(&mut writer, 0, vec![batch], &metrics()).unwrap();
+
+ let frames = pusher.frames.lock().unwrap();
+ assert!(frames.len() > 1);
+ let mut actual = Vec::new();
+ for (_, frame) in frames.iter() {
+ assert!(frame.len() <= 2_048);
+ let decoded = decode_frame(frame);
+ let structure = decoded
+ .column(0)
+ .as_any()
+ .downcast_ref::<StructArray>()
+ .unwrap();
+ let lists = structure
+ .column(0)
+ .as_any()
+ .downcast_ref::<ListArray>()
+ .unwrap();
+ assert_eq!(lists.value_offsets()[0], 0);
+ let dictionary = lists
+ .values()
+ .as_any()
+ .downcast_ref::<DictionaryArray<Int32Type>>()
+ .unwrap();
+ assert!(dictionary.values().len() < values.len());
+ let plain = cast(dictionary, &DataType::Utf8).unwrap();
+ actual.extend(
+ plain
+ .as_any()
+ .downcast_ref::<StringArray>()
+ .unwrap()
+ .iter()
+ .map(|value| value.unwrap().to_owned()),
+ );
+ }
+ assert_eq!(actual, values[7..23].to_vec());
+ }
+
+ #[test]
+ fn sliced_lists_and_maps_charge_only_their_live_nested_entries() {
+ let count = 16_384;
+ let values: ArrayRef = Arc::new(Int32Array::from_iter_values(0..count));
+ let offsets = OffsetBuffer::new((0..=count).collect::<Vec<i32>>().into());
+ let item = Arc::new(Field::new("item", DataType::Int32, false));
+ let list = ListArray::try_new(
+ Arc::clone(&item),
+ offsets.clone(),
+ Arc::clone(&values),
+ None,
+ )
+ .unwrap()
+ .slice(count as usize / 2, 1);
+ let list_batch = RecordBatch::try_new(
+ Arc::new(Schema::new(vec![Field::new(
+ "items",
+ DataType::List(item),
+ false,
+ )])),
+ vec![Arc::new(list)],
+ )
+ .unwrap();
+
+ let fields = vec![
+ Arc::new(Field::new("key", DataType::Int32, false)),
+ Arc::new(Field::new("value", DataType::Int32, false)),
+ ];
+ let entries = StructArray::try_new(
+ fields.clone().into(),
+ vec![Arc::clone(&values), Arc::clone(&values)],
+ None,
+ )
+ .unwrap();
+ let entry = Arc::new(Field::new(
+ "entries",
+ DataType::Struct(fields.into()),
+ false,
+ ));
+ let map = MapArray::try_new(Arc::clone(&entry), offsets, entries, None, false)
+ .unwrap()
+ .slice(count as usize / 2, 1);
+ let map_batch = RecordBatch::try_new(
+ Arc::new(Schema::new(vec![Field::new(
+ "entries",
+ DataType::Map(entry, false),
+ false,
+ )])),
+ vec![Arc::new(map)],
+ )
+ .unwrap();
+
+ let limit = 2_048;
+ for batch in [list_batch, map_batch] {
+ let pusher = Arc::new(ReservationRecordingPusher {
+ capacity: Some(64 * 1024),
+ max_reservation: Some(64 * 1024),
+ ..ReservationRecordingPusher::default()
+ });
+ let mut writer = writer(&batch, CompressionCodec::None, pusher.clone(), 1, limit);
+ finish_partition(&mut writer, 0, vec![batch.clone()], &metrics()).unwrap();
+ let frames = pusher.frames.lock().unwrap();
+ assert_eq!(frames.len(), 1);
+ assert_eq!(decode_frame(&frames[0].1), batch);
+ }
+ }
+
+ #[test]
+ fn input_failure_permanently_poisons_a_partially_pushed_map() {
+ let batch = sample_batch(0, 4);
+ let pusher = Arc::new(RecordingPusher::default());
+ let mut writer = writer(&batch, CompressionCodec::None, pusher.clone(), 1, 1_024);
+ let metrics = metrics();
+ let upstream_failure =
+ DataFusionError::External(Box::new(io::Error::other("shuffle input failed")));
+ let mut batches = [Ok(batch.clone()), Err(upstream_failure)].into_iter();
+
+ let error = writer.write(0, &mut batches, &metrics).unwrap_err();
+ let DataFusionError::External(cause) = error else {
+ panic!("upstream error was wrapped or stringified");
+ };
+ assert_eq!(cause.to_string(), "shuffle input failed");
+ assert_eq!(pusher.frames().len(), 1);
+ assert!(write_batches(&mut writer, 0, vec![batch], &metrics).is_err());
+ assert!(finish_partition(&mut writer, 0, vec![], &metrics).is_err());
+ assert!(writer.finish_all(&metrics).is_err());
+ assert_eq!(pusher.frames().len(), 1);
+ }
+
+ #[test]
+ fn legacy_jvm_frame_limit_uses_a_representable_java_encoding_reservation() {
+ struct IntegerLimitedPusher {
+ reservation: Mutex<Option<usize>>,
+ frames: Mutex<Vec<RecordedFrame>>,
+ }
+
+ impl ShufflePartitionPusher for IntegerLimitedPusher {
+ fn max_reservation_size(&self) -> usize {
+ i32::MAX as usize
+ }
+
+ fn reserve_partition_data(&self, bytes: usize) -> Result<()> {
+ assert!(i32::try_from(bytes).is_ok());
+ *self.reservation.lock().unwrap() = Some(bytes);
+ Ok(())
+ }
+
+ fn push_partition_data(&self, partition: i32, frame: &[u8]) -> Result<()> {
+ self.frames
+ .lock()
+ .unwrap()
+ .push((partition, frame.to_vec()));
+ Ok(())
+ }
+ }
+
+ let batch = sample_batch(0, 4);
+ let pusher = Arc::new(IntegerLimitedPusher {
+ reservation: Mutex::new(None),
+ frames: Mutex::new(Vec::new()),
+ });
+ let mut writer = writer(
+ &batch,
+ CompressionCodec::None,
+ pusher.clone(),
+ 1,
+ (i32::MAX - 8) as usize,
+ );
+
+ finish_partition(&mut writer, 0, vec![batch.clone()], &metrics()).unwrap();
+ let frames = pusher.frames.lock().unwrap();
+ assert_eq!(frames.len(), 1);
+ assert_eq!(decode_frame(&frames[0].1), batch);
+ let reservation = pusher.reservation.lock().unwrap().unwrap();
+ assert!(reservation >= frames[0].1.len() * 3);
+ assert!(
+ reservation < 64 * 1024,
+ "legacy callbacks should charge the actual batch instead of the JVM array maximum"
+ );
+ }
}
diff --git a/native/shuffle/src/writers/rss/rss_partition_writer.rs b/native/shuffle/src/writers/rss/rss_partition_writer.rs
index 3427136893..76503011db 100644
--- a/native/shuffle/src/writers/rss/rss_partition_writer.rs
+++ b/native/shuffle/src/writers/rss/rss_partition_writer.rs
@@ -18,11 +18,19 @@
use crate::metrics::ShufflePartitionerMetrics;
use crate::writers::partition_writer::PartitionWriter;
use crate::ShuffleBlockWriter;
-use arrow::array::RecordBatch;
+use arrow::array::cast::AsArray;
+use arrow::array::{
+ Array, ArrayRef, BinaryViewArray, FixedSizeListArray, LargeListArray, LargeListViewArray,
+ ListArray, ListViewArray, MapArray, RecordBatch, RunArray, StringViewArray, StructArray,
+ UnionArray,
+};
+use arrow::buffer::OffsetBuffer;
+use arrow::datatypes::{DataType, Field, Int16Type, Int32Type, Int64Type};
use arrow::ipc::writer::CompressionContext;
+use arrow_select::dictionary::garbage_collect_any_dictionary;
use datafusion::common::{DataFusionError, Result};
use datafusion_comet_jni_bridge::ShufflePartitionPusher;
-use std::io::Cursor;
+use std::io::{self, Cursor, Seek, SeekFrom, Write};
use std::sync::Arc;
/// Sends complete Comet shuffle blocks to a task-owned remote shuffle pusher.
@@ -42,9 +50,9 @@ pub(crate) struct RssPartitionWriter {
num_partitions: usize,
max_frame_size: usize,
compression_context: CompressionContext,
- frame: Vec<u8>,
next_partition_to_finish: usize,
finished: bool,
+ failed: bool,
}
impl RssPartitionWriter {
@@ -83,14 +91,14 @@ impl RssPartitionWriter {
num_partitions,
max_frame_size,
compression_context: CompressionContext::default(),
- frame: Vec::new(),
next_partition_to_finish: 0,
finished: false,
+ failed: false,
})
}
fn validate_writable_partition(&self, partition_id: usize) -> Result<i32> {
- if self.finished {
+ if self.finished || self.failed {
return Err(DataFusionError::Execution(
"Remote shuffle writer has already finished".to_string(),
));
@@ -125,35 +133,901 @@ impl RssPartitionWriter {
where
I: Iterator<Item = Result<RecordBatch>>,
{
- for batch in batches.by_ref() {
- let batch = batch?;
- self.frame.clear();
+ let result = (|| {
+ for batch in batches.by_ref() {
+ self.push_batch_within_limit(partition_id, &batch?, metrics)?;
+ }
+ Ok(())
+ })();
+ if result.is_err() {
+ // Earlier frames may have been accepted remotely; a partial map cannot be resumed or
+ // committed after its input, encoding, reservation, or callback fails.
+ self.failed = true;
+ }
+ result
+ }
- let encoded_size = self.block_writer.write_batch(
- &batch,
- &mut Cursor::new(&mut self.frame),
+ fn push_batch_within_limit(
+ &mut self,
+ partition_id: i32,
+ batch: &RecordBatch,
+ metrics: &ShufflePartitionerMetrics,
+ ) -> Result<()> {
+ if batch.num_rows() == 0 {
+ return Ok(());
+ }
+
+ let (metadata_scratch, planning_scratch) = Self::estimated_ipc_metadata_scratch(batch);
+ let codec_workspace = self.block_writer.rss_codec_workspace()?;
+ let reservation_limit = self.pusher.max_reservation_size();
+ if metadata_scratch
+ .saturating_add(codec_workspace)
+ .saturating_add(60)
+ > reservation_limit
+ {
+ // Neither schema descriptors nor codec workspace gets smaller when rows are split.
+ return Err(DataFusionError::Execution(format!(
+ "Remote shuffle frame schema and encoding workspace exceed the byte admission budget of {reservation_limit} bytes"
+ )));
+ }
+
+ // ArrayData and zero-copy nested slices still allocate descriptors. Admit that planning
+ // scratch before computing precise live-row sizes; the bound above only borrows the
+ // schema/arrays. Drop all planning temporaries and release this short-lived reservation
+ // before acquiring the full encoding bound, so there is no partial-reservation growth.
+ self.pusher.reserve_partition_data(planning_scratch)?;
+ let estimates = (|| {
+ Ok::<_, DataFusionError>((
+ Self::estimated_pre_compaction_ipc_data_size(batch)?,
+ Self::estimated_compaction_scratch(batch)?,
+ Self::estimated_minimum_compacted_ipc_data_size(batch)?,
+ ))
+ })();
+ let planning_released = self.pusher.release_partition_data_reservation();
+ let (original_size, compaction_scratch, minimum_size) = estimates?;
+ planning_released?;
+ if minimum_size > self.max_frame_size && batch.num_rows() > 1 {
+ return self.push_split_batch(partition_id, batch, metrics);
+ }
+
+ // Arrow IPC first materializes uncompressed message bodies even for a compressed frame.
+ // Those Vecs, schema/FlatBuffers scratch, compacted arrays, codec workspace, and the output
+ // buffer overlap. They must be admitted together, not only when the encoded output grows.
+ let ipc_scratch = original_size
+ .checked_mul(4)
+ .and_then(|bytes| bytes.checked_add(metadata_scratch))
+ .and_then(|bytes| bytes.checked_add(compaction_scratch))
+ .and_then(|bytes| bytes.checked_add(codec_workspace))
+ .ok_or_else(|| {
+ DataFusionError::Execution(
+ "Remote shuffle encoding workspace exceeds the native integer limit"
+ .to_string(),
+ )
+ })?;
+ let admitted_frame_limit =
+ (reservation_limit.saturating_sub(ipc_scratch) / 3).min(self.max_frame_size);
+ if admitted_frame_limit < 20 {
+ if batch.num_rows() > 1 {
+ return self.push_split_batch(partition_id, batch, metrics);
+ }
+ return Err(DataFusionError::Execution(format!(
+ "Remote shuffle frame for a single row and encoding workspace exceed the byte admission budget of {reservation_limit} bytes"
+ )));
+ }
+
+ // This controls only the first output capacity; metadata scratch is already charged above.
+ // A failed estimate drops every temporary allocation and releases its reservation before
+ // retrying, and no output buffer is allowed to grow beyond the admitted capacity.
+ const IPC_METADATA_RESERVATION_BYTES: usize = 4 * 1024;
+ let mut frame_bound = original_size
+ .saturating_add(IPC_METADATA_RESERVATION_BYTES)
+ .min(admitted_frame_limit);
+
+ loop {
+ // Native output capacity, its JNI byte array, and Celeborn's copied transport request
+ // overlap. Acquire all three plus the encoding workspace atomically, before either
+ // compaction or encoding, without ever clamping a calculated memory bound.
+ let reservation = frame_bound
+ .checked_mul(3)
+ .and_then(|bytes| bytes.checked_add(ipc_scratch))
+ .ok_or_else(|| {
+ DataFusionError::Execution(
+ "Remote shuffle frame-copy reservation exceeds the native integer limit"
+ .to_string(),
+ )
+ })?;
+ self.pusher.reserve_partition_data(reservation)?;
+
+ let compacted = match Self::compact_dictionary_columns(batch) {
+ Ok(compacted) => compacted,
+ Err(error) => {
+ self.pusher.release_partition_data_reservation()?;
+ return Err(error);
+ }
+ };
+ let compacted_batch = compacted.as_ref().unwrap_or(batch);
+ let encoded_estimate = match Self::estimated_ipc_data_size(compacted_batch) {
+ Ok(estimate) => estimate,
+ Err(error) => {
+ drop(compacted);
+ self.pusher.release_partition_data_reservation()?;
+ return Err(error);
+ }
+ };
+ if encoded_estimate > self.max_frame_size && compacted_batch.num_rows() > 1 {
+ drop(compacted);
+ self.pusher.release_partition_data_reservation()?;
+ return self.push_split_batch(partition_id, batch, metrics);
+ }
+
+ // Bound encoding by the admitted estimate, not only the configured maximum, so all
+ // three eventual frame copies fit the reservation that was acquired atomically.
+ let mut output = match BoundedBuffer::try_new(frame_bound) {
+ Ok(output) => output,
+ Err(error) => {
+ drop(compacted);
+ self.pusher.release_partition_data_reservation()?;
+ return Err(error);
+ }
+ };
+ if let Err(error) = self.block_writer.write_rss_batch(
+ compacted_batch,
+ &mut output,
&mut self.compression_context,
&metrics.encode_time,
- )?;
+ ) {
+ let exceeded = output.exceeded;
+ drop(output);
+ drop(compacted);
+ self.pusher.release_partition_data_reservation()?;
+ if !exceeded {
+ return Err(error);
+ }
+ if frame_bound < admitted_frame_limit {
+ frame_bound = frame_bound.saturating_mul(2).min(admitted_frame_limit);
+ continue;
+ }
+ if batch.num_rows() <= 1 {
+ return Err(DataFusionError::Execution(format!(
+ "Remote shuffle frame exceeds its configured maximum: a single row exceeds {} bytes",
+ admitted_frame_limit
+ )));
+ }
+ return self.push_split_batch(partition_id, batch, metrics);
+ }
+
+ // The Java callback must keep this reservation until both transport ownership and
+ // this native encoder's explicit acknowledgement have ended.
+ drop(compacted);
+ let frame = output.inner.get_ref();
+ if frame.is_empty() {
+ drop(output);
+ self.pusher.release_partition_data_reservation()?;
+ return Ok(());
+ }
+ let mut timer = metrics.write_time.timer();
+ let result = self.pusher.push_partition_data(partition_id, frame);
+ timer.stop();
+ // The JNI bridge has popped its local frame on return. Drop the native capacity
+ // before acknowledging completion even when the callback failed or completed inline.
+ drop(output);
+ let released = self.pusher.release_partition_data_reservation();
+ return result.and(released);
+ }
+ }
- if encoded_size == 0 {
- continue;
+ fn push_split_batch(
+ &mut self,
+ partition_id: i32,
+ batch: &RecordBatch,
+ metrics: &ShufflePartitionerMetrics,
+ ) -> Result<()> {
+ let midpoint = batch.num_rows() / 2;
+ self.push_batch_within_limit(partition_id, &batch.slice(0, midpoint), metrics)?;
+ self.push_batch_within_limit(
+ partition_id,
+ &batch.slice(midpoint, batch.num_rows() - midpoint),
+ metrics,
+ )
+ }
+
+ /// Bound the pinned Arrow IPC writer's schema, message, and array-descriptor allocations.
+ ///
+ /// Field/type/dictionary tables and metadata entries each fit in the 256-byte allowance;
+ /// record/dictionary messages add 16-byte field nodes and 16-byte buffer descriptors. Names,
+ /// metadata keys/values, and timestamp timezones are charged separately. Four copies cover
+ /// FlatBuffer/Vec geometric growth, realloc overlap, and the finished-message copy.
+ ///
+ /// ArrayData (136 bytes) and Buffer (24 bytes) descriptors are cloned while nested list/map
+ /// children are sliced for IPC. The depth-weighted allowance covers these simultaneously
+ /// retained trees, record-batch column vectors, and descriptor-Vec growth. This is independent
+ /// of the actual IPC body bytes and dictionary compaction scratch charged by the caller.
+ fn estimated_ipc_metadata_scratch(batch: &RecordBatch) -> (usize, usize) {
+ #[derive(Default)]
+ struct Metadata {
+ bytes: usize,
+ entries: usize,
+ nodes: usize,
+ dictionaries: usize,
+ depth: usize,
+ }
+
+ fn add_field(field: &Field, depth: usize, metadata: &mut Metadata) {
+ metadata.bytes = metadata.bytes.saturating_add(field.name().len());
+ metadata.entries = metadata.entries.saturating_add(field.metadata().len());
+ for (key, value) in field.metadata() {
+ metadata.bytes = metadata
+ .bytes
+ .saturating_add(key.len())
+ .saturating_add(value.len());
}
+ add_type(field.data_type(), depth, metadata);
+ }
- if encoded_size > self.max_frame_size {
- return Err(DataFusionError::Execution(format!(
- "Remote shuffle frame size {encoded_size} exceeds the configured maximum {}",
- self.max_frame_size
- )));
+ fn add_type(data_type: &DataType, depth: usize, metadata: &mut Metadata) {
+ metadata.nodes = metadata.nodes.saturating_add(1);
+ metadata.depth = metadata.depth.max(depth);
+ match data_type {
+ DataType::List(field)
+ | DataType::LargeList(field)
+ | DataType::ListView(field)
+ | DataType::LargeListView(field)
+ | DataType::FixedSizeList(field, _)
+ | DataType::Map(field, _) => add_field(field, depth + 1, metadata),
+ DataType::Struct(fields) => {
+ for field in fields {
+ add_field(field, depth + 1, metadata);
+ }
+ }
+ DataType::Union(fields, _) => {
+ for (_, field) in fields.iter() {
+ add_field(field, depth + 1, metadata);
+ }
+ }
+ DataType::RunEndEncoded(run_ends, values) => {
+ add_field(run_ends, depth + 1, metadata);
+ add_field(values, depth + 1, metadata);
+ }
+ DataType::Dictionary(_, values) => {
+ metadata.dictionaries = metadata.dictionaries.saturating_add(1);
+ add_type(values, depth + 1, metadata);
+ }
+ DataType::Timestamp(_, Some(timezone)) => {
+ metadata.bytes = metadata.bytes.saturating_add(timezone.len());
+ }
+ _ => {}
}
+ }
- let mut write_timer = metrics.write_time.timer();
- let result = self.pusher.push_partition_data(partition_id, &self.frame);
- write_timer.stop();
- result?;
+ // All ordinary Arrow layouts have at most three buffers including validity. View arrays
+ // add a variable number of backing buffers. Count those by borrowing concrete arrays,
+ // without ArrayData conversion or child slicing before the planning reservation exists.
+ fn view_backing_buffers(array: &dyn Array) -> usize {
+ if let Some(dictionary) = array.as_any_dictionary_opt() {
+ return view_backing_buffers(dictionary.values().as_ref());
+ }
+ match array.data_type() {
+ DataType::Utf8View => array
+ .as_any()
+ .downcast_ref::<StringViewArray>()
+ .unwrap()
+ .data_buffers()
+ .len(),
+ DataType::BinaryView => array
+ .as_any()
+ .downcast_ref::<BinaryViewArray>()
+ .unwrap()
+ .data_buffers()
+ .len(),
+ DataType::List(_) => view_backing_buffers(
+ array
+ .as_any()
+ .downcast_ref::<ListArray>()
+ .unwrap()
+ .values()
+ .as_ref(),
+ ),
+ DataType::LargeList(_) => view_backing_buffers(
+ array
+ .as_any()
+ .downcast_ref::<LargeListArray>()
+ .unwrap()
+ .values()
+ .as_ref(),
+ ),
+ DataType::ListView(_) => view_backing_buffers(
+ array
+ .as_any()
+ .downcast_ref::<ListViewArray>()
+ .unwrap()
+ .values()
+ .as_ref(),
+ ),
+ DataType::LargeListView(_) => view_backing_buffers(
+ array
+ .as_any()
+ .downcast_ref::<LargeListViewArray>()
+ .unwrap()
+ .values()
+ .as_ref(),
+ ),
+ DataType::FixedSizeList(_, _) => view_backing_buffers(
+ array
+ .as_any()
+ .downcast_ref::<FixedSizeListArray>()
+ .unwrap()
+ .values()
+ .as_ref(),
+ ),
+ DataType::Map(_, _) => view_backing_buffers(
+ array.as_any().downcast_ref::<MapArray>().unwrap().entries(),
+ ),
+ DataType::Struct(_) => array
+ .as_any()
+ .downcast_ref::<StructArray>()
+ .unwrap()
+ .columns()
+ .iter()
+ .fold(0usize, |total, column| {
+ total.saturating_add(view_backing_buffers(column.as_ref()))
+ }),
+ DataType::Union(fields, _) => {
+ let union = array.as_any().downcast_ref::<UnionArray>().unwrap();
+ fields.iter().fold(0usize, |total, (type_id, _)| {
+ total.saturating_add(view_backing_buffers(union.child(type_id).as_ref()))
+ })
+ }
+ DataType::RunEndEncoded(run_ends, _) => match run_ends.data_type() {
+ DataType::Int16 => view_backing_buffers(
+ array
+ .as_any()
+ .downcast_ref::<RunArray<Int16Type>>()
+ .unwrap()
+ .values()
+ .as_ref(),
+ ),
+ DataType::Int32 => view_backing_buffers(
+ array
+ .as_any()
+ .downcast_ref::<RunArray<Int32Type>>()
+ .unwrap()
+ .values()
+ .as_ref(),
+ ),
+ DataType::Int64 => view_backing_buffers(
+ array
+ .as_any()
+ .downcast_ref::<RunArray<Int64Type>>()
+ .unwrap()
+ .values()
+ .as_ref(),
+ ),
+ _ => unreachable!("Arrow run ends must use a supported signed integer"),
+ },
+ _ => 0,
+ }
}
- Ok(())
+ let schema = batch.schema();
+ let mut metadata = Metadata {
+ entries: schema.metadata().len(),
+ ..Metadata::default()
+ };
+ for (key, value) in schema.metadata() {
+ metadata.bytes = metadata
+ .bytes
+ .saturating_add(key.len())
+ .saturating_add(value.len());
+ }
+ for field in schema.fields() {
+ add_field(field, 1, &mut metadata);
+ }
+ let buffers = batch
+ .columns()
+ .iter()
+ .fold(metadata.nodes.saturating_mul(3), |total, column| {
+ total.saturating_add(view_backing_buffers(column.as_ref()))
+ });
+ let messages = metadata
+ .nodes
+ .saturating_add(metadata.dictionaries)
+ .saturating_add(metadata.entries)
+ .saturating_add(1);
+ let serialized = 1_024usize
+ .saturating_add(metadata.bytes)
+ .saturating_add(messages.saturating_mul(256))
+ .saturating_add(buffers.saturating_mul(32));
+ let descriptors = metadata
+ .nodes
+ .saturating_mul(512)
+ .saturating_add(buffers.saturating_mul(128))
+ .saturating_mul(metadata.depth.saturating_add(1));
+ let planning = descriptors.saturating_add(4 * 1024);
+ (
+ serialized.saturating_mul(4).saturating_add(planning),
+ planning,
+ )
+ }
+
+ fn estimated_ipc_data_size(batch: &RecordBatch) -> Result<usize> {
+ batch.columns().iter().try_fold(0usize, |total, column| {
+ let data = column.to_data();
+ let logical = data.get_slice_memory_size()?;
+ let additional = Self::additional_ipc_data_size(&data);
+ // IPC aligns every buffer to 64 bytes and may synthesize a validity buffer even
+ // when an array has no nulls. Include that padding before charging one frame.
+ let padding = Self::ipc_buffer_count(&data).saturating_mul(64);
+ Ok(total
+ .saturating_add(logical)
+ .saturating_add(additional)
+ .saturating_add(padding))
+ })
+ }
+
+ /// Estimate logical rows without charging unrelated list/map backing entries.
+ ///
+ /// `ArrayData::get_slice_memory_size` recurses through an entire list/map child even when the
+ /// parent offsets select only one entry. Their zero-copy child slices are sufficient here;
+ /// dictionary values remain fully charged because dense garbage collection traverses them.
+ fn estimated_pre_compaction_ipc_data_size(batch: &RecordBatch) -> Result<usize> {
+ batch.columns().iter().try_fold(0usize, |total, column| {
+ Ok(
+ total.saturating_add(Self::estimated_live_array_ipc_data_size(
+ column.as_ref(),
+ true,
+ )?),
+ )
+ })
+ }
+
+ /// Lower-bound the IPC buffers that must survive any dictionary-value compaction.
+ fn estimated_minimum_compacted_ipc_data_size(batch: &RecordBatch) -> Result<usize> {
+ batch.columns().iter().try_fold(0usize, |total, column| {
+ Ok(
+ total.saturating_add(Self::estimated_live_array_ipc_data_size(
+ column.as_ref(),
+ false,
+ )?),
+ )
+ })
+ }
+
+ fn estimated_live_array_ipc_data_size(
+ array: &dyn Array,
+ include_dictionary_values: bool,
+ ) -> Result<usize> {
+ let data = array.to_data();
+ let child_logical = data.child_data().iter().try_fold(0usize, |total, child| {
+ Ok::<_, DataFusionError>(total.saturating_add(child.get_slice_memory_size()?))
+ })?;
+ let logical = data.get_slice_memory_size()?.saturating_sub(child_logical);
+ let child_additional = data.child_data().iter().fold(0usize, |total, child| {
+ total.saturating_add(Self::additional_ipc_data_size(child))
+ });
+ let additional = Self::additional_ipc_data_size(&data).saturating_sub(child_additional);
+ let padding = data.buffers().len().saturating_add(1).saturating_mul(64);
+ let own = logical.saturating_add(additional).saturating_add(padding);
+
+ let children =
+ if let Some(dictionary) = array.as_any_dictionary_opt() {
+ if include_dictionary_values {
+ Self::estimated_live_array_ipc_data_size(
+ dictionary.values().as_ref(),
+ include_dictionary_values,
+ )?
+ } else {
+ 0
+ }
+ } else {
+ match array.data_type() {
+ DataType::List(_) => {
+ let list = array.as_any().downcast_ref::<ListArray>().unwrap();
+ let offsets = list.value_offsets();
+ let start = offsets[0] as usize;
+ let end = offsets[offsets.len() - 1] as usize;
+ let live_values = list.values().slice(start, end - start);
+ Self::estimated_live_array_ipc_data_size(
+ live_values.as_ref(),
+ include_dictionary_values,
+ )?
+ }
+ DataType::LargeList(_) => {
+ let list = array.as_any().downcast_ref::<LargeListArray>().unwrap();
+ let offsets = list.value_offsets();
+ let start = offsets[0] as usize;
+ let end = offsets[offsets.len() - 1] as usize;
+ let live_values = list.values().slice(start, end - start);
+ Self::estimated_live_array_ipc_data_size(
+ live_values.as_ref(),
+ include_dictionary_values,
+ )?
+ }
+ DataType::FixedSizeList(_, _) => {
+ let list = array.as_any().downcast_ref::<FixedSizeListArray>().unwrap();
+ Self::estimated_live_array_ipc_data_size(
+ list.values().as_ref(),
+ include_dictionary_values,
+ )?
+ }
+ DataType::Struct(_) => {
+ let structure = array.as_any().downcast_ref::<StructArray>().unwrap();
+ structure
+ .columns()
+ .iter()
+ .try_fold(0usize, |total, column| {
+ Ok::<_, DataFusionError>(total.saturating_add(
+ Self::estimated_live_array_ipc_data_size(
+ column.as_ref(),
+ include_dictionary_values,
+ )?,
+ ))
+ })?
+ }
+ DataType::Map(_, _) => {
+ let map = array.as_any().downcast_ref::<MapArray>().unwrap();
+ let offsets = map.value_offsets();
+ let start = offsets[0] as usize;
+ let end = offsets[offsets.len() - 1] as usize;
+ let live_entries = map.entries().slice(start, end - start);
+ Self::estimated_live_array_ipc_data_size(
+ &live_entries,
+ include_dictionary_values,
+ )?
+ }
+ _ => data.child_data().iter().try_fold(0usize, |total, child| {
+ let logical = child.get_slice_memory_size()?;
+ let additional = Self::additional_ipc_data_size(child);
+ let padding = Self::ipc_buffer_count(child).saturating_mul(64);
+ Ok::<_, DataFusionError>(total.saturating_add(
+ logical.saturating_add(additional).saturating_add(padding),
+ ))
+ })?,
+ }
+ };
+ Ok(own.saturating_add(children))
+ }
+
+ fn additional_ipc_data_size(data: &arrow::array::ArrayData) -> usize {
+ // Arrow IPC writes a validity bitmap even when the original array has no null buffer.
+ let validity = if data.nulls().is_none() {
+ data.len().div_ceil(8)
+ } else {
+ 0
+ };
+ // ArrayData's slice estimate charges one offset per row, while IPC writes the additional
+ // terminating offset. View arrays instead serialize every shared backing data buffer.
+ let own = if matches!(data.data_type(), DataType::BinaryView | DataType::Utf8View) {
+ data.buffers()
+ .iter()
+ .skip(1)
+ .fold(0usize, |total, buffer| total.saturating_add(buffer.len()))
+ } else {
+ match data.data_type() {
+ DataType::Binary | DataType::Utf8 => {
+ let temporary = if data.buffer::<i32>(0)[0] != 0 {
+ data.len().saturating_add(1).saturating_mul(4)
+ } else {
+ 0
+ };
+ 4usize.saturating_add(temporary)
+ }
+ DataType::LargeBinary | DataType::LargeUtf8 => {
+ let temporary = if data.buffer::<i64>(0)[0] != 0 {
+ data.len().saturating_add(1).saturating_mul(8)
+ } else {
+ 0
+ };
+ 8usize.saturating_add(temporary)
+ }
+ DataType::List(_) | DataType::Map(_, _) => 4,
+ DataType::LargeList(_) => 8,
+ _ => 0,
+ }
+ };
+ data.child_data()
+ .iter()
+ .fold(own.saturating_add(validity), |total, child| {
+ total.saturating_add(Self::additional_ipc_data_size(child))
+ })
+ }
+
+ fn ipc_buffer_count(data: &arrow::array::ArrayData) -> usize {
+ data.child_data()
+ .iter()
+ .fold(data.buffers().len().saturating_add(1), |total, child| {
+ total.saturating_add(Self::ipc_buffer_count(child))
+ })
+ }
+
+ /// Conservatively charge dense dictionary garbage collection and nested offset rebasing.
+ ///
+ /// This walks the original arrays without compacting, building occupancy masks, or allocating
+ /// per-row scratch. In particular, dense remaps scale with the *full* dictionary cardinality,
+ /// even when only one value is referenced by the current batch slice.
+ fn estimated_compaction_scratch(batch: &RecordBatch) -> Result<usize> {
+ batch.columns().iter().try_fold(0usize, |total, column| {
+ Ok(total.saturating_add(Self::estimated_array_compaction_scratch(column.as_ref())?))
+ })
+ }
+
+ fn estimated_array_compaction_scratch(array: &dyn Array) -> Result<usize> {
+ if let Some(dictionary) = array.as_any_dictionary_opt() {
+ let key_width = match dictionary.keys().data_type() {
+ DataType::Int8 | DataType::UInt8 => 1,
+ DataType::Int16 | DataType::UInt16 => 2,
+ DataType::Int32 | DataType::UInt32 => 4,
+ DataType::Int64 | DataType::UInt64 => 8,
+ _ => unreachable!("Arrow dictionary keys must have an integer type"),
+ };
+ let value_count = dictionary.values().len();
+ let key_count = dictionary.keys().len();
+ let occupancy = value_count.div_ceil(8).saturating_add(64);
+ let dense_remap = value_count.saturating_mul(key_width);
+ let copied_keys = key_count
+ .saturating_mul(key_width)
+ .saturating_add(key_count.div_ceil(8))
+ .saturating_add(64);
+ let filter_indices = key_count
+ .min(value_count)
+ .saturating_mul(std::mem::size_of::<(usize, usize)>());
+ let values = dictionary.values().to_data();
+ let copied_values = values
+ .get_slice_memory_size()?
+ .saturating_add(Self::additional_ipc_data_size(&values))
+ .saturating_add(Self::ipc_buffer_count(&values).saturating_mul(64));
+ let nested = Self::estimated_array_compaction_scratch(dictionary.values().as_ref())?;
+ return Ok(occupancy
+ .saturating_add(dense_remap)
+ .saturating_add(copied_keys)
+ .saturating_add(filter_indices)
+ .saturating_add(copied_values)
+ .saturating_add(nested));
+ }
+
+ match array.data_type() {
+ DataType::List(_) => {
+ let list = array.as_any().downcast_ref::<ListArray>().unwrap();
+ let offsets = list.value_offsets();
+ let start = offsets[0] as usize;
+ let end = offsets[offsets.len() - 1] as usize;
+ let normalized = start != 0 || end != list.values().len();
+ let rebased = if normalized {
+ list.len().saturating_add(1).saturating_mul(4)
+ } else {
+ 0
+ };
+ let live_values = list.values().slice(start, end - start);
+ Ok(
+ rebased.saturating_add(Self::estimated_array_compaction_scratch(
+ live_values.as_ref(),
+ )?),
+ )
+ }
+ DataType::LargeList(_) => {
+ let list = array.as_any().downcast_ref::<LargeListArray>().unwrap();
+ let offsets = list.value_offsets();
+ let start = offsets[0] as usize;
+ let end = offsets[offsets.len() - 1] as usize;
+ let normalized = start != 0 || end != list.values().len();
+ let rebased = if normalized {
+ list.len().saturating_add(1).saturating_mul(8)
+ } else {
+ 0
+ };
+ let live_values = list.values().slice(start, end - start);
+ Ok(
+ rebased.saturating_add(Self::estimated_array_compaction_scratch(
+ live_values.as_ref(),
+ )?),
+ )
+ }
+ DataType::FixedSizeList(_, _) => {
+ let list = array.as_any().downcast_ref::<FixedSizeListArray>().unwrap();
+ Self::estimated_array_compaction_scratch(list.values().as_ref())
+ }
+ DataType::Struct(_) => {
+ let structure = array.as_any().downcast_ref::<StructArray>().unwrap();
+ structure
+ .columns()
+ .iter()
+ .try_fold(0usize, |total, column| {
+ Ok(
+ total.saturating_add(Self::estimated_array_compaction_scratch(
+ column.as_ref(),
+ )?),
+ )
+ })
+ }
+ DataType::Map(_, _) => {
+ let map = array.as_any().downcast_ref::<MapArray>().unwrap();
+ let offsets = map.value_offsets();
+ let start = offsets[0] as usize;
+ let end = offsets[offsets.len() - 1] as usize;
+ let normalized = start != 0 || end != map.entries().len();
+ let rebased = if normalized {
+ map.len().saturating_add(1).saturating_mul(4)
+ } else {
+ 0
+ };
+ let live_entries = map.entries().slice(start, end - start);
+ Ok(
+ rebased
+ .saturating_add(Self::estimated_array_compaction_scratch(&live_entries)?),
+ )
+ }
+ _ => Ok(0),
+ }
+ }
+
+ fn compact_dictionary_columns(batch: &RecordBatch) -> Result<Option<RecordBatch>> {
+ let mut changed = false;
+ let columns = batch
+ .columns()
+ .iter()
+ .map(|column| -> Result<ArrayRef> {
+ let (compacted, column_changed) = Self::compact_array(column)?;
+ changed |= column_changed;
+ Ok(compacted)
+ })
+ .collect::<Result<Vec<_>>>()?;
+
+ if changed {
+ Ok(Some(RecordBatch::try_new(batch.schema(), columns)?))
+ } else {
+ Ok(None)
+ }
+ }
+
+ fn compact_array(array: &ArrayRef) -> Result<(ArrayRef, bool)> {
+ if let Some(dictionary) = array.as_any_dictionary_opt() {
+ let compacted = garbage_collect_any_dictionary(dictionary)?;
+ let compacted_dictionary = compacted.as_any_dictionary();
+ let dictionary_changed =
+ compacted_dictionary.values().len() != dictionary.values().len();
+ let (values, values_changed) = Self::compact_array(compacted_dictionary.values())?;
+ if values_changed {
+ return Ok((compacted_dictionary.with_values(values), true));
+ }
+ return Ok(if dictionary_changed {
+ (compacted, true)
+ } else {
+ (Arc::clone(array), false)
+ });
+ }
+
+ match array.data_type() {
+ DataType::List(field) => {
+ let list = array.as_any().downcast_ref::<ListArray>().unwrap();
+ let offsets = list.value_offsets();
+ let start = offsets[0] as usize;
+ let end = offsets[offsets.len() - 1] as usize;
+ let live_values = list.values().slice(start, end - start);
+ let (values, child_changed) = Self::compact_array(&live_values)?;
+ let normalized = start != 0 || end != list.values().len();
+ if !child_changed && !normalized {
+ return Ok((Arc::clone(array), false));
+ }
+ let rebased = if start == 0 {
+ list.offsets().clone()
+ } else {
+ OffsetBuffer::new(
+ offsets
+ .iter()
+ .map(|offset| offset - offsets[0])
+ .collect::<Vec<_>>()
+ .into(),
+ )
+ };
+ let rebuilt =
+ ListArray::try_new(Arc::clone(field), rebased, values, list.nulls().cloned())?;
+ Ok((Arc::new(rebuilt), true))
+ }
+ DataType::LargeList(field) => {
+ let list = array.as_any().downcast_ref::<LargeListArray>().unwrap();
+ let offsets = list.value_offsets();
+ let start = offsets[0] as usize;
+ let end = offsets[offsets.len() - 1] as usize;
+ let live_values = list.values().slice(start, end - start);
+ let (values, child_changed) = Self::compact_array(&live_values)?;
+ let normalized = start != 0 || end != list.values().len();
+ if !child_changed && !normalized {
+ return Ok((Arc::clone(array), false));
+ }
+ let rebased = if start == 0 {
+ list.offsets().clone()
+ } else {
+ OffsetBuffer::new(
+ offsets
+ .iter()
+ .map(|offset| offset - offsets[0])
+ .collect::<Vec<_>>()
+ .into(),
+ )
+ };
+ let rebuilt = LargeListArray::try_new(
+ Arc::clone(field),
+ rebased,
+ values,
+ list.nulls().cloned(),
+ )?;
+ Ok((Arc::new(rebuilt), true))
+ }
+ DataType::FixedSizeList(field, size) => {
+ let list = array.as_any().downcast_ref::<FixedSizeListArray>().unwrap();
+ let (values, changed) = Self::compact_array(list.values())?;
+ if !changed {
+ return Ok((Arc::clone(array), false));
+ }
+ let rebuilt = FixedSizeListArray::try_new_with_length(
+ Arc::clone(field),
+ *size,
+ values,
+ list.nulls().cloned(),
+ list.len(),
+ )?;
+ Ok((Arc::new(rebuilt), true))
+ }
+ DataType::Struct(fields) => {
+ let structure = array.as_any().downcast_ref::<StructArray>().unwrap();
+ let mut changed = false;
+ let columns = structure
+ .columns()
+ .iter()
+ .map(|column| {
+ let (compacted, child_changed) = Self::compact_array(column)?;
+ changed |= child_changed;
+ Ok(compacted)
+ })
+ .collect::<Result<Vec<_>>>()?;
+ if !changed {
+ return Ok((Arc::clone(array), false));
+ }
+ let rebuilt = StructArray::try_new_with_length(
+ fields.clone(),
+ columns,
+ structure.nulls().cloned(),
+ structure.len(),
+ )?;
+ Ok((Arc::new(rebuilt), true))
+ }
+ DataType::Map(field, ordered) => {
+ let map = array.as_any().downcast_ref::<MapArray>().unwrap();
+ let offsets = map.value_offsets();
+ let start = offsets[0] as usize;
+ let end = offsets[offsets.len() - 1] as usize;
+ let live_entries: ArrayRef = Arc::new(map.entries().slice(start, end - start));
+ let (entries, child_changed) = Self::compact_array(&live_entries)?;
+ let normalized = start != 0 || end != map.entries().len();
+ if !child_changed && !normalized {
+ return Ok((Arc::clone(array), false));
+ }
+ let rebased = if start == 0 {
+ map.offsets().clone()
+ } else {
+ OffsetBuffer::new(
+ offsets
+ .iter()
+ .map(|offset| offset - offsets[0])
+ .collect::<Vec<_>>()
+ .into(),
+ )
+ };
+ let entries = entries
+ .as_any()
+ .downcast_ref::<StructArray>()
+ .unwrap()
+ .clone();
+ let rebuilt = MapArray::try_new(
+ Arc::clone(field),
+ rebased,
+ entries,
+ map.nulls().cloned(),
+ *ordered,
+ )?;
+ Ok((Arc::new(rebuilt), true))
+ }
+ _ => Ok((Arc::clone(array), false)),
+ }
}
}
@@ -194,7 +1068,7 @@ impl PartitionWriter for RssPartitionWriter {
}
fn finish_all(&mut self, _metrics: &ShufflePartitionerMetrics) -> Result<()> {
- if self.finished {
+ if self.finished || self.failed {
return Err(DataFusionError::Execution(
"Remote shuffle writer has already finished".to_string(),
));
@@ -212,3 +1086,106 @@ impl PartitionWriter for RssPartitionWriter {
Ok(())
}
}
+
+/// A seekable encoder destination that rejects an oversized frame before growing its buffer.
+struct BoundedBuffer {
+ inner: Cursor<Vec<u8>>,
+ limit: usize,
+ exceeded: bool,
+}
+
+impl BoundedBuffer {
+ fn try_new(limit: usize) -> Result<Self> {
+ let mut bytes = Vec::new();
+ bytes.try_reserve_exact(limit).map_err(|error| {
+ DataFusionError::ResourcesExhausted(format!(
+ "Cannot allocate the admitted remote shuffle frame buffer: {error}"
+ ))
+ })?;
+ // Cursor<Vec<u8>> normally doubles capacity when its length grows. Preallocate the full
+ // admitted capacity so every permitted write is allocation-free instead.
+ if bytes.capacity() != limit {
+ return Err(DataFusionError::ResourcesExhausted(
+ "Remote shuffle allocator exceeded the admitted frame capacity".to_string(),
+ ));
+ }
+ Ok(Self {
+ inner: Cursor::new(bytes),
+ limit,
+ exceeded: false,
+ })
+ }
+
+ fn limit_error() -> io::Error {
+ io::Error::new(
+ io::ErrorKind::InvalidData,
+ "RSS frame exceeds its byte limit",
+ )
+ }
+}
+
+impl Write for BoundedBuffer {
+ fn write(&mut self, bytes: &[u8]) -> io::Result<usize> {
+ let end = self.inner.position().checked_add(bytes.len() as u64);
+ if end.is_none_or(|end| end > self.limit as u64) {
+ self.exceeded = true;
+ return Err(Self::limit_error());
+ }
+ self.inner.write(bytes)
+ }
+
+ fn flush(&mut self) -> io::Result<()> {
+ self.inner.flush()
+ }
+}
+
+impl Seek for BoundedBuffer {
+ fn seek(&mut self, position: SeekFrom) -> io::Result<u64> {
+ let previous = self.inner.position();
+ let next = self.inner.seek(position)?;
+ if next > self.limit as u64 {
+ self.inner.set_position(previous);
+ self.exceeded = true;
+ return Err(Self::limit_error());
+ }
+ Ok(next)
+ }
+}
+
+#[cfg(test)]
+mod buffer_tests {
+ use super::*;
+ use arrow::array::Int32Array;
+ use arrow::datatypes::Schema;
+ use datafusion::physical_plan::metrics::Time;
+
+ #[test]
+ fn large_uncompressed_frame_never_grows_past_its_admitted_capacity() {
+ let batch = RecordBatch::try_new(
+ Arc::new(Schema::new(vec![Field::new(
+ "value",
+ DataType::Int32,
+ false,
+ )])),
+ vec![Arc::new(Int32Array::from_iter_values(0..16_384))],
+ )
+ .unwrap();
+ let block_writer =
+ ShuffleBlockWriter::try_new(batch.schema().as_ref(), crate::CompressionCodec::None)
+ .unwrap();
+ let mut output = BoundedBuffer::try_new(67_996).unwrap();
+ block_writer
+ .write_rss_batch(
+ &batch,
+ &mut output,
+ &mut CompressionContext::default(),
+ &Time::default(),
+ )
+ .unwrap();
+ assert_eq!(output.inner.get_ref().len(), 67_996);
+ assert_eq!(output.inner.get_ref().capacity(), 67_996);
+ assert!(output.write(&[1]).is_err());
+ assert!(output.exceeded);
+ assert_eq!(output.inner.get_ref().capacity(), 67_996);
+ }
+}
diff --git a/native/shuffle/src/writers/shuffle_block_writer.rs b/native/shuffle/src/writers/shuffle_block_writer.rs
index 65e7bf4570..c398700342 100644
--- a/native/shuffle/src/writers/shuffle_block_writer.rs
+++ b/native/shuffle/src/writers/shuffle_block_writer.rs
@@ -44,9 +44,9 @@ pub enum CompressionCodec {
Snappy,
}
-/// How a writer encodes the Arrow IPC schema at the start of each block. The two arms are mutually
-/// exclusive by schema shape, decided once in [`ShuffleBlockWriter::try_new`], so the retained
-/// state never carries both a pre-encoded message and a schema at the same time.
+/// How a writer encodes the Arrow IPC schema at the start of each block. Local writers cache
+/// dictionary-free schemas; RSS serializes under admission. The retained state never carries
+/// both a pre-encoded message and a schema at the same time.
///
/// Both arms hold their payload behind an `Arc` because a `ShuffleBlockWriter` is cloned once per
/// output partition (see `LocalPartitionWriter`), and a shuffle can request millions of partitions.
@@ -57,8 +57,8 @@ enum SchemaEncoding {
/// Dictionary-free schema: the IPC schema message, pre-encoded once, is written verbatim at the
/// start of every block instead of being re-serialized per block.
Precoded(Arc<[u8]>),
- /// Schema containing dictionary types: the schema and record batch must share a dictionary
- /// tracker, so each block is encoded with `StreamWriter`, which re-serializes the schema.
+ /// Dictionary schemas and all RSS schemas: each block is encoded with `StreamWriter`, which
+ /// re-serializes the schema while sharing its dictionary tracker with the record batch.
Fallback(SchemaRef),
}
@@ -66,8 +66,8 @@ enum SchemaEncoding {
///
/// Each block is a self-contained Arrow IPC stream (schema message, dictionary messages, record
/// batch message, end-of-stream marker). For the common case of a schema with no dictionary types,
-/// the schema flatbuffer is encoded once in [`Self::try_new`] and written verbatim at the start of
-/// every block, rather than being re-serialized per block as `StreamWriter::try_new` would do.
+/// the schema flatbuffer for local writers is encoded once in [`Self::try_new`] and written verbatim
+/// at the start of every block, rather than re-serialized per block by `StreamWriter::try_new`.
/// Schemas that contain dictionary types fall back to `StreamWriter`, whose dictionary-id
/// bookkeeping ties schema and batch encoding together.
#[derive(Clone)]
@@ -83,7 +83,56 @@ pub struct ShuffleBlockWriter {
}
impl ShuffleBlockWriter {
+ /// Working memory for the RSS-only codec settings, excluding IPC and destination buffers.
+ ///
+ /// LZ4 uses a fixed 64 KiB input block, its compression-bound output, and a 16 KiB hash table.
+ /// Snappy uses a 64 KiB input block, a 76,490-byte output block, and its hash table. 256 KiB
+ /// conservatively covers either encoder. Zstd's streaming estimate includes the C-allocated
+ /// context, window, and input/output buffers; its Rust writer adds a fixed 32 KiB output Vec.
+ pub(crate) fn rss_codec_workspace(&self) -> Result<usize> {
+ match self.codec {
+ CompressionCodec::None => Ok(0),
+ CompressionCodec::Lz4Frame | CompressionCodec::Snappy => Ok(256 * 1024),
+ CompressionCodec::Zstd(level) => {
+ // The C estimator loops through the requested levels. Bound arbitrary user
+ // configuration before calling it (the encoder itself clamps to this range).
+ let levels = zstd::compression_level_range();
+ let level = level.clamp(*levels.start(), *levels.end());
+ // SAFETY: this pure estimator accepts an integer compression level and does not
+ // retain pointers. Our streaming encoder uses no dictionary or worker threads.
+ let estimate =
+ unsafe { zstd::zstd_safe::zstd_sys::ZSTD_estimateCStreamSize(level) };
+ // SAFETY: ZSTD_isError accepts every size_t returned by the estimator.
+ if unsafe { zstd::zstd_safe::zstd_sys::ZSTD_isError(estimate) } != 0 {
+ return Err(DataFusionError::Execution(
+ "Cannot estimate remote shuffle Zstd workspace".to_string(),
+ ));
+ }
+ estimate.checked_add(64 * 1024).ok_or_else(|| {
+ DataFusionError::Execution(
+ "Remote shuffle Zstd workspace exceeds the native integer limit"
+ .to_string(),
+ )
+ })
+ }
+ }
+ }
+
pub fn try_new(schema: &Schema, codec: CompressionCodec) -> Result<Self> {
+ Self::try_new_inner(schema, codec, None)
+ }
+
+ /// RSS cannot pre-encode a potentially large schema before its per-frame admission. Retain
+ /// the input's shared schema and serialize it inside the admitted encoding invocation instead.
+ pub(crate) fn try_new_rss(schema: SchemaRef, codec: CompressionCodec) -> Result<Self> {
+ Self::try_new_inner(schema.as_ref(), codec, Some(Arc::clone(&schema)))
+ }
+
+ fn try_new_inner(
+ schema: &Schema,
+ codec: CompressionCodec,
+ rss_schema: Option<SchemaRef>,
+ ) -> Result<Self> {
// Header layout: 8-byte block length placeholder + 8-byte field count (usize) + 4-byte
// codec tag = 20 bytes.
let mut header_bytes = Vec::with_capacity(20);
@@ -109,16 +158,16 @@ impl ShuffleBlockWriter {
// marker for V5. Alignment 64 and non-legacy framing match the arrow defaults.
let write_options = IpcWriteOptions::try_new(64, false, MetadataVersion::V5)?;
- // `flattened_fields` walks the full nested field tree, so this catches dictionary types
- // nested under any container arrow itself recurses into when emitting dictionaries.
- let has_dictionaries = schema
+ // For dictionary-free schemas, pre-encode the IPC schema message once so it does not have
+ // to be re-serialized per local block. RSS always delays serialization until admission.
+ // `flattened_fields` walks the full nested field tree for the local dictionary fallback.
+ let schema_encoding = if let Some(schema) = rss_schema {
+ SchemaEncoding::Fallback(schema)
+ } else if schema
.flattened_fields()
.iter()
- .any(|f| matches!(f.data_type(), DataType::Dictionary(_, _)));
-
- // For dictionary-free schemas, pre-encode the IPC schema message once so it does not have
- // to be re-serialized per block. Dictionary schemas use the `StreamWriter` fallback.
- let schema_encoding = if has_dictionaries {
+ .any(|f| matches!(f.data_type(), DataType::Dictionary(_, _)))
+ {
SchemaEncoding::Fallback(Arc::new(schema.clone()))
} else {
let data_gen = IpcDataGenerator::default();
@@ -186,6 +235,29 @@ impl ShuffleBlockWriter {
output: &mut W,
compression_context: &mut CompressionContext,
ipc_time: &Time,
+ ) -> Result<usize> {
+ self.write_batch_with_codec_limits(batch, output, compression_context, ipc_time, false)
+ }
+
+ /// Encode with the codec settings covered by [`Self::rss_codec_workspace`]. Local shuffle
+ /// retains its existing adaptive LZ4 block sizing through [`Self::write_batch`].
+ pub(crate) fn write_rss_batch<W: Write + Seek>(
+ &self,
+ batch: &RecordBatch,
+ output: &mut W,
+ compression_context: &mut CompressionContext,
+ ipc_time: &Time,
+ ) -> Result<usize> {
+ self.write_batch_with_codec_limits(batch, output, compression_context, ipc_time, true)
+ }
+
+ fn write_batch_with_codec_limits<W: Write + Seek>(
+ &self,
+ batch: &RecordBatch,
+ output: &mut W,
+ compression_context: &mut CompressionContext,
+ ipc_time: &Time,
+ bounded_rss_codec: bool,
) -> Result<usize> {
if batch.num_rows() == 0 {
return Ok(0);
@@ -202,7 +274,14 @@ impl ShuffleBlockWriter {
self.encode_ipc_stream(batch, output, compression_context)?;
}
CompressionCodec::Lz4Frame => {
- let mut wtr = lz4_flex::frame::FrameEncoder::new(&mut *output);
+ let frame_info = if bounded_rss_codec {
+ lz4_flex::frame::FrameInfo::new()
+ .block_size(lz4_flex::frame::BlockSize::Max64KB)
+ } else {
+ lz4_flex::frame::FrameInfo::default()
+ };
+ let mut wtr =
+ lz4_flex::frame::FrameEncoder::with_frame_info(frame_info, &mut *output);
self.encode_ipc_stream(batch, &mut wtr, compression_context)?;
wtr.finish().map_err(|e| {
DataFusionError::Execution(format!("lz4 compression error: {e}"))
@@ -248,6 +327,22 @@ mod tests {
use super::*;
use arrow::datatypes::{DataType, Field};
+ #[test]
+ fn rss_zstd_workspace_accounts_for_compression_level_without_unbounded_estimator_loops() {
+ let schema = Arc::new(Schema::empty());
+ let estimate = |level| {
+ ShuffleBlockWriter::try_new_rss(Arc::clone(&schema), CompressionCodec::Zstd(level))
+ .unwrap()
+ .rss_codec_workspace()
+ .unwrap()
+ };
+ let levels = zstd::compression_level_range();
+ assert!(estimate(1) > 64 * 1024);
+ assert!(estimate(*levels.end()) > estimate(1));
+ assert_eq!(estimate(i32::MAX), estimate(*levels.end()));
+ assert_eq!(estimate(i32::MIN), estimate(*levels.start()));
+ }
+
/// A `ShuffleBlockWriter` is cloned once per output partition (see `LocalPartitionWriter`), and
/// a shuffle can request millions of partitions (e.g. the SPARK-48037 test uses more than 16
/// million). Cloning must share the immutable header and pre-encoded schema buffers rather than
diff --git a/spark/src/main/java/org/apache/comet/shuffle/CelebornShufflePartitionPusher.java b/spark/src/main/java/org/apache/comet/shuffle/CelebornShufflePartitionPusher.java
index 132961ac44..7474e91742 100644
--- a/spark/src/main/java/org/apache/comet/shuffle/CelebornShufflePartitionPusher.java
+++ b/spark/src/main/java/org/apache/comet/shuffle/CelebornShufflePartitionPusher.java
@@ -20,26 +20,127 @@
package org.apache.comet.shuffle;
import java.io.IOException;
+import java.lang.reflect.Field;
import java.lang.reflect.InvocationTargetException;
import java.lang.reflect.Method;
import java.lang.reflect.Modifier;
import java.nio.ByteBuffer;
import java.nio.ByteOrder;
+import java.util.ArrayDeque;
+import java.util.ArrayList;
+import java.util.IdentityHashMap;
+import java.util.Iterator;
+import java.util.Map;
+import java.util.Optional;
+import java.util.concurrent.ScheduledFuture;
+import java.util.concurrent.ScheduledThreadPoolExecutor;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicLongArray;
+import java.util.concurrent.atomic.AtomicReference;
+import java.util.concurrent.atomic.LongAdder;
/** Adapts complete Comet shuffle frames to an existing, task-owned Celeborn shuffle client. */
public final class CelebornShufflePartitionPusher implements ShufflePartitionPusher {
private static final int CELEBORN_BATCH_HEADER_BYTES = 4 * Integer.BYTES;
private static final int MINIMUM_COMET_FRAME_BYTES = 2 * Long.BYTES;
+ private static final int MAX_JVM_ARRAY_BYTES = Integer.MAX_VALUE - 8;
+ private static final int DEFAULT_MAX_IN_FLIGHT_BYTES = 256 * 1024 * 1024;
+ private static final long RECONCILIATION_INTERVAL_MILLIS = 10;
+
+ // This daemon never owns client state: reconciliations are cancelled and removed once the
+ // associated task's transport requests complete.
+ private static final ScheduledThreadPoolExecutor COMPLETION_RECONCILER =
+ createCompletionReconciler();
private final Object shuffleClient;
private final Method pushOrMergeData;
private final Method computeBatchCRC;
+ private final Method mapperEnd;
+ private final Method cleanup;
+ private final Method getPushState;
+ private final Field clientPushStates;
+ private final Field inFlightRequestTracker;
+ private final Field totalInFlightRequests;
+ private final Field pushStateException;
+ private final Field cryptoHandler;
+ private final ExecutorShufflePushAdmission admission;
+ private final CelebornTransportCallbackTracker transportCallbacks;
private final int shuffleId;
private final int mapId;
private final int encodedAttemptId;
private final int numMappers;
private final int numPartitions;
+ private final int maxFrameBytes;
+ private final int maxReservationBytes;
+ private final AtomicLongArray partitionLengths;
+ private final Object lifecycleLock = new Object();
+ private final Object cleanupLock = new Object();
+ private final ThreadLocal<PushReservation> encodingReservation = new ThreadLocal<>();
+ private final ArrayDeque<PushReservation> pendingPushes = new ArrayDeque<>();
+ private final IdentityHashMap<Object, ObservedPushState> observedPushStates =
+ new IdentityHashMap<>();
+
+ private State state = State.OPEN;
+ private int activePushes;
+ private int activeClientPushes;
+ private int activeEncoders;
+ private ScheduledFuture<?> completionReconciliation;
+ private Thread mapperEndThread;
+ private boolean cleanupStarted;
+ private boolean cleanupAfterActivePushes;
+ private boolean finalCleanupStarted;
+ private IOException asynchronousFailure;
+
+ private enum State {
+ OPEN,
+ FINISHING,
+ FINISHED,
+ ABORTED
+ }
+
+ private static final class PushReservation {
+ private final int bytes;
+ private ObservedPushState pushState;
+ private CelebornTransportCallbackTracker.Push transportPush;
+ private boolean claimed;
+ private boolean submitted;
+ private boolean nativeReleased;
+ private boolean transportComplete;
+ private boolean released;
+
+ private PushReservation(int bytes, boolean nativeOwned) {
+ this.bytes = bytes;
+ this.nativeReleased = !nativeOwned;
+ }
+ }
+
+ private static final class ObservedPushState {
+ private final LongAdder inFlightRequests;
+ private final AtomicReference<?> exception;
+ private long submittedPushes;
+ private long releasedPushes;
+ private boolean terminalFailureCredited;
+
+ private ObservedPushState(LongAdder inFlightRequests, AtomicReference<?> exception) {
+ this.inFlightRequests = inFlightRequests;
+ this.exception = exception;
+ }
+ }
+
+ private static ScheduledThreadPoolExecutor createCompletionReconciler() {
+ ScheduledThreadPoolExecutor executor =
+ new ScheduledThreadPoolExecutor(
+ 1,
+ runnable -> {
+ Thread thread =
+ new Thread(runnable, "comet-celeborn-shuffle-push-admission-reconciler");
+ thread.setDaemon(true);
+ return thread;
+ });
+ executor.setRemoveOnCancelPolicy(true);
+ return executor;
+ }
/**
* Captures all shuffle and map-attempt identity before native worker threads invoke this pusher.
@@ -61,6 +162,47 @@ public final class CelebornShufflePartitionPusher implements ShufflePartitionPus
int encodedAttemptId,
int numMappers,
int numPartitions) {
+ this(
+ shuffleClient,
+ shuffleId,
+ mapId,
+ encodedAttemptId,
+ numMappers,
+ numPartitions,
+ MAX_JVM_ARRAY_BYTES,
+ DEFAULT_MAX_IN_FLIGHT_BYTES);
+ }
+
+ /** Binds one task to byte admission shared by its executor-side Celeborn client. */
+ public CelebornShufflePartitionPusher(
+ Object shuffleClient,
+ int shuffleId,
+ int mapId,
+ int encodedAttemptId,
+ int numMappers,
+ int numPartitions,
+ int maxInFlightBytes) {
+ this(
+ shuffleClient,
+ shuffleId,
+ mapId,
+ encodedAttemptId,
+ numMappers,
+ numPartitions,
+ MAX_JVM_ARRAY_BYTES,
+ maxInFlightBytes);
+ }
+
+ /** Binds one task to independently configured complete-frame and executor byte limits. */
+ public CelebornShufflePartitionPusher(
+ Object shuffleClient,
+ int shuffleId,
+ int mapId,
+ int encodedAttemptId,
+ int numMappers,
+ int numPartitions,
+ int configuredMaxFrameBytes,
+ int maxInFlightBytes) {
if (shuffleClient == null) {
throw new IllegalArgumentException("Celeborn shuffle client must not be null");
}
@@ -82,6 +224,13 @@ public final class CelebornShufflePartitionPusher implements ShufflePartitionPus
if (numPartitions <= 0) {
throw new IllegalArgumentException("Celeborn partition count must be positive");
}
+ if (configuredMaxFrameBytes < MINIMUM_COMET_FRAME_BYTES) {
+ throw new IllegalArgumentException("Celeborn maximum frame size must fit one Comet frame");
+ }
+ if (maxInFlightBytes < 3 * MINIMUM_COMET_FRAME_BYTES + CELEBORN_BATCH_HEADER_BYTES) {
+ throw new IllegalArgumentException(
+ "Celeborn in-flight byte limit must fit all copies of one Comet frame");
+ }
final Method pushMethod;
try {
@@ -150,19 +299,318 @@ public final class CelebornShufflePartitionPusher implements ShufflePartitionPus
"Celeborn integrity-accounting API must be an instance method returning void");
}
+ Method mapperEndMethod =
+ optionalLifecycleMethod(
+ shuffleClient, "mapperEnd", int.class, int.class, int.class, int.class, int.class);
+ if (mapperEndMethod == null) {
+ mapperEndMethod =
+ optionalLifecycleMethod(
+ shuffleClient, "mapperEnd", int.class, int.class, int.class, int.class);
+ }
+ Method cleanupMethod =
+ optionalLifecycleMethod(shuffleClient, "cleanup", int.class, int.class, int.class);
+
+ Method pushStateMethod = null;
+ Field pushStatesField = null;
+ Field trackerField = null;
+ Field requestCountField = null;
+ Field exceptionField = null;
+ try {
+ pushStateMethod = shuffleClient.getClass().getMethod("getPushState", String.class);
+ if (Modifier.isStatic(pushStateMethod.getModifiers())) {
+ throw new NoSuchMethodException("getPushState must be an instance method");
+ }
+ Class<?> pushStateClass = pushStateMethod.getReturnType();
+ trackerField = declaredField(pushStateClass, "inFlightRequestTracker");
+ requestCountField = declaredField(trackerField.getType(), "totalInflightReqs");
+ exceptionField = declaredField(pushStateClass, "exception");
+ if (requestCountField.getType() != LongAdder.class
+ || exceptionField.getType() != AtomicReference.class) {
+ throw new NoSuchFieldException("Celeborn PushState transport tracker is incompatible");
+ }
+ trackerField.setAccessible(true);
+ requestCountField.setAccessible(true);
+ exceptionField.setAccessible(true);
+ try {
+ pushStatesField = declaredField(shuffleClient.getClass(), "pushStates");
+ if (Map.class.isAssignableFrom(pushStatesField.getType())
+ && !Modifier.isStatic(pushStatesField.getModifiers())) {
+ pushStatesField.setAccessible(true);
+ } else {
+ pushStatesField = null;
+ }
+ } catch (NoSuchFieldException ignored) {
+ // Wrappers need not provide the optional side-effect-free state lookup.
+ }
+ } catch (NoSuchMethodException missing) {
+ // Compatibility-only clients can still push synchronously. Stock Celeborn 0.6 and 0.7
+ // expose getPushState and use completion-backed admission.
+ pushStateMethod = null;
+ pushStatesField = null;
+ trackerField = null;
+ requestCountField = null;
+ exceptionField = null;
+ } catch (ReflectiveOperationException | RuntimeException failure) {
+ throw new IllegalArgumentException(
+ "Celeborn shuffle client does not provide observable completion-backed push admission",
+ failure);
+ }
+
this.shuffleClient = shuffleClient;
this.pushOrMergeData = pushMethod;
this.computeBatchCRC = integrityMethod;
+ this.mapperEnd = mapperEndMethod;
+ this.cleanup = cleanupMethod;
+ this.getPushState = pushStateMethod;
+ this.clientPushStates = pushStatesField;
+ this.inFlightRequestTracker = trackerField;
+ this.totalInFlightRequests = requestCountField;
+ this.pushStateException = exceptionField;
+ this.cryptoHandler = optionalCryptoHandler(shuffleClient);
+ requireUnencryptedClient();
+ this.admission = ExecutorShufflePushAdmission.forClient(shuffleClient, maxInFlightBytes);
+ this.transportCallbacks =
+ getPushState == null ? null : admission.transportCallbacks(shuffleClient);
this.shuffleId = shuffleId;
this.mapId = mapId;
this.encodedAttemptId = encodedAttemptId;
this.numMappers = numMappers;
this.numPartitions = numPartitions;
+ this.maxReservationBytes = maxInFlightBytes - CELEBORN_BATCH_HEADER_BYTES;
+ this.maxFrameBytes =
+ Math.min(Math.min(configuredMaxFrameBytes, MAX_JVM_ARRAY_BYTES), maxReservationBytes / 3);
+ this.partitionLengths = new AtomicLongArray(numPartitions);
+ }
+
+ private static Method optionalLifecycleMethod(
+ Object shuffleClient, String name, Class<?>... parameterTypes) {
+ final Method method;
+ try {
+ method = shuffleClient.getClass().getMethod(name, parameterTypes);
+ } catch (NoSuchMethodException missing) {
+ return null;
+ } catch (SecurityException failure) {
+ throw new IllegalArgumentException("Cannot resolve the Celeborn " + name + " API", failure);
+ }
+ if (method.getReturnType() != void.class || Modifier.isStatic(method.getModifiers())) {
+ throw new IllegalArgumentException(
+ "Celeborn " + name + " API must be an instance method returning void");
+ }
+ return method;
+ }
+
+ private static Field declaredField(Class<?> owner, String name) throws NoSuchFieldException {
+ for (Class<?> current = owner; current != null; current = current.getSuperclass()) {
+ try {
+ return current.getDeclaredField(name);
+ } catch (NoSuchFieldException missing) {
+ // Application wrappers and compatibility clients can inherit these members.
+ }
+ }
+ throw new NoSuchFieldException(owner.getName() + "." + name);
+ }
+
+ private static Field optionalCryptoHandler(Object client) {
+ try {
+ Field field = declaredField(client.getClass(), "cryptoHandler");
+ if (Modifier.isStatic(field.getModifiers()) || field.getType() != Optional.class) {
+ throw new IllegalArgumentException("Cannot inspect the Celeborn client's encryption state");
+ }
+ field.setAccessible(true);
+ return field;
+ } catch (NoSuchFieldException absent) {
+ // Celeborn 0.6 has no encryption handler. Normal Spark construction is also guarded by
+ // spark.io.encryption.enabled in the factory, including legacy client APIs.
+ return null;
+ }
+ }
+
+ private void requireUnencryptedClient() {
+ if (cryptoHandler == null) {
+ return;
+ }
+ try {
+ Object handler = cryptoHandler.get(shuffleClient);
+ if (!(handler instanceof Optional) || ((Optional<?>) handler).isPresent()) {
+ // Stock SparkCryptoHandler retains a per-thread high-water buffer and creates additional
+ // encrypted copies. It has no bounded-workspace contract. Reject before encoding or raw
+ // submission; never remove the handler or downgrade the shared client's encryption.
+ throw new UnsupportedOperationException(
+ "Encrypted native Celeborn shuffle is not supported by bounded push admission; "
+ + "use ordinary Spark shuffle with encryption enabled");
+ }
+ } catch (IllegalAccessException failure) {
+ throw new IllegalStateException(
+ "Cannot inspect the Celeborn client's encryption state", failure);
+ }
+ }
+
+ @Override
+ public void reservePartitionData(int maxLength) throws IOException {
+ requireUnencryptedClient();
+ if (maxLength <= 0 || maxLength > maxReservationBytes) {
+ throw new IOException("Celeborn native frame reservation exceeds its byte limit");
+ }
+ if (encodingReservation.get() != null) {
+ throw new IOException("Celeborn native frame already has a reservation on this thread");
+ }
+
+ int requestBytes = maxLength + CELEBORN_BATCH_HEADER_BYTES;
+ admission.acquire(requestBytes, this::isAborted);
+ synchronized (lifecycleLock) {
+ if (state != State.OPEN) {
+ admission.release(requestBytes);
+ throw new IOException("Celeborn shuffle map attempt no longer accepts frame encoding");
+ }
+ activeEncoders++;
+ encodingReservation.set(new PushReservation(requestBytes, true));
+ }
+ }
+
+ @Override
+ public void releasePartitionDataReservation() {
+ PushReservation reservation = encodingReservation.get();
+ if (reservation == null) {
+ return;
+ }
+ encodingReservation.remove();
+ int released;
+ synchronized (lifecycleLock) {
+ reservation.nativeReleased = true;
+ if (!reservation.claimed) {
+ // Encoding failed (or was split) before entering the Java push path.
+ reservation.transportComplete = true;
+ }
+ released = releasableBytes(reservation);
+ activeEncoders--;
+ lifecycleLock.notifyAll();
+ }
+ admission.release(released);
}
/** Sends exactly one complete, already-compressed Comet shuffle frame to Celeborn. */
@Override
public void pushPartitionData(int partitionId, byte[] data, int length) throws IOException {
+ requireUnencryptedClient();
+ validateFrame(partitionId, data, length);
+ beginPush();
+
+ PushReservation reservation = null;
+ boolean registered = false;
+ boolean submitted = false;
+ boolean insideClient = false;
+ Throwable failure = null;
+ try {
+ reservation = claimEncodingReservation(length);
+ beginClientPush();
+ insideClient = true;
+
+ if (computeBatchCRC != null) {
+ computeBatchCRC.invoke(
+ shuffleClient, shuffleId, mapId, encodedAttemptId, partitionId, data, 0, length);
+ }
+
+ Object pushState = null;
+ ObservedPushState observed = null;
+ if (getPushState != null) {
+ pushState = getPushState.invoke(shuffleClient, mapKey());
+ if (pushState == null) {
+ throw new IOException("Celeborn returned a null push state for this map attempt");
+ }
+ observed = observePushState(pushState);
+ registerPendingPush(reservation, observed);
+ registered = true;
+ }
+
+ int accepted;
+ try (CelebornTransportCallbackTracker.Push transportPush =
+ transportCallbacks == null ? null : transportCallbacks.beginPush()) {
+ synchronized (lifecycleLock) {
+ reservation.transportPush = transportPush;
+ }
+ accepted =
+ (int)
+ pushOrMergeData.invoke(
+ shuffleClient,
+ shuffleId,
+ mapId,
+ encodedAttemptId,
+ partitionId,
+ data,
+ 0,
+ length,
+ numMappers,
+ numPartitions,
+ true,
+ true);
+ }
+ submitted = accepted > 0;
+ if (submitted && observed != null) {
+ if (clientPushStates != null) {
+ Object current = ((Map<?, ?>) clientPushStates.get(shuffleClient)).get(mapKey());
+ if (current != null && current != pushState) {
+ observed = observePushState(current);
+ }
+ }
+ markSubmitted(reservation, observed);
+ }
+
+ int minimumAccepted = length + CELEBORN_BATCH_HEADER_BYTES;
+ if (accepted < minimumAccepted) {
+ throw new IOException(
+ "Celeborn raw shuffle push accepted "
+ + accepted
+ + " bytes; expected at least "
+ + minimumAccepted
+ + " including its transport header");
+ }
+ if (accepted > reservation.bytes) {
+ throw new IOException(
+ "Celeborn encrypted shuffle request exceeds its reserved in-flight byte limit");
+ }
+ throwIfAsyncFailure();
+ if (isAborted()) {
+ throw new IOException("Celeborn shuffle map attempt was aborted during its push");
+ }
+ partitionLengths.addAndGet(partitionId, length);
+ } catch (IllegalAccessException cause) {
+ failure = new IOException("Cannot invoke the public Celeborn raw-push API", cause);
+ abortAndSuppress(failure);
+ throw (IOException) failure;
+ } catch (InvocationTargetException cause) {
+ failure = unwrapFailure("Celeborn raw shuffle push failed", cause);
+ abortAndSuppress(failure);
+ throwFailure(failure);
+ throw new AssertionError("unreachable");
+ } catch (IOException | RuntimeException | Error cause) {
+ failure = cause;
+ abortAndSuppress(cause);
+ throw cause;
+ } finally {
+ if (insideClient) {
+ endClientPush();
+ }
+ if (reservation != null) {
+ if (registered && !submitted) {
+ releaseUnsubmittedPush(reservation);
+ } else if (!registered) {
+ completeTransport(reservation);
+ }
+ }
+ try {
+ endPush();
+ } catch (IOException cleanupFailure) {
+ if (failure == null) {
+ throw cleanupFailure;
+ }
+ if (cleanupFailure != failure) {
+ failure.addSuppressed(cleanupFailure);
+ }
+ }
+ }
+ }
+
+ private void validateFrame(int partitionId, byte[] data, int length) throws IOException {
if (partitionId < 0 || partitionId >= numPartitions) {
throw new IOException("Celeborn output partition is outside this task's partition count");
}
@@ -175,8 +623,11 @@ public final class CelebornShufflePartitionPusher implements ShufflePartitionPus
if (length < MINIMUM_COMET_FRAME_BYTES || length > data.length) {
throw new IOException("Celeborn shuffle frame length must describe one complete frame");
}
+ if (length > maxFrameBytes) {
+ throw new IOException("Celeborn shuffle frame exceeds its configured maximum frame size");
+ }
- final long declaredBodyLength = ByteBuffer.wrap(data).order(ByteOrder.LITTLE_ENDIAN).getLong();
+ long declaredBodyLength = ByteBuffer.wrap(data).order(ByteOrder.LITTLE_ENDIAN).getLong();
if (declaredBodyLength != (long) length - Long.BYTES) {
throw new IOException(
"Celeborn shuffle frame declares "
@@ -184,53 +635,453 @@ public final class CelebornShufflePartitionPusher implements ShufflePartitionPus
+ " body bytes, but contains "
+ (length - Long.BYTES));
}
+ }
+
+ private PushReservation claimEncodingReservation(int frameBytes) throws IOException {
+ int required = Math.addExact(Math.multiplyExact(frameBytes, 3), CELEBORN_BATCH_HEADER_BYTES);
+ PushReservation reservation = encodingReservation.get();
+ if (reservation != null) {
+ if (required > reservation.bytes) {
+ throw new IOException("Celeborn shuffle push exceeds its native encoding reservation");
+ }
+ synchronized (lifecycleLock) {
+ if (reservation.claimed) {
+ throw new IOException("Celeborn native frame reservation has already been submitted");
+ }
+ reservation.claimed = true;
+ }
+ // The native Vec and JNI array still exist after this method and the Java push return.
+ // Keep the entire bound until native explicitly acknowledges their retirement. In
+ // particular, a fast network callback must not free bytes still owned by the encoder.
+ return reservation;
+ }
+ admission.acquire(required, this::isAborted);
+ return new PushReservation(required, false);
+ }
+
+ private ObservedPushState observePushState(Object pushState) throws IllegalAccessException {
+ synchronized (lifecycleLock) {
+ ObservedPushState observed = observedPushStates.get(pushState);
+ if (observed != null) {
+ return observed;
+ }
+ Object tracker = inFlightRequestTracker.get(pushState);
+ AtomicReference<?> exception = (AtomicReference<?>) pushStateException.get(pushState);
+ ObservedPushState created =
+ new ObservedPushState((LongAdder) totalInFlightRequests.get(tracker), exception);
+ observedPushStates.put(pushState, created);
+ return created;
+ }
+ }
- final int accepted;
+ private void beginPush() throws IOException {
+ synchronized (lifecycleLock) {
+ if (asynchronousFailure != null) {
+ throw asynchronousFailure;
+ }
+ if (state != State.OPEN) {
+ throw new IOException("Celeborn shuffle map attempt no longer accepts partition data");
+ }
+ activePushes++;
+ }
+ }
+
+ private void beginClientPush() throws IOException {
+ synchronized (lifecycleLock) {
+ if (state == State.ABORTED) {
+ throw new IOException("Celeborn shuffle map attempt was aborted before its client push");
+ }
+ activeClientPushes++;
+ }
+ }
+
+ private void registerPendingPush(PushReservation reservation, ObservedPushState pushState)
+ throws IOException {
+ synchronized (lifecycleLock) {
+ if (state == State.ABORTED) {
+ throw new IOException("Celeborn shuffle map attempt was aborted before submission");
+ }
+ reservation.pushState = pushState;
+ pendingPushes.addLast(reservation);
+ if (completionReconciliation == null || completionReconciliation.isDone()) {
+ completionReconciliation =
+ COMPLETION_RECONCILER.scheduleWithFixedDelay(
+ this::safelyReconcileAcceptedPushes,
+ RECONCILIATION_INTERVAL_MILLIS,
+ RECONCILIATION_INTERVAL_MILLIS,
+ TimeUnit.MILLISECONDS);
+ }
+ }
+ }
+
+ private void markSubmitted(PushReservation reservation, ObservedPushState pushState) {
+ synchronized (lifecycleLock) {
+ reservation.pushState = pushState;
+ reservation.submitted = true;
+ pushState.submittedPushes++;
+ }
+ reconcileAcceptedPushes();
+ }
+
+ // All reservation transitions are protected by lifecycleLock. Native retirement and transport
+ // completion may arrive in either order; neither alone proves that the overlapping copies died.
+ private int releasableBytes(PushReservation reservation) {
+ if (!reservation.released && reservation.nativeReleased && reservation.transportComplete) {
+ reservation.released = true;
+ return reservation.bytes;
+ }
+ return 0;
+ }
+
+ private void completeTransport(PushReservation reservation) {
+ int released;
+ synchronized (lifecycleLock) {
+ reservation.transportComplete = true;
+ released = releasableBytes(reservation);
+ }
+ admission.release(released);
+ }
+
+ private void safelyReconcileAcceptedPushes() {
try {
- if (computeBatchCRC != null) {
- computeBatchCRC.invoke(
- shuffleClient, shuffleId, mapId, encodedAttemptId, partitionId, data, 0, length);
+ reconcileAcceptedPushes();
+ } catch (RuntimeException | Error cause) {
+ IOException failure =
+ new IOException("Celeborn push completion reconciliation failed", cause);
+ synchronized (lifecycleLock) {
+ if (asynchronousFailure == null) {
+ asynchronousFailure = failure;
+ } else {
+ failure = asynchronousFailure;
+ }
}
+ abortAndSuppress(failure);
+ }
+ }
- accepted =
- (int)
- pushOrMergeData.invoke(
- shuffleClient,
- shuffleId,
- mapId,
- encodedAttemptId,
- partitionId,
- data,
- 0,
- length,
- numMappers,
- numPartitions,
- true,
- true);
- } catch (IllegalAccessException e) {
- throw new IOException("Cannot invoke the public Celeborn raw-push API", e);
- } catch (InvocationTargetException e) {
- Throwable cause = e.getCause();
- if (cause instanceof IOException) {
- throw (IOException) cause;
- }
- if (cause instanceof RuntimeException) {
- throw (RuntimeException) cause;
- }
- if (cause instanceof Error) {
- throw (Error) cause;
- }
- throw new IOException("Celeborn raw shuffle push failed", cause);
- }
-
- int minimumAccepted = length + CELEBORN_BATCH_HEADER_BYTES;
- if (accepted < minimumAccepted) {
- throw new IOException(
- "Celeborn raw shuffle push accepted "
- + accepted
- + " bytes; expected at least "
- + minimumAccepted
- + " including its transport header");
+ private void reconcileAcceptedPushes() {
+ ArrayList<PushReservation> completed = new ArrayList<>();
+ IOException detectedFailure = null;
+ synchronized (lifecycleLock) {
+ for (ObservedPushState observed : observedPushStates.values()) {
+ Object failure = observed.exception.get();
+ boolean cleanedUp =
+ failure instanceof IOException
+ && "Cleaned Up".equals(((IOException) failure).getMessage());
+ if (failure instanceof IOException && !cleanedUp) {
+ IOException observedFailure = (IOException) failure;
+ // Only the raw transport callback's pinned message proves request termination.
+ // PushState can also publish interruption/backpressure exceptions while an older
+ // request is still live; those failures must never free that request's admission.
+ if (!observed.terminalFailureCredited && isTerminalRawPushFailure(observedFailure)) {
+ observed.terminalFailureCredited = true;
+ }
+ if (asynchronousFailure == null) {
+ asynchronousFailure = observedFailure;
+ detectedFailure = asynchronousFailure;
+ }
+ }
+ if (transportCallbacks != null) {
+ // Stock counters can remain positive after cancellation. The transport path instead
+ // tracks each request through the raw call, transport callbacks, and queued retries.
+ continue;
+ }
+ long retainedRequests =
+ Math.max(
+ 0L, observed.inFlightRequests.sum() - (observed.terminalFailureCredited ? 1L : 0L));
+ long completions = Math.max(0L, observed.submittedPushes - retainedRequests);
+ while (observed.releasedPushes < completions) {
+ PushReservation reservation = smallestSubmittedReservation(observed);
+ if (reservation == null) {
+ break;
+ }
+ pendingPushes.remove(reservation);
+ observed.releasedPushes++;
+ completed.add(reservation);
+ }
+ }
+ if (transportCallbacks != null) {
+ Iterator<PushReservation> pending = pendingPushes.iterator();
+ while (pending.hasNext()) {
+ PushReservation reservation = pending.next();
+ if (reservation.submitted
+ && reservation.transportPush != null
+ && reservation.transportPush.isComplete()) {
+ pending.remove();
+ completed.add(reservation);
+ }
+ }
+ }
+ if (pendingPushes.isEmpty() && activeClientPushes == 0) {
+ stopCompletionReconciliation();
+ }
+ }
+ for (PushReservation reservation : completed) {
+ completeTransport(reservation);
+ }
+ if (detectedFailure != null) {
+ abortAndSuppress(detectedFailure);
+ }
+ }
+
+ private static boolean isTerminalRawPushFailure(IOException failure) {
+ String message = failure.getMessage();
+ return message != null
+ && message.startsWith("Push data to ")
+ && message.contains(" failed for shuffle ")
+ && message.contains(" batch ");
+ }
+
+ private PushReservation smallestSubmittedReservation(ObservedPushState observed) {
+ PushReservation smallest = null;
+ for (PushReservation reservation : pendingPushes) {
+ if (reservation.pushState == observed
+ && reservation.submitted
+ && (smallest == null || reservation.bytes < smallest.bytes)) {
+ smallest = reservation;
+ }
+ }
+ return smallest;
+ }
+
+ private void stopCompletionReconciliation() {
+ if (completionReconciliation != null) {
+ completionReconciliation.cancel(false);
+ completionReconciliation = null;
+ }
+ }
+
+ private void releaseUnsubmittedPush(PushReservation reservation) {
+ boolean removed;
+ synchronized (lifecycleLock) {
+ if (reservation.transportPush != null && !reservation.transportPush.isComplete()) {
+ // The raw call may throw after publishing transport or retry work. Keep its entire
+ // reservation until that work returns, even though the call never reported acceptance.
+ reservation.submitted = true;
+ removed = false;
+ } else {
+ removed = pendingPushes.remove(reservation);
+ }
+ if (pendingPushes.isEmpty() && activeClientPushes == 0) {
+ stopCompletionReconciliation();
+ }
+ }
+ if (removed) {
+ completeTransport(reservation);
+ }
+ }
+
+ private void endClientPush() {
+ synchronized (lifecycleLock) {
+ activeClientPushes--;
+ if (pendingPushes.isEmpty() && activeClientPushes == 0) {
+ stopCompletionReconciliation();
+ }
+ }
+ }
+
+ private void endPush() throws IOException {
+ boolean performFinalCleanup;
+ synchronized (lifecycleLock) {
+ activePushes--;
+ if (activePushes == 0) {
+ lifecycleLock.notifyAll();
+ }
+ performFinalCleanup =
+ activePushes == 0
+ && state == State.ABORTED
+ && cleanupAfterActivePushes
+ && !finalCleanupStarted;
+ if (performFinalCleanup) {
+ finalCleanupStarted = true;
+ }
+ }
+ if (performFinalCleanup) {
+ cleanupAttempt();
+ }
+ }
+
+ /** Drains asynchronous requests, commits this map, and returns Comet bytes per partition. */
+ public long[] finish() throws IOException {
+ try {
+ synchronized (lifecycleLock) {
+ if (asynchronousFailure != null) {
+ throw asynchronousFailure;
+ }
+ if (state == State.FINISHED) {
+ return snapshotPartitionLengths();
+ }
+ if (state != State.OPEN) {
+ throw new IOException("Celeborn shuffle map attempt is not available for completion");
+ }
+ state = State.FINISHING;
+ while ((activePushes != 0 || activeEncoders != 0) && state == State.FINISHING) {
+ lifecycleLock.wait();
+ }
+ if (asynchronousFailure != null) {
+ throw asynchronousFailure;
+ }
+ if (state != State.FINISHING) {
+ throw new IOException("Celeborn shuffle map attempt was aborted before completion");
+ }
+ if (mapperEnd == null) {
+ throw new IOException(
+ "Celeborn shuffle client does not provide the required public mapperEnd API");
+ }
+ mapperEndThread = Thread.currentThread();
+ }
+
+ try {
+ if (mapperEnd.getParameterCount() == 5) {
+ mapperEnd.invoke(
+ shuffleClient, shuffleId, mapId, encodedAttemptId, numMappers, numPartitions);
+ } else {
+ mapperEnd.invoke(shuffleClient, shuffleId, mapId, encodedAttemptId, numMappers);
+ }
+ } finally {
+ synchronized (lifecycleLock) {
+ mapperEndThread = null;
+ }
+ }
+ reconcileAcceptedPushes();
+ throwIfAsyncFailure();
+
+ synchronized (lifecycleLock) {
+ if (state != State.FINISHING) {
+ throw new IOException("Celeborn shuffle map attempt was aborted during completion");
+ }
+ state = State.FINISHED;
+ lifecycleLock.notifyAll();
+ }
+ return snapshotPartitionLengths();
+ } catch (InterruptedException cause) {
+ Thread.currentThread().interrupt();
+ IOException failure = new IOException("Interrupted while draining Celeborn pushes", cause);
+ abortAndSuppress(failure);
+ throw failure;
+ } catch (IllegalAccessException cause) {
+ IOException failure =
+ new IOException("Cannot invoke the public Celeborn mapperEnd API", cause);
+ abortAndSuppress(failure);
+ throw failure;
+ } catch (InvocationTargetException cause) {
+ Throwable failure = unwrapFailure("Celeborn shuffle map completion failed", cause);
+ abortAndSuppress(failure);
+ throwFailure(failure);
+ throw new AssertionError("unreachable");
+ } catch (IOException | RuntimeException | Error failure) {
+ abortAndSuppress(failure);
+ throw failure;
+ }
+ }
+
+ /** Cancels this map without releasing request bytes before actual transport completion. */
+ public void abort() throws IOException {
+ boolean shouldCleanup;
+ Thread completionThread;
+ synchronized (lifecycleLock) {
+ if (state == State.FINISHED) {
+ return;
+ }
+ state = State.ABORTED;
+ shouldCleanup = !cleanupStarted;
+ if (shouldCleanup) {
+ cleanupAfterActivePushes = activeClientPushes > 0;
+ }
+ cleanupStarted = true;
+ completionThread = mapperEndThread;
+ lifecycleLock.notifyAll();
+ }
+ if (completionThread != null && completionThread != Thread.currentThread()) {
+ completionThread.interrupt();
+ }
+ if (shouldCleanup) {
+ cleanupAttempt();
+ }
+ }
+
+ private void cleanupAttempt() throws IOException {
+ if (cleanup == null) {
+ return;
+ }
+ synchronized (cleanupLock) {
+ try {
+ cleanup.invoke(shuffleClient, shuffleId, mapId, encodedAttemptId);
+ } catch (IllegalAccessException failure) {
+ throw new IOException("Cannot invoke the public Celeborn cleanup API", failure);
+ } catch (InvocationTargetException failure) {
+ throwFailure(unwrapFailure("Celeborn shuffle map cleanup failed", failure));
+ }
+ }
+ }
+
+ private String mapKey() {
+ return shuffleId + "-" + mapId + "-" + encodedAttemptId;
+ }
+
+ private boolean isAborted() {
+ synchronized (lifecycleLock) {
+ return state == State.ABORTED;
+ }
+ }
+
+ private void throwIfAsyncFailure() throws IOException {
+ synchronized (lifecycleLock) {
+ if (asynchronousFailure != null) {
+ throw asynchronousFailure;
+ }
+ }
+ }
+
+ public int numPartitions() {
+ return numPartitions;
+ }
+
+ /** Largest configured frame whose native, JNI, and Celeborn copies fit shared admission. */
+ @Override
+ public int maxFrameBytes() {
+ return maxFrameBytes;
+ }
+
+ @Override
+ public int maxReservationBytes() {
+ return maxReservationBytes;
+ }
+
+ private long[] snapshotPartitionLengths() {
+ long[] sizes = new long[numPartitions];
+ for (int partition = 0; partition < numPartitions; partition++) {
+ sizes[partition] = partitionLengths.get(partition);
+ }
+ return sizes;
+ }
+
+ private void abortAndSuppress(Throwable failure) {
+ try {
+ abort();
+ } catch (Throwable cleanupFailure) {
+ if (cleanupFailure != failure) {
+ failure.addSuppressed(cleanupFailure);
+ }
+ }
+ }
+
+ private static Throwable unwrapFailure(String message, InvocationTargetException failure) {
+ Throwable cause = failure.getCause();
+ return cause instanceof IOException
+ || cause instanceof RuntimeException
+ || cause instanceof Error
+ ? cause
+ : new IOException(message, cause);
+ }
+
+ private static void throwFailure(Throwable failure) throws IOException {
+ if (failure instanceof IOException) {
+ throw (IOException) failure;
+ }
+ if (failure instanceof RuntimeException) {
+ throw (RuntimeException) failure;
}
+ throw (Error) failure;
}
}
diff --git a/spark/src/main/java/org/apache/comet/shuffle/CelebornTransportCallbackTracker.java b/spark/src/main/java/org/apache/comet/shuffle/CelebornTransportCallbackTracker.java
new file mode 100644
index 0000000000..cb4f35c0ca
--- /dev/null
+++ b/spark/src/main/java/org/apache/comet/shuffle/CelebornTransportCallbackTracker.java
@@ -0,0 +1,1000 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.comet.shuffle;
+
+import java.io.IOException;
+import java.lang.reflect.Array;
+import java.lang.reflect.Field;
+import java.lang.reflect.InvocationHandler;
+import java.lang.reflect.InvocationTargetException;
+import java.lang.reflect.Method;
+import java.lang.reflect.ParameterizedType;
+import java.lang.reflect.Proxy;
+import java.lang.reflect.Type;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.concurrent.AbstractExecutorService;
+import java.util.concurrent.Callable;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ConcurrentLinkedQueue;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Future;
+import java.util.concurrent.FutureTask;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.function.BiConsumer;
+import java.util.function.BiFunction;
+import java.util.function.Function;
+
+/** Tracks payload ownership across stock Celeborn transport callbacks and retry tasks. */
+final class CelebornTransportCallbackTracker {
+ private final Object shuffleClient;
+ private final Method getDataClientFactory;
+
+ private CelebornTransportCallbackTracker(Object shuffleClient, Method getDataClientFactory) {
+ this.shuffleClient = shuffleClient;
+ this.getDataClientFactory = getDataClientFactory;
+ }
+
+ static CelebornTransportCallbackTracker tryCreate(Object shuffleClient) {
+ try {
+ return new CelebornTransportCallbackTracker(
+ shuffleClient, shuffleClient.getClass().getMethod("getDataClientFactory"));
+ } catch (NoSuchMethodException ignored) {
+ // Compatibility clients without a transport factory retain their own completion tracking.
+ return null;
+ }
+ }
+
+ Push beginPush() throws IOException {
+ try {
+ Object factory = getDataClientFactory.invoke(shuffleClient);
+ if (factory == null) {
+ throw new IOException("Celeborn returned a null data transport factory");
+ }
+ synchronized (factory) {
+ Field bootstrapsField = field(factory.getClass(), "clientBootstraps");
+ Object bootstrapValue = bootstrapsField.get(factory);
+ if (!(bootstrapValue instanceof List<?>)) {
+ throw new IOException("Celeborn transport factory has no bootstrap list");
+ }
+ List<?> bootstraps = (List<?>) bootstrapValue;
+ FactoryHook hook = findHook(bootstraps);
+ if (hook == null) {
+ hook = new FactoryHook();
+ Class<?> bootstrapInterface = bootstrapInterface(bootstrapsField, factory);
+ ArrayList<Object> augmentedBootstraps = new ArrayList<>(bootstraps);
+ augmentedBootstraps.add(
+ Proxy.newProxyInstance(
+ bootstrapInterface.getClassLoader(), new Class<?>[] {bootstrapInterface}, hook));
+ // A client may already be iterating the old list. Do not mutate that iterator; the pool
+ // locks below wait for that client's construction before instrumenting its handler.
+ bootstrapsField.set(factory, augmentedBootstraps);
+ }
+ if (!hook.installed) {
+ hook.installExistingClients(factory);
+ hook.installed = true;
+ }
+ hook.installRetryPool(shuffleClient);
+ return new Push(hook);
+ }
+ } catch (ReflectiveOperationException | RuntimeException failure) {
+ throw new IOException("Cannot observe Celeborn shuffle payload ownership", failure);
+ }
+ }
+
+ private static FactoryHook findHook(List<?> bootstraps) {
+ for (Object bootstrap : bootstraps) {
+ if (bootstrap != null && Proxy.isProxyClass(bootstrap.getClass())) {
+ InvocationHandler handler = Proxy.getInvocationHandler(bootstrap);
+ if (handler instanceof FactoryHook) {
+ return (FactoryHook) handler;
+ }
+ }
+ }
+ return null;
+ }
+
+ private static Class<?> bootstrapInterface(Field bootstrapsField, Object factory)
+ throws ReflectiveOperationException {
+ Type genericType = bootstrapsField.getGenericType();
+ if (genericType instanceof ParameterizedType) {
+ Type elementType = ((ParameterizedType) genericType).getActualTypeArguments()[0];
+ if (elementType instanceof Class<?> && ((Class<?>) elementType).isInterface()) {
+ return (Class<?>) elementType;
+ }
+ }
+ return Class.forName(
+ "org.apache.celeborn.common.network.client.TransportClientBootstrap",
+ false,
+ factory.getClass().getClassLoader());
+ }
+
+ private static Field field(Class<?> type, String name) throws NoSuchFieldException {
+ for (Class<?> owner = type; owner != null; owner = owner.getSuperclass()) {
+ try {
+ Field found = owner.getDeclaredField(name);
+ found.setAccessible(true);
+ return found;
+ } catch (NoSuchFieldException ignored) {
+ // Stock clients and compatibility test clients can inherit these fields.
+ }
+ }
+ throw new NoSuchFieldException(type.getName() + "." + name);
+ }
+
+ /** A raw push plus every asynchronous owner created while handling that push. */
+ static final class Push implements AutoCloseable {
+ private final FactoryHook hook;
+ private final Push previous;
+ private final AtomicInteger owners = new AtomicInteger(1);
+ private final AtomicBoolean closed = new AtomicBoolean();
+ private final ConcurrentLinkedQueue<WriteOwnership> writes = new ConcurrentLinkedQueue<>();
+
+ private Push(FactoryHook hook) {
+ this.hook = hook;
+ this.previous = hook.current.get();
+ hook.current.set(this);
+ }
+
+ boolean isComplete() {
+ // A shutting-down event loop can reject listener notification after completing a promise.
+ // Reconciliation can still observe that write's completion without depending on notification.
+ for (WriteOwnership write : writes) {
+ write.completeIfDone();
+ }
+ return owners.get() == 0;
+ }
+
+ private Lease retain() {
+ int currentOwners;
+ do {
+ currentOwners = owners.get();
+ if (currentOwners == 0) {
+ throw new IllegalStateException("Cannot retain a completed Celeborn push");
+ }
+ } while (!owners.compareAndSet(currentOwners, currentOwners + 1));
+ return new Lease(this);
+ }
+
+ private Activation activate() {
+ return new Activation(this);
+ }
+
+ @Override
+ public void close() {
+ if (closed.compareAndSet(false, true)) {
+ if (hook.current.get() != this) {
+ throw new IllegalStateException(
+ "Celeborn raw push scopes must close in invocation order");
+ }
+ hook.restore(previous);
+ owners.decrementAndGet();
+ }
+ }
+ }
+
+ private static final class Lease implements AutoCloseable {
+ private final Push push;
+ private final AtomicBoolean closed = new AtomicBoolean();
+
+ private Lease(Push push) {
+ this.push = push;
+ }
+
+ @Override
+ public void close() {
+ if (closed.compareAndSet(false, true)) {
+ push.owners.decrementAndGet();
+ }
+ }
+ }
+
+ private static final class Activation implements AutoCloseable {
+ private final Push push;
+ private final Push previous;
+
+ private Activation(Push push) {
+ this.push = push;
+ previous = push.hook.current.get();
+ push.hook.current.set(push);
+ }
+
+ @Override
+ public void close() {
+ push.hook.restore(previous);
+ }
+ }
+
+ /** The factory owns its context; it is explicitly propagated at every asynchronous boundary. */
+ private static final class FactoryHook implements InvocationHandler {
+ // This is an invocation scope, not Spark TaskContext or native worker identity. A callback or
+ // retry can run on any thread; its wrapper restores this scope only while invoking its
+ // delegate.
+ private final ThreadLocal<Push> current = new ThreadLocal<>();
+ // Read and written only while holding the owning transport factory's monitor.
+ private boolean installed;
+
+ private void restore(Push previous) {
+ if (previous == null) {
+ current.remove();
+ } else {
+ current.set(previous);
+ }
+ }
+
+ private void installRetryPool(Object shuffleClient)
+ throws ReflectiveOperationException, IOException {
+ synchronized (shuffleClient) {
+ Field poolField = field(shuffleClient.getClass(), "pushDataRetryPool");
+ Object pool = poolField.get(shuffleClient);
+ if (pool instanceof TrackingRetryExecutor) {
+ if (((TrackingRetryExecutor) pool).hook != this) {
+ throw new IOException("Celeborn retry executor has a different callback tracker");
+ }
+ } else if (pool instanceof ExecutorService) {
+ poolField.set(shuffleClient, new TrackingRetryExecutor((ExecutorService) pool, this));
+ } else {
+ throw new IOException("Celeborn shuffle client has no compatible retry executor");
+ }
+ }
+ }
+
+ private void installExistingClients(Object factory)
+ throws ReflectiveOperationException, IOException {
+ Object pools = field(factory.getClass(), "connectionPool").get(factory);
+ if (!(pools instanceof Map<?, ?>)) {
+ throw new IOException("Celeborn transport factory has no connection pool");
+ }
+ for (Object pool : ((Map<?, ?>) pools).values()) {
+ Object clients = field(pool.getClass(), "clients").get(pool);
+ Object locks = field(pool.getClass(), "locks").get(pool);
+ if (clients == null
+ || locks == null
+ || !clients.getClass().isArray()
+ || !locks.getClass().isArray()
+ || Array.getLength(clients) != Array.getLength(locks)) {
+ throw new IOException("Celeborn connection pool has incompatible client slots");
+ }
+ for (int index = 0; index < Array.getLength(clients); index++) {
+ Object lock = Array.get(locks, index);
+ if (lock == null) {
+ throw new IOException("Celeborn connection pool has a null client lock");
+ }
+ synchronized (lock) {
+ Object client = Array.get(clients, index);
+ if (client != null) {
+ installClient(client);
+ }
+ }
+ }
+ }
+ }
+
+ private void installClient(Object client) throws ReflectiveOperationException, IOException {
+ Object handler = client.getClass().getMethod("getHandler").invoke(client);
+ if (handler == null) {
+ throw new IOException("Celeborn transport client has a null response handler");
+ }
+ synchronized (handler) {
+ installChannel(client);
+ Field requestsField = field(handler.getClass(), "outstandingPushes");
+ Object requests = requestsField.get(handler);
+ if (requests instanceof CallbackTrackingRequests) {
+ if (((CallbackTrackingRequests) requests).hook != this) {
+ throw new IOException("Celeborn response handler has a different callback tracker");
+ }
+ return;
+ }
+ if (!(requests instanceof ConcurrentHashMap<?, ?>)) {
+ throw new IOException("Celeborn response handler has incompatible push request storage");
+ }
+ @SuppressWarnings("unchecked")
+ ConcurrentHashMap<Object, Object> original = (ConcurrentHashMap<Object, Object>) requests;
+ // Keep the original map as the backing store: an unrelated task may already be using it.
+ requestsField.set(handler, new CallbackTrackingRequests(original, this));
+ }
+ }
+
+ private void installChannel(Object client) throws ReflectiveOperationException, IOException {
+ Field channelField = field(client.getClass(), "channel");
+ Object channel = channelField.get(client);
+ if (channel == null || !channelField.getType().isInterface()) {
+ throw new IOException("Celeborn transport client has no compatible channel interface");
+ }
+ if (Proxy.isProxyClass(channel.getClass())
+ && Proxy.getInvocationHandler(channel) instanceof TrackingChannel) {
+ if (((TrackingChannel) Proxy.getInvocationHandler(channel)).hook != this) {
+ throw new IOException("Celeborn channel has a different callback tracker");
+ }
+ return;
+ }
+ Class<?> channelInterface = channelField.getType();
+ TrackingChannel tracking = new TrackingChannel(channel, channelInterface, this);
+ channelField.set(
+ client,
+ Proxy.newProxyInstance(
+ channelInterface.getClassLoader(), new Class<?>[] {channelInterface}, tracking));
+ }
+
+ private PreparedRequest prepareRequest(Object requestInfo) {
+ Push push = current.get();
+ if (push == null) {
+ return null;
+ }
+ try {
+ Field callbackField = field(requestInfo.getClass(), "callback");
+ Object callback = callbackField.get(requestInfo);
+ if (callback == null) {
+ return null;
+ }
+ if (Proxy.isProxyClass(callback.getClass())
+ && Proxy.getInvocationHandler(callback) instanceof CompletionCallback) {
+ return null;
+ }
+ Class<?> callbackInterface = callbackField.getType();
+ if (!callbackInterface.isInterface()) {
+ throw new IllegalStateException("Celeborn push callback is not an interface");
+ }
+ CompletionCallback completion = new CompletionCallback(callback, push);
+ try {
+ Object wrapped =
+ Proxy.newProxyInstance(
+ callbackInterface.getClassLoader(),
+ new Class<?>[] {callbackInterface},
+ completion);
+ callbackField.set(requestInfo, wrapped);
+ return new PreparedRequest(requestInfo, callbackField, callback, wrapped, completion);
+ } catch (ReflectiveOperationException | RuntimeException | Error failure) {
+ completion.lease.close();
+ throw failure;
+ }
+ } catch (ReflectiveOperationException failure) {
+ throw new IllegalStateException("Cannot observe a Celeborn push callback", failure);
+ }
+ }
+
+ @Override
+ public Object invoke(Object proxy, Method method, Object[] args) throws Throwable {
+ if (method.getDeclaringClass() == Object.class) {
+ return objectMethod(proxy, method, args);
+ }
+ if ("doBootstrap".equals(method.getName()) && args != null && args.length == 1) {
+ installClient(args[0]);
+ return null;
+ }
+ throw new UnsupportedOperationException(
+ "Unsupported Celeborn transport bootstrap: " + method);
+ }
+ }
+
+ /** Pins the actual write future independently of the response handler's timeout bookkeeping. */
+ private static final class WriteOwnership implements InvocationHandler {
+ private final Push push;
+ private final Lease lease;
+ private final Method isDone;
+ private volatile Object future;
+
+ private WriteOwnership(Push push, Method isDone) {
+ this.push = push;
+ this.lease = push.retain();
+ this.isDone = isDone;
+ push.writes.add(this);
+ }
+
+ private void finish() {
+ lease.close();
+ push.writes.remove(this);
+ }
+
+ private void completeIfDone() {
+ Object actual = future;
+ if (actual != null) {
+ try {
+ if ((boolean) isDone.invoke(actual)) {
+ finish();
+ }
+ } catch (ReflectiveOperationException failure) {
+ throw new IllegalStateException(
+ "Cannot inspect Celeborn outbound write completion", failure);
+ }
+ }
+ }
+
+ @Override
+ public Object invoke(Object proxy, Method method, Object[] args) {
+ if (method.getDeclaringClass() == Object.class) {
+ return objectMethod(proxy, method, args);
+ }
+ if ("operationComplete".equals(method.getName())) {
+ finish();
+ return null;
+ }
+ throw new UnsupportedOperationException("Unsupported Celeborn write listener: " + method);
+ }
+ }
+
+ private static final class TrackingChannel implements InvocationHandler {
+ private final Object channel;
+ private final FactoryHook hook;
+ private final Class<?> futureInterface;
+ private final Method addListener;
+ private final Method isDone;
+
+ private TrackingChannel(Object channel, Class<?> channelInterface, FactoryHook hook)
+ throws NoSuchMethodException {
+ this.channel = channel;
+ this.hook = hook;
+ futureInterface = channelInterface.getMethod("writeAndFlush", Object.class).getReturnType();
+ if (!futureInterface.isInterface()) {
+ throw new NoSuchMethodException("Celeborn channel writes must return a future interface");
+ }
+ isDone = futureInterface.getMethod("isDone");
+ Method listenerMethod = null;
+ for (Method candidate : futureInterface.getMethods()) {
+ if ("addListener".equals(candidate.getName())
+ && candidate.getParameterCount() == 1
+ && candidate.getParameterTypes()[0].isInterface()) {
+ listenerMethod = candidate;
+ break;
+ }
+ }
+ if (listenerMethod == null) {
+ throw new NoSuchMethodException(
+ "Celeborn channel futures must support completion listeners");
+ }
+ addListener = listenerMethod;
+ }
+
+ @Override
+ public Object invoke(Object proxy, Method method, Object[] args) throws Throwable {
+ if (method.getDeclaringClass() == Object.class) {
+ return objectMethod(proxy, method, args);
+ }
+ Push push = hook.current.get();
+ if (push == null || !"writeAndFlush".equals(method.getName())) {
+ try {
+ Object result = method.invoke(channel, args);
+ return result == channel ? proxy : result;
+ } catch (InvocationTargetException failure) {
+ throw failure.getCause();
+ }
+ }
+ if (args == null || args.length != 1) {
+ throw new IOException("Celeborn owned writes must create their own transport promise");
+ }
+ WriteOwnership ownership = new WriteOwnership(push, isDone);
+ final Object future;
+ try {
+ future = method.invoke(channel, args);
+ } catch (InvocationTargetException failure) {
+ ownership.finish();
+ throw failure.getCause();
+ } catch (ReflectiveOperationException | RuntimeException | Error failure) {
+ ownership.finish();
+ throw failure;
+ }
+ ownership.future = future;
+ Class<?> listenerInterface = addListener.getParameterTypes()[0];
+ Object listener =
+ Proxy.newProxyInstance(
+ listenerInterface.getClassLoader(), new Class<?>[] {listenerInterface}, ownership);
+ try {
+ addListener.invoke(future, listener);
+ } catch (InvocationTargetException failure) {
+ ownership.completeIfDone();
+ throw failure.getCause();
+ }
+ ownership.completeIfDone();
+ return Proxy.newProxyInstance(
+ futureInterface.getClassLoader(),
+ new Class<?>[] {futureInterface},
+ new UncancellableWriteFuture(future));
+ }
+ }
+
+ /** A timeout may finish the response callback, but must not mark an unwritten body complete. */
+ private static final class UncancellableWriteFuture implements InvocationHandler {
+ private final Object future;
+
+ private UncancellableWriteFuture(Object future) {
+ this.future = future;
+ }
+
+ @Override
+ public Object invoke(Object proxy, Method method, Object[] args) throws Throwable {
+ if (method.getDeclaringClass() == Object.class) {
+ return objectMethod(proxy, method, args);
+ }
+ if ("cancel".equals(method.getName()) || "isCancellable".equals(method.getName())) {
+ // Celeborn already tolerates cancellation failing once Netty has flushed a write. Keep
+ // that behavior for queued Comet writes too, so only actual write retirement releases it.
+ return false;
+ }
+ try {
+ Object result = method.invoke(future, args);
+ // Stock TransportClient stores writeAndFlush(...).addListener(...), not the initial
+ // future. Preserve this wrapper through all fluent future operations.
+ return result == future ? proxy : result;
+ } catch (InvocationTargetException failure) {
+ throw failure.getCause();
+ }
+ }
+ }
+
+ private static final class PreparedRequest {
+ private final Object requestInfo;
+ private final Field callbackField;
+ private final Object original;
+ private final Object wrapped;
+ private final CompletionCallback completion;
+
+ private PreparedRequest(
+ Object requestInfo,
+ Field callbackField,
+ Object original,
+ Object wrapped,
+ CompletionCallback completion) {
+ this.requestInfo = requestInfo;
+ this.callbackField = callbackField;
+ this.original = original;
+ this.wrapped = wrapped;
+ this.completion = completion;
+ }
+
+ private void discard() {
+ if (completion.delivered.compareAndSet(false, true)) {
+ try {
+ if (callbackField.get(requestInfo) == wrapped) {
+ callbackField.set(requestInfo, original);
+ }
+ } catch (IllegalAccessException failure) {
+ throw new IllegalStateException(
+ "Cannot restore an unpublished Celeborn callback", failure);
+ } finally {
+ completion.callback = null;
+ completion.lease.close();
+ }
+ }
+ }
+ }
+
+ private static final class CompletionCallback implements InvocationHandler {
+ private volatile Object callback;
+ private final Push push;
+ private final Lease lease;
+ private final AtomicBoolean delivered = new AtomicBoolean();
+
+ private CompletionCallback(Object callback, Push push) {
+ this.callback = callback;
+ this.push = push;
+ this.lease = push.retain();
+ }
+
+ private Object invokeCallback(Method method, Object[] args) throws Throwable {
+ try {
+ return method.invoke(callback, args);
+ } catch (InvocationTargetException failure) {
+ throw failure.getCause();
+ }
+ }
+
+ @Override
+ public Object invoke(Object proxy, Method method, Object[] args) throws Throwable {
+ if (method.getDeclaringClass() == Object.class) {
+ return objectMethod(proxy, method, args);
+ }
+ if (!delivered.compareAndSet(false, true)) {
+ return null;
+ }
+ try (Activation ignored = push.activate()) {
+ return invokeCallback(method, args);
+ } finally {
+ // Request-map removal is not completion. Retry submission during the original callback
+ // has already retained its own lease before this transport owner's lease is released.
+ // The handler may retain RequestInfo after invoking it. Clear its payload-owning delegate
+ // only after the invocation frame has returned, and before publishing completion.
+ callback = null;
+ lease.close();
+ }
+ }
+ }
+
+ private interface OwnedRetry {
+ void discard();
+ }
+
+ /** A running retry retains its payload even if its Future is cancelled or reports completion. */
+ private static final class RetryOwnership {
+ private static final int QUEUED = 0;
+ private static final int RUNNING = 1;
+ private static final int FINISHED = 2;
+
+ private final Push push;
+ private final Lease lease;
+ private final AtomicInteger state = new AtomicInteger(QUEUED);
+
+ private RetryOwnership(Push push) {
+ this.push = push;
+ lease = push.retain();
+ }
+
+ private boolean start() {
+ return state.compareAndSet(QUEUED, RUNNING);
+ }
+
+ private void finish() {
+ state.set(FINISHED);
+ lease.close();
+ }
+
+ private boolean discard() {
+ return state.compareAndSet(QUEUED, FINISHED);
+ }
+ }
+
+ private static final class TrackingRetryRunnable implements Runnable, OwnedRetry {
+ private volatile Runnable delegate;
+ private final RetryOwnership ownership;
+
+ private TrackingRetryRunnable(Runnable delegate, Push push) {
+ this.delegate = delegate;
+ ownership = new RetryOwnership(push);
+ }
+
+ private void invokeDelegate() {
+ delegate.run();
+ }
+
+ @Override
+ public void run() {
+ if (!ownership.start()) {
+ return;
+ }
+ try (Activation ignored = ownership.push.activate()) {
+ invokeDelegate();
+ } finally {
+ // An executor or shutdown caller can retain the completed wrapper. The delegate's own
+ // invocation frame must have returned before dropping its payload reference and lease.
+ delegate = null;
+ ownership.finish();
+ }
+ }
+
+ @Override
+ public void discard() {
+ if (ownership.discard()) {
+ delegate = null;
+ ownership.finish();
+ }
+ }
+ }
+
+ private static final class TrackingRetryFuture<T> extends FutureTask<T> implements OwnedRetry {
+ private final RetryOwnership ownership;
+
+ private TrackingRetryFuture(Callable<T> callable, Push push) {
+ super(callable);
+ ownership = new RetryOwnership(push);
+ }
+
+ private TrackingRetryFuture(Runnable runnable, T result, Push push) {
+ super(runnable, result);
+ ownership = new RetryOwnership(push);
+ }
+
+ @Override
+ public void run() {
+ if (!ownership.start()) {
+ return;
+ }
+ try (Activation ignored = ownership.push.activate()) {
+ super.run();
+ } finally {
+ ownership.finish();
+ }
+ }
+
+ @Override
+ public boolean cancel(boolean mayInterruptIfRunning) {
+ boolean cancelled = super.cancel(mayInterruptIfRunning);
+ if (cancelled) {
+ // FutureTask can become cancelled while user code is still running. Only a queued task
+ // may release here; a running task releases from run's finally block after it returns.
+ if (ownership.discard()) {
+ ownership.finish();
+ }
+ }
+ return cancelled;
+ }
+
+ @Override
+ public void discard() {
+ cancel(false);
+ }
+ }
+
+ private static final class TrackingRetryExecutor extends AbstractExecutorService {
+ private final ExecutorService delegate;
+ private final FactoryHook hook;
+
+ private TrackingRetryExecutor(ExecutorService delegate, FactoryHook hook) {
+ this.delegate = delegate;
+ this.hook = hook;
+ }
+
+ @Override
+ public void execute(Runnable command) {
+ if (command == null) {
+ throw new NullPointerException("command");
+ }
+ Push push = hook.current.get();
+ if (push == null) {
+ delegate.execute(command);
+ } else {
+ submitOwned(new TrackingRetryRunnable(command, push));
+ }
+ }
+
+ @Override
+ public Future<?> submit(Runnable task) {
+ return submit(task, null);
+ }
+
+ @Override
+ public <T> Future<T> submit(Runnable task, T result) {
+ Push push = hook.current.get();
+ if (push == null) {
+ return delegate.submit(task, result);
+ }
+ TrackingRetryFuture<T> future = new TrackingRetryFuture<>(task, result, push);
+ submitOwned(future);
+ return future;
+ }
+
+ @Override
+ public <T> Future<T> submit(Callable<T> task) {
+ Push push = hook.current.get();
+ if (push == null) {
+ return delegate.submit(task);
+ }
+ TrackingRetryFuture<T> future = new TrackingRetryFuture<>(task, push);
+ submitOwned(future);
+ return future;
+ }
+
+ private <T extends Runnable & OwnedRetry> void submitOwned(T task) {
+ try {
+ delegate.execute(task);
+ } catch (RuntimeException | Error failure) {
+ task.discard();
+ throw failure;
+ }
+ }
+
+ @Override
+ public void shutdown() {
+ // Stock Celeborn uses graceful shutdown: queued retries retain ownership until they drain.
+ delegate.shutdown();
+ }
+
+ @Override
+ public List<Runnable> shutdownNow() {
+ List<Runnable> queued = delegate.shutdownNow();
+ for (Runnable task : queued) {
+ if (task instanceof OwnedRetry) {
+ ((OwnedRetry) task).discard();
+ }
+ }
+ // Return cancelled wrappers, as with FutureTasks, without reviving the payload-owning
+ // original Runnable after its reservation was released. Unowned tasks are unchanged.
+ return queued;
+ }
+
+ @Override
+ public boolean isShutdown() {
+ return delegate.isShutdown();
+ }
+
+ @Override
+ public boolean isTerminated() {
+ return delegate.isTerminated();
+ }
+
+ @Override
+ public boolean awaitTermination(long timeout, TimeUnit unit) throws InterruptedException {
+ return delegate.awaitTermination(timeout, unit);
+ }
+ }
+
+ private static Object objectMethod(Object proxy, Method method, Object[] args) {
+ switch (method.getName()) {
+ case "equals":
+ return proxy == args[0];
+ case "hashCode":
+ return System.identityHashCode(proxy);
+ case "toString":
+ return "Comet Celeborn transport callback observer";
+ default:
+ throw new UnsupportedOperationException(method.toString());
+ }
+ }
+
+ /**
+ * Intercepts stock transport's put-based request publication. Other map operations delegate to
+ * the original backing store without introducing ownership for values that might not be inserted.
+ */
+ private static final class CallbackTrackingRequests extends ConcurrentHashMap<Object, Object> {
+ private final ConcurrentHashMap<Object, Object> delegate;
+ private final FactoryHook hook;
+
+ private CallbackTrackingRequests(ConcurrentHashMap<Object, Object> delegate, FactoryHook hook) {
+ this.delegate = delegate;
+ this.hook = hook;
+ }
+
+ @Override
+ public Object put(Object key, Object value) {
+ PreparedRequest prepared = hook.prepareRequest(value);
+ boolean published = false;
+ try {
+ Object previous = delegate.put(key, value);
+ published = true;
+ return previous;
+ } finally {
+ if (!published && prepared != null) {
+ prepared.discard();
+ }
+ }
+ }
+
+ @Override
+ public Object putIfAbsent(Object key, Object value) {
+ PreparedRequest prepared = hook.prepareRequest(value);
+ boolean published = false;
+ try {
+ Object previous = delegate.putIfAbsent(key, value);
+ published = previous == null;
+ return previous;
+ } finally {
+ if (!published && prepared != null) {
+ prepared.discard();
+ }
+ }
+ }
+
+ @Override
+ public void putAll(Map<?, ?> values) {
+ values.forEach(this::put);
+ }
+
+ @Override
+ public Object get(Object key) {
+ return delegate.get(key);
+ }
+
+ @Override
+ public Object getOrDefault(Object key, Object defaultValue) {
+ return delegate.getOrDefault(key, defaultValue);
+ }
+
+ @Override
+ public boolean containsKey(Object key) {
+ return delegate.containsKey(key);
+ }
+
+ @Override
+ public boolean containsValue(Object value) {
+ return delegate.containsValue(value);
+ }
+
+ @Override
+ public Object remove(Object key) {
+ return delegate.remove(key);
+ }
+
+ @Override
+ public boolean remove(Object key, Object value) {
+ return delegate.remove(key, value);
+ }
+
+ @Override
+ public Object replace(Object key, Object value) {
+ return delegate.replace(key, value);
+ }
+
+ @Override
+ public boolean replace(Object key, Object oldValue, Object newValue) {
+ return delegate.replace(key, oldValue, newValue);
+ }
+
+ @Override
+ public Object computeIfAbsent(Object key, Function<? super Object, ?> mappingFunction) {
+ return delegate.computeIfAbsent(key, mappingFunction);
+ }
+
+ @Override
+ public Object computeIfPresent(
+ Object key, BiFunction<? super Object, ? super Object, ?> remappingFunction) {
+ return delegate.computeIfPresent(key, remappingFunction);
+ }
+
+ @Override
+ public Object compute(
+ Object key, BiFunction<? super Object, ? super Object, ?> remappingFunction) {
+ return delegate.compute(key, remappingFunction);
+ }
+
+ @Override
+ public Object merge(
+ Object key, Object value, BiFunction<? super Object, ? super Object, ?> remappingFunction) {
+ return delegate.merge(key, value, remappingFunction);
+ }
+
+ @Override
+ public void replaceAll(BiFunction<? super Object, ? super Object, ?> function) {
+ delegate.replaceAll(function);
+ }
+
+ @Override
+ public void clear() {
+ delegate.clear();
+ }
+
+ @Override
+ public int size() {
+ return delegate.size();
+ }
+
+ @Override
+ public boolean isEmpty() {
+ return delegate.isEmpty();
+ }
+
+ @Override
+ public long mappingCount() {
+ return delegate.mappingCount();
+ }
+
+ @Override
+ public KeySetView<Object, Object> keySet() {
+ return delegate.keySet();
+ }
+
+ @Override
+ public Set<Map.Entry<Object, Object>> entrySet() {
+ return delegate.entrySet();
+ }
+
+ @Override
+ public Collection<Object> values() {
+ return delegate.values();
+ }
+
+ @Override
+ public void forEach(BiConsumer<? super Object, ? super Object> action) {
+ delegate.forEach(action);
+ }
+ }
+}
diff --git a/spark/src/main/java/org/apache/comet/shuffle/ExecutorShufflePushAdmission.java b/spark/src/main/java/org/apache/comet/shuffle/ExecutorShufflePushAdmission.java
new file mode 100644
index 0000000000..71d05b71e4
--- /dev/null
+++ b/spark/src/main/java/org/apache/comet/shuffle/ExecutorShufflePushAdmission.java
@@ -0,0 +1,103 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.comet.shuffle;
+
+import java.io.IOException;
+import java.util.IdentityHashMap;
+import java.util.concurrent.Semaphore;
+import java.util.concurrent.TimeUnit;
+import java.util.function.BooleanSupplier;
+
+/** Fair, executor-wide byte admission shared by tasks using the same Celeborn client. */
+final class ExecutorShufflePushAdmission {
+ // At most one entry exists for each application-owned Celeborn client. The shuffle manager
+ // explicitly removes its client on shutdown, so this registry never owns an application after
+ // the application's own client has been released.
+ private static final IdentityHashMap<Object, ExecutorShufflePushAdmission> CLIENTS =
+ new IdentityHashMap<>();
+
+ private final int limit;
+ private final Semaphore available;
+ private CelebornTransportCallbackTracker transportCallbacks;
+ private boolean transportCallbacksInitialized;
+
+ private ExecutorShufflePushAdmission(int limit) {
+ this.limit = limit;
+ this.available = new Semaphore(limit, true);
+ }
+
+ static ExecutorShufflePushAdmission forClient(Object client, int limit) {
+ synchronized (CLIENTS) {
+ ExecutorShufflePushAdmission admission = CLIENTS.get(client);
+ if (admission == null) {
+ admission = new ExecutorShufflePushAdmission(limit);
+ CLIENTS.put(client, admission);
+ } else if (admission.limit != limit) {
+ throw new IllegalArgumentException(
+ "Tasks sharing a Celeborn client must use the same in-flight byte limit");
+ }
+ return admission;
+ }
+ }
+
+ static void releaseClient(Object client) {
+ synchronized (CLIENTS) {
+ CLIENTS.remove(client);
+ }
+ }
+
+ synchronized CelebornTransportCallbackTracker transportCallbacks(Object client) {
+ if (!transportCallbacksInitialized) {
+ transportCallbacks = CelebornTransportCallbackTracker.tryCreate(client);
+ transportCallbacksInitialized = true;
+ }
+ return transportCallbacks;
+ }
+
+ void acquire(int bytes, BooleanSupplier cancelled) throws IOException {
+ if (bytes <= 0 || bytes > limit) {
+ throw new IOException(
+ "Celeborn push requires "
+ + bytes
+ + " bytes, exceeding the executor in-flight byte limit of "
+ + limit);
+ }
+ try {
+ while (!available.tryAcquire(bytes, 25, TimeUnit.MILLISECONDS)) {
+ if (cancelled.getAsBoolean()) {
+ throw new IOException("Celeborn shuffle map attempt was cancelled during admission");
+ }
+ }
+ if (cancelled.getAsBoolean()) {
+ available.release(bytes);
+ throw new IOException("Celeborn shuffle map attempt was cancelled during admission");
+ }
+ } catch (InterruptedException failure) {
+ Thread.currentThread().interrupt();
+ throw new IOException("Interrupted while waiting for Celeborn push admission", failure);
+ }
+ }
+
+ void release(int bytes) {
+ if (bytes > 0) {
+ available.release(bytes);
+ }
+ }
+}
diff --git a/spark/src/main/java/org/apache/comet/shuffle/ShufflePartitionPusher.java b/spark/src/main/java/org/apache/comet/shuffle/ShufflePartitionPusher.java
index 05771e921a..3de42b1e51 100644
--- a/spark/src/main/java/org/apache/comet/shuffle/ShufflePartitionPusher.java
+++ b/spark/src/main/java/org/apache/comet/shuffle/ShufflePartitionPusher.java
@@ -30,6 +30,26 @@ import java.io.IOException;
@FunctionalInterface
public interface ShufflePartitionPusher {
+ /** Reserves an upper bound before a native worker starts encoding a shuffle frame. */
+ default void reservePartitionData(int maxLength) throws IOException {}
+
+ /**
+ * Acknowledges that native encoding buffers and JNI local references have been released, on both
+ * success and failure. Implementations may retain admission until asynchronous pushes also
+ * finish.
+ */
+ default void releasePartitionDataReservation() {}
+
+ /** Returns the largest encoding reservation this callback can admit before allocating buffers. */
+ default int maxReservationBytes() {
+ return Integer.MAX_VALUE;
+ }
+
+ /** Returns the largest complete frame that this callback can safely accept. */
+ default int maxFrameBytes() {
+ return Integer.MAX_VALUE - 8;
+ }
+
/** Pushes one complete, length-prefixed Arrow IPC block for the given output partition. */
void pushPartitionData(int partitionId, byte[] data, int length) throws IOException;
}
diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala
index d8fe5b6989..7654b8949d 100644
--- a/spark/src/main/scala/org/apache/comet/CometConf.scala
+++ b/spark/src/main/scala/org/apache/comet/CometConf.scala
@@ -607,6 +607,34 @@ object CometConf extends ShimCometConf {
.checkValue(v => v >= 0, "Must not be negative")
.createWithDefault(0)
+ val COMET_SHUFFLE_RSS_MAX_FRAME_BYTES: ConfigEntry[Long] =
+ conf("spark.comet.shuffle.rss.maxFrameBytes")
+ .category(CATEGORY_SHUFFLE)
+ .doc("Maximum encoded size of one complete native shuffle frame sent to a remote " +
+ "shuffle service. Frames are never split across remote push requests.")
+ .bytesConf(ByteUnit.BYTE)
+ .checkValue(
+ value => value >= 20 && value <= Int.MaxValue - 16,
+ "Remote shuffle frame size must fit a complete Comet frame and a Celeborn request")
+ .createWithDefault(64L * 1024 * 1024)
+
+ val COMET_SHUFFLE_RSS_MAX_IN_FLIGHT_BYTES: ConfigEntry[Long] =
+ conf("spark.comet.shuffle.rss.maxInFlightBytes")
+ .category(CATEGORY_SHUFFLE)
+ .doc(
+ "Maximum shuffle bytes admitted concurrently by native Comet map attempts sharing " +
+ "an executor-side remote shuffle client. Admission includes native encoding " +
+ "scratch and overlapping native, JNI, and remote shuffle frame copies. " +
+ "A frame must fit its codec and Arrow workspace as well as its encoded bytes; " +
+ "too-small limits fail before encoding. Encrypted native RSS is not supported; " +
+ "use ordinary Spark shuffle when spark.io.encryption.enabled is true.")
+ .bytesConf(ByteUnit.BYTE)
+ .checkValue(
+ value => value >= 76 && value <= Int.MaxValue,
+ "Remote shuffle in-flight byte limit must fit three complete frame copies and a " +
+ "Celeborn request header")
+ .createWithDefault(256L * 1024 * 1024)
+
val COMET_DEBUG_ENABLED: ConfigEntry[Boolean] =
conf("spark.comet.debug.enabled")
.category(CATEGORY_EXEC)
diff --git a/spark/src/main/scala/org/apache/comet/shuffle/CelebornShufflePusherFactory.scala b/spark/src/main/scala/org/apache/comet/shuffle/CelebornShufflePusherFactory.scala
new file mode 100644
index 0000000000..37d67f55ca
--- /dev/null
+++ b/spark/src/main/scala/org/apache/comet/shuffle/CelebornShufflePusherFactory.scala
@@ -0,0 +1,380 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.comet.shuffle
+
+import java.io.IOException
+import java.lang.reflect.InvocationTargetException
+import java.util.Optional
+
+import org.apache.spark.{ShuffleDependency, SparkConf, TaskContext}
+import org.apache.spark.shuffle.ShuffleHandle
+import org.apache.spark.storage.BlockManagerId
+
+import org.apache.comet.CometConf
+import org.apache.comet.util.ClassLoaders
+
+/** Resolves an application's optional Celeborn client without adding a Celeborn dependency. */
+object CelebornShufflePusherFactory {
+
+ private val CELEBORN_SHUFFLE_HANDLE = "org.apache.spark.shuffle.celeborn.CelebornShuffleHandle"
+ private val CELEBORN_SPARK_UTILS = "org.apache.spark.shuffle.celeborn.SparkUtils"
+ private val CELEBORN_SPARK_COMMON_UTILS = "org.apache.spark.shuffle.celeborn.SparkCommonUtils"
+ private val CELEBORN_SHUFFLE_CLIENT = "org.apache.celeborn.client.ShuffleClient"
+ private val CELEBORN_CONF = "org.apache.celeborn.common.CelebornConf"
+ private val CELEBORN_USER_IDENTIFIER = "org.apache.celeborn.common.identity.UserIdentifier"
+ private val MAX_STAGE_ATTEMPTS = 1 << 15
+ private val MAX_TASK_ATTEMPTS = 1 << 16
+
+ private[shuffle] def encodeAttemptNumber(stageAttempt: Int, taskAttempt: Int): Int = {
+ require(
+ stageAttempt >= 0 && stageAttempt < MAX_STAGE_ATTEMPTS,
+ s"Celeborn stage attempt must be between 0 and ${MAX_STAGE_ATTEMPTS - 1}: $stageAttempt")
+ require(
+ taskAttempt >= 0 && taskAttempt < MAX_TASK_ATTEMPTS,
+ s"Celeborn task attempt must be between 0 and ${MAX_TASK_ATTEMPTS - 1}: $taskAttempt")
+ (stageAttempt << 16) | taskAttempt
+ }
+
+ /** Bind one already-resolved Celeborn generation to one Spark map attempt. */
+ def create(
+ conf: SparkConf,
+ client: AnyRef,
+ celebornShuffleId: Int,
+ numMappers: Int,
+ numPartitions: Int,
+ taskContext: TaskContext): CelebornShufflePartitionPusher = {
+ requireUnencryptedNativeShuffle(conf)
+ val frameEntry = CometConf.COMET_SHUFFLE_RSS_MAX_FRAME_BYTES
+ val maxFrameBytes = conf.getSizeAsBytes(frameEntry.key, frameEntry.defaultValue.get.toString)
+ require(
+ maxFrameBytes >= 20 && maxFrameBytes <= Int.MaxValue - 16,
+ "Celeborn frame bytes must fit a complete Comet frame and a request header")
+ val limitEntry = CometConf.COMET_SHUFFLE_RSS_MAX_IN_FLIGHT_BYTES
+ val maxInFlightBytes =
+ conf.getSizeAsBytes(limitEntry.key, limitEntry.defaultValue.get.toString)
+ require(
+ maxInFlightBytes >= 76 && maxInFlightBytes <= Int.MaxValue,
+ "Celeborn executor in-flight bytes must fit three complete frames and a request header")
+
+ new CelebornShufflePartitionPusher(
+ client,
+ celebornShuffleId,
+ taskContext.partitionId(),
+ encodeAttemptNumber(taskContext.stageAttemptNumber(), taskContext.attemptNumber()),
+ numMappers,
+ numPartitions,
+ maxFrameBytes.toInt,
+ maxInFlightBytes.toInt)
+ }
+
+ /**
+ * Reuse the Celeborn client described by Spark's actual Celeborn handle. The acquisition hook
+ * runs before generation resolution so manager shutdown can release clients even if setup
+ * fails.
+ */
+ def createFromHandle(
+ conf: SparkConf,
+ handle: ShuffleHandle,
+ taskContext: TaskContext,
+ onClientAcquired: AnyRef => Unit,
+ onShuffleGenerationResolved: (Int, Int) => Unit,
+ onShuffleGenerationInvalidated: (Int, Int) => Unit = (_, _) => (),
+ onShuffleGenerationInvalidationUnsafe: (Int, Int) => Boolean = (_, _) => false)
+ : ResolvedCelebornShufflePusher = {
+ requireUnencryptedNativeShuffle(conf)
+ try {
+ val handleClass = ClassLoaders.loadClass(CELEBORN_SHUFFLE_HANDLE)
+ require(
+ handleClass.isInstance(handle),
+ "Native Comet shuffle requires an actual Celeborn shuffle handle; " +
+ s"received ${handle.getClass.getName}")
+
+ val sparkUtilsClass = ClassLoaders.loadClass(CELEBORN_SPARK_UTILS)
+ val shuffleClientClass = ClassLoaders.loadClass(CELEBORN_SHUFFLE_CLIENT)
+ val celebornConfClass = ClassLoaders.loadClass(CELEBORN_CONF)
+ val userIdentifierClass = ClassLoaders.loadClass(CELEBORN_USER_IDENTIFIER)
+ val celebornConf =
+ sparkUtilsClass.getMethod("fromSparkConf", classOf[SparkConf]).invoke(null, conf)
+
+ def handleValue(name: String): AnyRef = handleClass.getMethod(name).invoke(handle)
+
+ val client = acquireClient(
+ resolveClient(
+ conf,
+ shuffleClientClass,
+ celebornConfClass,
+ userIdentifierClass,
+ Array[AnyRef](
+ handleValue("appUniqueId"),
+ handleValue("lifecycleManagerHost"),
+ handleValue("lifecycleManagerPort"),
+ celebornConf,
+ handleValue("userIdentifier"),
+ handleValue("extension")),
+ sparkConf =>
+ ClassLoaders
+ .loadClass(CELEBORN_SPARK_COMMON_UTILS)
+ .getMethod("getCryptoHandler", classOf[SparkConf])
+ .invoke(null, sparkConf)
+ .asInstanceOf[Optional[_]]),
+ onClientAcquired)
+
+ if (!handleValue("stageRerunEnabled").asInstanceOf[Boolean]) {
+ throw new IllegalStateException(
+ "Native Celeborn shuffle requires stage reruns to recover ambiguous map attempts")
+ }
+
+ registerBarrierFailureListener(
+ sparkUtilsClass,
+ shuffleClientClass,
+ handleClass,
+ client,
+ taskContext,
+ handle)
+
+ val celebornShuffleId = sparkUtilsClass
+ .getMethod(
+ "celebornShuffleId",
+ shuffleClientClass,
+ handleClass,
+ classOf[TaskContext],
+ classOf[java.lang.Boolean])
+ .invoke(null, client, handle, taskContext, java.lang.Boolean.TRUE)
+ .asInstanceOf[Int]
+ val numMappers = handleValue("numMappers").asInstanceOf[Int]
+ onShuffleGenerationResolved(celebornShuffleId, numMappers)
+
+ if (taskContext.attemptNumber() > 0) {
+ rejectRetriedAttempt(
+ client,
+ handle.shuffleId,
+ celebornShuffleId,
+ taskContext,
+ authorizedToCommit = true,
+ () => onShuffleGenerationInvalidated(handle.shuffleId, celebornShuffleId),
+ () => onShuffleGenerationInvalidationUnsafe(handle.shuffleId, celebornShuffleId))
+ }
+
+ val numPartitions = handleValue("dependency")
+ .asInstanceOf[ShuffleDependency[_, _, _]]
+ .partitioner
+ .numPartitions
+ ResolvedCelebornShufflePusher(
+ create(conf, client, celebornShuffleId, numMappers, numPartitions, taskContext),
+ client,
+ celebornShuffleId)
+ } catch {
+ case failure: InvocationTargetException =>
+ throw new IllegalStateException(
+ "Could not resolve the application-owned Celeborn shuffle client",
+ Option(failure.getCause).getOrElse(failure))
+ case failure: ReflectiveOperationException =>
+ throw new IllegalStateException(
+ "The Celeborn client does not expose the required native shuffle handle API",
+ failure)
+ }
+ }
+
+ private def requireUnencryptedNativeShuffle(conf: SparkConf): Unit = {
+ require(
+ !conf.getBoolean("spark.io.encryption.enabled", false),
+ "Encrypted native Celeborn shuffle is not supported by bounded push admission; " +
+ "use ordinary Spark shuffle when spark.io.encryption.enabled is true")
+ }
+
+ /** Preserve Spark's shuffle encryption when Celeborn exposes its crypto-aware client API. */
+ private[shuffle] def resolveClient(
+ conf: SparkConf,
+ shuffleClientClass: Class[_],
+ celebornConfClass: Class[_],
+ userIdentifierClass: Class[_],
+ arguments: Array[AnyRef],
+ resolveCryptoHandler: SparkConf => Optional[_]): AnyRef = {
+ val argumentTypes: Array[Class[_]] = Array(
+ classOf[String],
+ classOf[String],
+ java.lang.Integer.TYPE,
+ celebornConfClass,
+ userIdentifierClass,
+ classOf[Array[Byte]])
+
+ val cryptoAwareMethod =
+ try {
+ Some(shuffleClientClass.getMethod("get", (argumentTypes :+ classOf[Optional[_]]): _*))
+ } catch {
+ case _: NoSuchMethodException => None
+ }
+
+ cryptoAwareMethod
+ .map { method =>
+ method.invoke(null, (arguments :+ resolveCryptoHandler(conf)): _*)
+ }
+ .getOrElse(
+ shuffleClientClass.getMethod("get", argumentTypes: _*).invoke(null, arguments: _*))
+ .asInstanceOf[AnyRef]
+ }
+
+ private[shuffle] def acquireClient(
+ resolveClient: => AnyRef,
+ onClientAcquired: AnyRef => Unit): AnyRef = {
+ val client = resolveClient
+ onClientAcquired(client)
+ client
+ }
+
+ private[shuffle] def registerBarrierFailureListener(
+ sparkUtilsClass: Class[_],
+ shuffleClientClass: Class[_],
+ handleClass: Class[_],
+ client: AnyRef,
+ taskContext: TaskContext,
+ handle: ShuffleHandle): Unit = {
+ sparkUtilsClass
+ .getMethod(
+ "addFailureListenerIfBarrierTask",
+ shuffleClientClass,
+ classOf[TaskContext],
+ handleClass)
+ .invoke(null, client, taskContext, handle)
+ }
+
+ /** Check driver-side task liveness before invalidating a Celeborn generation. */
+ def shouldReportShuffleFetchFailure(taskAttemptId: Long): Boolean =
+ ClassLoaders
+ .loadClass(CELEBORN_SPARK_UTILS)
+ .getMethod("shouldReportShuffleFetchFailure", java.lang.Long.TYPE)
+ .invoke(null, Long.box(taskAttemptId))
+ .asInstanceOf[Boolean]
+
+ private[shuffle] def rejectRetriedAttempt(
+ client: AnyRef,
+ sparkShuffleId: Int,
+ celebornShuffleId: Int,
+ taskContext: TaskContext,
+ authorizedToCommit: Boolean,
+ onGenerationInvalidated: () => Unit = () => (),
+ onGenerationInvalidationUnsafe: () => Boolean = () => false): Unit = {
+ if (!authorizedToCommit || onGenerationInvalidationUnsafe()) {
+ throw commitDenied(taskContext)
+ }
+
+ val invalidated = client.getClass
+ .getMethod(
+ "reportShuffleFetchFailure",
+ java.lang.Integer.TYPE,
+ java.lang.Integer.TYPE,
+ java.lang.Long.TYPE)
+ .invoke(
+ client,
+ Int.box(sparkShuffleId),
+ Int.box(celebornShuffleId),
+ Long.box(taskContext.taskAttemptId()))
+ .asInstanceOf[Boolean]
+
+ if (!invalidated) {
+ throw new IOException(
+ "Could not invalidate the Celeborn shuffle generation for a retried map attempt")
+ }
+ onGenerationInvalidated()
+
+ val failure = ClassLoaders
+ .loadClass("org.apache.spark.shuffle.FetchFailedException")
+ .getConstructor(
+ classOf[BlockManagerId],
+ java.lang.Integer.TYPE,
+ java.lang.Long.TYPE,
+ java.lang.Integer.TYPE,
+ java.lang.Integer.TYPE,
+ classOf[String],
+ classOf[Throwable])
+ .newInstance(
+ null,
+ Int.box(sparkShuffleId),
+ Long.box(-1L),
+ Int.box(-1),
+ Int.box(-1),
+ "Retried Celeborn map attempt requires a new shuffle generation: " +
+ s"$sparkShuffleId/$celebornShuffleId",
+ null)
+ .asInstanceOf[Throwable]
+ throw failure
+ }
+
+ /** Preserve Spark's non-counted failure classification for speculative losers. */
+ def commitDenied(taskContext: TaskContext): Throwable =
+ ClassLoaders
+ .loadClass("org.apache.spark.executor.CommitDeniedException")
+ .getConstructor(
+ classOf[String],
+ java.lang.Integer.TYPE,
+ java.lang.Integer.TYPE,
+ java.lang.Integer.TYPE)
+ .newInstance(
+ "Another Spark map attempt already owns the Celeborn shuffle commit",
+ Int.box(taskContext.stageId()),
+ Int.box(taskContext.partitionId()),
+ Int.box(taskContext.attemptNumber()))
+ .asInstanceOf[Throwable]
+
+ /** Remove task-independent state for one unregistered Celeborn generation. */
+ def cleanupShuffle(client: AnyRef, celebornShuffleId: Int): Unit = {
+ try {
+ client.getClass
+ .getMethod("cleanupShuffle", java.lang.Integer.TYPE)
+ .invoke(client, Int.box(celebornShuffleId))
+ } catch {
+ case failure: InvocationTargetException =>
+ throw new IllegalStateException(
+ "Could not clean up the Celeborn shuffle generation",
+ Option(failure.getCause).getOrElse(failure))
+ case failure: ReflectiveOperationException =>
+ throw new IllegalStateException("Celeborn shuffle cleanup is unavailable", failure)
+ }
+ }
+
+ /** Release only application clients that this composite manager actually acquired. */
+ def releaseClient(client: AnyRef): Unit = {
+ try {
+ val shuffleClientClass = ClassLoaders.loadClass(CELEBORN_SHUFFLE_CLIENT)
+ try {
+ shuffleClientClass.getMethod("removeInstance", shuffleClientClass).invoke(null, client)
+ } catch {
+ case _: NoSuchMethodException =>
+ // Celeborn 0.6 owns one shared application client and cannot remove a single instance.
+ // Its global reset would also invalidate ordinary Spark shuffle, so leave it untouched.
+ }
+ } catch {
+ case failure: InvocationTargetException =>
+ throw new IllegalStateException(
+ "Could not release the application-owned Celeborn shuffle client",
+ Option(failure.getCause).getOrElse(failure))
+ case failure: ReflectiveOperationException =>
+ throw new IllegalStateException("Celeborn shuffle client cleanup is unavailable", failure)
+ } finally {
+ ExecutorShufflePushAdmission.releaseClient(client)
+ }
+ }
+}
+
+/** A task-owned pusher, its application-owned client, and the resolved shuffle generation. */
+final case class ResolvedCelebornShufflePusher(
+ pusher: CelebornShufflePartitionPusher,
+ client: AnyRef,
+ celebornShuffleId: Int)
diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleManager.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleManager.scala
index a3759df11e..edafa8d9a3 100644
--- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleManager.scala
+++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleManager.scala
@@ -20,18 +20,27 @@
package org.apache.spark.sql.comet.execution.shuffle
import java.lang.reflect.InvocationTargetException
+import java.util.concurrent.ConcurrentHashMap
-import org.apache.spark.{ShuffleDependency, SparkConf, TaskContext}
+import scala.collection.mutable
+import scala.jdk.CollectionConverters._
+
+import org.apache.spark.{ShuffleDependency, SparkConf, SparkEnv, TaskContext}
+import org.apache.spark.rpc.{RpcCallContext, RpcEndpointRef, RpcEnv, ThreadSafeRpcEndpoint}
+import org.apache.spark.scheduler.OutputCommitCoordinator
import org.apache.spark.shuffle.{BaseShuffleHandle, ShuffleBlockResolver, ShuffleHandle, ShuffleManager, ShuffleReader, ShuffleReadMetricsReporter, ShuffleWriteMetricsReporter, ShuffleWriter}
+import org.apache.spark.util.RpcUtils
+import org.apache.comet.CometConf
+import org.apache.comet.shuffle.{CelebornShufflePartitionPusher, CelebornShufflePusherFactory}
import org.apache.comet.util.ClassLoaders
/**
- * Preserves the application's existing Celeborn shuffle manager for ordinary Spark shuffles.
+ * Preserves the application's existing Celeborn shuffle manager for ordinary Spark shuffles while
+ * binding native Comet map output to Celeborn's application-owned client.
*
* Celeborn is an optional, application-provided dependency, so its shuffle manager is loaded
- * reflectively. Comet shuffle dependencies remain unsupported until subsequent changes add their
- * remote map-side lifecycle and reduce-side reader.
+ * reflectively. Native reduce-side reads and JVM Comet shuffles remain unsupported.
*/
class CometCelebornShuffleManager private[shuffle] (
conf: SparkConf,
@@ -46,11 +55,31 @@ class CometCelebornShuffleManager private[shuffle] (
private val backend = Option(backendFactory(conf, isDriver)).getOrElse {
throw new IllegalStateException("Celeborn Spark shuffle manager factory returned null")
}
+ private val nativeShuffleClients =
+ new ConcurrentHashMap[Int, ConcurrentHashMap[Int, AnyRef]]()
+ private val ownedNativeClients = new ConcurrentHashMap[AnyRef, java.lang.Boolean]()
+ @volatile private var nativeGenerationCoordinator: CelebornShuffleGenerationCoordinator = _
+ @volatile private var nativeGenerationEndpoint: RpcEndpointRef = _
override def registerShuffle[K, V, C](
shuffleId: Int,
dependency: ShuffleDependency[K, V, C]): ShuffleHandle = {
dependency match {
+ case native: CometShuffleDependency[_, _, _] if native.shuffleType == CometNativeShuffle =>
+ val handle = backend.registerShuffle(shuffleId, dependency)
+ if (!CometCelebornShuffleManager.isCelebornHandle(handle)) {
+ val failure = new UnsupportedOperationException(
+ "Native Comet shuffle cannot use Celeborn's local fallback writer")
+ try backend.unregisterShuffle(shuffleId)
+ catch {
+ case cleanupFailure: Throwable => failure.addSuppressed(cleanupFailure)
+ }
+ throw failure
+ }
+ if (isDriver) {
+ initializeNativeGenerationCoordinator()
+ }
+ handle
case _: CometShuffleDependency[_, _, _] => rejectCometShuffle()
case _ => backend.registerShuffle(shuffleId, dependency)
}
@@ -61,8 +90,76 @@ class CometCelebornShuffleManager private[shuffle] (
mapId: Long,
context: TaskContext,
metrics: ShuffleWriteMetricsReporter): ShuffleWriter[K, V] = {
- rejectCometHandle(handle)
- backend.getWriter(handle, mapId, context, metrics)
+ nativeDependency(handle) match {
+ case Some(dependency) =>
+ val earlyClaim = claimNativeShuffleAttempt(handle.shuffleId, context)
+ if (!earlyClaim.authorized && context.attemptNumber() > 0 &&
+ !earlyClaim.requiresGenerationResolution) {
+ throw CelebornShufflePusherFactory.commitDenied(context)
+ }
+ var preparedClaim = earlyClaim
+ val resolved = CelebornShufflePusherFactory.createFromHandle(
+ conf,
+ handle,
+ context,
+ client => ownedNativeClients.put(client, java.lang.Boolean.TRUE),
+ (celebornShuffleId, numMappers) =>
+ preparedClaim = prepareNativeShuffleGeneration(
+ handle.shuffleId,
+ celebornShuffleId,
+ numMappers,
+ context,
+ earlyClaim),
+ (sparkShuffleId, celebornShuffleId) =>
+ invalidateNativeShuffleGeneration(
+ sparkShuffleId,
+ celebornShuffleId,
+ context,
+ preparedClaim),
+ (sparkShuffleId, celebornShuffleId) =>
+ abandonNativeShuffleAttempt(
+ sparkShuffleId,
+ celebornShuffleId,
+ context,
+ preparedClaim))
+ nativeShuffleClients
+ .computeIfAbsent(handle.shuffleId, _ => new ConcurrentHashMap[Int, AnyRef]())
+ .put(resolved.celebornShuffleId, resolved.client)
+
+ new CometNativeShuffleWriter[K, V](
+ dependency.nativeShuffleSpec.getOrElse {
+ throw new IllegalStateException("Native Comet shuffle has no execution plan")
+ },
+ dependency.outputPartitioning.getOrElse {
+ throw new IllegalStateException("Native Comet shuffle has no output partitioning")
+ },
+ dependency.outputAttributes,
+ dependency.shuffleWriteMetrics,
+ dependency.numParts,
+ dependency.shuffleId,
+ mapId,
+ context,
+ metrics,
+ dependency.rangePartitionBounds,
+ Some(
+ CelebornNativeShuffleDestination(
+ resolved.pusher,
+ CometCelebornShuffleManager.maxNativeFrameBytes(
+ CometConf.COMET_SHUFFLE_RSS_MAX_FRAME_BYTES.get().toInt,
+ resolved.pusher),
+ dependency.partitioner.numPartitions,
+ commitAuthorized = true,
+ commitValidator = () =>
+ validateNativeShuffleAttempt(
+ handle.shuffleId,
+ resolved.celebornShuffleId,
+ context,
+ preparedClaim))))
+
+ case None =>
+ rejectCometHandle(handle)
+ backend.getWriter(handle, mapId, context, metrics)
+ }
}
// Spark's final all-mapper reader overload delegates to this mapper-range overload.
@@ -74,6 +171,9 @@ class CometCelebornShuffleManager private[shuffle] (
endPartition: Int,
context: TaskContext,
metrics: ShuffleReadMetricsReporter): ShuffleReader[K, C] = {
+ if (nativeDependency(handle).nonEmpty) {
+ rejectCometShuffle()
+ }
rejectCometHandle(handle)
backend.getReader(
handle,
@@ -87,9 +187,160 @@ class CometCelebornShuffleManager private[shuffle] (
override def shuffleBlockResolver: ShuffleBlockResolver = backend.shuffleBlockResolver
- override def unregisterShuffle(shuffleId: Int): Boolean = backend.unregisterShuffle(shuffleId)
+ override def unregisterShuffle(shuffleId: Int): Boolean = {
+ Option(nativeShuffleClients.remove(shuffleId)).foreach { generations =>
+ generations.forEach { (celebornShuffleId, client) =>
+ CelebornShufflePusherFactory.cleanupShuffle(client, celebornShuffleId)
+ }
+ }
+ if (isDriver) {
+ Option(nativeGenerationCoordinator).foreach(_.unregisterShuffle(shuffleId))
+ }
+ backend.unregisterShuffle(shuffleId)
+ }
+
+ override def stop(): Unit = {
+ try backend.stop()
+ finally {
+ try {
+ if (isDriver) {
+ Option(nativeGenerationEndpoint).foreach { endpoint =>
+ SparkEnv.get.rpcEnv.stop(endpoint)
+ }
+ }
+ } finally {
+ ownedNativeClients.keySet().asScala.foreach(CelebornShufflePusherFactory.releaseClient)
+ ownedNativeClients.clear()
+ nativeShuffleClients.clear()
+ }
+ }
+ }
+
+ private def initializeNativeGenerationCoordinator(): Unit = synchronized {
+ if (nativeGenerationEndpoint == null) {
+ val env = Option(SparkEnv.get).getOrElse {
+ throw new IllegalStateException("Spark environment is unavailable for native shuffle")
+ }
+ val coordinator = new CelebornShuffleGenerationCoordinator(
+ env.outputCommitCoordinator,
+ CelebornShufflePusherFactory.shouldReportShuffleFetchFailure)
+ val endpoint = env.rpcEnv.setupEndpoint(
+ CometCelebornShuffleManager.GENERATION_COORDINATOR_ENDPOINT,
+ new CelebornShuffleGenerationEndpoint(env.rpcEnv, coordinator))
+ nativeGenerationCoordinator = coordinator
+ nativeGenerationEndpoint = endpoint
+ }
+ }
- override def stop(): Unit = backend.stop()
+ private def generationEndpoint: RpcEndpointRef =
+ Option(nativeGenerationEndpoint).getOrElse {
+ synchronized {
+ if (nativeGenerationEndpoint == null) {
+ if (isDriver) {
+ initializeNativeGenerationCoordinator()
+ } else {
+ nativeGenerationEndpoint = RpcUtils.makeDriverRef(
+ CometCelebornShuffleManager.GENERATION_COORDINATOR_ENDPOINT,
+ conf,
+ SparkEnv.get.rpcEnv)
+ }
+ }
+ nativeGenerationEndpoint
+ }
+ }
+
+ private def claimNativeShuffleAttempt(
+ shuffleId: Int,
+ taskContext: TaskContext): CelebornMapAttemptClaim =
+ generationEndpoint.askSync[CelebornMapAttemptClaim](
+ ClaimCelebornMapAttempt(
+ shuffleId,
+ taskContext.stageId(),
+ taskContext.stageAttemptNumber(),
+ taskContext.partitionId(),
+ taskContext.attemptNumber()))
+
+ private def prepareNativeShuffleGeneration(
+ shuffleId: Int,
+ celebornShuffleId: Int,
+ numMappers: Int,
+ taskContext: TaskContext,
+ earlyClaim: CelebornMapAttemptClaim): CelebornMapAttemptClaim = {
+ val prepared = generationEndpoint.askSync[CelebornMapAttemptClaim](
+ PrepareCelebornShuffleGeneration(
+ shuffleId,
+ celebornShuffleId,
+ taskContext.stageId(),
+ taskContext.stageAttemptNumber(),
+ numMappers,
+ taskContext.partitionId(),
+ taskContext.attemptNumber(),
+ earlyClaim.epoch,
+ earlyClaim.authorized))
+ if (!prepared.authorized) {
+ throw CelebornShufflePusherFactory.commitDenied(taskContext)
+ }
+ prepared
+ }
+
+ private def validateNativeShuffleAttempt(
+ shuffleId: Int,
+ celebornShuffleId: Int,
+ taskContext: TaskContext,
+ claim: CelebornMapAttemptClaim): Boolean =
+ generationEndpoint.askSync[Boolean](
+ ValidateCelebornMapAttempt(
+ shuffleId,
+ celebornShuffleId,
+ taskContext.stageId(),
+ taskContext.stageAttemptNumber(),
+ taskContext.partitionId(),
+ taskContext.attemptNumber(),
+ claim.epoch))
+
+ private def invalidateNativeShuffleGeneration(
+ shuffleId: Int,
+ celebornShuffleId: Int,
+ taskContext: TaskContext,
+ claim: CelebornMapAttemptClaim): Unit = {
+ generationEndpoint.askSync[Boolean](
+ InvalidateCelebornShuffleGeneration(
+ shuffleId,
+ celebornShuffleId,
+ taskContext.stageId(),
+ taskContext.stageAttemptNumber(),
+ claim.epoch))
+ ()
+ }
+
+ private def abandonNativeShuffleAttempt(
+ shuffleId: Int,
+ celebornShuffleId: Int,
+ taskContext: TaskContext,
+ claim: CelebornMapAttemptClaim): Boolean =
+ generationEndpoint.askSync[Boolean](
+ AbandonCelebornMapAttempt(
+ shuffleId,
+ celebornShuffleId,
+ taskContext.stageId(),
+ taskContext.stageAttemptNumber(),
+ taskContext.partitionId(),
+ taskContext.attemptNumber(),
+ claim.epoch,
+ taskContext.taskAttemptId()))
+
+ private def nativeDependency(handle: ShuffleHandle): Option[CometShuffleDependency[_, _, _]] = {
+ if (!CometCelebornShuffleManager.isCelebornHandle(handle)) {
+ None
+ } else {
+ handle.asInstanceOf[BaseShuffleHandle[_, _, _]].dependency match {
+ case dependency: CometShuffleDependency[_, _, _]
+ if dependency.shuffleType == CometNativeShuffle =>
+ Some(dependency)
+ case _ => None
+ }
+ }
+ }
private def rejectCometHandle(handle: ShuffleHandle): Unit = handle match {
case _: CometNativeShuffleHandle[_, _] => rejectCometShuffle()
@@ -103,13 +354,24 @@ class CometCelebornShuffleManager private[shuffle] (
private def rejectCometShuffle(): Nothing = {
throw new UnsupportedOperationException(
- "Comet shuffle over Celeborn is not supported until native shuffle integration is complete")
+ "Comet shuffle over Celeborn is not supported for JVM shuffle, local handles, or reads")
}
}
private[shuffle] object CometCelebornShuffleManager {
+ private[shuffle] val GENERATION_COORDINATOR_ENDPOINT =
+ "CometCelebornShuffleGenerationCoordinator"
private val CELEBORN_MANAGER_CLASS = "org.apache.spark.shuffle.celeborn.SparkShuffleManager"
+ private val CELEBORN_HANDLE_CLASS = "org.apache.spark.shuffle.celeborn.CelebornShuffleHandle"
+
+ private[shuffle] def maxNativeFrameBytes(
+ configuredMaxFrameBytes: Int,
+ pusher: CelebornShufflePartitionPusher): Int =
+ math.min(configuredMaxFrameBytes, pusher.maxFrameBytes())
+
+ private[shuffle] def isCelebornHandle(handle: ShuffleHandle): Boolean =
+ handle != null && handle.getClass.getName == CELEBORN_HANDLE_CLASS
private[shuffle] def createBackend(conf: SparkConf, isDriver: Boolean): ShuffleManager = {
try {
@@ -143,3 +405,378 @@ private[shuffle] object CometCelebornShuffleManager {
}
}
}
+
+private[shuffle] final case class PrepareCelebornShuffleGeneration(
+ shuffleId: Int,
+ celebornShuffleId: Int,
+ stageId: Int,
+ stageAttempt: Int,
+ numMappers: Int,
+ mapId: Int = -1,
+ taskAttempt: Int = -1,
+ claimEpoch: Long = -1L,
+ claimAuthorized: Boolean = false)
+ extends Serializable
+
+private[shuffle] final case class ClaimCelebornMapAttempt(
+ shuffleId: Int,
+ stageId: Int,
+ stageAttempt: Int,
+ mapId: Int,
+ taskAttempt: Int)
+ extends Serializable
+
+private[shuffle] final case class CelebornMapAttemptClaim(
+ authorized: Boolean,
+ epoch: Long,
+ requiresGenerationResolution: Boolean = false)
+ extends Serializable
+
+private[shuffle] final case class ValidateCelebornMapAttempt(
+ shuffleId: Int,
+ celebornShuffleId: Int,
+ stageId: Int,
+ stageAttempt: Int,
+ mapId: Int,
+ taskAttempt: Int,
+ claimEpoch: Long)
+ extends Serializable
+
+private[shuffle] final case class InvalidateCelebornShuffleGeneration(
+ shuffleId: Int,
+ celebornShuffleId: Int,
+ stageId: Int,
+ stageAttempt: Int,
+ claimEpoch: Long)
+ extends Serializable
+
+private[shuffle] final case class AbandonCelebornMapAttempt(
+ shuffleId: Int,
+ celebornShuffleId: Int,
+ stageId: Int,
+ stageAttempt: Int,
+ mapId: Int,
+ taskAttempt: Int,
+ claimEpoch: Long,
+ taskAttemptId: Long)
+ extends Serializable
+
+/**
+ * Keeps Comet map admission aligned with Spark task failures and Celeborn shuffle generations.
+ */
+private[shuffle] final class CelebornShuffleGenerationCoordinator(
+ outputCommitCoordinator: OutputCommitCoordinator,
+ shouldReportShuffleFetchFailure: Long => Boolean = _ => true) {
+
+ private val generations = mutable.HashMap.empty[Int, PrepareCelebornShuffleGeneration]
+ private val invalidatedGenerations = mutable.HashSet.empty[Int]
+ private val generationEpochs = mutable.HashMap.empty[Int, Long]
+ private val claimOwners = mutable.HashMap.empty[(Int, Int, Int, Int), (Int, Long)]
+ private val deniedAttempts =
+ mutable.HashMap.empty[(Int, Int, Int, Int), mutable.HashSet[Int]]
+
+ private def currentEpoch(shuffleId: Int): Long = generationEpochs.getOrElse(shuffleId, 0L)
+
+ private def ownerKey(
+ shuffleId: Int,
+ stageId: Int,
+ stageAttempt: Int,
+ mapId: Int): (Int, Int, Int, Int) =
+ (shuffleId, stageId, stageAttempt, mapId)
+
+ private def authorize(
+ shuffleId: Int,
+ stageId: Int,
+ stageAttempt: Int,
+ mapId: Int,
+ taskAttempt: Int): CelebornMapAttemptClaim = {
+ val epoch = currentEpoch(shuffleId)
+ val key = ownerKey(shuffleId, stageId, stageAttempt, mapId)
+ claimOwners.get(key).foreach { case (owner, ownerEpoch) =>
+ if (ownerEpoch != epoch || !attemptCanRun(stageId, stageAttempt, mapId, owner)) {
+ claimOwners.remove(key)
+ }
+ }
+ val authorized =
+ attemptCanRun(stageId, stageAttempt, mapId, taskAttempt) &&
+ !deniedAttempts.get(key).exists(_.contains(taskAttempt)) &&
+ claimOwners.get(key).forall(_ == ((taskAttempt, epoch)))
+ if (authorized) {
+ claimOwners.update(key, (taskAttempt, epoch))
+ }
+ CelebornMapAttemptClaim(authorized, epoch)
+ }
+
+ private def invalidateOwners(shuffleId: Int): Unit = {
+ generationEpochs.update(shuffleId, currentEpoch(shuffleId) + 1L)
+ claimOwners.retain { case ((ownerShuffleId, _, _, _), _) => ownerShuffleId != shuffleId }
+ deniedAttempts.retain { case ((ownerShuffleId, _, _, _), _) => ownerShuffleId != shuffleId }
+ }
+
+ private def attemptCanRun(
+ stageId: Int,
+ stageAttempt: Int,
+ mapId: Int,
+ taskAttempt: Int): Boolean =
+ outputCommitCoordinator.synchronized {
+ // Spark records task failures here even when no file-output commit was requested. Read that
+ // record without acquiring its commit slot: Spark 3.4 aborts the whole stage if an authorized
+ // file committer fails, including executor loss before native map output has been written.
+ val stageStatesField = classOf[OutputCommitCoordinator].getDeclaredField("stageStates")
+ stageStatesField.setAccessible(true)
+ val stageStates = stageStatesField
+ .get(outputCommitCoordinator)
+ .asInstanceOf[mutable.Map[Int, AnyRef]]
+ stageStates.get(stageId).exists { stageState =>
+ val owners = stageState.getClass
+ .getMethod("authorizedCommitters")
+ .invoke(stageState)
+ .asInstanceOf[Array[AnyRef]]
+ val failures = stageState.getClass
+ .getMethod("failures")
+ .invoke(stageState)
+ .asInstanceOf[mutable.Map[Int, mutable.Set[AnyRef]]]
+ mapId >= 0 && mapId < owners.length && !failures.get(mapId).exists { attempts =>
+ attempts.exists { attempt =>
+ attempt.getClass.getMethod("stageAttempt").invoke(attempt).asInstanceOf[Int] ==
+ stageAttempt &&
+ attempt.getClass.getMethod("taskAttempt").invoke(attempt).asInstanceOf[Int] ==
+ taskAttempt
+ }
+ }
+ }
+ }
+
+ def claimMapAttempt(claim: ClaimCelebornMapAttempt): CelebornMapAttemptClaim = synchronized {
+ val previousGeneration = generations.get(claim.shuffleId)
+ val stale = previousGeneration.exists { generation =>
+ generation.stageId == claim.stageId &&
+ (generation.stageAttempt > claim.stageAttempt ||
+ (generation.stageAttempt == claim.stageAttempt &&
+ invalidatedGenerations.contains(claim.shuffleId)))
+ }
+ if (stale) {
+ CelebornMapAttemptClaim(false, currentEpoch(claim.shuffleId))
+ } else {
+ val key = ownerKey(claim.shuffleId, claim.stageId, claim.stageAttempt, claim.mapId)
+ val epoch = currentEpoch(claim.shuffleId)
+ if (claimOwners.get(key).contains((claim.taskAttempt, epoch))) {
+ return authorize(
+ claim.shuffleId,
+ claim.stageId,
+ claim.stageAttempt,
+ claim.mapId,
+ claim.taskAttempt)
+ }
+
+ // Spark resolves input before asking for a writer. Reserve an earlier live attempt even if
+ // its speculative copy reaches this manager first; rejected attempts must not become owners.
+ val alreadyDenied = deniedAttempts.getOrElse(key, mutable.HashSet.empty[Int])
+ var candidate = 0
+ while (candidate < claim.taskAttempt) {
+ if (!alreadyDenied.contains(candidate) &&
+ !claimOwners.get(key).contains((candidate, epoch))) {
+ val earlier =
+ authorize(claim.shuffleId, claim.stageId, claim.stageAttempt, claim.mapId, candidate)
+ if (earlier.authorized) {
+ deniedAttempts.getOrElseUpdate(key, mutable.HashSet.empty[Int]).add(claim.taskAttempt)
+ return CelebornMapAttemptClaim(false, epoch)
+ }
+ }
+ candidate += 1
+ }
+
+ val result = authorize(
+ claim.shuffleId,
+ claim.stageId,
+ claim.stageAttempt,
+ claim.mapId,
+ claim.taskAttempt)
+ if (!result.authorized && claim.taskAttempt > 0) {
+ val requiresResolution = previousGeneration.exists { generation =>
+ generation.stageId == claim.stageId && generation.stageAttempt < claim.stageAttempt
+ } && !claimOwners.get(key).exists { case (_, ownerEpoch) => ownerEpoch == epoch }
+ if (!requiresResolution) {
+ deniedAttempts.getOrElseUpdate(key, mutable.HashSet.empty[Int]).add(claim.taskAttempt)
+ }
+ result.copy(requiresGenerationResolution = requiresResolution)
+ } else {
+ result
+ }
+ }
+ }
+
+ def prepareGeneration(generation: PrepareCelebornShuffleGeneration): Boolean = synchronized {
+ require(generation.numMappers > 0, "Celeborn shuffle mapper count must be positive")
+
+ generations.get(generation.shuffleId) match {
+ case Some(previous)
+ if invalidatedGenerations.contains(generation.shuffleId) &&
+ previous.celebornShuffleId == generation.celebornShuffleId =>
+ false
+ case Some(previous)
+ if previous.stageId == generation.stageId &&
+ previous.stageAttempt > generation.stageAttempt =>
+ false
+ case Some(previous)
+ if previous.stageId == generation.stageId &&
+ previous.stageAttempt == generation.stageAttempt =>
+ previous.celebornShuffleId == generation.celebornShuffleId
+ case Some(previous) if previous.celebornShuffleId != generation.celebornShuffleId =>
+ invalidateOwners(generation.shuffleId)
+ invalidatedGenerations.remove(generation.shuffleId)
+ generations.update(generation.shuffleId, generation)
+ true
+ case Some(_) =>
+ generations.update(generation.shuffleId, generation)
+ true
+ case None =>
+ generations.update(generation.shuffleId, generation)
+ true
+ }
+ }
+
+ def prepareGenerationAndClaim(
+ generation: PrepareCelebornShuffleGeneration): CelebornMapAttemptClaim = synchronized {
+ if (!prepareGeneration(generation)) {
+ deniedAttempts
+ .getOrElseUpdate(
+ ownerKey(
+ generation.shuffleId,
+ generation.stageId,
+ generation.stageAttempt,
+ generation.mapId),
+ mutable.HashSet.empty[Int])
+ .add(generation.taskAttempt)
+ return CelebornMapAttemptClaim(false, currentEpoch(generation.shuffleId))
+ }
+
+ val epoch = currentEpoch(generation.shuffleId)
+ val expectedOwner = claimOwners.get(
+ ownerKey(
+ generation.shuffleId,
+ generation.stageId,
+ generation.stageAttempt,
+ generation.mapId))
+ if (generation.claimAuthorized && generation.claimEpoch == epoch &&
+ expectedOwner.contains((generation.taskAttempt, epoch)) &&
+ attemptCanRun(
+ generation.stageId,
+ generation.stageAttempt,
+ generation.mapId,
+ generation.taskAttempt)) {
+ CelebornMapAttemptClaim(true, epoch)
+ } else {
+ val claim = authorize(
+ generation.shuffleId,
+ generation.stageId,
+ generation.stageAttempt,
+ generation.mapId,
+ generation.taskAttempt)
+ if (!claim.authorized) {
+ deniedAttempts
+ .getOrElseUpdate(
+ ownerKey(
+ generation.shuffleId,
+ generation.stageId,
+ generation.stageAttempt,
+ generation.mapId),
+ mutable.HashSet.empty[Int])
+ .add(generation.taskAttempt)
+ }
+ claim
+ }
+ }
+
+ def validateMapAttempt(validation: ValidateCelebornMapAttempt): Boolean = synchronized {
+ generations.get(validation.shuffleId).exists { generation =>
+ !invalidatedGenerations.contains(validation.shuffleId) &&
+ generation.celebornShuffleId == validation.celebornShuffleId &&
+ generation.stageId == validation.stageId &&
+ generation.stageAttempt == validation.stageAttempt &&
+ currentEpoch(validation.shuffleId) == validation.claimEpoch &&
+ attemptCanRun(
+ validation.stageId,
+ validation.stageAttempt,
+ validation.mapId,
+ validation.taskAttempt) &&
+ claimOwners
+ .get(
+ ownerKey(
+ validation.shuffleId,
+ validation.stageId,
+ validation.stageAttempt,
+ validation.mapId))
+ .contains((validation.taskAttempt, validation.claimEpoch))
+ }
+ }
+
+ def invalidateGeneration(invalidation: InvalidateCelebornShuffleGeneration): Boolean =
+ synchronized {
+ val current = generations.get(invalidation.shuffleId).exists { generation =>
+ generation.celebornShuffleId == invalidation.celebornShuffleId &&
+ generation.stageId == invalidation.stageId &&
+ generation.stageAttempt == invalidation.stageAttempt &&
+ currentEpoch(invalidation.shuffleId) == invalidation.claimEpoch
+ }
+ if (current) {
+ invalidatedGenerations.add(invalidation.shuffleId)
+ invalidateOwners(invalidation.shuffleId)
+ }
+ current
+ }
+
+ def abandonMapAttempt(abandoned: AbandonCelebornMapAttempt): Boolean = synchronized {
+ val current = validateMapAttempt(
+ ValidateCelebornMapAttempt(
+ abandoned.shuffleId,
+ abandoned.celebornShuffleId,
+ abandoned.stageId,
+ abandoned.stageAttempt,
+ abandoned.mapId,
+ abandoned.taskAttempt,
+ abandoned.claimEpoch))
+ if (!current) {
+ // This result is the factory's invalidation-unsafe gate. A stale or already-failed claim
+ // must stop before the external report RPC, without touching a replacement map owner.
+ return true
+ }
+ val abandon = !shouldReportShuffleFetchFailure(abandoned.taskAttemptId)
+ if (abandon) {
+ val key =
+ ownerKey(abandoned.shuffleId, abandoned.stageId, abandoned.stageAttempt, abandoned.mapId)
+ claimOwners.remove(key)
+ deniedAttempts.getOrElseUpdate(key, mutable.HashSet.empty[Int]).add(abandoned.taskAttempt)
+ }
+ abandon
+ }
+
+ def unregisterShuffle(shuffleId: Int): Unit = synchronized {
+ generations.remove(shuffleId)
+ invalidatedGenerations.remove(shuffleId)
+ generationEpochs.remove(shuffleId)
+ claimOwners.retain { case ((ownerShuffleId, _, _, _), _) => ownerShuffleId != shuffleId }
+ deniedAttempts.retain { case ((ownerShuffleId, _, _, _), _) => ownerShuffleId != shuffleId }
+ }
+}
+
+private[shuffle] final class CelebornShuffleGenerationEndpoint(
+ override val rpcEnv: RpcEnv,
+ coordinator: CelebornShuffleGenerationCoordinator)
+ extends ThreadSafeRpcEndpoint {
+
+ override def receiveAndReply(context: RpcCallContext): PartialFunction[Any, Unit] = {
+ case claim: ClaimCelebornMapAttempt =>
+ context.reply(coordinator.claimMapAttempt(claim))
+ case generation: PrepareCelebornShuffleGeneration if generation.mapId < 0 =>
+ context.reply(coordinator.prepareGeneration(generation))
+ case generation: PrepareCelebornShuffleGeneration =>
+ context.reply(coordinator.prepareGenerationAndClaim(generation))
+ case validation: ValidateCelebornMapAttempt =>
+ context.reply(coordinator.validateMapAttempt(validation))
+ case invalidation: InvalidateCelebornShuffleGeneration =>
+ context.reply(coordinator.invalidateGeneration(invalidation))
+ case abandoned: AbandonCelebornMapAttempt =>
+ context.reply(coordinator.abandonMapAttempt(abandoned))
+ }
+}
diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleWriter.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleWriter.scala
index e8c4b3f0c9..89be100dbd 100644
--- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleWriter.scala
+++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleWriter.scala
@@ -21,6 +21,8 @@ package org.apache.spark.sql.comet.execution.shuffle
import java.nio.{ByteBuffer, ByteOrder}
import java.nio.file.{Files, Paths}
+import java.util.concurrent.{ScheduledFuture, TimeUnit}
+import java.util.concurrent.atomic.AtomicBoolean
import scala.collection.mutable
import scala.jdk.CollectionConverters._
@@ -35,11 +37,13 @@ import org.apache.spark.sql.catalyst.plans.physical.{HashPartitioning, Partition
import org.apache.spark.sql.comet.{CometExec, CometMetricNode, CometScalarSubquery, PlanDataInjector}
import org.apache.spark.sql.execution.metric.SQLMetric
import org.apache.spark.sql.types.StructField
+import org.apache.spark.util.{ThreadUtils, Utils}
import org.apache.comet.{CometConf, CometExecIterator}
import org.apache.comet.serde.{OperatorOuterClass, PartitioningOuterClass, QueryPlanSerde}
import org.apache.comet.serde.OperatorOuterClass.{CompressionCodec, Operator}
import org.apache.comet.serde.operator.schema2Proto
+import org.apache.comet.shuffle.{CelebornShufflePartitionPusher, CelebornShufflePusherFactory, ShufflePartitionPusher}
/**
* Drives the native shuffle write in a single [[CometExecIterator]] per partition. The plan is
@@ -60,7 +64,8 @@ class CometNativeShuffleWriter[K, V](
mapId: Long,
context: TaskContext,
metricsReporter: ShuffleWriteMetricsReporter,
- rangePartitionBounds: Option[Seq[InternalRow]] = None)
+ rangePartitionBounds: Option[Seq[InternalRow]] = None,
+ remoteDestination: Option[CelebornNativeShuffleDestination] = None)
extends ShuffleWriter[K, V]
with Logging {
@@ -68,16 +73,55 @@ class CometNativeShuffleWriter[K, V](
var partitionLengths: Array[Long] = _
var mapStatus: MapStatus = _
+ private var stopped = false
+ private lazy val effectivePartitionCount =
+ remoteDestination.map(_.numPartitions).getOrElse(outputPartitioning.numPartitions)
+
+ private val cancellationWatch: Option[ScheduledFuture[_]] = remoteDestination.flatMap {
+ destination =>
+ Option(context).map { taskContext =>
+ val watcher = CelebornNativeShuffleDestination.watchForCancellation(
+ taskContext,
+ destination.pusher,
+ failure =>
+ logWarning("Could not abort an interrupted Celeborn shuffle map attempt", failure))
+ taskContext.addTaskCompletionListener[Unit] { _ =>
+ watcher.cancel(false)
+ destination.pusher.abort()
+ }
+ watcher
+ }
+ }
override def write(inputs: Iterator[Product2[K, V]]): Unit = {
- val shuffleBlockResolver =
- SparkEnv.get.shuffleManager.shuffleBlockResolver.asInstanceOf[IndexShuffleBlockResolver]
- val dataFile = shuffleBlockResolver.getDataFile(shuffleId, mapId)
- val indexFile = shuffleBlockResolver.getIndexFile(shuffleId, mapId)
- val tempDataFilename = dataFile.getPath.replace(".data", ".data.tmp")
- val tempIndexFilename = indexFile.getPath.replace(".index", ".index.tmp")
- val tempDataFilePath = Paths.get(tempDataFilename)
- val tempIndexFilePath = Paths.get(tempIndexFilename)
+ try {
+ writeInternal(inputs)
+ } catch {
+ case failure: Throwable =>
+ remoteDestination.foreach { destination =>
+ try destination.pusher.abort()
+ catch {
+ case cleanupFailure: Throwable => failure.addSuppressed(cleanupFailure)
+ }
+ }
+ throw failure
+ }
+ }
+
+ private def writeInternal(inputs: Iterator[Product2[K, V]]): Unit = {
+ val localOutput = if (remoteDestination.isEmpty) {
+ val resolver =
+ SparkEnv.get.shuffleManager.shuffleBlockResolver.asInstanceOf[IndexShuffleBlockResolver]
+ val dataFile = resolver.getDataFile(shuffleId, mapId)
+ val indexFile = resolver.getIndexFile(shuffleId, mapId)
+ Some(
+ LocalShuffleOutput(
+ resolver,
+ dataFile.getPath.replace(".data", ".data.tmp"),
+ indexFile.getPath.replace(".index", ".index.tmp")))
+ } else {
+ None
+ }
// The dep's _rdd is always a CometNativeShuffleInputRDD on this path. Pattern-match instead
// of asInstanceOf so a future RDD-layering change produces a clear error here rather than a
@@ -94,7 +138,10 @@ class CometNativeShuffleWriter[K, V](
val inputObjects = shuffleInputIter.inputObjects
val shuffleBlockIters = shuffleInputIter.shuffleBlockIterators
- val unifiedPlan = buildUnifiedPlan(tempDataFilename, tempIndexFilename)
+ val unifiedPlan = localOutput match {
+ case Some(output) => buildUnifiedPlan(output.dataFile, output.indexFile)
+ case None => buildUnifiedPlan("", "")
+ }
val ctx = spec.execContext
val finalNativePlan = if (ctx.commonByKey.nonEmpty) {
// This partition's plan-data slice rides on the input iterator's Partition object (populated
@@ -147,7 +194,8 @@ class CometNativeShuffleWriter[K, V](
partitionIdx,
ctx.broadcastedHadoopConfForEncryption,
ctx.encryptedFilePaths,
- shuffleBlockIters)
+ shuffleBlockIters,
+ shufflePartitionPusher = remoteDestination.map(_.callback))
// Register subqueries against the iterator id so native callbacks resolve them to values.
ctx.subqueries.foreach { sub =>
@@ -161,51 +209,77 @@ class CometNativeShuffleWriter[K, V](
}
}
- while (cometIter.hasNext) {
- cometIter.next()
+ CometNativeShuffleWriter.drainAndClose(cometIter, () => cometIter.close())
+
+ remoteDestination match {
+ case Some(destination) =>
+ val completionStart = System.nanoTime()
+ val authorized = if (destination.commitAuthorized) {
+ destination.commitValidator()
+ } else {
+ SparkEnv.get.outputCommitCoordinator.canCommit(
+ context.stageId(),
+ context.stageAttemptNumber(),
+ context.partitionId(),
+ context.attemptNumber())
+ }
+ if (!authorized) {
+ throw CelebornShufflePusherFactory.commitDenied(context)
+ }
+ partitionLengths = destination.pusher.finish()
+ if (destination.commitAuthorized && !destination.commitValidator()) {
+ throw CelebornShufflePusherFactory.commitDenied(context)
+ }
+ metricsReporter.incBytesWritten(partitionLengths.sum)
+ metricsReporter.incWriteTime(System.nanoTime() - completionStart)
+ mapStatus = MapStatus.apply(
+ SparkEnv.get.blockManager.shuffleServerId,
+ partitionLengths,
+ context.taskAttemptId())
+
+ case None =>
+ val output = localOutput.get
+ val tempDataFilePath = Paths.get(output.dataFile)
+ val tempIndexFilePath = Paths.get(output.indexFile)
+
+ var offset = 0L
+ partitionLengths = Files
+ .readAllBytes(tempIndexFilePath)
+ .grouped(OFFSET_LENGTH)
+ .drop(1)
+ .map(indexBytes => {
+ val partitionOffset =
+ ByteBuffer.wrap(indexBytes).order(ByteOrder.LITTLE_ENDIAN).getLong
+ val partitionLength = partitionOffset - offset
+ offset = partitionOffset
+ partitionLength
+ })
+ .toArray
+ Files.delete(tempIndexFilePath)
+
+ metricsReporter.incBytesWritten(Files.size(tempDataFilePath))
+ output.resolver.writeMetadataFileAndCommit(
+ shuffleId,
+ mapId,
+ partitionLengths,
+ Array.empty,
+ tempDataFilePath.toFile)
+ mapStatus =
+ MapStatus.apply(SparkEnv.get.blockManager.shuffleServerId, partitionLengths, mapId)
}
- cometIter.close()
-
- // get partition lengths from shuffle write output index file
- var offset = 0L
- partitionLengths = Files
- .readAllBytes(tempIndexFilePath)
- .grouped(OFFSET_LENGTH)
- .drop(1) // first partition offset is always 0
- .map(indexBytes => {
- val partitionOffset =
- ByteBuffer.wrap(indexBytes).order(ByteOrder.LITTLE_ENDIAN).getLong
- val partitionLength = partitionOffset - offset
- offset = partitionOffset
- partitionLength
- })
- .toArray
- Files.delete(tempIndexFilePath)
-
- // Total written bytes at native
- metricsReporter.incBytesWritten(Files.size(tempDataFilePath))
+
metricsReporter.incRecordsWritten(metricsOutputRows.value)
metricsReporter.incWriteTime(metricsWriteTime.value)
-
- // commit
- shuffleBlockResolver.writeMetadataFileAndCommit(
- shuffleId,
- mapId,
- partitionLengths,
- Array.empty, // TODO: add checksums
- tempDataFilePath.toFile)
- mapStatus =
- MapStatus.apply(SparkEnv.get.blockManager.shuffleServerId, partitionLengths, mapId)
}
private def isSinglePartitioning(p: Partitioning): Boolean = p match {
case SinglePartition => true
- case rp: RangePartitioning =>
+ case _: RangePartitioning =>
// Spark sometimes generates RangePartitioning schemes with numPartitions == 1,
// or the computed bounds results in a single target partition.
// In this case Comet just serializes a SinglePartition scheme to native.
- rp.numPartitions == 1 || rangePartitionBounds.forall(_.isEmpty)
- case hp: HashPartitioning => hp.numPartitions == 1
+ effectivePartitionCount == 1 || rangePartitionBounds.forall(_.isEmpty)
+ case _: HashPartitioning => effectivePartitionCount == 1
case _ => false
}
@@ -213,22 +287,30 @@ class CometNativeShuffleWriter[K, V](
* Build the unified `ShuffleWriter(child = childNativeOp)` plan with the partitioning serde,
* compression settings, and output file paths.
*/
- private def buildUnifiedPlan(dataFile: String, indexFile: String): Operator = {
+ private[shuffle] def buildUnifiedPlan(dataFile: String, indexFile: String): Operator = {
val shuffleWriterBuilder = OperatorOuterClass.ShuffleWriter.newBuilder()
- // Keep the legacy output fields populated so older native libraries can still consume local
- // shuffle plans while newer libraries use the explicit partition-writer destination.
- shuffleWriterBuilder.setOutputDataFile(dataFile)
- shuffleWriterBuilder.setOutputIndexFile(indexFile)
- shuffleWriterBuilder.setPartitionWriter(
- OperatorOuterClass.PartitionWriter
- .newBuilder()
- .setLocal(
- OperatorOuterClass.LocalPartitionWriter
+ remoteDestination match {
+ case Some(_) =>
+ shuffleWriterBuilder.setPartitionWriter(
+ OperatorOuterClass.PartitionWriter
.newBuilder()
- .setOutputDataFile(dataFile)
- .setOutputIndexFile(indexFile)
+ .setRss(OperatorOuterClass.RssPartitionWriter.getDefaultInstance)
.build())
- .build())
+ case None =>
+ // Keep legacy paths for older native libraries while newer libraries use the destination.
+ shuffleWriterBuilder.setOutputDataFile(dataFile)
+ shuffleWriterBuilder.setOutputIndexFile(indexFile)
+ shuffleWriterBuilder.setPartitionWriter(
+ OperatorOuterClass.PartitionWriter
+ .newBuilder()
+ .setLocal(
+ OperatorOuterClass.LocalPartitionWriter
+ .newBuilder()
+ .setOutputDataFile(dataFile)
+ .setOutputIndexFile(indexFile)
+ .build())
+ .build())
+ }
if (SparkEnv.get.conf.getBoolean("spark.shuffle.compress", true)) {
val codec = CometConf.COMET_SHUFFLE_COMPRESSION_CODEC.get() match {
@@ -255,7 +337,7 @@ class CometNativeShuffleWriter[K, V](
case _: HashPartitioning =>
val hashPartitioning = outputPartitioning.asInstanceOf[HashPartitioning]
val partitioning = PartitioningOuterClass.HashPartition.newBuilder()
- partitioning.setNumPartitions(outputPartitioning.numPartitions)
+ partitioning.setNumPartitions(effectivePartitionCount)
val partitionExprs = hashPartitioning.expressions
.flatMap(e => QueryPlanSerde.exprToProto(e, outputAttributes))
@@ -273,7 +355,7 @@ class CometNativeShuffleWriter[K, V](
case _: RangePartitioning =>
val rangePartitioning = outputPartitioning.asInstanceOf[RangePartitioning]
val partitioning = PartitioningOuterClass.RangePartition.newBuilder()
- partitioning.setNumPartitions(outputPartitioning.numPartitions)
+ partitioning.setNumPartitions(effectivePartitionCount)
// Detect duplicates by tracking expressions directly, similar to DataFusion's LexOrdering
// DataFusion will deduplicate identical sort expressions in LexOrdering,
@@ -337,7 +419,7 @@ class CometNativeShuffleWriter[K, V](
case _: RoundRobinPartitioning =>
val partitioning = PartitioningOuterClass.RoundRobinPartition.newBuilder()
- partitioning.setNumPartitions(outputPartitioning.numPartitions)
+ partitioning.setNumPartitions(effectivePartitionCount)
partitioning.setMaxHashColumns(
CometConf.COMET_SHUFFLE_NATIVE_ROUND_ROBIN_PARTITIONING_MAX_HASH_COLUMNS.get())
@@ -366,12 +448,112 @@ class CometNativeShuffleWriter[K, V](
}
override def stop(success: Boolean): Option[MapStatus] = {
- if (success) {
- Some(mapStatus)
- } else {
- None
+ remoteDestination match {
+ case None => if (success) Some(mapStatus) else None
+ case Some(destination) =>
+ synchronized {
+ if (stopped) {
+ None
+ } else {
+ stopped = true
+ try {
+ if (success) {
+ if (mapStatus == null) {
+ throw new IllegalStateException(
+ "Cannot complete a Celeborn shuffle map task before writing its data")
+ }
+ Some(mapStatus)
+ } else {
+ None
+ }
+ } finally {
+ cancellationWatch.foreach(_.cancel(false))
+ destination.pusher.abort()
+ }
+ }
+ }
}
}
override def getPartitionLengths(): Array[Long] = partitionLengths
+
+ private final case class LocalShuffleOutput(
+ resolver: IndexShuffleBlockResolver,
+ dataFile: String,
+ indexFile: String)
+}
+
+private[shuffle] object CometNativeShuffleWriter {
+ def drainAndClose(iterator: Iterator[_], close: () => Unit): Unit = {
+ Utils.tryWithSafeFinally {
+ while (iterator.hasNext) {
+ iterator.next()
+ }
+ } {
+ // In particular, cleanup must not replace a typed FetchFailedException that Spark needs
+ // to trigger stage recovery. The original failure retains cleanup errors as suppressed.
+ close()
+ }
+ }
+}
+
+/** A task-owned remote destination whose callback never crosses map-attempt boundaries. */
+private[shuffle] final case class CelebornNativeShuffleDestination(
+ pusher: CelebornShufflePartitionPusher,
+ maxFrameBytes: Int,
+ numPartitions: Int,
+ commitAuthorized: Boolean = false,
+ commitValidator: () => Boolean = () => true) {
+ require(pusher != null, "The Celeborn shuffle partition pusher must not be null")
+ require(maxFrameBytes >= 20, "The Celeborn shuffle frame limit must fit a complete frame")
+ require(
+ maxFrameBytes <= pusher.maxFrameBytes(),
+ "The Celeborn shuffle destination cannot exceed its task-owned pusher's frame limit")
+ require(
+ numPartitions == pusher.numPartitions(),
+ "The Celeborn shuffle destination must use its actual reducer partition count")
+
+ // SQL execution configuration can impose a tighter bound than the executor SparkConf used to
+ // create the pusher. Native planning reads the bound from its registered task-owned callback.
+ private[shuffle] val callback: ShufflePartitionPusher = new ShufflePartitionPusher {
+ override def pushPartitionData(partitionId: Int, data: Array[Byte], length: Int): Unit =
+ pusher.pushPartitionData(partitionId, data, length)
+
+ override def reservePartitionData(bytes: Int): Unit = pusher.reservePartitionData(bytes)
+
+ override def releasePartitionDataReservation(): Unit =
+ pusher.releasePartitionDataReservation()
+
+ override def maxFrameBytes(): Int = CelebornNativeShuffleDestination.this.maxFrameBytes
+
+ override def maxReservationBytes(): Int = pusher.maxReservationBytes()
+ }
+}
+
+private[shuffle] object CelebornNativeShuffleDestination {
+ // One daemon bounds cancellation-polling threads per executor; each task cancels its scheduled
+ // future on completion, so the process-lifetime scheduler does not retain completed callbacks.
+ private val cancellationWatcher =
+ ThreadUtils.newDaemonSingleThreadScheduledExecutor("comet-celeborn-task-cancellation")
+
+ private[shuffle] def watchForCancellation(
+ taskContext: TaskContext,
+ pusher: CelebornShufflePartitionPusher,
+ reportFailure: Throwable => Unit): ScheduledFuture[_] = {
+ val handled = new AtomicBoolean(false)
+ cancellationWatcher.scheduleWithFixedDelay(
+ new Runnable {
+ override def run(): Unit = {
+ if (taskContext.isInterrupted() && handled.compareAndSet(false, true)) {
+ try pusher.abort()
+ catch {
+ case failure: Throwable => reportFailure(failure)
+ }
+ }
+ }
+ },
+ 0L,
+ 25L,
+ TimeUnit.MILLISECONDS)
+ }
}
diff --git a/spark/src/test/java/org/apache/comet/shuffle/CelebornTransportOwnershipTestHelper.java b/spark/src/test/java/org/apache/comet/shuffle/CelebornTransportOwnershipTestHelper.java
new file mode 100644
index 0000000000..09c6fb524a
--- /dev/null
+++ b/spark/src/test/java/org/apache/comet/shuffle/CelebornTransportOwnershipTestHelper.java
@@ -0,0 +1,382 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.comet.shuffle;
+
+import java.io.IOException;
+import java.lang.reflect.Field;
+import java.lang.reflect.Proxy;
+import java.nio.ByteBuffer;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.concurrent.AbstractExecutorService;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.RejectedExecutionException;
+import java.util.concurrent.TimeUnit;
+
+import io.netty.buffer.ByteBuf;
+import io.netty.buffer.Unpooled;
+import io.netty.channel.Channel;
+import io.netty.channel.ChannelFuture;
+import io.netty.channel.ChannelHandlerContext;
+import io.netty.channel.ChannelOutboundBuffer;
+import io.netty.channel.ChannelOutboundHandlerAdapter;
+import io.netty.channel.ChannelPromise;
+import io.netty.channel.embedded.EmbeddedChannel;
+
+/** Real Netty write-lifetime checks invoked by CelebornShufflePartitionPusherSuite. */
+public final class CelebornTransportOwnershipTestHelper {
+ private CelebornTransportOwnershipTestHelper() {}
+
+ public static void assertTimedOutWriteRetainsOwnership(boolean existingClient) throws Exception {
+ ShuffleClient shuffle = new ShuffleClient();
+ Client client = existingClient ? shuffle.factory.createClient() : null;
+ CelebornTransportCallbackTracker.Push push =
+ CelebornTransportCallbackTracker.tryCreate(shuffle).beginPush();
+ if (client == null) {
+ client = shuffle.factory.createClient();
+ }
+ ByteBuf body = Unpooled.wrappedBuffer(new byte[8860]);
+ try {
+ Request request = new Request();
+ client.handler.outstandingPushes.put(1L, request);
+ // Stock TransportClient stores the result of this fluent addListener call.
+ ChannelFuture write = client.getChannel().writeAndFlush(body).addListener(ignored -> {});
+ push.close();
+ check(!write.cancel(true), "owned writes must not complete via timeout cancellation");
+ client.handler.outstandingPushes.remove(1L);
+ request.callback.onFailure(new IOException("Cleaned Up"));
+ check(!write.isDone(), "the blocked outbound write must still be pending");
+ check(body.refCnt() == 1, "the actual outbound buffer must still own the body");
+ check(!push.isComplete(), "callback completion must not release a pending outbound write");
+
+ client.underlying.retireWrite();
+ check(write.isDone(), "write retirement must complete its real future");
+ check(body.refCnt() == 0, "write retirement must release the actual body");
+ check(push.isComplete(), "admission may finish after both callback and write retire");
+ } finally {
+ push.close();
+ client.underlying.close();
+ shuffle.pushDataRetryPool.shutdownNow();
+ }
+ }
+
+ public static void assertUnflushedWriteCannotBeCancelledEarly() throws Exception {
+ ShuffleClient shuffle = new ShuffleClient();
+ Client client = shuffle.factory.createClient();
+ HoldingWriteHandler queued = new HoldingWriteHandler();
+ client.underlying.pipeline().addLast(queued);
+ CelebornTransportCallbackTracker.Push push =
+ CelebornTransportCallbackTracker.tryCreate(shuffle).beginPush();
+ ByteBuf body = Unpooled.wrappedBuffer(new byte[8860]);
+ try {
+ Request request = new Request();
+ client.handler.outstandingPushes.put(1L, request);
+ ChannelFuture write = client.getChannel().writeAndFlush(body).addListener(ignored -> {});
+ push.close();
+ check(!write.isCancellable(), "the exposed future must consistently refuse cancellation");
+ check(!write.cancel(true), "timeout must not cancel a still-owned queued write");
+ client.handler.outstandingPushes.remove(1L);
+ request.callback.onFailure(new IOException("Cleaned Up"));
+ check(body.refCnt() == 1, "the queued write still owns its payload");
+ check(!push.isComplete(), "a queued body must stay charged after callback failure");
+
+ queued.failWrite();
+ check(body.refCnt() == 0, "failing the queued write must release its body");
+ check(push.isComplete(), "the failed write and callback have both returned ownership");
+ } finally {
+ push.close();
+ queued.failWrite();
+ client.underlying.close();
+ shuffle.pushDataRetryPool.shutdownNow();
+ }
+ }
+
+ public static void assertUnownedWritesPreserveCancellation() throws Exception {
+ ShuffleClient shuffle = new ShuffleClient();
+ Client client = shuffle.factory.createClient();
+ CelebornTransportCallbackTracker.Push push =
+ CelebornTransportCallbackTracker.tryCreate(shuffle).beginPush();
+ push.close();
+ HoldingWriteHandler queued = new HoldingWriteHandler();
+ client.underlying.pipeline().addLast(queued);
+ ByteBuf body = Unpooled.wrappedBuffer(new byte[16]);
+ try {
+ ChannelFuture write = client.getChannel().writeAndFlush(body);
+ check(write.cancel(true), "ordinary Spark writes must keep their normal cancellation API");
+ } finally {
+ queued.failWrite();
+ client.underlying.close();
+ shuffle.pushDataRetryPool.shutdownNow();
+ }
+ }
+
+ public static void assertCompletedCallbacksForgetPayloads() throws Exception {
+ for (boolean throwFromCallback : new boolean[] {false, true}) {
+ ShuffleClient shuffle = new ShuffleClient();
+ Client client = shuffle.factory.createClient();
+ CelebornTransportCallbackTracker.Push push =
+ CelebornTransportCallbackTracker.tryCreate(shuffle).beginPush();
+ final byte[] payload = new byte[8860];
+ final RuntimeException callbackFailure = new RuntimeException("callback test failure");
+ Request retained = new Request();
+ retained.callback =
+ new Callback() {
+ @Override
+ public void onSuccess(ByteBuffer response) {
+ check(payload.length == 8860, "the active callback must still own its payload");
+ check(!push.isComplete(), "the callback must stay charged while running");
+ if (throwFromCallback) {
+ throw callbackFailure;
+ }
+ }
+
+ @Override
+ public void onFailure(Throwable failure) {
+ onSuccess(null);
+ }
+ };
+ try {
+ client.handler.outstandingPushes.put(1L, retained);
+ push.close();
+ client.handler.outstandingPushes.remove(1L);
+ try {
+ retained.callback.onFailure(new IOException("Cleaned Up"));
+ check(!throwFromCallback, "the callback exception must be preserved");
+ } catch (RuntimeException failure) {
+ check(failure == callbackFailure && throwFromCallback, "unexpected callback exception");
+ }
+ // Keep the request and proxy reachable, modeling a handler paused after invocation.
+ checkDelegateCleared(Proxy.getInvocationHandler(retained.callback), "callback");
+ check(push.isComplete(), "only a payload-free callback wrapper may release admission");
+ } finally {
+ push.close();
+ client.underlying.close();
+ shuffle.pushDataRetryPool.shutdownNow();
+ }
+ }
+ }
+
+ public static void assertCompletedRetriesForgetPayloads() throws Exception {
+ for (int mode = 0; mode < 3; mode++) {
+ QueuedExecutor executor = new QueuedExecutor();
+ ShuffleClient shuffle = new ShuffleClient(executor);
+ CelebornTransportCallbackTracker.Push push =
+ CelebornTransportCallbackTracker.tryCreate(shuffle).beginPush();
+ final byte[] payload = new byte[8860];
+ final boolean throwFromRetry = mode == 1;
+ final RuntimeException retryFailure = new RuntimeException("retry test failure");
+ try {
+ shuffle.pushDataRetryPool.execute(
+ () -> {
+ check(payload.length == 8860, "the active retry must still own its payload");
+ check(!push.isComplete(), "the retry must stay charged while running");
+ if (throwFromRetry) {
+ throw retryFailure;
+ }
+ });
+ Runnable retained = executor.tasks.get(0);
+ push.close();
+ check(!push.isComplete(), "queued retries must retain admission");
+ if (mode == 2) {
+ List<Runnable> discarded = shuffle.pushDataRetryPool.shutdownNow();
+ check(discarded.get(0) == retained, "shutdown must return the discarded wrapper");
+ retained.run();
+ } else {
+ executor.tasks.clear();
+ try {
+ retained.run();
+ check(!throwFromRetry, "the retry exception must be preserved");
+ } catch (RuntimeException failure) {
+ check(failure == retryFailure && throwFromRetry, "unexpected retry exception");
+ }
+ }
+ // A queue, executor worker, or shutdown caller may keep the wrapper after completion.
+ checkDelegateCleared(retained, "delegate");
+ check(push.isComplete(), "only a payload-free retry wrapper may release admission");
+ } finally {
+ push.close();
+ shuffle.pushDataRetryPool.shutdownNow();
+ }
+ }
+ }
+
+ private static void checkDelegateCleared(Object owner, String name) throws Exception {
+ Field delegate = owner.getClass().getDeclaredField(name);
+ delegate.setAccessible(true);
+ check(delegate.get(owner) == null, "completed ownership must not retain its " + name);
+ }
+
+ private static void check(boolean condition, String message) {
+ if (!condition) {
+ throw new AssertionError(message);
+ }
+ }
+
+ public interface Bootstrap {
+ void doBootstrap(Client client);
+ }
+
+ public interface Callback {
+ void onSuccess(ByteBuffer response);
+
+ void onFailure(Throwable failure);
+ }
+
+ public static final class Request {
+ public Callback callback =
+ new Callback() {
+ @Override
+ public void onSuccess(ByteBuffer response) {}
+
+ @Override
+ public void onFailure(Throwable failure) {}
+ };
+ }
+
+ public static final class ResponseHandler {
+ private final ConcurrentHashMap<Long, Request> outstandingPushes = new ConcurrentHashMap<>();
+ }
+
+ public static final class Client {
+ private final HoldingChannel underlying = new HoldingChannel();
+ private final Channel channel = underlying;
+ private final ResponseHandler handler = new ResponseHandler();
+
+ public ResponseHandler getHandler() {
+ return handler;
+ }
+
+ public Channel getChannel() {
+ return channel;
+ }
+ }
+
+ public static final class Pool {
+ private final Client[] clients = new Client[1];
+ private final Object[] locks = {new Object()};
+ }
+
+ public static final class Factory {
+ private final List<Bootstrap> clientBootstraps = new ArrayList<>();
+ private final ConcurrentHashMap<String, Pool> connectionPool = new ConcurrentHashMap<>();
+
+ private Client createClient() {
+ Pool pool = connectionPool.computeIfAbsent("worker", ignored -> new Pool());
+ synchronized (pool.locks[0]) {
+ Client client = new Client();
+ for (Bootstrap bootstrap : clientBootstraps) {
+ bootstrap.doBootstrap(client);
+ }
+ pool.clients[0] = client;
+ return client;
+ }
+ }
+ }
+
+ public static final class ShuffleClient {
+ private final Factory factory = new Factory();
+ private final ExecutorService pushDataRetryPool;
+
+ private ShuffleClient() {
+ this(Executors.newSingleThreadExecutor());
+ }
+
+ private ShuffleClient(ExecutorService retryPool) {
+ pushDataRetryPool = retryPool;
+ }
+
+ public Factory getDataClientFactory() {
+ return factory;
+ }
+ }
+
+ private static final class QueuedExecutor extends AbstractExecutorService {
+ private final List<Runnable> tasks = new ArrayList<>();
+ private boolean shutdown;
+
+ @Override
+ public void execute(Runnable command) {
+ if (shutdown) {
+ throw new RejectedExecutionException("test executor is shut down");
+ }
+ tasks.add(command);
+ }
+
+ @Override
+ public void shutdown() {
+ shutdown = true;
+ }
+
+ @Override
+ public List<Runnable> shutdownNow() {
+ shutdown = true;
+ List<Runnable> discarded = new ArrayList<>(tasks);
+ tasks.clear();
+ return discarded;
+ }
+
+ @Override
+ public boolean isShutdown() {
+ return shutdown;
+ }
+
+ @Override
+ public boolean isTerminated() {
+ return shutdown && tasks.isEmpty();
+ }
+
+ @Override
+ public boolean awaitTermination(long timeout, TimeUnit unit) {
+ return isTerminated();
+ }
+ }
+
+ private static final class HoldingChannel extends EmbeddedChannel {
+ @Override
+ protected void doWrite(ChannelOutboundBuffer outbound) {
+ // Leave the real, flushed Netty outbound buffer blocked until the test retires it.
+ }
+
+ private void retireWrite() {
+ unsafe().outboundBuffer().remove(new IOException("write retired by test"));
+ }
+ }
+
+ private static final class HoldingWriteHandler extends ChannelOutboundHandlerAdapter {
+ private ByteBuf body;
+ private ChannelPromise promise;
+
+ @Override
+ public void write(ChannelHandlerContext context, Object message, ChannelPromise writePromise) {
+ body = (ByteBuf) message;
+ promise = writePromise;
+ }
+
+ private void failWrite() {
+ if (body != null) {
+ body.release();
+ body = null;
+ promise.tryFailure(new IOException("queued write retired by test"));
+ }
+ }
+ }
+}
diff --git a/spark/src/test/scala/org/apache/comet/CometConfSuite.scala b/spark/src/test/scala/org/apache/comet/CometConfSuite.scala
index a5b1dc5d5f..d6db22b282 100644
--- a/spark/src/test/scala/org/apache/comet/CometConfSuite.scala
+++ b/spark/src/test/scala/org/apache/comet/CometConfSuite.scala
@@ -157,6 +157,41 @@ class CometConfSuite extends AnyFunSuite {
assert(CometConf.COMET_EXPLAIN_FALLBACK_ENABLED.get(conf))
}
+ test("remote shuffle frame and admission limits have bounded defaults") {
+ val conf = new SQLConf
+
+ assert(CometConf.COMET_SHUFFLE_RSS_MAX_FRAME_BYTES.get(conf) == 64L * 1024 * 1024)
+ assert(CometConf.COMET_SHUFFLE_RSS_MAX_IN_FLIGHT_BYTES.get(conf) == 256L * 1024 * 1024)
+ }
+
+ test("remote shuffle frame limit rejects incomplete frames and oversized JVM requests") {
+ val conf = new SQLConf
+ val entry = CometConf.COMET_SHUFFLE_RSS_MAX_FRAME_BYTES
+
+ conf.setConfString(entry.key, "19b")
+ assertThrows[IllegalArgumentException](entry.get(conf))
+
+ conf.setConfString(entry.key, s"${Int.MaxValue - 15}b")
+ assertThrows[IllegalArgumentException](entry.get(conf))
+
+ conf.setConfString(entry.key, "20b")
+ assert(entry.get(conf) == 20)
+ }
+
+ test("remote shuffle admission limit reserves complete native, JNI, and client frames") {
+ val conf = new SQLConf
+ val entry = CometConf.COMET_SHUFFLE_RSS_MAX_IN_FLIGHT_BYTES
+
+ conf.setConfString(entry.key, "75b")
+ assertThrows[IllegalArgumentException](entry.get(conf))
+
+ conf.setConfString(entry.key, s"${Int.MaxValue.toLong + 1}b")
+ assertThrows[IllegalArgumentException](entry.get(conf))
+
+ conf.setConfString(entry.key, "76b")
+ assert(entry.get(conf) == 76)
+ }
+
test(
"COMET_EXPLAIN_FALLBACK_LOG_ENABLED reads deprecated logFallbackReasons.enabled as alias") {
val conf = new SQLConf
diff --git a/spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala
index ee76189bd8..83d185b3d9 100644
--- a/spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala
+++ b/spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala
@@ -114,9 +114,20 @@ class CometNativeShuffleSuite extends CometTestBase with AdaptiveSparkPlanHelper
StructType(Nil),
"native-rss-callback-test")
val captured = mutable.ArrayBuffer.empty[(Int, Int, Int, Long)]
+ val ownership = mutable.ArrayBuffer.empty[String]
val pusher = new ShufflePartitionPusher {
+ override def maxReservationBytes(): Int = 1024 * 1024
+
+ override def reservePartitionData(bytes: Int): Unit = {
+ assert(bytes > 0 && bytes <= maxReservationBytes())
+ ownership += "reserve"
+ }
+
+ override def releasePartitionDataReservation(): Unit = ownership += "release"
+
override def pushPartitionData(partitionId: Int, data: Array[Byte], length: Int)
: Unit = {
+ ownership += "push"
val encodedLength = ByteBuffer.wrap(data).order(ByteOrder.LITTLE_ENDIAN).getLong
captured.synchronized {
captured += ((partitionId, data.length, length, encodedLength))
@@ -137,6 +148,7 @@ class CometNativeShuffleSuite extends CometTestBase with AdaptiveSparkPlanHelper
while (iterator.hasNext) {
iterator.next().close()
}
+ assert(ownership.takeRight(3).toSeq == Seq("reserve", "push", "release"))
captured.iterator
} finally {
iterator.close()
@@ -164,9 +176,17 @@ class CometNativeShuffleSuite extends CometTestBase with AdaptiveSparkPlanHelper
StructType(Nil),
"native-rss-callback-failure-test")
val expected = new IOException("remote shuffle callback sentinel")
+ val ownership = mutable.ArrayBuffer.empty[String]
val pusher = new ShufflePartitionPusher {
- override def pushPartitionData(partitionId: Int, data: Array[Byte], length: Int): Unit =
+ override def reservePartitionData(bytes: Int): Unit = ownership += "reserve"
+
+ override def releasePartitionDataReservation(): Unit = ownership += "release"
+
+ override def pushPartitionData(partitionId: Int, data: Array[Byte], length: Int)
+ : Unit = {
+ ownership += "push"
throw expected
+ }
}
val iterator = new CometExecIterator(
CometExec.newIterId,
@@ -183,7 +203,9 @@ class CometNativeShuffleSuite extends CometTestBase with AdaptiveSparkPlanHelper
iterator.hasNext
Iterator.single((false, "callback unexpectedly succeeded"))
} catch {
- case actual: IOException => Iterator.single((actual eq expected, actual.getMessage))
+ case actual: IOException =>
+ assert(ownership.takeRight(3).toSeq == Seq("reserve", "push", "release"))
+ Iterator.single((actual eq expected, actual.getMessage))
}
} finally {
iterator.close()
@@ -194,6 +216,50 @@ class CometNativeShuffleSuite extends CometTestBase with AdaptiveSparkPlanHelper
assert(failures.sameElements(Array((true, "remote shuffle callback sentinel"))))
}
+ test("native RSS checks the callback's admission limit before allocating or reserving") {
+ val planBytes = rssShufflePlanBytes
+ val failures = spark.sparkContext
+ .parallelize(Seq(17), 1)
+ .mapPartitions { rowCounts =>
+ val batch = new ColumnarBatch(Array.empty[ColumnVector], rowCounts.next())
+ val inputs = CometArrowStream.inputObjects(
+ Iterator.single(batch),
+ StructType(Nil),
+ "native-rss-admission-limit-test")
+ var callbacks = 0
+ val pusher = new ShufflePartitionPusher {
+ override def maxReservationBytes(): Int = 1
+
+ override def reservePartitionData(bytes: Int): Unit = callbacks += 1
+
+ override def pushPartitionData(partitionId: Int, data: Array[Byte], length: Int): Unit =
+ callbacks += 1
+ }
+ val iterator = new CometExecIterator(
+ CometExec.newIterId,
+ inputs,
+ 0,
+ planBytes,
+ CometMetricNode(Map.empty),
+ 1,
+ 0,
+ shufflePartitionPusher = Some(pusher))
+ try {
+ try {
+ iterator.hasNext
+ Iterator.single((false, callbacks))
+ } catch {
+ case failure: Exception =>
+ Iterator.single((failure.getMessage.contains("admission budget"), callbacks))
+ }
+ } finally {
+ iterator.close()
+ }
+ }
+ .collect()
+ assert(failures.sameElements(Array((true, 0))))
+ }
+
test(
"failed RSS callback registration closes every Arrow and shuffle input and preserves its error") {
val planBytes = rssShufflePlanBytes
diff --git a/spark/src/test/scala/org/apache/comet/shuffle/CelebornShufflePartitionPusherSuite.scala b/spark/src/test/scala/org/apache/comet/shuffle/CelebornShufflePartitionPusherSuite.scala
index 3a661ecbc9..84cecddea0 100644
--- a/spark/src/test/scala/org/apache/comet/shuffle/CelebornShufflePartitionPusherSuite.scala
+++ b/spark/src/test/scala/org/apache/comet/shuffle/CelebornShufflePartitionPusherSuite.scala
@@ -21,14 +21,17 @@ package org.apache.comet.shuffle
import java.io.IOException
import java.nio.{ByteBuffer, ByteOrder}
-import java.util.Arrays
-import java.util.concurrent.atomic.AtomicReference
+import java.util.{ArrayList => JArrayList, Arrays, List => JList, Optional}
+import java.util.concurrent.{AbstractExecutorService, ConcurrentHashMap, CountDownLatch, ExecutorService, LinkedBlockingQueue, RejectedExecutionException, TimeUnit}
+import java.util.concurrent.atomic.{AtomicInteger, AtomicLong, AtomicReference, LongAdder}
import java.util.zip.CRC32
import scala.collection.mutable
import org.scalatest.funsuite.AnyFunSuite
+import org.apache.spark.SparkConf
+
class CelebornShufflePartitionPusherSuite extends AnyFunSuite {
private val encodedAttemptId = (4 << 16) | 7
@@ -36,6 +39,14 @@ class CelebornShufflePartitionPusherSuite extends AnyFunSuite {
private def pusher(client: AnyRef): CelebornShufflePartitionPusher =
new CelebornShufflePartitionPusher(client, 19, 3, encodedAttemptId, 12, 9)
+ private def clientArguments: Array[AnyRef] = Array[AnyRef](
+ "native-celeborn-application",
+ "localhost",
+ Int.box(9097),
+ new RecordingCelebornClientConf,
+ new RecordingCelebornUserIdentifier,
+ Array[Byte](1, 2, 3))
+
private def frame(bodyLength: Int = java.lang.Long.BYTES): Array[Byte] = {
val buffer = ByteBuffer
.allocate(java.lang.Long.BYTES + bodyLength)
@@ -170,36 +181,124 @@ class CelebornShufflePartitionPusherSuite extends AnyFunSuite {
}
}
- test("raw push accepts frames expanded by Spark IO encryption without pushing twice") {
+ test("encrypted native clients are rejected before encoding or raw allocation") {
val client = new RecordingCelebornPushClient
val cryptoHandler = new RecordingCelebornCryptoHandler
- val bytes = frame()
- client.cryptoHandler = Some(cryptoHandler)
+ client.cryptoHandler = Optional.of(cryptoHandler)
- pusher(client).pushPartitionData(6, bytes, bytes.length)
+ val failure = intercept[UnsupportedOperationException] {
+ pusher(client)
+ }
- assert(client.pushCount == 1)
+ assert(failure.getMessage.contains("Encrypted native Celeborn shuffle is not supported"))
+ assert(client.pushCount == 0)
+ assert(cryptoHandler.encryptionCount == 0)
+ assert(client.cryptoHandler.get() eq cryptoHandler)
+ }
+
+ test("Spark IO encryption is rejected before creating a native task pusher") {
+ val client = new RecordingCelebornPushClient
+ val conf = new SparkConf(false).set("spark.io.encryption.enabled", "true")
+ val failure = intercept[IllegalArgumentException] {
+ CelebornShufflePusherFactory.create(conf, client, 19, 12, 9, null)
+ }
+ assert(failure.getMessage.contains("Encrypted native Celeborn shuffle is not supported"))
+ assert(client.pushCount == 0)
+ }
+
+ test("crypto-aware client acquisition preserves encryption for ordinary Spark shuffle") {
+ val conf = new SparkConf(false).set("spark.io.encryption.enabled", "true")
+ val cryptoHandler = new RecordingCelebornCryptoHandler
+ val bytes = frame()
+ RecordingCryptoAwareCelebornClientFactory.reset()
+ var observedConf: SparkConf = null
+
+ val client = CelebornShufflePusherFactory
+ .resolveClient(
+ conf,
+ classOf[RecordingCryptoAwareCelebornClientFactory],
+ classOf[RecordingCelebornClientConf],
+ classOf[RecordingCelebornUserIdentifier],
+ clientArguments,
+ sparkConf => {
+ observedConf = sparkConf
+ Optional.of(cryptoHandler)
+ })
+ .asInstanceOf[RecordingCelebornPushClient]
+
+ // The shared application client must keep its real crypto handler even though native RSS
+ // cannot currently bound that handler's extra allocations and retained high-water buffer.
+ client.pushOrMergeData(19, 3, encodedAttemptId, 6, bytes, 0, bytes.length, 12, 9, true, true)
+
+ assert(observedConf eq conf)
+ assert(RecordingCryptoAwareCelebornClientFactory.cryptoAwareCalls.get() == 1)
+ assert(RecordingCryptoAwareCelebornClientFactory.legacyCalls.get() == 0)
assert(cryptoHandler.encryptionCount == 1)
assert(cryptoHandler.plaintext.sameElements(bytes))
assert(cryptoHandler.encryptedLength == bytes.length + 20)
- assert(client.lastPush.bytes eq bytes)
- assert(client.lastPush.skipCompress)
}
- test("integrity accounting covers plaintext before Spark IO encryption expands the frame") {
+ test("older Celeborn clients retain their six-argument client acquisition API") {
+ val conf = new SparkConf(false)
+ var cryptoHandlerResolved = false
+ RecordingLegacyCelebornClientFactory.calls.set(0)
+
+ val client = CelebornShufflePusherFactory.resolveClient(
+ conf,
+ classOf[RecordingLegacyCelebornClientFactory],
+ classOf[RecordingCelebornClientConf],
+ classOf[RecordingCelebornUserIdentifier],
+ clientArguments,
+ _ => {
+ cryptoHandlerResolved = true
+ Optional.empty[AnyRef]()
+ })
+
+ assert(client.isInstanceOf[RecordingCelebornPushClient])
+ assert(RecordingLegacyCelebornClientFactory.calls.get() == 1)
+ assert(!cryptoHandlerResolved)
+ }
+
+ test("crypto-handler failures never fall back to an unencrypted Celeborn client") {
+ val conf = new SparkConf(false).set("spark.io.encryption.enabled", "true")
+ val expected = new IllegalStateException("Spark shuffle encryption key was unavailable")
+ RecordingCryptoAwareCelebornClientFactory.reset()
+
+ val actual = intercept[IllegalStateException] {
+ CelebornShufflePusherFactory.resolveClient(
+ conf,
+ classOf[RecordingCryptoAwareCelebornClientFactory],
+ classOf[RecordingCelebornClientConf],
+ classOf[RecordingCelebornUserIdentifier],
+ clientArguments,
+ _ => throw expected)
+ }
+
+ assert(actual eq expected)
+ assert(RecordingCryptoAwareCelebornClientFactory.cryptoAwareCalls.get() == 0)
+ assert(RecordingCryptoAwareCelebornClientFactory.legacyCalls.get() == 0)
+ }
+
+ test(
+ "encryption enabled after binding is rejected before reservation or integrity accounting") {
val client = new IntegrityCheckingCelebornPushClient
val cryptoHandler = new RecordingCelebornCryptoHandler
val bytes = frame(13)
- client.cryptoHandler = Some(cryptoHandler)
+ val adapter = pusher(client)
+ client.cryptoHandler = Optional.of(cryptoHandler)
- pusher(client).pushPartitionData(4, bytes, bytes.length)
+ intercept[UnsupportedOperationException] {
+ adapter.reservePartitionData(3 * bytes.length)
+ }
+ intercept[UnsupportedOperationException] {
+ adapter.pushPartitionData(4, bytes, bytes.length)
+ }
- assert(client.partitionCrc(4) == combinedCrc(bytes))
- assert(client.partitionBytes(4) == bytes.length)
- assert(client.invocationOrder.toSeq == Seq("crc:4", "push:4"))
- assert(cryptoHandler.plaintext.sameElements(bytes))
- assert(cryptoHandler.encryptedLength == bytes.length + 20)
- assert(client.pushCount == 1)
+ assert(client.accountedFrames.isEmpty)
+ assert(client.invocationOrder.isEmpty)
+ assert(cryptoHandler.encryptionCount == 0)
+ assert(client.pushCount == 0)
+ assert(client.cryptoHandler.get() eq cryptoHandler)
}
test("raw push preserves the exact IOException thrown by the Celeborn client") {
@@ -360,14 +459,796 @@ class CelebornShufflePartitionPusherSuite extends AnyFunSuite {
assert(client.lastPush.attemptId == encodedAttemptId)
assert(client.lastPush.partitionId == 4)
}
+
+ test("configured complete-frame and three-copy executor bounds are both enforced") {
+ val client = new RecordingCelebornPushClient
+ val bounded = new CelebornShufflePartitionPusher(client, 19, 3, 7, 12, 9, 24, 160)
+ assert(bounded.maxFrameBytes() == 24)
+ assert(bounded.numPartitions() == 9)
+
+ val oversized = frame(17)
+ val failure = intercept[IOException] {
+ bounded.pushPartitionData(0, oversized, oversized.length)
+ }
+ assert(failure.getMessage.contains("configured maximum"))
+ assert(client.pushCount == 0)
+
+ val budgetClient = new RecordingCelebornPushClient
+ val admissionBounded =
+ new CelebornShufflePartitionPusher(budgetClient, 19, 3, 7, 12, 9, 1024, 76)
+ assert(admissionBounded.maxFrameBytes() == 20)
+
+ intercept[IllegalArgumentException] {
+ new CelebornShufflePartitionPusher(new RecordingCelebornPushClient, 19, 3, 7, 12, 9, 15, 76)
+ }
+ intercept[IllegalArgumentException] {
+ new CelebornShufflePartitionPusher(new RecordingCelebornPushClient, 19, 3, 7, 12, 9, 20, 63)
+ }
+ }
+
+ test("Celeborn 0.7 mapperEnd receives all five task fields and reports Comet partition bytes") {
+ val client = new RecordingCelebornPushClient
+ val adapter = pusher(client)
+ val first = frame(8)
+ val second = frame(11)
+ adapter.pushPartitionData(2, first, first.length)
+ adapter.pushPartitionData(6, second, second.length)
+
+ val sizes = adapter.finish()
+ assert(sizes.length == 9)
+ assert(sizes(2) == first.length)
+ assert(sizes(6) == second.length)
+ assert(
+ sizes.zipWithIndex
+ .filterNot { case (_, index) => index == 2 || index == 6 }
+ .forall(_._1 == 0))
+ assert(client.mapperEndCalls.get() == 1)
+ assert(client.lastMapperEnd == ((19, 3, encodedAttemptId, 12, 9)))
+ assert(adapter.finish().sameElements(sizes))
+ assert(client.mapperEndCalls.get() == 1)
+
+ intercept[IOException] {
+ adapter.pushPartitionData(2, first, first.length)
+ }
+ }
+
+ test("Celeborn 0.6 four-argument mapperEnd remains supported") {
+ val client = new LegacyMapperEndCelebornPushClient
+ val adapter = new CelebornShufflePartitionPusher(client, 19, 3, encodedAttemptId, 12, 9)
+ val bytes = frame()
+ adapter.pushPartitionData(5, bytes, bytes.length)
+
+ val sizes = adapter.finish()
+ assert(sizes(5) == bytes.length)
+ assert(client.mapperEndCalls.get() == 1)
+ assert(client.lastMapperEnd == ((19, 3, encodedAttemptId, 12)))
+ }
+
+ test("mapperEnd preserves its original failure and cancels the failed map exactly once") {
+ val client = new RecordingCelebornPushClient
+ val expected = new IOException("Celeborn asynchronous worker push failed")
+ client.mapperEndFailure = expected
+
+ val actual = intercept[IOException] {
+ pusher(client).finish()
+ }
+ assert(actual eq expected)
+ assert(client.mapperEndCalls.get() == 1)
+ assert(client.cleanupCalls.get() == 1)
+ assert(client.lastCleanup == ((19, 3, encodedAttemptId)))
+ }
+
+ test("executor admission remains charged until stock Celeborn completes its async request") {
+ val client = new AsyncRecordingCelebornPushClient
+ val adapter = new CelebornShufflePartitionPusher(client, 19, 3, encodedAttemptId, 12, 9, 64)
+ val bytes = frame()
+ adapter.pushPartitionData(2, bytes, bytes.length)
+ val state = client.currentState(19, 3, encodedAttemptId)
+ assert(state.inFlightRequestTracker.totalInflightReqs.sum() == 1)
+
+ val failure = new AtomicReference[Throwable]()
+ val second = new Thread(() => {
+ try adapter.pushPartitionData(3, bytes, bytes.length)
+ catch { case error: Throwable => failure.set(error) }
+ })
+ second.start()
+ second.join(100)
+ assert(second.isAlive, "the first async request must retain shared byte admission")
+ assert(client.pushCount == 1)
+
+ client.complete(state)
+ second.join(5000)
+ assert(!second.isAlive)
+ assert(failure.get() == null)
+ assert(client.pushCount == 2)
+ client.complete(state)
+ val sizes = adapter.finish()
+ assert(sizes(2) == bytes.length)
+ assert(sizes(3) == bytes.length)
+ }
+
+ test("cleanup does not release pinned request bytes before the cancelled transport completes") {
+ val client = new AsyncRecordingCelebornPushClient
+ val first = new CelebornShufflePartitionPusher(client, 19, 3, encodedAttemptId, 12, 9, 64)
+ val bytes = frame()
+ first.pushPartitionData(0, bytes, bytes.length)
+ val cancelledState = client.currentState(19, 3, encodedAttemptId)
+ first.abort()
+ first.abort()
+ assert(client.cleanupCalls.get() == 1)
+ assert(cancelledState.exception.get().getMessage == "Cleaned Up")
+ assert(cancelledState.inFlightRequestTracker.totalInflightReqs.sum() == 1)
+
+ val retry =
+ new CelebornShufflePartitionPusher(client, 19, 3, encodedAttemptId + 1, 12, 9, 64)
+ val failure = new AtomicReference[Throwable]()
+ val worker = new Thread(() => {
+ try retry.pushPartitionData(1, bytes, bytes.length)
+ catch { case error: Throwable => failure.set(error) }
+ })
+ worker.start()
+ worker.join(100)
+ assert(worker.isAlive, "cleanup alone must not free an unfinished network request")
+ assert(client.pushCount == 1)
+
+ client.complete(cancelledState)
+ worker.join(5000)
+ assert(!worker.isAlive)
+ assert(failure.get() == null)
+ val retryState = client.currentState(19, 3, encodedAttemptId + 1)
+ client.complete(retryState)
+ assert(retry.finish()(1) == bytes.length)
+ }
+
+ test("failed transport after cleanup releases only its completed shared admission") {
+ val client = new TransportRecordingCelebornPushClient
+ val first = new CelebornShufflePartitionPusher(client, 19, 3, encodedAttemptId, 12, 9, 64)
+ val bytes = frame()
+ first.pushPartitionData(0, bytes, bytes.length)
+ val cancelledState = client.currentState(19, 3, encodedAttemptId)
+ val unrelatedState = new RecordingCelebornPushState
+ client.dataClientFactory.handler.add(unrelatedState)
+ first.abort()
+
+ assert(cancelledState.exception.get().getMessage == "Cleaned Up")
+ assert(cancelledState.inFlightRequestTracker.totalInflightReqs.sum() == 1)
+
+ val retry =
+ new CelebornShufflePartitionPusher(client, 19, 3, encodedAttemptId + 1, 12, 9, 64)
+ val failure = new AtomicReference[Throwable]()
+ val worker = new Thread(() => {
+ try retry.pushPartitionData(1, bytes, bytes.length)
+ catch { case error: Throwable => failure.set(error) }
+ })
+ worker.start()
+ worker.join(100)
+ assert(worker.isAlive, "cleanup must not release a held transport request")
+
+ client.failTransport(unrelatedState, new IOException("unrelated connection closed"))
+ worker.join(100)
+ assert(worker.isAlive, "another map's completed request must not release this map's bytes")
+
+ client.failTransport(cancelledState, new IOException("cancelled connection closed"))
+ worker.join(5000)
+ assert(!worker.isAlive, "failed transport completion after cleanup must release admission")
+ assert(failure.get() == null)
+ assert(cancelledState.inFlightRequestTracker.totalInflightReqs.sum() == 1)
+
+ val retryState = client.currentState(19, 3, encodedAttemptId + 1)
+ client.complete(retryState)
+ assert(retry.finish()(1) == bytes.length)
+ }
+
+ test("failure callbacks removed before the raw push returns remain observable after cleanup") {
+ Seq(false, true).foreach { openConnection =>
+ val client = new TransportRecordingCelebornPushClient
+ client.openConnectionBeforePush = openConnection
+ val first =
+ new CelebornShufflePartitionPusher(client, 19, 3, encodedAttemptId, 12, 9, 604, 2416)
+ val bytes = frame(596)
+ val removed = new CountDownLatch(1)
+ val resumeCallback = new CountDownLatch(1)
+ val callbackFailure = new AtomicReference[Throwable]()
+ var callbackThread: Thread = null
+
+ client.beforePushReturns = state => {
+ val request = client.dataClientFactory.handler.remove(state)
+ callbackThread = new Thread(() => {
+ try {
+ removed.countDown()
+ assert(resumeCallback.await(5, TimeUnit.SECONDS))
+ request.callback.onFailure(new IOException("connection closed after cancellation"))
+ } catch {
+ case failure: Throwable => callbackFailure.set(failure)
+ }
+ })
+ callbackThread.start()
+ assert(removed.await(5, TimeUnit.SECONDS))
+ }
+
+ var worker: Thread = null
+ try {
+ first.reservePartitionData(1812)
+ first.pushPartitionData(0, bytes, bytes.length)
+ first.releasePartitionDataReservation()
+ client.beforePushReturns = _ => ()
+ val state = client.currentState(19, 3, encodedAttemptId)
+ assert(client.dataClientFactory.handler.outstandingPushes.isEmpty)
+ first.abort()
+ assert(state.exception.get().getMessage == "Cleaned Up")
+ assert(state.inFlightRequestTracker.totalInflightReqs.sum() == 1)
+
+ val retry = new CelebornShufflePartitionPusher(
+ client,
+ 19,
+ 3,
+ encodedAttemptId + 1,
+ 12,
+ 9,
+ 604,
+ 2416)
+ val failure = new AtomicReference[Throwable]()
+ worker = new Thread(() => {
+ try {
+ retry.reservePartitionData(1812)
+ try retry.pushPartitionData(1, bytes, bytes.length)
+ finally retry.releasePartitionDataReservation()
+ } catch {
+ case error: Throwable => failure.set(error)
+ }
+ })
+ worker.start()
+ worker.join(100)
+ assert(worker.isAlive, "removal and cleanup must not release the paused callback's bytes")
+
+ resumeCallback.countDown()
+ callbackThread.join(5000)
+ assert(!callbackThread.isAlive)
+ assert(callbackFailure.get() == null)
+ worker.join(5000)
+ assert(!worker.isAlive, "completion must remain observable after transport-map removal")
+ assert(failure.get() == null)
+ assert(state.inFlightRequestTracker.totalInflightReqs.sum() == 1)
+
+ val retryState = client.currentState(19, 3, encodedAttemptId + 1)
+ client.complete(retryState)
+ assert(retry.finish()(1) == bytes.length)
+ } finally {
+ resumeCallback.countDown()
+ if (callbackThread != null) {
+ callbackThread.join(5000)
+ }
+ if (worker != null) {
+ worker.interrupt()
+ worker.join(5000)
+ }
+ }
+ }
+ }
+
+ private def assertNextMapWaitsForCompletion(client: TransportRecordingCelebornPushClient)(
+ completePending: => Unit): Unit = {
+ val bytes = frame(596)
+ val next =
+ new CelebornShufflePartitionPusher(client, 19, 3, encodedAttemptId + 1, 12, 9, 604, 2416)
+ val started = new CountDownLatch(1)
+ val failure = new AtomicReference[Throwable]()
+ val worker = new Thread(() => {
+ try {
+ started.countDown()
+ next.reservePartitionData(1812)
+ try next.pushPartitionData(1, bytes, bytes.length)
+ finally next.releasePartitionDataReservation()
+ } catch {
+ case error: Throwable => failure.set(error)
+ }
+ })
+ try {
+ worker.start()
+ assert(started.await(5, TimeUnit.SECONDS))
+ worker.join(100)
+ assert(worker.isAlive, "unfinished callbacks or retries must retain their admission")
+ completePending
+ worker.join(5000)
+ assert(!worker.isAlive, "completed callbacks and retries must release shared admission")
+ assert(failure.get() == null)
+ client.complete(client.currentState(19, 3, encodedAttemptId + 1))
+ assert(next.finish()(1) == bytes.length)
+ } finally {
+ worker.interrupt()
+ worker.join(5000)
+ }
+ }
+
+ test("cleanup during raw submission keeps callbacks on a recreated push state observable") {
+ val client = new TransportRecordingCelebornPushClient
+ val first =
+ new CelebornShufflePartitionPusher(client, 19, 3, encodedAttemptId, 12, 9, 604, 2416)
+ val original = client.getPushState(s"19-3-$encodedAttemptId")
+ val recreated = new AtomicReference[RecordingCelebornPushState]()
+ client.beforePushBegins = () => first.abort()
+ client.beforePushReturns = state => recreated.set(state)
+ val bytes = frame(596)
+ intercept[IOException] {
+ first.pushPartitionData(0, bytes, bytes.length)
+ }
+ client.beforePushBegins = () => ()
+ client.beforePushReturns = _ => ()
+ val state = recreated.get()
+ assert(state != null && (state ne original))
+ assert(state.exception.get().getMessage == "Cleaned Up")
+
+ assertNextMapWaitsForCompletion(client) {
+ client.failTransport(state, new IOException("connection failed after final cleanup"))
+ assert(state.inFlightRequestTracker.totalInflightReqs.sum() == 1)
+ }
+ }
+
+ test("a raw invocation that throws after publication retains its outstanding callback") {
+ val client = new TransportRecordingCelebornPushClient
+ val first =
+ new CelebornShufflePartitionPusher(client, 19, 3, encodedAttemptId, 12, 9, 604, 2416)
+ val published = new AtomicReference[RecordingCelebornPushState]()
+ val expected = new IOException("raw push failed after publishing the request")
+ client.beforePushReturns = state => {
+ published.set(state)
+ throw expected
+ }
+ val bytes = frame(596)
+ val failure = intercept[IOException] {
+ first.pushPartitionData(0, bytes, bytes.length)
+ }
+ assert(failure eq expected)
+ client.beforePushReturns = _ => ()
+
+ assertNextMapWaitsForCompletion(client) {
+ client.failTransport(published.get(), new IOException("cancelled transport failed"))
+ }
+ }
+
+ test("queued retry completion after cleanup releases admission without a transport callback") {
+ val client = new TransportRecordingCelebornPushClient
+ client.retriesBeforeFailure = 1
+ val first =
+ new CelebornShufflePartitionPusher(client, 19, 3, encodedAttemptId, 12, 9, 604, 2416)
+ val bytes = frame(596)
+ first.pushPartitionData(0, bytes, bytes.length)
+ val state = client.currentState(19, 3, encodedAttemptId)
+ client.failTransport(state, new IOException("retry the failed transport"))
+ assert(client.retryExecutor.pendingCount == 1)
+ assert(client.dataClientFactory.handler.outstandingPushes.isEmpty)
+ first.abort()
+
+ assertNextMapWaitsForCompletion(client) {
+ // Stock retries retain the original callback, not the proxy in PushRequestInfo.
+ client.retryExecutor.runNext()
+ assert(client.retryExecutor.pendingCount == 0)
+ assert(state.inFlightRequestTracker.totalInflightReqs.sum() == 1)
+ }
+ }
+
+ test("retry handoff to another transport keeps admission until its callback returns") {
+ val client = new TransportRecordingCelebornPushClient
+ client.retriesBeforeFailure = 1
+ val first =
+ new CelebornShufflePartitionPusher(client, 19, 3, encodedAttemptId, 12, 9, 604, 2416)
+ val bytes = frame(596)
+ first.pushPartitionData(0, bytes, bytes.length)
+ val state = client.currentState(19, 3, encodedAttemptId)
+ client.retryCallback = callback => client.dataClientFactory.handler.add(state, callback)
+ client.failTransport(state, new IOException("retry at another worker"))
+ assert(client.retryExecutor.pendingCount == 1)
+ client.retryExecutor.runNext()
+ assert(client.retryExecutor.pendingCount == 0)
+ assert(client.dataClientFactory.handler.outstandingPushes.size() == 1)
+ first.abort()
+
+ assertNextMapWaitsForCompletion(client) {
+ client.failTransport(state, new IOException("retried transport failed after cancellation"))
+ assert(state.inFlightRequestTracker.totalInflightReqs.sum() == 1)
+ }
+ }
+
+ test("cancelling a queued retry releases ownership without running its body") {
+ val client = new TransportRecordingCelebornPushClient
+ val tracker = CelebornTransportCallbackTracker.tryCreate(client)
+ val push = tracker.beginPush()
+ val calls = new AtomicInteger()
+ val retry =
+ try {
+ client.pushDataRetryPool.submit(new Runnable {
+ override def run(): Unit = { calls.incrementAndGet(); () }
+ })
+ } finally {
+ push.close()
+ }
+ assert(!push.isComplete)
+ assert(retry.cancel(true))
+ assert(push.isComplete)
+ client.retryExecutor.runNext()
+ assert(calls.get() == 0)
+ }
+
+ test("cancelling a running retry retains ownership until its body actually returns") {
+ Seq(false, true).foreach { interrupt =>
+ val client = new TransportRecordingCelebornPushClient
+ val tracker = CelebornTransportCallbackTracker.tryCreate(client)
+ val started = new CountDownLatch(1)
+ val resume = new CountDownLatch(1)
+ val push = tracker.beginPush()
+ val retry =
+ try {
+ client.pushDataRetryPool.submit(new Runnable {
+ override def run(): Unit = {
+ started.countDown()
+ try {
+ assert(resume.await(5, TimeUnit.SECONDS))
+ } catch {
+ // A cancelled Future is done even when the retry ignores interruption.
+ case _: InterruptedException => assert(resume.await(5, TimeUnit.SECONDS))
+ }
+ }
+ })
+ } finally {
+ push.close()
+ }
+ val worker = new Thread(() => client.retryExecutor.runNext())
+ try {
+ worker.start()
+ assert(started.await(5, TimeUnit.SECONDS))
+ assert(retry.cancel(interrupt))
+ assert(retry.isDone)
+ assert(!push.isComplete, "Future cancellation is not completion of a running retry")
+ resume.countDown()
+ worker.join(5000)
+ assert(!worker.isAlive)
+ assert(push.isComplete)
+ } finally {
+ resume.countDown()
+ worker.join(5000)
+ }
+ }
+ }
+
+ test("rejected retry submissions do not leak ownership") {
+ val client = new TransportRecordingCelebornPushClient
+ val tracker = CelebornTransportCallbackTracker.tryCreate(client)
+ val push = tracker.beginPush()
+ client.retryExecutor.shutdown()
+ try {
+ intercept[RejectedExecutionException] {
+ client.pushDataRetryPool.submit(new Runnable {
+ override def run(): Unit = fail("A rejected retry must not run")
+ })
+ }
+ } finally {
+ push.close()
+ }
+ assert(push.isComplete)
+ }
+
+ test("cancelled successful transport does not release another live transport twice") {
+ val client = new TransportRecordingCelebornPushClient
+ // The complete reservations are 112 + 64 bytes. Neither individual completion can admit
+ // the next map's 160-byte reservation, regardless of transport callback ordering.
+ val first = new CelebornShufflePartitionPusher(client, 19, 3, encodedAttemptId, 12, 9, 176)
+ val large = frame(24)
+ val small = frame()
+ val retryBytes = frame(40)
+ first.pushPartitionData(0, large, large.length)
+ first.pushPartitionData(1, small, small.length)
+ val cancelledState = client.currentState(19, 3, encodedAttemptId)
+ first.abort()
+ assert(cancelledState.inFlightRequestTracker.totalInflightReqs.sum() == 2)
+
+ val retry =
+ new CelebornShufflePartitionPusher(client, 19, 3, encodedAttemptId + 1, 12, 9, 176)
+ val failure = new AtomicReference[Throwable]()
+ val worker = new Thread(() => {
+ try retry.pushPartitionData(2, retryBytes, retryBytes.length)
+ catch { case error: Throwable => failure.set(error) }
+ })
+ worker.start()
+ worker.join(100)
+ assert(worker.isAlive)
+
+ client.complete(cancelledState)
+ worker.join(100)
+ assert(worker.isAlive, "one success must not also release the other live request")
+ assert(cancelledState.inFlightRequestTracker.totalInflightReqs.sum() == 1)
+
+ client.failTransport(cancelledState, new IOException("cancelled connection closed"))
+ worker.join(5000)
+ assert(!worker.isAlive)
+ assert(failure.get() == null)
+ assert(cancelledState.inFlightRequestTracker.totalInflightReqs.sum() == 1)
+
+ val retryState = client.currentState(19, 3, encodedAttemptId + 1)
+ client.complete(retryState)
+ assert(retry.finish()(2) == retryBytes.length)
+ }
+
+ test(
+ "terminal async failure without removeBatch preserves its exception and releases one request") {
+ val client = new AsyncRecordingCelebornPushClient
+ val adapter = new CelebornShufflePartitionPusher(client, 19, 3, encodedAttemptId, 12, 9, 64)
+ val bytes = frame()
+ adapter.pushPartitionData(0, bytes, bytes.length)
+ val failedState = client.currentState(19, 3, encodedAttemptId)
+ val expected = new IOException(
+ "Push data to worker-1 failed for shuffle 19 map 3 attempt 7 partition 0 batch 1.")
+ client.failWithoutRemovingRequest(failedState, expected)
+
+ // mapperEnd observes the exception immediately, before the polling reconciler has sampled it.
+ val actual = intercept[IOException] {
+ adapter.finish()
+ }
+ assert(actual eq expected)
+ assert(failedState.inFlightRequestTracker.totalInflightReqs.sum() == 1)
+
+ val retry =
+ new CelebornShufflePartitionPusher(client, 19, 3, encodedAttemptId + 1, 12, 9, 64)
+ val failure = new AtomicReference[Throwable]()
+ val worker = new Thread(() => {
+ try retry.pushPartitionData(4, bytes, bytes.length)
+ catch { case error: Throwable => failure.set(error) }
+ })
+ worker.start()
+ worker.join(5000)
+ assert(!worker.isAlive, "the one proven terminal failure must release its shared admission")
+ assert(failure.get() == null)
+ val retryState = client.currentState(19, 3, encodedAttemptId + 1)
+ client.complete(retryState)
+ assert(retry.finish()(4) == bytes.length)
+ }
+
+ test("encoding reservations enforce one thread-owned claim and release unused capacity") {
+ val client = new RecordingCelebornPushClient
+ val adapter = new CelebornShufflePartitionPusher(client, 19, 3, encodedAttemptId, 12, 9, 64)
+ val bytes = frame()
+ adapter.reservePartitionData(48)
+ intercept[IOException] {
+ adapter.reservePartitionData(48)
+ }
+ adapter.releasePartitionDataReservation()
+ adapter.reservePartitionData(48)
+ adapter.pushPartitionData(2, bytes, bytes.length)
+ adapter.releasePartitionDataReservation()
+ assert(adapter.finish()(2) == bytes.length)
+ }
+
+ test("native and JNI frame admission waits for both retirement and transport completion") {
+ Seq(false, true).foreach { nativeRetiresFirst =>
+ val client = new TransportRecordingCelebornPushClient
+ val bytes = frame(8852)
+ val budget = bytes.length * 3 + 16
+ val first =
+ new CelebornShufflePartitionPusher(client, 19, 3, encodedAttemptId, 12, 9, 8860, budget)
+ val next =
+ new CelebornShufflePartitionPusher(client, 19, 4, encodedAttemptId, 12, 9, 8860, budget)
+ assert(first.maxReservationBytes() == budget - 16)
+ first.reservePartitionData(budget - 16)
+ first.pushPartitionData(0, bytes, bytes.length)
+ val state = client.currentState(19, 3, encodedAttemptId)
+ val waiting = new CountDownLatch(1)
+ val acquired = new CountDownLatch(1)
+ val failure = new AtomicReference[Throwable]()
+ val worker = new Thread(() => {
+ try {
+ waiting.countDown()
+ next.reservePartitionData(budget - 16)
+ try acquired.countDown()
+ finally next.releasePartitionDataReservation()
+ } catch {
+ case error: Throwable => failure.set(error)
+ }
+ })
+ try {
+ worker.start()
+ assert(waiting.await(5, TimeUnit.SECONDS))
+ if (nativeRetiresFirst) first.releasePartitionDataReservation()
+ else client.complete(state)
+ assert(!acquired.await(100, TimeUnit.MILLISECONDS))
+ if (nativeRetiresFirst) client.complete(state)
+ else first.releasePartitionDataReservation()
+ assert(acquired.await(5, TimeUnit.SECONDS))
+ worker.join(5000)
+ assert(failure.get() == null)
+ assert(first.finish()(0) == bytes.length)
+ } finally {
+ first.releasePartitionDataReservation()
+ worker.interrupt()
+ worker.join(5000)
+ }
+ }
+ }
+
+ test("timed-out writes retain ownership until Netty releases their outbound buffer") {
+ CelebornTransportOwnershipTestHelper.assertTimedOutWriteRetainsOwnership(true)
+ CelebornTransportOwnershipTestHelper.assertTimedOutWriteRetainsOwnership(false)
+ }
+
+ test("native queued writes cannot be cancelled before their buffers are released") {
+ CelebornTransportOwnershipTestHelper.assertUnflushedWriteCannotBeCancelledEarly()
+ }
+
+ test("ordinary Spark transport writes preserve their cancellation behavior") {
+ CelebornTransportOwnershipTestHelper.assertUnownedWritesPreserveCancellation()
+ }
+
+ test("completed transport callbacks no longer retain captured payloads") {
+ CelebornTransportOwnershipTestHelper.assertCompletedCallbacksForgetPayloads()
+ }
+
+ test("completed and discarded retry wrappers no longer retain captured payloads") {
+ CelebornTransportOwnershipTestHelper.assertCompletedRetriesForgetPayloads()
+ }
+
+ test("failed push keeps native frame admission until the JNI caller retires its buffers") {
+ Seq(false, true).foreach { publishedTransport =>
+ val client = new TransportRecordingCelebornPushClient
+ val bytes = frame()
+ val expected = new IOException("submission failed while native and JNI buffers are live")
+ val publishedState = new AtomicReference[RecordingCelebornPushState]()
+ if (publishedTransport) {
+ client.beforePushReturns = state => {
+ publishedState.set(state)
+ throw expected
+ }
+ } else client.failure = expected
+ val first =
+ new CelebornShufflePartitionPusher(client, 19, 3, encodedAttemptId, 12, 9, 64)
+ val next =
+ new CelebornShufflePartitionPusher(client, 19, 4, encodedAttemptId, 12, 9, 64)
+ first.reservePartitionData(48)
+ assert(intercept[IOException] {
+ first.pushPartitionData(0, bytes, bytes.length)
+ } eq expected)
+ if (publishedTransport) {
+ client.failTransport(publishedState.get(), new IOException("cancelled request failed"))
+ }
+ val waiting = new CountDownLatch(1)
+ val acquired = new CountDownLatch(1)
+ val failure = new AtomicReference[Throwable]()
+ val worker = new Thread(() => {
+ try {
+ waiting.countDown()
+ next.reservePartitionData(48)
+ try acquired.countDown()
+ finally next.releasePartitionDataReservation()
+ } catch {
+ case error: Throwable => failure.set(error)
+ }
+ })
+ try {
+ worker.start()
+ assert(waiting.await(5, TimeUnit.SECONDS))
+ assert(!acquired.await(100, TimeUnit.MILLISECONDS))
+ first.releasePartitionDataReservation()
+ assert(acquired.await(5, TimeUnit.SECONDS))
+ worker.join(5000)
+ assert(failure.get() == null)
+ } finally {
+ first.releasePartitionDataReservation()
+ worker.interrupt()
+ worker.join(5000)
+ }
+ }
+ }
+
+ test("interruption failures never release another still-live async transport request") {
+ val client = new AsyncRecordingCelebornPushClient
+ val first = new CelebornShufflePartitionPusher(client, 19, 3, encodedAttemptId, 12, 9, 64)
+ val bytes = frame()
+ first.pushPartitionData(0, bytes, bytes.length)
+ val heldState = client.currentState(19, 3, encodedAttemptId)
+ val interruption = new IOException(
+ "Interrupted while limiting Celeborn in-flight requests",
+ new InterruptedException("task was interrupted"))
+ client.failWithoutRemovingRequest(heldState, interruption)
+
+ val observed = intercept[IOException] {
+ first.finish()
+ }
+ assert(observed eq interruption)
+
+ val retry =
+ new CelebornShufflePartitionPusher(client, 19, 3, encodedAttemptId + 1, 12, 9, 64)
+ val failure = new AtomicReference[Throwable]()
+ val worker = new Thread(() => {
+ try retry.pushPartitionData(2, bytes, bytes.length)
+ catch { case error: Throwable => failure.set(error) }
+ })
+ worker.start()
+ worker.join(100)
+ assert(worker.isAlive, "an interrupted limiter does not prove its older request completed")
+ assert(heldState.inFlightRequestTracker.totalInflightReqs.sum() == 1)
+
+ client.complete(heldState)
+ worker.join(5000)
+ assert(!worker.isAlive)
+ assert(failure.get() == null)
+ val retryState = client.currentState(19, 3, encodedAttemptId + 1)
+ client.complete(retryState)
+ assert(retry.finish()(2) == bytes.length)
+ }
+}
+
+/** Stand-ins for optional Celeborn types used by reflective client-acquisition tests. */
+final class RecordingCelebornClientConf
+final class RecordingCelebornUserIdentifier
+
+/** Mirrors Celeborn 0.7's modern and legacy static client-acquisition overloads. */
+final class RecordingCryptoAwareCelebornClientFactory
+
+object RecordingCryptoAwareCelebornClientFactory {
+ val cryptoAwareCalls: AtomicInteger = new AtomicInteger()
+ val legacyCalls: AtomicInteger = new AtomicInteger()
+
+ def reset(): Unit = {
+ cryptoAwareCalls.set(0)
+ legacyCalls.set(0)
+ }
+
+ def get(
+ appUniqueId: String,
+ lifecycleManagerHost: String,
+ lifecycleManagerPort: Int,
+ conf: RecordingCelebornClientConf,
+ userIdentifier: RecordingCelebornUserIdentifier,
+ extension: Array[Byte],
+ cryptoHandler: Optional[_]): RecordingCelebornPushClient = {
+ cryptoAwareCalls.incrementAndGet()
+ val client = new RecordingCelebornPushClient
+ if (cryptoHandler.isPresent) {
+ client.cryptoHandler =
+ Optional.of(cryptoHandler.get().asInstanceOf[RecordingCelebornCryptoHandler])
+ }
+ client
+ }
+
+ def get(
+ appUniqueId: String,
+ lifecycleManagerHost: String,
+ lifecycleManagerPort: Int,
+ conf: RecordingCelebornClientConf,
+ userIdentifier: RecordingCelebornUserIdentifier,
+ extension: Array[Byte]): RecordingCelebornPushClient = {
+ legacyCalls.incrementAndGet()
+ new RecordingCelebornPushClient
+ }
+}
+
+/** Mirrors Celeborn 0.6's original static client-acquisition API. */
+final class RecordingLegacyCelebornClientFactory
+
+object RecordingLegacyCelebornClientFactory {
+ val calls: AtomicInteger = new AtomicInteger()
+
+ def get(
+ appUniqueId: String,
+ lifecycleManagerHost: String,
+ lifecycleManagerPort: Int,
+ conf: RecordingCelebornClientConf,
+ userIdentifier: RecordingCelebornUserIdentifier,
+ extension: Array[Byte]): RecordingCelebornPushClient = {
+ calls.incrementAndGet()
+ new RecordingCelebornPushClient
+ }
}
/** Public so the adapter can resolve and invoke the optional client's API using reflection. */
class RecordingCelebornPushClient {
+ val mapperEndCalls: AtomicInteger = new AtomicInteger()
+ val cleanupCalls: AtomicInteger = new AtomicInteger()
@volatile var acceptedBytes: Option[Int] = None
- @volatile var cryptoHandler: Option[RecordingCelebornCryptoHandler] = None
+ @volatile var cryptoHandler: Optional[RecordingCelebornCryptoHandler] = Optional.empty()
@volatile var failure: Throwable = _
+ @volatile var mapperEndFailure: Throwable = _
+ @volatile var cleanupFailure: Throwable = _
+ @volatile var lastMapperEnd: (Int, Int, Int, Int, Int) = _
+ @volatile var lastCleanup: (Int, Int, Int) = _
@volatile var lastPush: RecordedCelebornPush = _
@volatile var pushCount: Int = 0
@@ -402,9 +1283,381 @@ class RecordingCelebornPushClient {
throw failure
}
val transportPayloadLength =
- cryptoHandler.fold(length)(handler => handler.encrypt(bytes, offset, length).length)
+ if (cryptoHandler.isPresent) cryptoHandler.get().encrypt(bytes, offset, length).length
+ else length
acceptedBytes.getOrElse(transportPayloadLength + 16)
}
+
+ @throws[IOException]
+ def mapperEnd(
+ shuffleId: Int,
+ mapId: Int,
+ attemptId: Int,
+ numMappers: Int,
+ numPartitions: Int): Unit = {
+ mapperEndCalls.incrementAndGet()
+ lastMapperEnd = (shuffleId, mapId, attemptId, numMappers, numPartitions)
+ if (mapperEndFailure != null) {
+ throw mapperEndFailure
+ }
+ }
+
+ @throws[IOException]
+ def cleanup(shuffleId: Int, mapId: Int, attemptId: Int): Unit = {
+ cleanupCalls.incrementAndGet()
+ lastCleanup = (shuffleId, mapId, attemptId)
+ if (cleanupFailure != null) {
+ throw cleanupFailure
+ }
+ }
+}
+
+/** Mirrors stock Celeborn's private request tracker without requiring its optional dependency. */
+final class RecordingCelebornInFlightTracker {
+ val totalInflightReqs: LongAdder = new LongAdder()
+}
+
+/** Mirrors the public PushState failure slot and its stock private tracker member. */
+final class RecordingCelebornPushState {
+ val inFlightRequestTracker: RecordingCelebornInFlightTracker =
+ new RecordingCelebornInFlightTracker()
+ val exception: AtomicReference[IOException] = new AtomicReference[IOException]()
+}
+
+/** Exposes the same lifecycle and completion state as the public Apache Celeborn client. */
+class AsyncRecordingCelebornPushClient extends RecordingCelebornPushClient {
+ val pushStates: ConcurrentHashMap[String, RecordingCelebornPushState] =
+ new ConcurrentHashMap[String, RecordingCelebornPushState]()
+
+ def getPushState(mapKey: String): RecordingCelebornPushState =
+ pushStates.computeIfAbsent(mapKey, _ => new RecordingCelebornPushState())
+
+ def currentState(shuffleId: Int, mapId: Int, attemptId: Int): RecordingCelebornPushState =
+ pushStates.get(s"$shuffleId-$mapId-$attemptId")
+
+ def complete(state: RecordingCelebornPushState): Unit = state.synchronized {
+ state.inFlightRequestTracker.totalInflightReqs.decrement()
+ state.notifyAll()
+ }
+
+ def failWithoutRemovingRequest(state: RecordingCelebornPushState, failure: IOException): Unit =
+ state.synchronized {
+ state.exception.compareAndSet(null, failure)
+ state.notifyAll()
+ }
+
+ @throws[IOException]
+ override def pushOrMergeData(
+ shuffleId: Int,
+ mapId: Int,
+ attemptId: Int,
+ partitionId: Int,
+ bytes: Array[Byte],
+ offset: Int,
+ length: Int,
+ numMappers: Int,
+ numPartitions: Int,
+ doPush: Boolean,
+ skipCompress: Boolean): Int = {
+ val state = getPushState(s"$shuffleId-$mapId-$attemptId")
+ state.inFlightRequestTracker.totalInflightReqs.increment()
+ super.pushOrMergeData(
+ shuffleId,
+ mapId,
+ attemptId,
+ partitionId,
+ bytes,
+ offset,
+ length,
+ numMappers,
+ numPartitions,
+ doPush,
+ skipCompress)
+ }
+
+ @throws[IOException]
+ override def mapperEnd(
+ shuffleId: Int,
+ mapId: Int,
+ attemptId: Int,
+ numMappers: Int,
+ numPartitions: Int): Unit = {
+ super.mapperEnd(shuffleId, mapId, attemptId, numMappers, numPartitions)
+ val key = s"$shuffleId-$mapId-$attemptId"
+ val state = pushStates.get(key)
+ if (state != null) {
+ state.synchronized {
+ while (state.exception.get() == null &&
+ state.inFlightRequestTracker.totalInflightReqs.sum() > 0) {
+ state.wait(25)
+ }
+ val failure = state.exception.get()
+ if (failure != null) {
+ throw failure
+ }
+ }
+ pushStates.remove(key, state)
+ }
+ }
+
+ @throws[IOException]
+ override def cleanup(shuffleId: Int, mapId: Int, attemptId: Int): Unit = {
+ super.cleanup(shuffleId, mapId, attemptId)
+ val removed = pushStates.remove(s"$shuffleId-$mapId-$attemptId")
+ if (removed != null) {
+ removed.synchronized {
+ removed.exception.compareAndSet(null, new IOException("Cleaned Up"))
+ removed.notifyAll()
+ }
+ }
+ }
+}
+
+/** Reproduces stock Celeborn's client-factory, handler, and callback completion boundaries. */
+final class TransportRecordingCelebornPushClient extends AsyncRecordingCelebornPushClient {
+ val dataClientFactory: RecordingCelebornTransportClientFactory =
+ new RecordingCelebornTransportClientFactory
+ val retryExecutor = new RecordingCelebornRetryExecutor
+ val pushDataRetryPool: ExecutorService = retryExecutor
+
+ def getDataClientFactory: RecordingCelebornTransportClientFactory = dataClientFactory
+
+ var openConnectionBeforePush: Boolean = false
+ var beforePushBegins: () => Unit = () => ()
+ var beforePushReturns: RecordingCelebornPushState => Unit = (_: RecordingCelebornPushState) =>
+ ()
+ var retriesBeforeFailure: Int = 0
+ var retryCallback: RecordingCelebornTransportCallbackApi => Unit =
+ callback => callback.onFailure(new IOException("revive failed"))
+
+ @throws[IOException]
+ override def pushOrMergeData(
+ shuffleId: Int,
+ mapId: Int,
+ attemptId: Int,
+ partitionId: Int,
+ bytes: Array[Byte],
+ offset: Int,
+ length: Int,
+ numMappers: Int,
+ numPartitions: Int,
+ doPush: Boolean,
+ skipCompress: Boolean): Int = {
+ beforePushBegins()
+ val accepted = super.pushOrMergeData(
+ shuffleId,
+ mapId,
+ attemptId,
+ partitionId,
+ bytes,
+ offset,
+ length,
+ numMappers,
+ numPartitions,
+ doPush,
+ skipCompress)
+ if (openConnectionBeforePush) {
+ openConnectionBeforePush = false
+ dataClientFactory.openConnection()
+ }
+ val state = currentState(shuffleId, mapId, attemptId)
+ dataClientFactory.handler.add(
+ state,
+ new RecordingCelebornTransportCallback(
+ state,
+ retriesBeforeFailure,
+ callback => {
+ pushDataRetryPool.submit(new Runnable {
+ override def run(): Unit = retryCallback(callback)
+ })
+ ()
+ }))
+ beforePushReturns(state)
+ accepted
+ }
+
+ override def complete(state: RecordingCelebornPushState): Unit = {
+ dataClientFactory.handler.remove(state).callback.onSuccess(ByteBuffer.allocate(0))
+ }
+
+ def failTransport(state: RecordingCelebornPushState, failure: IOException): Unit = {
+ val request = dataClientFactory.handler.remove(state)
+ request.callback.onFailure(failure)
+ }
+}
+
+final class RecordingCelebornRetryExecutor extends AbstractExecutorService {
+ private val pending = new LinkedBlockingQueue[Runnable]()
+ @volatile private var stopped = false
+
+ def pendingCount: Int = pending.size()
+
+ def runNext(): Unit = {
+ val task = pending.poll(5, TimeUnit.SECONDS)
+ require(task != null, "Expected a queued Celeborn retry")
+ task.run()
+ }
+
+ override def execute(command: Runnable): Unit = {
+ if (stopped) {
+ throw new RejectedExecutionException("Celeborn retry executor is stopped")
+ }
+ pending.add(command)
+ }
+
+ override def shutdown(): Unit = stopped = true
+
+ override def shutdownNow(): JList[Runnable] = {
+ stopped = true
+ val tasks = new JArrayList[Runnable]()
+ pending.drainTo(tasks)
+ tasks
+ }
+
+ override def isShutdown: Boolean = stopped
+
+ override def isTerminated: Boolean = stopped && pending.isEmpty
+
+ override def awaitTermination(timeout: Long, unit: TimeUnit): Boolean = isTerminated
+}
+
+final class RecordingCelebornTransportClientFactory {
+ var handler: RecordingCelebornTransportResponseHandler =
+ new RecordingCelebornTransportResponseHandler
+ val clientBootstraps: JList[RecordingCelebornTransportClientBootstrap] =
+ new JArrayList[RecordingCelebornTransportClientBootstrap]()
+ val connectionPool: ConcurrentHashMap[String, RecordingCelebornTransportClientPool] =
+ new ConcurrentHashMap[String, RecordingCelebornTransportClientPool]()
+ connectionPool.put("worker", new RecordingCelebornTransportClientPool(handler))
+
+ def openConnection(): Unit = {
+ handler = new RecordingCelebornTransportResponseHandler
+ val pool = new RecordingCelebornTransportClientPool(handler)
+ val bootstraps = clientBootstraps.iterator()
+ while (bootstraps.hasNext) {
+ bootstraps.next().doBootstrap(pool.clients(0))
+ }
+ connectionPool.put("worker", pool)
+ }
+}
+
+trait RecordingCelebornTransportClientBootstrap {
+ def doBootstrap(client: RecordingCelebornTransportClient): Unit
+}
+
+final class RecordingCelebornTransportClientPool(
+ handler: RecordingCelebornTransportResponseHandler) {
+ val clients: Array[RecordingCelebornTransportClient] =
+ Array(new RecordingCelebornTransportClient(handler))
+ val locks: Array[Object] = Array(new Object)
+}
+
+final class RecordingCelebornTransportClient(handler: RecordingCelebornTransportResponseHandler) {
+ val channel: io.netty.channel.Channel = new io.netty.channel.embedded.EmbeddedChannel()
+ def getChannel: io.netty.channel.Channel = channel
+ def getHandler: RecordingCelebornTransportResponseHandler = handler
+}
+
+final class RecordingCelebornTransportResponseHandler {
+ private val nextRequestId = new AtomicLong()
+ val outstandingPushes: ConcurrentHashMap[java.lang.Long, RecordingCelebornTransportRequest] =
+ new ConcurrentHashMap[java.lang.Long, RecordingCelebornTransportRequest]()
+
+ def add(pushState: RecordingCelebornPushState): Unit =
+ add(pushState, new RecordingCelebornTransportCallback(pushState))
+
+ def add(
+ pushState: RecordingCelebornPushState,
+ callback: RecordingCelebornTransportCallbackApi): Unit = {
+ outstandingPushes.put(
+ Long.box(nextRequestId.incrementAndGet()),
+ new RecordingCelebornTransportRequest(pushState, callback))
+ }
+
+ def remove(pushState: RecordingCelebornPushState): RecordingCelebornTransportRequest = {
+ val entries = outstandingPushes.entrySet().iterator()
+ while (entries.hasNext) {
+ val entry = entries.next()
+ if (entry.getValue.pushState eq pushState) {
+ val removed = outstandingPushes.remove(entry.getKey)
+ if (removed != null) {
+ return removed
+ }
+ }
+ }
+ throw new IllegalStateException("Celeborn transport request is no longer outstanding")
+ }
+}
+
+final class RecordingCelebornTransportRequest(
+ val pushState: RecordingCelebornPushState,
+ var callback: RecordingCelebornTransportCallbackApi)
+
+trait RecordingCelebornTransportCallbackApi {
+ def onSuccess(response: ByteBuffer): Unit
+ def onFailure(failure: Throwable): Unit
+}
+
+final class RecordingCelebornTransportCallback(
+ val pushState: RecordingCelebornPushState,
+ private var retriesRemaining: Int = 0,
+ submitRetry: RecordingCelebornTransportCallbackApi => Unit = _ => ())
+ extends RecordingCelebornTransportCallbackApi {
+ override def onSuccess(response: ByteBuffer): Unit = pushState.synchronized {
+ pushState.inFlightRequestTracker.totalInflightReqs.decrement()
+ pushState.notifyAll()
+ }
+
+ override def onFailure(failure: Throwable): Unit = {
+ if (pushState.exception.get() == null) {
+ if (retriesRemaining > 0) {
+ retriesRemaining -= 1
+ submitRetry(this)
+ } else {
+ pushState.exception.compareAndSet(null, new IOException(failure))
+ }
+ }
+ }
+}
+
+/** Implements the older public Celeborn 0.6 four-argument mapper-completion API. */
+final class LegacyMapperEndCelebornPushClient {
+ private val delegate = new RecordingCelebornPushClient
+ val mapperEndCalls: AtomicInteger = new AtomicInteger()
+ @volatile var lastMapperEnd: (Int, Int, Int, Int) = _
+
+ @throws[IOException]
+ def pushOrMergeData(
+ shuffleId: Int,
+ mapId: Int,
+ attemptId: Int,
+ partitionId: Int,
+ bytes: Array[Byte],
+ offset: Int,
+ length: Int,
+ numMappers: Int,
+ numPartitions: Int,
+ doPush: Boolean,
+ skipCompress: Boolean): Int =
+ delegate.pushOrMergeData(
+ shuffleId,
+ mapId,
+ attemptId,
+ partitionId,
+ bytes,
+ offset,
+ length,
+ numMappers,
+ numPartitions,
+ doPush,
+ skipCompress)
+
+ def mapperEnd(shuffleId: Int, mapId: Int, attemptId: Int, numMappers: Int): Unit = {
+ mapperEndCalls.incrementAndGet()
+ lastMapperEnd = (shuffleId, mapId, attemptId, numMappers)
+ }
+
+ def cleanup(shuffleId: Int, mapId: Int, attemptId: Int): Unit = ()
}
/** Models the Spark crypto wire format's minimum 4-byte length plus 16-byte IV overhead. */
diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornNativeShuffleWriterSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornNativeShuffleWriterSuite.scala
new file mode 100644
index 0000000000..c3f291af85
--- /dev/null
+++ b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornNativeShuffleWriterSuite.scala
@@ -0,0 +1,481 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.spark.sql.comet.execution.shuffle
+
+import java.io.IOException
+import java.util.concurrent.{CountDownLatch, TimeUnit}
+
+import org.apache.spark.{ShuffleDependency, SparkConf, SparkEnv, TaskContext}
+import org.apache.spark.executor.CommitDeniedException
+import org.apache.spark.shuffle.{BaseShuffleHandle, FetchFailedException, IndexShuffleBlockResolver, ShuffleBlockResolver, ShuffleHandle, ShuffleManager, ShuffleReader, ShuffleReadMetricsReporter, ShuffleWriteMetricsReporter, ShuffleWriter}
+import org.apache.spark.sql.CometTestBase
+import org.apache.spark.sql.vectorized.ColumnarBatch
+
+import org.apache.comet.CometConf
+import org.apache.comet.shuffle.{CelebornShufflePartitionPusher, CelebornShufflePusherFactory, RecordingCelebornPushClient}
+
+/** Exercises real Spark map tasks, native RSS planning, and the Celeborn map lifecycle. */
+class CometCelebornNativeShuffleWriterSuite extends CometTestBase {
+
+ import testImplicits._
+
+ test("native iteration preserves fetch failures when closing also fails") {
+ Seq(false, true).foreach { failOnHasNext =>
+ val expected = new FetchFailedException(null, 19, 0L, 0, 0, "fetch failure sentinel", null)
+ val closeFailure = new IOException("native iterator close sentinel")
+ var closed = false
+ val iterator = new Iterator[Int] {
+ override def hasNext: Boolean = if (failOnHasNext) throw expected else true
+
+ override def next(): Int = throw expected
+ }
+ val actual = intercept[FetchFailedException] {
+ CometNativeShuffleWriter.drainAndClose(
+ iterator,
+ () => {
+ closed = true
+ throw closeFailure
+ })
+ }
+ assert(actual eq expected)
+ assert(actual.getSuppressed.toSeq == Seq(closeFailure))
+ assert(closed)
+ }
+ }
+
+ test("successful native iteration closes once and propagates standalone close failures") {
+ var rows = 0
+ var closes = 0
+ CometNativeShuffleWriter.drainAndClose(
+ Iterator(1, 2, 3).map { value =>
+ rows += 1
+ value
+ },
+ () => closes += 1)
+ assert(rows == 3)
+ assert(closes == 1)
+
+ val expected = new IOException("standalone native iterator close failure")
+ val actual = intercept[IOException] {
+ CometNativeShuffleWriter.drainAndClose(Iterator.empty, () => throw expected)
+ }
+ assert(actual eq expected)
+ }
+
+ private def withNativeShuffleDependency(rangePartitioning: Boolean = false)(
+ run: CometShuffleDependency[Int, ColumnarBatch, ColumnarBatch] => Unit): Unit = {
+ withSQLConf(
+ CometConf.COMET_EXEC_ENABLED.key -> "true",
+ CometConf.COMET_SHUFFLE_MODE.key -> "native",
+ CometConf.COMET_SHUFFLE_ENABLED.key -> "true",
+ "spark.sql.adaptive.enabled" -> "false") {
+ val rows = (0 until 24).map { value =>
+ val key = if (rangePartitioning) value % 2 else value
+ (key, s"row-$value")
+ }
+ withParquetTable(rows, "celeborn_native_rows") {
+ val input = sql("SELECT * FROM celeborn_native_rows")
+ val shuffled =
+ if (rangePartitioning) input.repartitionByRange(10, $"_1")
+ else input.repartition(3, $"_1")
+ val exchange = shuffled.queryExecution.executedPlan
+ .collectFirst { case value: CometShuffleExchangeExec => value }
+ .getOrElse(fail("Expected a native Comet shuffle exchange"))
+ run(
+ exchange.shuffleDependency
+ .asInstanceOf[CometShuffleDependency[Int, ColumnarBatch, ColumnarBatch]])
+ }
+ }
+ }
+
+ test("configured frame limits and executor admission both constrain native RSS callbacks") {
+ val context = TaskContext.empty()
+ val configured = new SparkConf(false)
+ .set(CometConf.COMET_SHUFFLE_RSS_MAX_FRAME_BYTES.key, "128")
+ .set(CometConf.COMET_SHUFFLE_RSS_MAX_IN_FLIGHT_BYTES.key, "4096")
+ val limitedFrame = CelebornShufflePusherFactory.create(
+ configured,
+ new RecordingCelebornPushClient,
+ 89,
+ 1,
+ 1,
+ context)
+ assert(limitedFrame.maxFrameBytes() == 128)
+ val tighterDestination = CelebornNativeShuffleDestination(limitedFrame, 64, 1)
+ assert(tighterDestination.callback.maxFrameBytes() == 64)
+ assert(tighterDestination.callback.maxReservationBytes() == 4096 - 16)
+ tighterDestination.callback.reservePartitionData(64)
+ tighterDestination.callback.releasePartitionDataReservation()
+
+ val constrained =
+ configured.clone().set(CometConf.COMET_SHUFFLE_RSS_MAX_IN_FLIGHT_BYTES.key, "112")
+ val limitedAdmission = CelebornShufflePusherFactory.create(
+ constrained,
+ new RecordingCelebornPushClient,
+ 90,
+ 1,
+ 1,
+ context)
+ assert(limitedAdmission.maxFrameBytes() == 32)
+ limitedFrame.abort()
+ limitedAdmission.abort()
+ }
+
+ test("task interruption proactively aborts its task-owned Celeborn callback") {
+ val interrupted = new CountDownLatch(1)
+ val client = new RecordingCelebornPushClient {
+ override def cleanup(shuffleId: Int, mapId: Int, attemptId: Int): Unit = {
+ super.cleanup(shuffleId, mapId, attemptId)
+ interrupted.countDown()
+ }
+ }
+ val context = TaskContext.empty()
+ val pusher = new CelebornShufflePartitionPusher(client, 90, 0, 0, 1, 1)
+ val watcher = CelebornNativeShuffleDestination.watchForCancellation(
+ context,
+ pusher,
+ failure => throw failure)
+
+ try {
+ context.markInterrupted("Spark cancelled its native shuffle map attempt")
+ assert(interrupted.await(5, TimeUnit.SECONDS))
+ assert(client.cleanupCalls.get() == 1)
+ } finally {
+ watcher.cancel(false)
+ pusher.abort()
+ }
+ }
+
+ test("native map attempts push complete frames and commit remotely without local files") {
+ withNativeShuffleDependency() { dependency =>
+ val numMappers = dependency.rdd.getNumPartitions
+ val results = spark.sparkContext.runJob(
+ dependency.rdd,
+ (context: TaskContext, inputs: Iterator[Product2[Int, ColumnarBatch]]) => {
+ val client = new RecordingCelebornPushClient
+ val pusher = CelebornShufflePusherFactory.create(
+ SparkEnv.get.conf,
+ client,
+ 91,
+ numMappers,
+ dependency.partitioner.numPartitions,
+ context)
+ assert(
+ SparkEnv.get.outputCommitCoordinator.canCommit(
+ context.stageId(),
+ context.stageAttemptNumber(),
+ context.partitionId(),
+ context.attemptNumber()))
+ var validations = 0
+ val writer = CometCelebornNativeShuffleWriterSuite.newWriter(
+ dependency,
+ context,
+ pusher,
+ commitAuthorized = true,
+ commitValidator = () => {
+ validations += 1
+ true
+ })
+ val plan = writer.buildUnifiedPlan("", "").getShuffleWriter
+ assert(plan.getPartitionWriter.hasRss)
+ assert(plan.getOutputDataFile.isEmpty)
+ assert(plan.getOutputIndexFile.isEmpty)
+
+ writer.write(inputs)
+ val status = writer.stop(success = true).get
+ val localDataFile = SparkEnv.get.shuffleManager.shuffleBlockResolver
+ .asInstanceOf[IndexShuffleBlockResolver]
+ .getDataFile(dependency.shuffleId, context.taskAttemptId())
+
+ (
+ writer.getPartitionLengths(),
+ status.mapId,
+ context.taskAttemptId(),
+ client.mapperEndCalls.get(),
+ client.cleanupCalls.get(),
+ Option(client.lastPush).map(_.length).getOrElse(0),
+ localDataFile.exists(),
+ validations,
+ writer.stop(success = false).isEmpty)
+ })
+
+ assert(results.nonEmpty)
+ assert(results.forall(_._1.length == 3))
+ assert(results.exists(_._1.sum > 0))
+ assert(results.forall(result => result._2 == result._3))
+ assert(results.forall(_._4 == 1))
+ assert(results.forall(_._5 == 0))
+ assert(results.exists(_._6 >= 20))
+ assert(results.forall(!_._7))
+ assert(results.forall(_._8 == 2))
+ assert(results.forall(_._9))
+ }
+ }
+
+ test("mapper completion failures abort remote attempts without publishing MapStatus") {
+ withNativeShuffleDependency() { dependency =>
+ val numMappers = dependency.rdd.getNumPartitions
+ val results = spark.sparkContext.runJob(
+ dependency.rdd,
+ (context: TaskContext, inputs: Iterator[Product2[Int, ColumnarBatch]]) => {
+ val client = new RecordingCelebornPushClient
+ val expected = new IOException("Celeborn rejected completed map output")
+ client.mapperEndFailure = expected
+ val pusher = CelebornShufflePusherFactory.create(
+ SparkEnv.get.conf,
+ client,
+ 92,
+ numMappers,
+ dependency.partitioner.numPartitions,
+ context)
+ val writer =
+ CometCelebornNativeShuffleWriterSuite.newWriter(dependency, context, pusher)
+ val failure =
+ try {
+ writer.write(inputs)
+ null
+ } catch {
+ case error: IOException => error
+ }
+
+ (
+ failure eq expected,
+ client.mapperEndCalls.get(),
+ client.cleanupCalls.get(),
+ writer.mapStatus == null,
+ writer.stop(success = false).isEmpty)
+ })
+
+ assert(results.nonEmpty)
+ assert(results.forall(_._1))
+ assert(results.forall(_._2 == 1))
+ assert(results.forall(_._3 == 1))
+ assert(results.forall(_._4))
+ assert(results.forall(_._5))
+ }
+ }
+
+ test("low-cardinality range shuffle plans use the actual reducer partition count") {
+ withNativeShuffleDependency(rangePartitioning = true) { dependency =>
+ val actualPartitions = dependency.partitioner.numPartitions
+ assert(actualPartitions < dependency.outputPartitioning.get.numPartitions)
+ val numMappers = dependency.rdd.getNumPartitions
+ val results = spark.sparkContext.runJob(
+ dependency.rdd,
+ (context: TaskContext, inputs: Iterator[Product2[Int, ColumnarBatch]]) => {
+ val client = new RecordingCelebornPushClient
+ val pusher = CelebornShufflePusherFactory.create(
+ SparkEnv.get.conf,
+ client,
+ 93,
+ numMappers,
+ actualPartitions,
+ context)
+ val writer =
+ CometCelebornNativeShuffleWriterSuite.newWriter(dependency, context, pusher)
+ writer.write(inputs)
+ writer.stop(success = true).get
+
+ (
+ writer.getPartitionLengths().length,
+ Option(client.lastPush).map(_.numPartitions).getOrElse(actualPartitions))
+ })
+
+ assert(results.nonEmpty)
+ assert(results.forall(_._1 == actualPartitions))
+ assert(results.forall(_._2 == actualPartitions))
+ }
+ }
+
+ test("speculative losing attempts cannot finalize or publish remote shuffle output") {
+ withNativeShuffleDependency() { dependency =>
+ val numMappers = dependency.rdd.getNumPartitions
+ val results = spark.sparkContext.runJob(
+ dependency.rdd,
+ (context: TaskContext, inputs: Iterator[Product2[Int, ColumnarBatch]]) => {
+ val client = new RecordingCelebornPushClient
+ val pusher = CelebornShufflePusherFactory.create(
+ SparkEnv.get.conf,
+ client,
+ 94,
+ numMappers,
+ dependency.partitioner.numPartitions,
+ context)
+ val writer =
+ CometCelebornNativeShuffleWriterSuite.newWriter(dependency, context, pusher)
+ assert(
+ SparkEnv.get.outputCommitCoordinator.canCommit(
+ context.stageId(),
+ context.stageAttemptNumber(),
+ context.partitionId(),
+ context.attemptNumber() + 1))
+ val failure =
+ try {
+ writer.write(inputs)
+ null
+ } catch {
+ case denied: CommitDeniedException => denied
+ }
+
+ (
+ failure != null && !failure.toTaskCommitDeniedReason.countTowardsTaskFailures,
+ client.mapperEndCalls.get(),
+ client.cleanupCalls.get(),
+ writer.mapStatus == null,
+ writer.stop(success = false).isEmpty)
+ })
+
+ assert(results.nonEmpty)
+ assert(results.forall(_._1))
+ assert(results.forall(_._2 == 0))
+ assert(results.forall(_._3 == 1))
+ assert(results.forall(_._4))
+ assert(results.forall(_._5))
+ }
+ }
+
+ test("generation invalidation during completion cannot publish a stale MapStatus") {
+ withNativeShuffleDependency() { dependency =>
+ val numMappers = dependency.rdd.getNumPartitions
+ val results = spark.sparkContext.runJob(
+ dependency.rdd,
+ (context: TaskContext, inputs: Iterator[Product2[Int, ColumnarBatch]]) => {
+ val client = new RecordingCelebornPushClient
+ val pusher = CelebornShufflePusherFactory.create(
+ SparkEnv.get.conf,
+ client,
+ 95,
+ numMappers,
+ dependency.partitioner.numPartitions,
+ context)
+ var validations = 0
+ val writer = CometCelebornNativeShuffleWriterSuite.newWriter(
+ dependency,
+ context,
+ pusher,
+ commitAuthorized = true,
+ commitValidator = () => {
+ validations += 1
+ validations == 1
+ })
+ val failure =
+ try {
+ writer.write(inputs)
+ null
+ } catch {
+ case denied: CommitDeniedException => denied
+ }
+
+ (
+ failure != null && !failure.toTaskCommitDeniedReason.countTowardsTaskFailures,
+ validations,
+ client.mapperEndCalls.get(),
+ client.cleanupCalls.get(),
+ writer.mapStatus == null,
+ writer.stop(success = false).isEmpty)
+ })
+
+ assert(results.nonEmpty)
+ assert(results.forall(_._1))
+ assert(results.forall(_._2 == 2))
+ assert(results.forall(_._3 == 1))
+ assert(results.forall(_._4 == 0))
+ assert(results.forall(_._5))
+ assert(results.forall(_._6))
+ }
+ }
+
+ test("native registration rejects and cleans up Celeborn's delegated local fallback") {
+ withNativeShuffleDependency() { dependency =>
+ val backend = new CometCelebornNativeShuffleWriterSuite.LocalFallbackShuffleManager
+ val manager =
+ new CometCelebornShuffleManager(spark.sparkContext.getConf, true, (_, _) => backend)
+ val failure = intercept[UnsupportedOperationException] {
+ manager.registerShuffle(dependency.shuffleId, dependency)
+ }
+
+ assert(failure.getMessage.contains("local fallback"))
+ assert(backend.unregisteredShuffle.contains(dependency.shuffleId))
+ }
+ }
+}
+
+private[shuffle] object CometCelebornNativeShuffleWriterSuite {
+
+ def newWriter(
+ dependency: CometShuffleDependency[Int, ColumnarBatch, ColumnarBatch],
+ context: TaskContext,
+ pusher: CelebornShufflePartitionPusher,
+ commitAuthorized: Boolean = false,
+ commitValidator: () => Boolean = () => true): CometNativeShuffleWriter[Int, ColumnarBatch] =
+ new CometNativeShuffleWriter[Int, ColumnarBatch](
+ dependency.nativeShuffleSpec.get,
+ dependency.outputPartitioning.get,
+ dependency.outputAttributes,
+ dependency.shuffleWriteMetrics,
+ dependency.numParts,
+ dependency.shuffleId,
+ context.taskAttemptId(),
+ context,
+ context.taskMetrics().shuffleWriteMetrics,
+ dependency.rangePartitionBounds,
+ Some(
+ CelebornNativeShuffleDestination(
+ pusher,
+ pusher.maxFrameBytes(),
+ dependency.partitioner.numPartitions,
+ commitAuthorized,
+ commitValidator)))
+
+ final class LocalFallbackShuffleManager extends ShuffleManager {
+ var unregisteredShuffle: Option[Int] = None
+
+ override def registerShuffle[K, V, C](
+ shuffleId: Int,
+ dependency: ShuffleDependency[K, V, C]): ShuffleHandle =
+ new BaseShuffleHandle(shuffleId, dependency)
+
+ override def getWriter[K, V](
+ handle: ShuffleHandle,
+ mapId: Long,
+ context: TaskContext,
+ metrics: ShuffleWriteMetricsReporter): ShuffleWriter[K, V] =
+ throw new AssertionError("Native Comet data must not reach a local fallback writer")
+
+ override def getReader[K, C](
+ handle: ShuffleHandle,
+ startMapIndex: Int,
+ endMapIndex: Int,
+ startPartition: Int,
+ endPartition: Int,
+ context: TaskContext,
+ metrics: ShuffleReadMetricsReporter): ShuffleReader[K, C] =
+ throw new AssertionError("Native Comet data must not reach a local fallback reader")
+
+ override def shuffleBlockResolver: ShuffleBlockResolver = null
+
+ override def unregisterShuffle(shuffleId: Int): Boolean = {
+ unregisteredShuffle = Some(shuffleId)
+ true
+ }
+
+ override def stop(): Unit = ()
+ }
+}
diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleManagerSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleManagerSuite.scala
index 486aa1e481..7f1725f7a5 100644
--- a/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleManagerSuite.scala
+++ b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleManagerSuite.scala
@@ -19,9 +19,12 @@
package org.apache.spark.sql.comet.execution.shuffle
+import scala.collection.mutable
+
import org.scalatest.funsuite.AnyFunSuite
-import org.apache.spark.{ShuffleDependency, SparkConf, TaskContext}
+import org.apache.spark.{ExecutorLostFailure, ShuffleDependency, SparkConf, TaskContext, TaskEndReason, UnknownReason}
+import org.apache.spark.scheduler.OutputCommitCoordinator
import org.apache.spark.shuffle.{BaseShuffleHandle, ShuffleBlockResolver, ShuffleHandle, ShuffleManager, ShuffleReader, ShuffleReadMetricsReporter, ShuffleWriteMetricsReporter, ShuffleWriter}
class CometCelebornShuffleManagerSuite extends AnyFunSuite {
@@ -98,6 +101,250 @@ class CometCelebornShuffleManagerSuite extends AnyFunSuite {
new CometCelebornShuffleManager(conf, isDriver, (_, _) => backend)
}
+ private def startStage(
+ coordinator: OutputCommitCoordinator,
+ stageId: Int,
+ numMappers: Int): Unit = {
+ coordinator.getClass
+ .getMethod("stageStart", java.lang.Integer.TYPE, java.lang.Integer.TYPE)
+ .invoke(coordinator, Int.box(stageId), Int.box(numMappers - 1))
+ }
+
+ private def completeFailedAttempt(
+ coordinator: OutputCommitCoordinator,
+ stageId: Int,
+ stageAttempt: Int,
+ mapId: Int,
+ taskAttempt: Int,
+ reason: TaskEndReason = UnknownReason): Unit = {
+ coordinator.getClass
+ .getMethod(
+ "taskCompleted",
+ java.lang.Integer.TYPE,
+ java.lang.Integer.TYPE,
+ java.lang.Integer.TYPE,
+ java.lang.Integer.TYPE,
+ classOf[TaskEndReason])
+ .invoke(
+ coordinator,
+ Int.box(stageId),
+ Int.box(stageAttempt),
+ Int.box(mapId),
+ Int.box(taskAttempt),
+ reason)
+ }
+
+ private def sparkCommitOwners(
+ coordinator: OutputCommitCoordinator,
+ stageId: Int): Array[AnyRef] = coordinator.synchronized {
+ val statesField = classOf[OutputCommitCoordinator].getDeclaredField("stageStates")
+ statesField.setAccessible(true)
+ val state = statesField.get(coordinator).asInstanceOf[mutable.Map[Int, AnyRef]](stageId)
+ state.getClass.getMethod("authorizedCommitters").invoke(state).asInstanceOf[Array[AnyRef]]
+ }
+
+ test("replacement shuffle generations invalidate stale Comet map ownership") {
+ val sparkCoordinator = new OutputCommitCoordinator(new SparkConf(false), true)
+ val coordinator = new CelebornShuffleGenerationCoordinator(sparkCoordinator)
+ val stageId = 71
+ startStage(sparkCoordinator, stageId, numMappers = 2)
+
+ val first = PrepareCelebornShuffleGeneration(7, 91, stageId, 0, 2)
+ assert(coordinator.prepareGeneration(first))
+ val owner = coordinator.claimMapAttempt(ClaimCelebornMapAttempt(7, stageId, 0, 0, 0))
+ assert(owner.authorized)
+ assert(sparkCommitOwners(sparkCoordinator, stageId).forall(_ == null))
+ assert(
+ coordinator.validateMapAttempt(
+ ValidateCelebornMapAttempt(7, 91, stageId, 0, 0, 0, owner.epoch)))
+ assert(!coordinator.claimMapAttempt(ClaimCelebornMapAttempt(7, stageId, 0, 0, 1)).authorized)
+
+ val replacement = first.copy(celebornShuffleId = 92, stageAttempt = 1)
+ assert(coordinator.prepareGeneration(replacement))
+ val replacementOwner =
+ coordinator.claimMapAttempt(ClaimCelebornMapAttempt(7, stageId, 1, 0, 0))
+ assert(replacementOwner.authorized)
+ assert(replacementOwner.epoch > owner.epoch)
+ assert(
+ !coordinator.validateMapAttempt(
+ ValidateCelebornMapAttempt(7, 91, stageId, 0, 0, 0, owner.epoch)))
+ assert(
+ coordinator.validateMapAttempt(
+ ValidateCelebornMapAttempt(7, 92, stageId, 1, 0, 0, replacementOwner.epoch)))
+ }
+
+ test("speculation arriving first reserves the original Spark map attempt") {
+ val sparkCoordinator = new OutputCommitCoordinator(new SparkConf(false), true)
+ val coordinator = new CelebornShuffleGenerationCoordinator(sparkCoordinator)
+ val stageId = 72
+ startStage(sparkCoordinator, stageId, numMappers = 1)
+
+ val speculative = coordinator.claimMapAttempt(ClaimCelebornMapAttempt(8, stageId, 0, 0, 1))
+ assert(!speculative.authorized)
+ assert(!speculative.requiresGenerationResolution)
+
+ val original = coordinator.claimMapAttempt(ClaimCelebornMapAttempt(8, stageId, 0, 0, 0))
+ assert(original.authorized)
+ assert(original.epoch == speculative.epoch)
+ }
+
+ test("abandoning an unsafe owner releases Comet admission without taking Spark commit locks") {
+ val sparkCoordinator = new OutputCommitCoordinator(new SparkConf(false), true)
+ val coordinator = new CelebornShuffleGenerationCoordinator(sparkCoordinator, _ => false)
+ val stageId = 73
+ startStage(sparkCoordinator, stageId, numMappers = 1)
+
+ val generation = PrepareCelebornShuffleGeneration(9, 93, stageId, 0, 1)
+ assert(coordinator.prepareGeneration(generation))
+ val original = coordinator.claimMapAttempt(ClaimCelebornMapAttempt(9, stageId, 0, 0, 0))
+ assert(original.authorized)
+ val prepared = coordinator.prepareGenerationAndClaim(
+ generation
+ .copy(mapId = 0, taskAttempt = 0, claimEpoch = original.epoch, claimAuthorized = true))
+ assert(prepared.authorized)
+ assert(sparkCommitOwners(sparkCoordinator, stageId).forall(_ == null))
+
+ assert(
+ coordinator.abandonMapAttempt(
+ AbandonCelebornMapAttempt(9, 93, stageId, 0, 0, 0, prepared.epoch, 101L)))
+ val retry = coordinator.claimMapAttempt(ClaimCelebornMapAttempt(9, stageId, 0, 0, 1))
+ assert(retry.authorized)
+ assert(sparkCommitOwners(sparkCoordinator, stageId).forall(_ == null))
+ assert(
+ coordinator.abandonMapAttempt(
+ AbandonCelebornMapAttempt(9, 93, stageId, 0, 0, 0, prepared.epoch, 101L)))
+ assert(coordinator.claimMapAttempt(ClaimCelebornMapAttempt(9, stageId, 0, 0, 1)).authorized)
+ }
+
+ test("stale and recorded-failed claims cannot pass the generation invalidation gate") {
+ val sparkCoordinator = new OutputCommitCoordinator(new SparkConf(false), true)
+ var livenessChecks = 0
+ val coordinator = new CelebornShuffleGenerationCoordinator(
+ sparkCoordinator,
+ _ => {
+ livenessChecks += 1
+ true
+ })
+ val stageId = 77
+ startStage(sparkCoordinator, stageId, numMappers = 1)
+ val firstGeneration = PrepareCelebornShuffleGeneration(13, 100, stageId, 0, 1)
+ assert(coordinator.prepareGeneration(firstGeneration))
+ val first = coordinator.claimMapAttempt(ClaimCelebornMapAttempt(13, stageId, 0, 0, 0))
+ assert(first.authorized)
+
+ completeFailedAttempt(sparkCoordinator, stageId, 0, 0, 0)
+ assert(
+ coordinator.abandonMapAttempt(
+ AbandonCelebornMapAttempt(13, 100, stageId, 0, 0, 0, first.epoch, 201L)))
+ assert(livenessChecks == 0)
+
+ val replacement = firstGeneration.copy(celebornShuffleId = 101, stageAttempt = 1)
+ assert(coordinator.prepareGeneration(replacement))
+ val current = coordinator.claimMapAttempt(ClaimCelebornMapAttempt(13, stageId, 1, 0, 0))
+ assert(current.authorized)
+ assert(
+ coordinator.abandonMapAttempt(
+ AbandonCelebornMapAttempt(13, 100, stageId, 0, 0, 0, first.epoch, 201L)))
+ assert(livenessChecks == 0)
+ assert(
+ coordinator.validateMapAttempt(
+ ValidateCelebornMapAttempt(13, 101, stageId, 1, 0, 0, current.epoch)))
+ assert(
+ !coordinator.abandonMapAttempt(
+ AbandonCelebornMapAttempt(13, 101, stageId, 1, 0, 0, current.epoch, 202L)))
+ assert(livenessChecks == 1)
+ }
+
+ Seq[(String, TaskEndReason)](
+ "original task failure" -> UnknownReason,
+ "executor loss" -> ExecutorLostFailure(
+ "lost-executor",
+ exitCausedByApp = false,
+ reason = Some("executor disconnected before map completion"))).foreach {
+ case (description, reason) =>
+ test(
+ s"$description releases Comet admission without becoming a Spark file-commit failure") {
+ val sparkCoordinator = new OutputCommitCoordinator(new SparkConf(false), true)
+ val coordinator = new CelebornShuffleGenerationCoordinator(sparkCoordinator)
+ val stageId = 76
+ startStage(sparkCoordinator, stageId, numMappers = 1)
+ assert(
+ coordinator.prepareGeneration(PrepareCelebornShuffleGeneration(12, 98, stageId, 0, 1)))
+ val original = coordinator.claimMapAttempt(ClaimCelebornMapAttempt(12, stageId, 0, 0, 0))
+ assert(original.authorized)
+
+ // Spark 3.4 calls stageFailed when taskCompleted matches an authorized file committer.
+ // Keeping that slot empty also covers executor loss, where no task-side cleanup can run.
+ assert(sparkCommitOwners(sparkCoordinator, stageId).forall(_ == null))
+ completeFailedAttempt(sparkCoordinator, stageId, 0, 0, 0, reason)
+ assert(sparkCommitOwners(sparkCoordinator, stageId).forall(_ == null))
+ assert(
+ !coordinator.validateMapAttempt(
+ ValidateCelebornMapAttempt(12, 98, stageId, 0, 0, 0, original.epoch)))
+ assert(
+ !coordinator.claimMapAttempt(ClaimCelebornMapAttempt(12, stageId, 0, 0, 0)).authorized)
+
+ val retry = coordinator.claimMapAttempt(ClaimCelebornMapAttempt(12, stageId, 0, 0, 1))
+ assert(retry.authorized)
+ assert(
+ !coordinator.claimMapAttempt(ClaimCelebornMapAttempt(12, stageId, 0, 0, 2)).authorized)
+ assert(sparkCommitOwners(sparkCoordinator, stageId).forall(_ == null))
+ }
+ }
+
+ test("replacement generations preserve another partition's previously recorded failure") {
+ val sparkCoordinator = new OutputCommitCoordinator(new SparkConf(false), true)
+ val coordinator = new CelebornShuffleGenerationCoordinator(sparkCoordinator)
+ val stageId = 74
+ startStage(sparkCoordinator, stageId, numMappers = 2)
+
+ val original = PrepareCelebornShuffleGeneration(10, 94, stageId, 0, 2)
+ assert(coordinator.prepareGeneration(original))
+ assert(coordinator.claimMapAttempt(ClaimCelebornMapAttempt(10, stageId, 0, 0, 0)).authorized)
+ completeFailedAttempt(sparkCoordinator, stageId, 1, mapId = 0, taskAttempt = 0)
+
+ val firstReplacement =
+ coordinator.claimMapAttempt(ClaimCelebornMapAttempt(10, stageId, 1, 1, 0))
+ assert(firstReplacement.authorized)
+ val replacement = original.copy(celebornShuffleId = 95, stageAttempt = 1)
+ assert(
+ coordinator
+ .prepareGenerationAndClaim(
+ replacement.copy(
+ mapId = 1,
+ taskAttempt = 0,
+ claimEpoch = firstReplacement.epoch,
+ claimAuthorized = true))
+ .authorized)
+
+ val failedPartitionRetry =
+ coordinator.claimMapAttempt(ClaimCelebornMapAttempt(10, stageId, 1, 0, 1))
+ assert(failedPartitionRetry.authorized)
+ }
+
+ test("invalidated generations reject stale owners and permit a fresh replacement") {
+ val sparkCoordinator = new OutputCommitCoordinator(new SparkConf(false), true)
+ val coordinator = new CelebornShuffleGenerationCoordinator(sparkCoordinator)
+ val stageId = 75
+ startStage(sparkCoordinator, stageId, numMappers = 1)
+
+ val generation = PrepareCelebornShuffleGeneration(11, 96, stageId, 0, 1)
+ assert(coordinator.prepareGeneration(generation))
+ val owner = coordinator.claimMapAttempt(ClaimCelebornMapAttempt(11, stageId, 0, 0, 0))
+ assert(owner.authorized)
+ assert(
+ coordinator.invalidateGeneration(
+ InvalidateCelebornShuffleGeneration(11, 96, stageId, 0, owner.epoch)))
+ assert(
+ !coordinator.validateMapAttempt(
+ ValidateCelebornMapAttempt(11, 96, stageId, 0, 0, 0, owner.epoch)))
+ assert(!coordinator.prepareGeneration(generation))
+
+ val replacement = generation.copy(celebornShuffleId = 97, stageAttempt = 1)
+ assert(coordinator.prepareGeneration(replacement))
+ assert(coordinator.claimMapAttempt(ClaimCelebornMapAttempt(11, stageId, 1, 0, 0)).authorized)
+ }
+
test("manager preserves the application Spark configuration and driver identity") {
val conf = new SparkConf(false).set("spark.app.name", "existing-celeborn-application")
val backend = new RecordingShuffleManager
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org