sunchao commented on code in PR #5513:
URL: https://github.com/apache/datafusion-comet/pull/5513#discussion_r3877264838
##########
native/shuffle/src/writers/rss/rss_partition_writer.rs:
##########
@@ -125,35 +132,633 @@ 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
+ }
+
+ fn push_batch_within_limit(
+ &mut self,
+ partition_id: i32,
+ batch: &RecordBatch,
+ metrics: &ShufflePartitionerMetrics,
+ ) -> Result<()> {
+ if batch.num_rows() == 0 {
+ return Ok(());
+ }
+
+ // Estimate only live nested rows before allocation. Dictionary values
remain charged
+ // because Arrow's dense garbage collection scans their complete value
tables.
+ let original_size =
Self::estimated_pre_compaction_ipc_data_size(batch)?;
+ let compaction_scratch = Self::estimated_compaction_scratch(batch)?;
+ let minimum_size =
Self::estimated_minimum_compacted_ipc_data_size(batch)?;
+ if minimum_size > self.max_frame_size && batch.num_rows() > 1 {
+ return self.push_split_batch(partition_id, batch, metrics);
+ }
+
+ // Arrow's array estimates omit the schema, FlatBuffers messages, and
compression-frame
+ // overhead. A small metadata allowance avoids most retries without
charging every tiny
+ // batch for the configured maximum frame. Unusually large metadata is
handled below by
+ // releasing the entire reservation and retrying with a larger,
still-bounded estimate.
+ const IPC_METADATA_RESERVATION_BYTES: usize = 4 * 1024;
+ let mut frame_bound = original_size
+ .saturating_add(IPC_METADATA_RESERVATION_BYTES)
+ .min(self.max_frame_size);
- let encoded_size = self.block_writer.write_batch(
- &batch,
- &mut Cursor::new(&mut self.frame),
+ loop {
+ // Native IPC, its JNI byte array, and Celeborn's copied transport
request overlap.
+ // Acquire all three copies before compaction/encoding to avoid
allocation before
+ // admission and deadlocks caused by growing another task's
partial reservation.
+ let overlapping_copies = frame_bound.checked_mul(3).ok_or_else(|| {
+ DataFusionError::Execution(
+ "Remote shuffle frame-copy reservation exceeds the native
integer limit"
+ .to_string(),
+ )
+ })?;
+ // Legacy Java callbacks retain the full JVM-soft frame limit but
inherit a no-op
+ // reservation method whose argument is still a signed Java int.
Concrete bounded
+ // pushers advertise a smaller frame limit, so their actual
three-copy reservation
+ // fits; capping here preserves legacy callbacks and their failure
propagation.
+ let reservation = overlapping_copies
+ .max(original_size.saturating_add(compaction_scratch))
+ .min(self.pusher.max_reservation_size());
+ self.pusher.reserve_partition_data(reservation)?;
Review Comment:
[P2] Charge codec workspace and allocated capacity before encoding
At exact head `be5f4af7d7876e862af9be172b95e2e704acf049`, the new
`spark.comet.shuffle.rss.maxInFlightBytes` contract includes native encoding
scratch, but this reservation covers estimated IPC/frame lengths and compaction
only. The actual `ShuffleBlockWriter` also creates codec working buffers, and
`BoundedBuffer` bounds its Vec's length rather than its allocated capacity.
Consequently, a valid budget can be exceeded even by one admitted encoder,
before transport completion or the separate callback-lifetime issue matters.
I independently compiled the exact current/base RSS writer sources with the
unchanged block writer and measured successful live Rust heap allocation
requests only during encoding, excluding input/schema/metric construction. With
a 16,384-byte budget, the real Java pusher accepts a one-Int32-row LZ4
reservation of 12,687 bytes (12,703 including the Celeborn header; effective
frame cap 5,456), but encoding peaks at 154,497 bytes and produces a 238-byte
frame. The uncompressed control peaks at 1,160 bytes. This is not only a
tiny-budget codec case: with 16,384 Int32 rows and a 262,144-byte budget, the
accepted reservation is 215,424 + 16, while LZ4 peaks at 285,321 bytes. Without
compression, that batch's 67,996-byte frame has 135,976 bytes of Vec capacity;
adding its JNI and Celeborn copies requires at least 271,984 bytes, also above
the budget. The Java reservation checks were exercised with stock Celeborn
0.6.3 and 0.7.0.
Base `eabb5d4773091b983d8fce713f0e34b1cf93f877` has the same codec
allocation peaks but no encoding-admission protocol or byte-budget setting. The
introduced defect is the new limit admitting work without accounting for memory
it explicitly promises to cover; this is not a claim that the codec allocations
themselves were introduced here. Please include codec workspace and actual
buffer capacities in the bound, or reject budgets that cannot accommodate them,
before allowing encoding to begin.
##########
spark/src/main/java/org/apache/comet/shuffle/CelebornShufflePartitionPusher.java:
##########
@@ -175,62 +570,449 @@ public void pushPartitionData(int partitionId, byte[]
data, int length) throws I
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 "
+ declaredBodyLength
+ " body bytes, but contains "
+ (length - Long.BYTES));
}
+ }
- final int accepted;
+ 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");
+ }
+ encodingReservation.remove();
+ synchronized (lifecycleLock) {
+ activeEncoders--;
+ lifecycleLock.notifyAll();
+ }
+ if (required < reservation.bytes) {
+ admission.release(reservation.bytes - required);
+ reservation.bytes = required;
+ }
+ return reservation;
+ }
+ admission.acquire(required, this::isAborted);
+ return new PushReservation(required);
+ }
+
+ 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;
+ }
+ }
+
+ 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, int
acceptedBytes) {
+ int released;
+ synchronized (lifecycleLock) {
+ reservation.pushState = pushState;
+ reservation.submitted = true;
+ pushState.submittedPushes++;
+ int retained = Math.min(reservation.bytes, acceptedBytes);
+ released = reservation.bytes - retained;
+ reservation.bytes = retained;
Review Comment:
[P2] Retain native and JNI frame charges until their buffers are released
At exact head `be5f4af7d7876e862af9be172b95e2e704acf049`, `markSubmitted`
shrinks admission to the Celeborn transport byte count while the caller still
owns both the native output Vec and the JNI byte array. A fast transport
completion can then release the remainder. The JNI local frame is popped only
after `JavaShufflePartitionPusher::push_partition_data` returns from its
callback, and the native RSS writer drops its output afterward. Another map can
therefore acquire the released permits while those two frame copies remain live.
I reproduced this using the actual Rust RSS writer, real
JavaShufflePartitionPusher JNI bridge, unmodified Java pusher, and stock
Celeborn 0.6.3/0.7.0 request registration/removal/success callbacks. Only
network delivery and scheduling before JNI return were controlled. Two maps
each encode 2,048 Int32 rows without compression, producing 8,860-byte frames,
with `maxFrameBytes=8860` and the valid shared budget `maxInFlightBytes=26596`.
After the first production Java push returns but before JNI releases its
array/native frame, all 26,596 permits are available. The second encoder
enters; the two paused calls then own at least 35,440 bytes of native/JNI frame
data alone, excluding capacity and all transport/allocator overhead. A control
deferring semaphore release until the first native writer returns leaves zero
permits, blocks the second encoder until those buffers are gone, and then lets
both complete.
The exact base `eabb5d4773091b983d8fce713f0e34b1cf93f877` has these same
frame lifetimes but no shared admission protocol; the faulty release transition
is new in this PR. This is separate from the fixed cleanup/transport-callback
leak and from codec workspace accounting. Please keep the native/JNI portion
charged until those owners actually finish rather than releasing it when the
raw Celeborn call returns or its transport callback completes.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
