sunchao commented on code in PR #5513:
URL: https://github.com/apache/datafusion-comet/pull/5513#discussion_r3876101388
##########
spark/src/main/java/org/apache/comet/shuffle/CelebornShufflePartitionPusher.java:
##########
@@ -207,30 +479,503 @@ public void pushPartitionData(int partitionId, byte[]
data, int length) throws I
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;
+ 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, accepted);
}
- if (cause instanceof RuntimeException) {
- throw (RuntimeException) 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");
+ }
+ if (accepted > reservation.bytes) {
+ throw new IOException(
+ "Celeborn encrypted shuffle request exceeds its reserved in-flight
byte limit");
}
- if (cause instanceof Error) {
- throw (Error) cause;
+ 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) {
+ admission.release(reservation.bytes);
+ }
+ }
+ try {
+ endPush();
+ } catch (IOException cleanupFailure) {
+ if (failure == null) {
+ throw cleanupFailure;
+ }
+ if (cleanupFailure != failure) {
+ failure.addSuppressed(cleanupFailure);
+ }
}
- throw new IOException("Celeborn raw shuffle push failed", cause);
}
+ }
- int minimumAccepted = length + CELEBORN_BATCH_HEADER_BYTES;
- if (accepted < minimumAccepted) {
+ 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");
+ }
+ if (data == null) {
+ throw new IOException("Celeborn shuffle frame must not be null");
+ }
+ if (length > Integer.MAX_VALUE - CELEBORN_BATCH_HEADER_BYTES) {
+ throw new IOException("Celeborn shuffle frame and transport header
exceed the byte limit");
+ }
+ 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");
+ }
+
+ long declaredBodyLength =
ByteBuffer.wrap(data).order(ByteOrder.LITTLE_ENDIAN).getLong();
+ if (declaredBodyLength != (long) length - Long.BYTES) {
throw new IOException(
- "Celeborn raw shuffle push accepted "
- + accepted
- + " bytes; expected at least "
- + minimumAccepted
- + " including its transport header");
+ "Celeborn shuffle frame declares "
+ + declaredBodyLength
+ + " 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");
+ }
+ 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);
+ ObservedPushState created =
+ new ObservedPushState(
+ (LongAdder) totalInFlightRequests.get(tracker),
+ (AtomicReference<?>) pushStateException.get(pushState));
+ 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;
+ }
+ if (released > 0) {
+ admission.release(released);
+ }
+ reconcileAcceptedPushes();
+ }
+
+ private void safelyReconcileAcceptedPushes() {
+ try {
+ 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);
+ }
+ }
+
+ private void reconcileAcceptedPushes() {
+ ArrayList<PushReservation> completed = new ArrayList<>();
+ IOException detectedFailure = null;
+ synchronized (lifecycleLock) {
+ for (ObservedPushState observed : observedPushStates.values()) {
+ Object failure = observed.exception.get();
+ if (failure instanceof IOException
+ && !"Cleaned Up".equals(((IOException) failure).getMessage())) {
+ 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;
+ }
+ }
+ long retainedRequests =
+ Math.max(
+ 0L, observed.inFlightRequests.sum() -
(observed.terminalFailureCredited ? 1L : 0L));
+ long completions = Math.max(0L, observed.submittedPushes -
retainedRequests);
Review Comment:
[P2] Track failed transport completion after task cleanup
On `bd439e12137b12b81276b40907f7130371ef4c14`, cancelling a map with an
outstanding push can permanently consume this executor-wide admission budget.
Stock Celeborn 0.6.3 and 0.7.0 `cleanup()` sets the captured
`PushState.exception` to `"Cleaned Up"`; if that request subsequently fails
(for example, its connection closes or it times out), the stock raw-push
callback returns immediately because the exception is already set, without
calling `removeBatch()`. Thus the counter used here remains positive even after
the transport has completed. The reconciler never releases the reservation, and
later native maps sharing the client can wait indefinitely in `acquire()`.
I reproduced this with the unmodified Comet pusher and real Celeborn
`ShuffleClientImpl`, `PushState`, cleanup, and raw callbacks, replacing only
network delivery. A real native `ShuffleBlockWriter` frame containing 32 Int32
rows is 604 bytes. With valid `maxFrameBytes=604` and `maxInFlightBytes=2416`,
reserve 1812 bytes, push the frame, abort the map, and deliver the held stock
callback's `onFailure`. All original transport callbacks have completed, but
the stock counter stays at 1 and available admission stays at 1796; another
map's required 1828-byte admission blocks until interrupted. The head's
successful-callback control releases all 2416 bytes, and the same
failed-callback sequence on base `eabb5d4773091b983d8fce713f0e34b1cf93f877`
does not block another map.
Please account for actual failed transport completion after cleanup rather
than relying on this counter to eventually fall. The existing cancellation test
only completes the cancelled request successfully, so it misses this failure
path.
--
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]
