ziting-openai commented on code in PR #5513:
URL: https://github.com/apache/datafusion-comet/pull/5513#discussion_r3876294379
##########
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] This still reproduces at 16d71677 despite the new transport snapshot.
Celeborn 0.6/0.7 removes a failed request from outstandingPushes before
invoking its callback. If the transport thread pauses between removal and
callback, abort() snapshots no request, cleanup() sets Cleaned Up, and the
resumed callback returns without removeBatch(); cancelledTransportBaseline is
never initialized, so executor-wide admission remains permanently reserved. The
new regression only fails the transport after abort has already captured the
request. Please cover the removal-before-snapshot/callback-after-cleanup
interleaving and track completion independently of transport-map membership.
--
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]
