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 7474e917420..f168df4b51a 100644 --- a/spark/src/main/java/org/apache/comet/shuffle/CelebornShufflePartitionPusher.java +++ b/spark/src/main/java/org/apache/comet/shuffle/CelebornShufflePartitionPusher.java @@ -66,6 +66,7 @@ public final class CelebornShufflePartitionPusher implements ShufflePartitionPus private final Field cryptoHandler; private final ExecutorShufflePushAdmission admission; private final CelebornTransportCallbackTracker transportCallbacks; + private final Object submissionLock = new Object(); private final int shuffleId; private final int mapId; private final int encodedAttemptId; @@ -86,10 +87,10 @@ public final class CelebornShufflePartitionPusher implements ShufflePartitionPus private int activeClientPushes; private int activeEncoders; private ScheduledFuture completionReconciliation; + private Thread activeClientThread; private Thread mapperEndThread; - private boolean cleanupStarted; - private boolean cleanupAfterActivePushes; - private boolean finalCleanupStarted; + private boolean cleanupRequested; + private boolean cleanupClaimed; private IOException asynchronousFailure; private enum State { @@ -103,11 +104,13 @@ private static final class PushReservation { private final int bytes; private ObservedPushState pushState; private CelebornTransportCallbackTracker.Push transportPush; + private Object clientPushState; private boolean claimed; private boolean submitted; private boolean nativeReleased; private boolean transportComplete; private boolean released; + private boolean pushStateFallback; private PushReservation(int bytes, boolean nativeOwned) { this.bytes = bytes; @@ -118,9 +121,8 @@ private PushReservation(int bytes, boolean nativeOwned) { private static final class ObservedPushState { private final LongAdder inFlightRequests; private final AtomicReference exception; - private long submittedPushes; - private long releasedPushes; - private boolean terminalFailureCredited; + private long fallbackSubmittedPushes; + private long fallbackReleasedPushes; private ObservedPushState(LongAdder inFlightRequests, AtomicReference exception) { this.inFlightRequests = inFlightRequests; @@ -494,16 +496,37 @@ public void pushPartitionData(int partitionId, byte[] data, int length) throws I requireUnencryptedClient(); validateFrame(partitionId, data, length); beginPush(); + PushReservation reservation; + try { + // Admission can wait for an older request's completion. Never hold the submission lock + // while waiting; callbacks and retries must remain able to make progress. + reservation = claimEncodingReservation(length); + } catch (IOException | RuntimeException | Error failure) { + abortAndSuppress(failure); + try { + endPush(); + } catch (IOException cleanupFailure) { + if (cleanupFailure != failure) { + failure.addSuppressed(cleanupFailure); + } + } + throw failure; + } + synchronized (submissionLock) { + pushClaimedPartitionData(partitionId, data, length, reservation); + } + } - PushReservation reservation = null; + private void pushClaimedPartitionData( + int partitionId, byte[] data, int length, PushReservation reservation) throws IOException { boolean registered = false; boolean submitted = false; boolean insideClient = false; Throwable failure = null; try { - reservation = claimEncodingReservation(length); beginClientPush(); insideClient = true; + awaitPushStateFallbackSlot(); if (computeBatchCRC != null) { computeBatchCRC.invoke( @@ -518,7 +541,7 @@ public void pushPartitionData(int partitionId, byte[] data, int length) throws I throw new IOException("Celeborn returned a null push state for this map attempt"); } observed = observePushState(pushState); - registerPendingPush(reservation, observed); + registerPendingPush(reservation, pushState, observed); registered = true; } @@ -527,6 +550,10 @@ public void pushPartitionData(int partitionId, byte[] data, int length) throws I transportCallbacks == null ? null : transportCallbacks.beginPush()) { synchronized (lifecycleLock) { reservation.transportPush = transportPush; + if (state == State.ABORTED) { + throw new IOException( + "Celeborn shuffle map attempt was aborted before its client invocation"); + } } accepted = (int) @@ -549,10 +576,11 @@ public void pushPartitionData(int partitionId, byte[] data, int length) throws I if (clientPushStates != null) { Object current = ((Map) clientPushStates.get(shuffleClient)).get(mapKey()); if (current != null && current != pushState) { + pushState = current; observed = observePushState(current); } } - markSubmitted(reservation, observed); + markSubmitted(reservation, pushState, observed); } int minimumAccepted = length + CELEBORN_BATCH_HEADER_BYTES; @@ -590,23 +618,38 @@ public void pushPartitionData(int partitionId, byte[] data, int length) throws I if (insideClient) { endClientPush(); } + IOException deferredCleanupFailure = null; if (reservation != null) { - if (registered && !submitted) { - releaseUnsubmittedPush(reservation); - } else if (!registered) { - completeTransport(reservation); + try { + if (registered && !submitted) { + releaseUnsubmittedPush(reservation); + } else if (!registered) { + completeTransport(reservation); + } + } catch (IOException cleanupFailure) { + if (failure == null) { + deferredCleanupFailure = cleanupFailure; + } else if (cleanupFailure != failure) { + failure.addSuppressed(cleanupFailure); + } } } try { endPush(); } catch (IOException cleanupFailure) { if (failure == null) { - throw cleanupFailure; - } - if (cleanupFailure != failure) { + if (deferredCleanupFailure == null) { + deferredCleanupFailure = cleanupFailure; + } else if (cleanupFailure != deferredCleanupFailure) { + deferredCleanupFailure.addSuppressed(cleanupFailure); + } + } else if (cleanupFailure != failure) { failure.addSuppressed(cleanupFailure); } } + if (deferredCleanupFailure != null) { + throw deferredCleanupFailure; + } } } @@ -692,15 +735,35 @@ private void beginClientPush() throws IOException { throw new IOException("Celeborn shuffle map attempt was aborted before its client push"); } activeClientPushes++; + activeClientThread = Thread.currentThread(); } } - private void registerPendingPush(PushReservation reservation, ObservedPushState pushState) + private void awaitPushStateFallbackSlot() throws IOException { + synchronized (lifecycleLock) { + while (state != State.ABORTED && hasPendingPushStateFallback()) { + try { + lifecycleLock.wait(); + } catch (InterruptedException failure) { + Thread.currentThread().interrupt(); + throw new IOException( + "Interrupted while waiting for a prior Celeborn fallback push", failure); + } + } + if (state == State.ABORTED) { + throw new IOException("Celeborn shuffle map attempt was aborted before its client push"); + } + } + } + + private void registerPendingPush( + PushReservation reservation, Object clientPushState, ObservedPushState pushState) throws IOException { synchronized (lifecycleLock) { if (state == State.ABORTED) { throw new IOException("Celeborn shuffle map attempt was aborted before submission"); } + reservation.clientPushState = clientPushState; reservation.pushState = pushState; pendingPushes.addLast(reservation); if (completionReconciliation == null || completionReconciliation.isDone()) { @@ -714,15 +777,27 @@ private void registerPendingPush(PushReservation reservation, ObservedPushState } } - private void markSubmitted(PushReservation reservation, ObservedPushState pushState) { + private void markSubmitted( + PushReservation reservation, Object clientPushState, ObservedPushState pushState) + throws IOException { synchronized (lifecycleLock) { + reservation.clientPushState = clientPushState; reservation.pushState = pushState; reservation.submitted = true; - pushState.submittedPushes++; + registerPushStateFallback(reservation); } reconcileAcceptedPushes(); } + private void registerPushStateFallback(PushReservation reservation) { + if (!reservation.pushStateFallback + && (reservation.transportPush == null + || !reservation.transportPush.usesTransportOwnership())) { + reservation.pushStateFallback = true; + reservation.pushState.fallbackSubmittedPushes++; + } + } + // 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) { @@ -745,9 +820,11 @@ private void completeTransport(PushReservation reservation) { private void safelyReconcileAcceptedPushes() { try { reconcileAcceptedPushes(); - } catch (RuntimeException | Error cause) { + } catch (IOException | RuntimeException | Error cause) { IOException failure = - new IOException("Celeborn push completion reconciliation failed", cause); + cause instanceof IOException + ? (IOException) cause + : new IOException("Celeborn push completion reconciliation failed", cause); synchronized (lifecycleLock) { if (asynchronousFailure == null) { asynchronousFailure = failure; @@ -759,10 +836,22 @@ private void safelyReconcileAcceptedPushes() { } } - private void reconcileAcceptedPushes() { + private void reconcileAcceptedPushes() throws IOException { ArrayList completed = new ArrayList<>(); IOException detectedFailure = null; synchronized (lifecycleLock) { + Iterator pending = pendingPushes.iterator(); + while (pending.hasNext()) { + PushReservation reservation = pending.next(); + if (!reservation.submitted) { + continue; + } + registerPushStateFallback(reservation); + if (!reservation.pushStateFallback && reservation.transportPush.isComplete()) { + pending.remove(); + completed.add(reservation); + } + } for (ObservedPushState observed : observedPushStates.values()) { Object failure = observed.exception.get(); boolean cleanedUp = @@ -770,51 +859,41 @@ private void reconcileAcceptedPushes() { && "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. + if (hasPendingTransportOwnership(observed)) { + // A stock counter cannot distinguish a fallback request from a transport-owned request. + // Wait for precise requests on this push state before crediting fallback completions. 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) { + boolean terminalFailure = + failure instanceof IOException + && !cleanedUp + && isTerminalRawPushFailure((IOException) failure); + long counterCredits = terminalFailure ? 1L : 0L; + long retainedRequests = Math.max(0L, observed.inFlightRequests.sum() - counterCredits); + long completions = Math.max(0L, observed.fallbackSubmittedPushes - retainedRequests); + while (observed.fallbackReleasedPushes < completions) { + PushReservation reservation = smallestFallbackReservation(observed); + if (reservation == null + || (reservation.transportPush != null + && !reservation.transportPush.retainedTransportOwnershipComplete())) { break; } pendingPushes.remove(reservation); - observed.releasedPushes++; + observed.fallbackReleasedPushes++; completed.add(reservation); } } - if (transportCallbacks != null) { - Iterator 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(); } + if (!completed.isEmpty()) { + lifecycleLock.notifyAll(); + } } for (PushReservation reservation : completed) { completeTransport(reservation); @@ -822,6 +901,7 @@ private void reconcileAcceptedPushes() { if (detectedFailure != null) { abortAndSuppress(detectedFailure); } + performDeferredCleanupIfReady(); } private static boolean isTerminalRawPushFailure(IOException failure) { @@ -832,11 +912,23 @@ private static boolean isTerminalRawPushFailure(IOException failure) { && message.contains(" batch "); } - private PushReservation smallestSubmittedReservation(ObservedPushState observed) { + private boolean hasPendingTransportOwnership(ObservedPushState observed) { + for (PushReservation reservation : pendingPushes) { + if (reservation.pushState == observed + && reservation.submitted + && !reservation.pushStateFallback) { + return true; + } + } + return false; + } + + private PushReservation smallestFallbackReservation(ObservedPushState observed) { PushReservation smallest = null; for (PushReservation reservation : pendingPushes) { if (reservation.pushState == observed && reservation.submitted + && reservation.pushStateFallback && (smallest == null || reservation.bytes < smallest.bytes)) { smallest = reservation; } @@ -851,14 +943,50 @@ private void stopCompletionReconciliation() { } } - private void releaseUnsubmittedPush(PushReservation reservation) { + private void refreshPendingPushState(PushReservation reservation) throws IOException { + if (clientPushStates == null) { + return; + } + final Object current; + try { + Object states = clientPushStates.get(shuffleClient); + if (!(states instanceof Map)) { + throw new IOException("Celeborn returned incompatible push-state storage"); + } + current = ((Map) states).get(mapKey()); + if (current == null || current == reservation.clientPushState) { + return; + } + ObservedPushState observed = observePushState(current); + synchronized (lifecycleLock) { + reservation.clientPushState = current; + reservation.pushState = observed; + } + } catch (IllegalAccessException failure) { + throw new IOException("Cannot inspect Celeborn push-state storage", failure); + } + } + + private void releaseUnsubmittedPush(PushReservation reservation) throws IOException { + // The raw call may recreate its PushState during cleanup before throwing. Resolve that state + // again before deciding whether counter-backed completion owns the published request. + refreshPendingPushState(reservation); boolean removed; synchronized (lifecycleLock) { - if (reservation.transportPush != null && !reservation.transportPush.isComplete()) { + if (reservation.transportPush != null + && reservation.transportPush.usesTransportOwnership() + && !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 if (reservation.transportPush == null + || !reservation.transportPush.usesTransportOwnership()) { + // The raw call may have published an uninstrumented request before throwing. Counter + // reconciliation releases it immediately when no request was actually retained. + reservation.submitted = true; + registerPushStateFallback(reservation); + removed = false; } else { removed = pendingPushes.remove(reservation); } @@ -873,6 +1001,9 @@ private void releaseUnsubmittedPush(PushReservation reservation) { private void endClientPush() { synchronized (lifecycleLock) { + if (activeClientThread == Thread.currentThread()) { + activeClientThread = null; + } activeClientPushes--; if (pendingPushes.isEmpty() && activeClientPushes == 0) { stopCompletionReconciliation(); @@ -881,24 +1012,13 @@ private void endClientPush() { } 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(); } + performDeferredCleanupIfReady(); } /** Drains asynchronous requests, commits this map, and returns Comet bytes per partition. */ @@ -977,29 +1097,76 @@ public long[] finish() throws IOException { /** 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; + cleanupRequested = true; + if (activeClientThread != null && activeClientThread != Thread.currentThread()) { + // Interrupt while ownership is protected by lifecycleLock. Otherwise the raw call could + // return and this Thread could resume unrelated work before the interrupt is delivered. + activeClientThread.interrupt(); } - cleanupStarted = true; completionThread = mapperEndThread; lifecycleLock.notifyAll(); } if (completionThread != null && completionThread != Thread.currentThread()) { completionThread.interrupt(); } - if (shouldCleanup) { + performDeferredCleanupIfReady(); + } + + private void performDeferredCleanupIfReady() throws IOException { + boolean performCleanup; + synchronized (lifecycleLock) { + // Stock Celeborn cleanup records "Cleaned Up" in PushState. A later failure callback then + // returns before retiring the stock request counter. Wait for raw calls to be classified and + // for counter-only callbacks to complete. Seal exact owners against a later global downgrade + // before cleanup prevents their callbacks and queued retries from invoking stock code. + performCleanup = + cleanupRequested + && !cleanupClaimed + && activePushes == 0 + && sealPendingTransportOwnershipForCleanup(); + if (performCleanup) { + cleanupClaimed = true; + } + } + if (performCleanup) { cleanupAttempt(); } } + private boolean hasPendingPushStateFallback() { + for (PushReservation reservation : pendingPushes) { + registerPushStateFallback(reservation); + if (reservation.pushStateFallback) { + return true; + } + } + return false; + } + + private boolean sealPendingTransportOwnershipForCleanup() { + ArrayList exactPushes = new ArrayList<>(); + for (PushReservation reservation : pendingPushes) { + registerPushStateFallback(reservation); + if (reservation.pushStateFallback || reservation.transportPush == null) { + return false; + } + exactPushes.add(reservation.transportPush); + } + if (CelebornTransportCallbackTracker.Push.trySealCohortForCancellation(exactPushes)) { + return true; + } + for (PushReservation reservation : pendingPushes) { + registerPushStateFallback(reservation); + } + return false; + } + private void cleanupAttempt() throws IOException { if (cleanup == null) { return; diff --git a/spark/src/main/java/org/apache/comet/shuffle/CelebornTransportCallbackTracker.java b/spark/src/main/java/org/apache/comet/shuffle/CelebornTransportCallbackTracker.java index cb4f35c0caa..f139f0efaa3 100644 --- a/spark/src/main/java/org/apache/comet/shuffle/CelebornTransportCallbackTracker.java +++ b/spark/src/main/java/org/apache/comet/shuffle/CelebornTransportCallbackTracker.java @@ -47,10 +47,18 @@ import java.util.function.BiFunction; import java.util.function.Function; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + /** Tracks payload ownership across stock Celeborn transport callbacks and retry tasks. */ final class CelebornTransportCallbackTracker { + private static final Logger LOG = LoggerFactory.getLogger(CelebornTransportCallbackTracker.class); + private final Object shuffleClient; private final Method getDataClientFactory; + private FactoryHook factoryHook; + private boolean initialized; + private boolean closed; private CelebornTransportCallbackTracker(Object shuffleClient, Method getDataClientFactory) { this.shuffleClient = shuffleClient; @@ -67,7 +75,28 @@ static CelebornTransportCallbackTracker tryCreate(Object shuffleClient) { } } - Push beginPush() throws IOException { + synchronized Push beginPush() throws IOException { + if (closed) { + return null; + } + if (!initialized) { + initialized = true; + factoryHook = installFactoryHook(); + } + return factoryHook == null ? null : factoryHook.beginPush(); + } + + synchronized void close() { + if (closed) { + return; + } + closed = true; + if (factoryHook != null) { + factoryHook.releaseRegistration(); + } + } + + private FactoryHook installFactoryHook() { try { Object factory = getDataClientFactory.invoke(shuffleClient); if (factory == null) { @@ -81,26 +110,61 @@ Push beginPush() throws IOException { } List bootstraps = (List) bootstrapValue; FactoryHook hook = findHook(bootstraps); + boolean addBootstrap = false; + if (hook == null) { + hook = findRetryHook(shuffleClient, factory); + addBootstrap = true; + } if (hook == null) { - hook = new FactoryHook(); Class bootstrapInterface = bootstrapInterface(bootstrapsField, factory); - ArrayList augmentedBootstraps = new ArrayList<>(bootstraps); - augmentedBootstraps.add( + hook = new FactoryHook(factory, bootstrapsField); + hook.bootstrap = 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); + bootstrapInterface.getClassLoader(), new Class[] {bootstrapInterface}, hook); } - if (!hook.installed) { - hook.installExistingClients(factory); - hook.installed = true; + hook.retainRegistration(); + try { + if (addBootstrap) { + ArrayList augmentedBootstraps = new ArrayList<>(bootstraps); + augmentedBootstraps.add(hook.bootstrap); + // 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 it. + bootstrapsField.set(factory, augmentedBootstraps); + } + try { + if (hook.acceptsTransportOwnership() && !hook.installed) { + hook.installExistingClients(factory); + hook.installed = true; + } + if (hook.acceptsTransportOwnership()) { + hook.installRetryPool(shuffleClient); + } + } catch (ReflectiveOperationException | IOException | RuntimeException failure) { + Error invocationError = invocationError(failure); + if (invocationError != null) { + throw invocationError; + } + hook.disableTransportOwnership(failure); + } + return hook; + } catch (ReflectiveOperationException | RuntimeException failure) { + hook.releaseRegistration(); + throw failure; + } catch (Error failure) { + rollbackFatalInitialization(hook, failure); + throw failure; } - hook.installRetryPool(shuffleClient); - return new Push(hook); } - } catch (ReflectiveOperationException | RuntimeException failure) { - throw new IOException("Cannot observe Celeborn shuffle payload ownership", failure); + } catch (ReflectiveOperationException | IOException | RuntimeException failure) { + Error invocationError = invocationError(failure); + if (invocationError != null) { + throw invocationError; + } + LOG.warn( + "Cannot install Celeborn transport ownership tracking; falling back to push-state " + + "completion", + failure); + return null; } } @@ -116,6 +180,53 @@ private static FactoryHook findHook(List bootstraps) { return null; } + private static Error invocationError(Throwable failure) { + if (failure instanceof InvocationTargetException + && ((InvocationTargetException) failure).getCause() instanceof Error) { + return (Error) ((InvocationTargetException) failure).getCause(); + } + return failure instanceof Error ? (Error) failure : null; + } + + private static void rollbackFatalInitialization(FactoryHook hook, Error failure) { + try { + hook.disableTransportOwnership(failure); + } catch (Throwable cleanupFailure) { + if (cleanupFailure != failure) { + failure.addSuppressed(cleanupFailure); + } + } + try { + hook.releaseRegistration(); + } catch (Throwable cleanupFailure) { + if (cleanupFailure != failure) { + failure.addSuppressed(cleanupFailure); + } + } + } + + private static void disableAfterFatalBootstrap(FactoryHook hook, Error failure) { + try { + hook.disableTransportOwnership(failure); + } catch (Throwable cleanupFailure) { + if (cleanupFailure != failure) { + failure.addSuppressed(cleanupFailure); + } + } + } + + private static FactoryHook findRetryHook(Object shuffleClient, Object factory) + throws ReflectiveOperationException { + synchronized (shuffleClient) { + Object retryPool = field(shuffleClient.getClass(), "pushDataRetryPool").get(shuffleClient); + if (retryPool instanceof TrackingRetryExecutor) { + FactoryHook hook = ((TrackingRetryExecutor) retryPool).hook; + return hook.factory == factory ? hook : null; + } + return null; + } + } + private static Class bootstrapInterface(Field bootstrapsField, Object factory) throws ReflectiveOperationException { Type genericType = bootstrapsField.getGenericType(); @@ -151,6 +262,15 @@ static final class Push implements AutoCloseable { private final AtomicInteger owners = new AtomicInteger(1); private final AtomicBoolean closed = new AtomicBoolean(); private final ConcurrentLinkedQueue writes = new ConcurrentLinkedQueue<>(); + private final Set ownedRetries = ConcurrentHashMap.newKeySet(); + private OwnershipMode ownershipMode = OwnershipMode.TRANSPORT_OWNED; + private int activeDelegates; + + private enum OwnershipMode { + TRANSPORT_OWNED, + PUSH_STATE_FALLBACK, + CANCELLATION_SEALED + } private Push(FactoryHook hook) { this.hook = hook; @@ -159,6 +279,13 @@ private Push(FactoryHook hook) { } boolean isComplete() { + if (!usesTransportOwnership()) { + return false; + } + return retainedTransportOwnershipComplete() && usesTransportOwnership(); + } + + boolean retainedTransportOwnershipComplete() { // 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) { @@ -167,6 +294,83 @@ boolean isComplete() { return owners.get() == 0; } + boolean usesTransportOwnership() { + synchronized (hook) { + return ownershipMode != OwnershipMode.PUSH_STATE_FALLBACK; + } + } + + private void fallBackToPushState() { + synchronized (hook) { + if (ownershipMode == OwnershipMode.TRANSPORT_OWNED) { + ownershipMode = OwnershipMode.PUSH_STATE_FALLBACK; + } + } + } + + static boolean trySealCohortForCancellation(List pushes) { + if (pushes.isEmpty()) { + return true; + } + FactoryHook hook = pushes.get(0).hook; + ArrayList retries; + synchronized (hook) { + for (Push push : pushes) { + if (push.hook != hook + || !push.closed.get() + || push.activeDelegates != 0 + || push.ownershipMode == OwnershipMode.PUSH_STATE_FALLBACK) { + return false; + } + } + retries = new ArrayList<>(); + for (Push push : pushes) { + push.ownershipMode = OwnershipMode.CANCELLATION_SEALED; + retries.addAll(push.ownedRetries); + } + } + for (OwnedRetry retry : retries) { + retry.discard(); + } + return true; + } + + private boolean beginDelegateUnlessCancelled() { + synchronized (hook) { + if (ownershipMode == OwnershipMode.CANCELLATION_SEALED) { + return false; + } + activeDelegates++; + return true; + } + } + + private void endDelegate() { + synchronized (hook) { + activeDelegates--; + } + } + + private void registerRetry(OwnedRetry retry) { + ownedRetries.add(retry); + synchronized (hook) { + if (ownershipMode != OwnershipMode.CANCELLATION_SEALED) { + return; + } + } + retry.discard(); + } + + private void unregisterRetry(OwnedRetry retry) { + ownedRetries.remove(retry); + } + + private void releaseOwner() { + if (owners.decrementAndGet() == 0) { + hook.activePushes.remove(this); + } + } + private Lease retain() { int currentOwners; do { @@ -190,7 +394,7 @@ public void close() { "Celeborn raw push scopes must close in invocation order"); } hook.restore(previous); - owners.decrementAndGet(); + releaseOwner(); } } } @@ -206,7 +410,7 @@ private Lease(Push push) { @Override public void close() { if (closed.compareAndSet(false, true)) { - push.owners.decrementAndGet(); + push.releaseOwner(); } } } @@ -233,9 +437,102 @@ private static final class FactoryHook implements InvocationHandler { // retry can run on any thread; its wrapper restores this scope only while invoking its // delegate. private final ThreadLocal current = new ThreadLocal<>(); + private final Set activePushes = ConcurrentHashMap.newKeySet(); + private final Object factory; + private final Field bootstrapsField; + private final AtomicBoolean acceptingTransportOwnership = new AtomicBoolean(true); + private volatile boolean active; + private Object bootstrap; + // Read and written only while holding the owning transport factory's monitor. + private int registrations; // Read and written only while holding the owning transport factory's monitor. private boolean installed; + private FactoryHook(Object factory, Field bootstrapsField) { + this.factory = factory; + this.bootstrapsField = bootstrapsField; + } + + private synchronized Push beginPush() { + if (!active || !acceptingTransportOwnership.get()) { + return null; + } + Push push = new Push(this); + activePushes.add(push); + return push; + } + + private boolean acceptsTransportOwnership() { + return active && acceptingTransportOwnership.get(); + } + + private Push currentPush() { + Push push = currentObservedPush(); + return push != null && push.usesTransportOwnership() ? push : null; + } + + private Push currentObservedPush() { + Push push = current.get(); + return active ? push : null; + } + + private synchronized void disableTransportOwnership(Throwable failure) { + for (Push push : activePushes) { + push.fallBackToPushState(); + } + if (!acceptingTransportOwnership.compareAndSet(true, false)) { + return; + } + LOG.warn( + "Cannot instrument a Celeborn transport client; falling back to push-state completion", + failure); + } + + private synchronized void retainRegistration() { + registrations++; + active = true; + } + + private void releaseRegistration() { + synchronized (factory) { + if (registrations == 0) { + return; + } + registrations--; + if (registrations != 0) { + return; + } + installed = false; + synchronized (this) { + active = false; + for (Push push : activePushes) { + push.fallBackToPushState(); + } + } + try { + Object value = bootstrapsField.get(factory); + if (!(value instanceof List)) { + LOG.warn("Cannot remove the released Celeborn transport ownership hook"); + return; + } + ArrayList retained = new ArrayList<>(); + boolean removed = false; + for (Object candidate : (List) value) { + if (candidate == bootstrap) { + removed = true; + } else { + retained.add(candidate); + } + } + if (removed) { + bootstrapsField.set(factory, retained); + } + } catch (ReflectiveOperationException | RuntimeException failure) { + LOG.warn("Cannot remove the released Celeborn transport ownership hook", failure); + } + } + } + private void restore(Push previous) { if (previous == null) { current.remove(); @@ -339,7 +636,7 @@ private void installChannel(Object client) throws ReflectiveOperationException, } private PreparedRequest prepareRequest(Object requestInfo) { - Push push = current.get(); + Push push = currentPush(); if (push == null) { return null; } @@ -381,8 +678,35 @@ public Object invoke(Object proxy, Method method, Object[] args) throws Throwabl return objectMethod(proxy, method, args); } if ("doBootstrap".equals(method.getName()) && args != null && args.length == 1) { - installClient(args[0]); - return null; + synchronized (this) { + // Serialize the active check with the zero-registration transition. Once the last + // owner has returned from releaseRegistration, no stale copy-on-write bootstrap + // snapshot may instrument or reject a later ordinary client. + if (!active) { + return null; + } + if (!acceptingTransportOwnership.get()) { + Push push = current.get(); + if (push != null) { + push.fallBackToPushState(); + } + return null; + } + try { + installClient(args[0]); + } catch (ReflectiveOperationException | IOException | RuntimeException failure) { + Error invocationError = invocationError(failure); + if (invocationError != null) { + disableAfterFatalBootstrap(this, invocationError); + throw invocationError; + } + disableTransportOwnership(failure); + } catch (Error failure) { + disableAfterFatalBootstrap(this, failure); + throw failure; + } + return null; + } } throw new UnsupportedOperationException( "Unsupported Celeborn transport bootstrap: " + method); @@ -472,7 +796,7 @@ public Object invoke(Object proxy, Method method, Object[] args) throws Throwabl if (method.getDeclaringClass() == Object.class) { return objectMethod(proxy, method, args); } - Push push = hook.current.get(); + Push push = hook.currentObservedPush(); if (push == null || !"writeAndFlush".equals(method.getName())) { try { Object result = method.invoke(channel, args); @@ -608,9 +932,18 @@ public Object invoke(Object proxy, Method method, Object[] args) throws Throwabl if (!delivered.compareAndSet(false, true)) { return null; } - try (Activation ignored = push.activate()) { - return invokeCallback(method, args); + boolean invokeDelegate = push.beginDelegateUnlessCancelled(); + try { + if (invokeDelegate) { + try (Activation ignored = push.activate()) { + return invokeCallback(method, args); + } + } + return null; } finally { + if (invokeDelegate) { + push.endDelegate(); + } // 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 @@ -672,13 +1005,22 @@ public void run() { if (!ownership.start()) { return; } - try (Activation ignored = ownership.push.activate()) { - invokeDelegate(); + boolean invokeDelegate = ownership.push.beginDelegateUnlessCancelled(); + try { + if (invokeDelegate) { + try (Activation ignored = ownership.push.activate()) { + invokeDelegate(); + } + } } finally { + if (invokeDelegate) { + ownership.push.endDelegate(); + } // 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(); + ownership.push.unregisterRetry(this); } } @@ -687,6 +1029,7 @@ public void discard() { if (ownership.discard()) { delegate = null; ownership.finish(); + ownership.push.unregisterRetry(this); } } } @@ -709,10 +1052,21 @@ public void run() { if (!ownership.start()) { return; } - try (Activation ignored = ownership.push.activate()) { - super.run(); + boolean invokeDelegate = ownership.push.beginDelegateUnlessCancelled(); + try { + if (invokeDelegate) { + try (Activation ignored = ownership.push.activate()) { + super.run(); + } + } else { + super.cancel(false); + } } finally { + if (invokeDelegate) { + ownership.push.endDelegate(); + } ownership.finish(); + ownership.push.unregisterRetry(this); } } @@ -724,6 +1078,7 @@ public boolean cancel(boolean mayInterruptIfRunning) { // may release here; a running task releases from run's finally block after it returns. if (ownership.discard()) { ownership.finish(); + ownership.push.unregisterRetry(this); } } return cancelled; @@ -749,11 +1104,11 @@ public void execute(Runnable command) { if (command == null) { throw new NullPointerException("command"); } - Push push = hook.current.get(); + Push push = hook.currentPush(); if (push == null) { delegate.execute(command); } else { - submitOwned(new TrackingRetryRunnable(command, push)); + submitOwned(new TrackingRetryRunnable(command, push), push); } } @@ -764,27 +1119,28 @@ public Future submit(Runnable task) { @Override public Future submit(Runnable task, T result) { - Push push = hook.current.get(); + Push push = hook.currentPush(); if (push == null) { return delegate.submit(task, result); } TrackingRetryFuture future = new TrackingRetryFuture<>(task, result, push); - submitOwned(future); + submitOwned(future, push); return future; } @Override public Future submit(Callable task) { - Push push = hook.current.get(); + Push push = hook.currentPush(); if (push == null) { return delegate.submit(task); } TrackingRetryFuture future = new TrackingRetryFuture<>(task, push); - submitOwned(future); + submitOwned(future, push); return future; } - private void submitOwned(T task) { + private void submitOwned(T task, Push push) { + push.registerRetry(task); try { delegate.execute(task); } catch (RuntimeException | Error failure) { diff --git a/spark/src/main/java/org/apache/comet/shuffle/ExecutorShufflePushAdmission.java b/spark/src/main/java/org/apache/comet/shuffle/ExecutorShufflePushAdmission.java index 71d05b71e49..199d057e533 100644 --- a/spark/src/main/java/org/apache/comet/shuffle/ExecutorShufflePushAdmission.java +++ b/spark/src/main/java/org/apache/comet/shuffle/ExecutorShufflePushAdmission.java @@ -37,6 +37,7 @@ final class ExecutorShufflePushAdmission { private final Semaphore available; private CelebornTransportCallbackTracker transportCallbacks; private boolean transportCallbacksInitialized; + private boolean closed; private ExecutorShufflePushAdmission(int limit) { this.limit = limit; @@ -58,12 +59,19 @@ static ExecutorShufflePushAdmission forClient(Object client, int limit) { } static void releaseClient(Object client) { + ExecutorShufflePushAdmission admission; synchronized (CLIENTS) { - CLIENTS.remove(client); + admission = CLIENTS.remove(client); + } + if (admission != null) { + admission.close(); } } synchronized CelebornTransportCallbackTracker transportCallbacks(Object client) { + if (closed) { + return null; + } if (!transportCallbacksInitialized) { transportCallbacks = CelebornTransportCallbackTracker.tryCreate(client); transportCallbacksInitialized = true; @@ -71,6 +79,16 @@ synchronized CelebornTransportCallbackTracker transportCallbacks(Object client) return transportCallbacks; } + private synchronized void close() { + if (closed) { + return; + } + closed = true; + if (transportCallbacks != null) { + transportCallbacks.close(); + } + } + void acquire(int bytes, BooleanSupplier cancelled) throws IOException { if (bytes <= 0 || bytes > limit) { throw new IOException( diff --git a/spark/src/test/java/org/apache/comet/shuffle/CelebornTransportOwnershipTestHelper.java b/spark/src/test/java/org/apache/comet/shuffle/CelebornTransportOwnershipTestHelper.java index 09c6fb524a3..404f491e36e 100644 --- a/spark/src/test/java/org/apache/comet/shuffle/CelebornTransportOwnershipTestHelper.java +++ b/spark/src/test/java/org/apache/comet/shuffle/CelebornTransportOwnershipTestHelper.java @@ -27,10 +27,12 @@ import java.util.List; import java.util.concurrent.AbstractExecutorService; import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.CountDownLatch; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; import java.util.concurrent.RejectedExecutionException; import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicReference; import io.netty.buffer.ByteBuf; import io.netty.buffer.Unpooled; @@ -46,6 +48,265 @@ public final class CelebornTransportOwnershipTestHelper { private CelebornTransportOwnershipTestHelper() {} + @SuppressWarnings({"rawtypes", "unchecked"}) + public static void assertBootstrapFailureDoesNotEscapeClientCreation() throws Exception { + ShuffleClient shuffle = new ShuffleClient(); + ExecutorShufflePushAdmission admission = ExecutorShufflePushAdmission.forClient(shuffle, 64); + Client existing = shuffle.factory.createClient(true); + Client incompatible = null; + Client ordinary = null; + CelebornTransportCallbackTracker.Push push = null; + try { + CelebornTransportCallbackTracker tracker = admission.transportCallbacks(shuffle); + push = tracker.beginPush(); + check(push != null, "the compatible factory must initially support transport ownership"); + + // Celeborn treats bootstrap exceptions as fatal to its shared client factory. A Comet-only + // reflection mismatch must instead disable tracking and let ordinary clients keep opening. + incompatible = shuffle.factory.createClient(false); + check(!push.usesTransportOwnership(), "the active push must switch to counter fallback"); + + // An existing request map may already be instrumented. Once the push falls back, publishing + // a request with an unknown callback shape must delegate without further reflection. + ConcurrentHashMap rawRequests = existing.handler.outstandingPushes; + Object unknownRequest = new Object(); + rawRequests.put(1L, unknownRequest); + check( + rawRequests.get(1L) == unknownRequest, + "fallback publication must not inspect an unknown request shape"); + + push.close(); + push = null; + check(tracker.beginPush() == null, "a failed bootstrap must persistently select fallback"); + ordinary = shuffle.factory.createClient(true); + } finally { + if (push != null) { + push.close(); + } + existing.underlying.close(); + if (incompatible != null) { + incompatible.underlying.close(); + } + if (ordinary != null) { + ordinary.underlying.close(); + } + ExecutorShufflePushAdmission.releaseClient(shuffle); + shuffle.pushDataRetryPool.shutdownNow(); + } + } + + public static void assertBootstrapHookFollowsClientLifetime() throws Exception { + Factory factory = new Factory(); + Bootstrap before = ignored -> {}; + Bootstrap after = ignored -> {}; + factory.clientBootstraps.add(before); + factory.clientBootstraps.add(after); + ShuffleClient first = new ShuffleClient(factory, Executors.newSingleThreadExecutor()); + ShuffleClient second = new ShuffleClient(factory, Executors.newSingleThreadExecutor()); + try { + ExecutorShufflePushAdmission firstAdmission = + ExecutorShufflePushAdmission.forClient(first, 64); + CelebornTransportCallbackTracker.Push firstPush = + firstAdmission.transportCallbacks(first).beginPush(); + firstPush.close(); + + ExecutorShufflePushAdmission secondAdmission = + ExecutorShufflePushAdmission.forClient(second, 64); + CelebornTransportCallbackTracker.Push secondPush = + secondAdmission.transportCallbacks(second).beginPush(); + secondPush.close(); + + check(factory.clientBootstraps.size() == 3, "shared factories need one Comet bootstrap"); + Bootstrap cometBootstrap = factory.clientBootstraps.get(2); + check(Proxy.isProxyClass(cometBootstrap.getClass()), "the Comet bootstrap must be a proxy"); + + ExecutorShufflePushAdmission.releaseClient(first); + check(factory.clientBootstraps.size() == 3, "the first owner must retain the shared hook"); + ExecutorShufflePushAdmission.releaseClient(first); + check(factory.clientBootstraps.size() == 3, "client release must be idempotent"); + + ExecutorShufflePushAdmission.releaseClient(second); + check(factory.clientBootstraps.size() == 2, "the last owner must remove the Comet hook"); + check( + factory.clientBootstraps.get(0) == before && factory.clientBootstraps.get(1) == after, + "hook removal must preserve unrelated bootstraps and their order"); + + Client gapClient = factory.createClient(); + check( + !Proxy.isProxyClass(gapClient.channel.getClass()), + "clients opened without a Comet owner must remain ordinary"); + check( + gapClient.handler.outstandingPushes.getClass() == ConcurrentHashMap.class, + "clients opened without a Comet owner must keep their ordinary request map"); + + // A factory can already be iterating the old copy-on-write list when the last owner exits. + // Its captured proxy must become inert even though that old list remains reachable. + Client staleClient = new Client(); + cometBootstrap.doBootstrap(staleClient); + check( + !Proxy.isProxyClass(staleClient.channel.getClass()), + "a released bootstrap snapshot must not wrap later channels"); + check( + staleClient.handler.outstandingPushes.getClass() == ConcurrentHashMap.class, + "a released bootstrap snapshot must not wrap later request maps"); + staleClient.underlying.close(); + + ExecutorShufflePushAdmission reboundAdmission = + ExecutorShufflePushAdmission.forClient(first, 64); + CelebornTransportCallbackTracker.Push reboundPush = + reboundAdmission.transportCallbacks(first).beginPush(); + check(reboundPush != null, "a surviving client must reuse its compatible transport hook"); + reboundPush.close(); + check( + factory.clientBootstraps.size() == 3 && factory.clientBootstraps.get(2) == cometBootstrap, + "reacquiring a surviving client must restore the same bootstrap proxy"); + check( + Proxy.isProxyClass(gapClient.channel.getClass()), + "reacquisition must instrument clients opened while the hook was absent"); + check( + gapClient.handler.outstandingPushes.getClass() != ConcurrentHashMap.class, + "reacquisition must instrument request maps opened while the hook was absent"); + ExecutorShufflePushAdmission.releaseClient(first); + check(factory.clientBootstraps.size() == 2, "the rebound hook must remain releasable"); + gapClient.underlying.close(); + } finally { + ExecutorShufflePushAdmission.releaseClient(first); + ExecutorShufflePushAdmission.releaseClient(second); + first.pushDataRetryPool.shutdownNow(); + second.pushDataRetryPool.shutdownNow(); + } + } + + public static void assertFatalBootstrapErrorsDoNotLeakRegistrations() throws Exception { + Factory initializationFactory = new Factory(); + Client existing = initializationFactory.createClient(); + ShuffleClient initializationShuffle = + new ShuffleClient(initializationFactory, Executors.newSingleThreadExecutor()); + AssertionError initializationError = new AssertionError("fatal existing-client mismatch"); + existing.handlerFailure = initializationError; + try { + CelebornTransportCallbackTracker tracker = + ExecutorShufflePushAdmission.forClient(initializationShuffle, 64) + .transportCallbacks(initializationShuffle); + try { + tracker.beginPush(); + throw new AssertionError("a fatal reflected bootstrap error must escape unchanged"); + } catch (AssertionError actual) { + check(actual == initializationError, "the reflected fatal error must stay unchanged"); + } + check( + initializationFactory.clientBootstraps.isEmpty(), + "fatal initialization must remove the unowned bootstrap registration"); + existing.handlerFailure = null; + Client ordinary = initializationFactory.createClient(); + ordinary.underlying.close(); + } finally { + existing.handlerFailure = null; + existing.underlying.close(); + ExecutorShufflePushAdmission.releaseClient(initializationShuffle); + initializationShuffle.pushDataRetryPool.shutdownNow(); + } + + ShuffleClient activeShuffle = new ShuffleClient(); + try { + CelebornTransportCallbackTracker tracker = + ExecutorShufflePushAdmission.forClient(activeShuffle, 64) + .transportCallbacks(activeShuffle); + CelebornTransportCallbackTracker.Push push = tracker.beginPush(); + push.close(); + Bootstrap bootstrap = activeShuffle.factory.clientBootstraps.get(0); + Client fatalClient = new Client(); + AssertionError callbackError = new AssertionError("fatal new-client mismatch"); + fatalClient.handlerFailure = callbackError; + try { + bootstrap.doBootstrap(fatalClient); + throw new AssertionError("a fatal bootstrap callback error must escape unchanged"); + } catch (AssertionError actual) { + check(actual == callbackError, "the callback fatal error must stay unchanged"); + } finally { + fatalClient.underlying.close(); + } + Client ordinary = activeShuffle.factory.createClient(); + ordinary.underlying.close(); + check(tracker.beginPush() == null, "fatal callback errors must leave the shared hook inert"); + } finally { + ExecutorShufflePushAdmission.releaseClient(activeShuffle); + activeShuffle.pushDataRetryPool.shutdownNow(); + } + } + + public static void assertBootstrapReleaseWaitsForActiveInvocation() throws Exception { + ShuffleClient shuffle = new ShuffleClient(); + Thread bootstrapThread = null; + Thread releaseThread = null; + Client racingClient = new Client(); + Client staleClient = new Client(); + CountDownLatch handlerEntered = new CountDownLatch(1); + CountDownLatch resumeHandler = new CountDownLatch(1); + AtomicReference bootstrapFailure = new AtomicReference<>(); + racingClient.handlerEntered = handlerEntered; + racingClient.resumeHandler = resumeHandler; + try { + CelebornTransportCallbackTracker tracker = + ExecutorShufflePushAdmission.forClient(shuffle, 64).transportCallbacks(shuffle); + CelebornTransportCallbackTracker.Push push = tracker.beginPush(); + push.close(); + Bootstrap bootstrap = shuffle.factory.clientBootstraps.get(0); + + bootstrapThread = + new Thread( + () -> { + try { + bootstrap.doBootstrap(racingClient); + } catch (Throwable failure) { + bootstrapFailure.set(failure); + } + }); + bootstrapThread.start(); + check(handlerEntered.await(5, TimeUnit.SECONDS), "the bootstrap must enter client setup"); + + releaseThread = new Thread(() -> ExecutorShufflePushAdmission.releaseClient(shuffle)); + releaseThread.start(); + releaseThread.join(100); + check(releaseThread.isAlive(), "last-owner release must wait for an active bootstrap"); + + resumeHandler.countDown(); + bootstrapThread.join(5000); + releaseThread.join(5000); + check(!bootstrapThread.isAlive(), "the active bootstrap must finish"); + check(!releaseThread.isAlive(), "last-owner release must finish after the bootstrap"); + check(bootstrapFailure.get() == null, "the active bootstrap must remain valid"); + check( + Proxy.isProxyClass(racingClient.channel.getClass()), + "a bootstrap started before release may finish instrumentation"); + check( + shuffle.factory.clientBootstraps.isEmpty(), + "last-owner release must remove the bootstrap before returning"); + + bootstrap.doBootstrap(staleClient); + check( + !Proxy.isProxyClass(staleClient.channel.getClass()), + "a stale bootstrap must remain inert after release returns"); + check( + staleClient.handler.outstandingPushes.getClass() == ConcurrentHashMap.class, + "a stale bootstrap must not instrument request maps after release returns"); + } finally { + resumeHandler.countDown(); + if (bootstrapThread != null) { + bootstrapThread.interrupt(); + bootstrapThread.join(5000); + } + if (releaseThread != null) { + releaseThread.interrupt(); + releaseThread.join(5000); + } + ExecutorShufflePushAdmission.releaseClient(shuffle); + racingClient.underlying.close(); + staleClient.underlying.close(); + shuffle.pushDataRetryPool.shutdownNow(); + } + } + public static void assertTimedOutWriteRetainsOwnership(boolean existingClient) throws Exception { ShuffleClient shuffle = new ShuffleClient(); Client client = existingClient ? shuffle.factory.createClient() : null; @@ -258,10 +519,35 @@ public static final class ResponseHandler { public static final class Client { private final HoldingChannel underlying = new HoldingChannel(); - private final Channel channel = underlying; + private final Channel channel; private final ResponseHandler handler = new ResponseHandler(); + private Error handlerFailure; + private CountDownLatch handlerEntered; + private CountDownLatch resumeHandler; + + private Client() { + this(true); + } + + private Client(boolean compatible) { + channel = compatible ? underlying : null; + } public ResponseHandler getHandler() { + if (handlerEntered != null) { + handlerEntered.countDown(); + try { + if (!resumeHandler.await(5, TimeUnit.SECONDS)) { + throw new IllegalStateException("timed out waiting to resume client setup"); + } + } catch (InterruptedException failure) { + Thread.currentThread().interrupt(); + throw new IllegalStateException("client setup was interrupted", failure); + } + } + if (handlerFailure != null) { + throw handlerFailure; + } return handler; } @@ -280,9 +566,13 @@ public static final class Factory { private final ConcurrentHashMap connectionPool = new ConcurrentHashMap<>(); private Client createClient() { + return createClient(true); + } + + private Client createClient(boolean compatible) { Pool pool = connectionPool.computeIfAbsent("worker", ignored -> new Pool()); synchronized (pool.locks[0]) { - Client client = new Client(); + Client client = new Client(compatible); for (Bootstrap bootstrap : clientBootstraps) { bootstrap.doBootstrap(client); } @@ -293,14 +583,19 @@ private Client createClient() { } public static final class ShuffleClient { - private final Factory factory = new Factory(); + private final Factory factory; private final ExecutorService pushDataRetryPool; private ShuffleClient() { - this(Executors.newSingleThreadExecutor()); + this(new Factory(), Executors.newSingleThreadExecutor()); } private ShuffleClient(ExecutorService retryPool) { + this(new Factory(), retryPool); + } + + private ShuffleClient(Factory factory, ExecutorService retryPool) { + this.factory = factory; pushDataRetryPool = retryPool; } 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 84cecddea06..ab9d341ba4a 100644 --- a/spark/src/test/scala/org/apache/comet/shuffle/CelebornShufflePartitionPusherSuite.scala +++ b/spark/src/test/scala/org/apache/comet/shuffle/CelebornShufflePartitionPusherSuite.scala @@ -575,8 +575,8 @@ class CelebornShufflePartitionPusherSuite extends AnyFunSuite { val cancelledState = client.currentState(19, 3, encodedAttemptId) first.abort() first.abort() - assert(client.cleanupCalls.get() == 1) - assert(cancelledState.exception.get().getMessage == "Cleaned Up") + assert(client.cleanupCalls.get() == 0) + assert(cancelledState.exception.get() == null) assert(cancelledState.inFlightRequestTracker.totalInflightReqs.sum() == 1) val retry = @@ -595,6 +595,7 @@ class CelebornShufflePartitionPusherSuite extends AnyFunSuite { worker.join(5000) assert(!worker.isAlive) assert(failure.get() == null) + awaitCleanup(client) val retryState = client.currentState(19, 3, encodedAttemptId + 1) client.complete(retryState) assert(retry.finish()(1) == bytes.length) @@ -726,6 +727,228 @@ class CelebornShufflePartitionPusherSuite extends AnyFunSuite { } } + test("cancelled exact ownership remains sealed after a later bootstrap downgrade") { + val client = new TransportRecordingCelebornPushClient + 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) + val originalHandler = client.dataClientFactory.handler + + first.abort() + assert(state.exception.get().getMessage == "Cleaned Up") + client.dataClientFactory.openUninstrumentableConnection() + + assertNextMapWaitsForCompletion(client) { + originalHandler + .remove(state) + .callback + .onFailure(new IOException("sealed callback failed after global downgrade")) + } + } + + test("failed bootstrap instrumentation falls back to push-state completion") { + val client = new TransportRecordingCelebornPushClient + client.openUninstrumentableConnectionBeforePush = true + val bytes = frame() + val first = + new CelebornShufflePartitionPusher(client, 19, 3, encodedAttemptId, 12, 9, 64) + + first.pushPartitionData(0, bytes, bytes.length) + val firstState = client.currentState(19, 3, encodedAttemptId) + + def startPush(attemptId: Int) + : (CelebornShufflePartitionPusher, AtomicReference[Throwable], CountDownLatch, Thread) = { + val next = new CelebornShufflePartitionPusher(client, 19, 3, attemptId, 12, 9, 64) + val failure = new AtomicReference[Throwable]() + val started = new CountDownLatch(1) + val worker = new Thread(() => { + try { + started.countDown() + next.pushPartitionData(1, bytes, bytes.length) + } catch { case error: Throwable => failure.set(error) } + }) + worker.start() + (next, failure, started, worker) + } + + val (second, secondFailure, secondStarted, secondWorker) = startPush(encodedAttemptId + 1) + try { + assert(secondStarted.await(5, TimeUnit.SECONDS)) + secondWorker.join(100) + assert(secondWorker.isAlive, "an uninstrumented push must retain counter-backed admission") + client.complete(firstState) + secondWorker.join(5000) + assert(!secondWorker.isAlive) + assert(secondFailure.get() == null) + + val secondState = client.currentState(19, 3, encodedAttemptId + 1) + val (third, thirdFailure, thirdStarted, thirdWorker) = startPush(encodedAttemptId + 2) + try { + assert(thirdStarted.await(5, TimeUnit.SECONDS)) + thirdWorker.join(100) + assert(thirdWorker.isAlive, "counter fallback must persist for later pushes") + client.complete(secondState) + thirdWorker.join(5000) + assert(!thirdWorker.isAlive) + assert(thirdFailure.get() == null) + val thirdState = client.currentState(19, 3, encodedAttemptId + 2) + client.complete(thirdState) + assert(third.finish()(1) == bytes.length) + } finally { + thirdWorker.interrupt() + thirdWorker.join(5000) + } + assert(second.finish()(1) == bytes.length) + assert(first.finish()(0) == bytes.length) + } finally { + secondWorker.interrupt() + secondWorker.join(5000) + } + } + + test("aborted fallback push releases admission after its failed transport completes") { + val client = new TransportRecordingCelebornPushClient + client.openUninstrumentableConnectionBeforePush = true + 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) + first.abort() + assert(client.cleanupCalls.get() == 0) + assert(state.exception.get() == null) + + val terminalFailure = new IOException( + s"Push data to worker failed for shuffle 19 map 3 attempt $encodedAttemptId " + + "partition 0 batch 1.") + assertNextMapWaitsForCompletion(client) { + client.failTransport(state, terminalFailure) + // Stock Celeborn's failure path leaves this counter charged. The terminal failure itself + // is the completion signal for the cancelled fallback request. + assert(state.inFlightRequestTracker.totalInflightReqs.sum() == 1) + } + awaitCleanup(client) + assert(state.exception.get() eq terminalFailure) + } + + test("aborted fallback push keeps its state until a published request completes") { + val client = new TransportRecordingCelebornPushClient + client.openUninstrumentableConnectionBeforePush = true + val first = + new CelebornShufflePartitionPusher(client, 19, 3, encodedAttemptId, 12, 9, 604, 2416) + val original = client.getPushState(s"19-3-$encodedAttemptId") + val published = new AtomicReference[RecordingCelebornPushState]() + val expected = new IOException("raw fallback push failed after publishing the request") + client.beforePushBegins = () => first.abort() + 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.beforePushBegins = () => () + client.beforePushReturns = _ => () + assert(published.get() eq original) + assert(client.cleanupCalls.get() == 0) + assert(original.exception.get() == null) + + assertNextMapWaitsForCompletion(client) { + client.complete(published.get()) + } + awaitCleanup(client) + } + + test("fallback submissions wait for prior transport completion") { + val client = new TransportRecordingCelebornPushClient + val bytes = frame(596) + val admissionBytes = 5484 + + // Disable precise ownership once, then verify later fallback submissions on the same pusher. + client.openUninstrumentableConnectionBeforePush = true + val initializer = + new CelebornShufflePartitionPusher( + client, + 19, + 3, + encodedAttemptId, + 12, + 9, + 604, + admissionBytes) + initializer.pushPartitionData(0, bytes, bytes.length) + client.complete(client.currentState(19, 3, encodedAttemptId)) + assert(initializer.finish()(0) == bytes.length) + + val attemptId = encodedAttemptId + 1 + val pusher = + new CelebornShufflePartitionPusher(client, 19, 3, attemptId, 12, 9, 604, admissionBytes) + val rawEntries = new AtomicInteger() + val firstPublished = new CountDownLatch(1) + val resumeFirst = new CountDownLatch(1) + val firstFailure = new AtomicReference[Throwable]() + val secondFailure = new AtomicReference[Throwable]() + client.beforePushBegins = () => rawEntries.incrementAndGet() + client.beforePushReturns = _ => { + if (firstPublished.getCount > 0) { + firstPublished.countDown() + if (!resumeFirst.await(5, TimeUnit.SECONDS)) { + throw new IOException("timed out waiting to resume the first fallback submission") + } + } + } + val first = new Thread(() => { + try pusher.pushPartitionData(0, bytes, bytes.length) + catch { case failure: Throwable => firstFailure.set(failure) } + }) + val second = new Thread(() => { + try pusher.pushPartitionData(1, bytes, bytes.length) + catch { case failure: Throwable => secondFailure.set(failure) } + }) + try { + first.start() + assert(firstPublished.await(5, TimeUnit.SECONDS)) + second.start() + second.join(100) + assert(second.isAlive) + assert( + rawEntries.get() == 1, + "a second raw submission must not race the first submission's PushState resolution") + + resumeFirst.countDown() + first.join(5000) + assert(!first.isAlive) + assert(firstFailure.get() == null) + second.join(100) + assert(second.isAlive, "a second fallback push must wait for the first callback") + assert(rawEntries.get() == 1) + + val state = client.currentState(19, 3, attemptId) + client.complete(state) + second.join(5000) + assert(!second.isAlive) + assert(secondFailure.get() == null) + assert(rawEntries.get() == 2) + + client.complete(state) + assert(pusher.finish().take(2).sum == 2L * bytes.length) + } finally { + resumeFirst.countDown() + first.interrupt() + second.interrupt() + first.join(5000) + second.join(5000) + client.beforePushBegins = () => () + client.beforePushReturns = _ => () + } + } + private def assertNextMapWaitsForCompletion(client: TransportRecordingCelebornPushClient)( completePending: => Unit): Unit = { val bytes = frame(596) @@ -760,7 +983,15 @@ class CelebornShufflePartitionPusherSuite extends AnyFunSuite { } } - test("cleanup during raw submission keeps callbacks on a recreated push state observable") { + private def awaitCleanup(client: RecordingCelebornPushClient): Unit = { + val deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(5) + while (client.cleanupCalls.get() == 0 && System.nanoTime() < deadline) { + Thread.sleep(10) + } + assert(client.cleanupCalls.get() == 1) + } + + test("cleanup during raw submission waits until callback ownership is registered") { val client = new TransportRecordingCelebornPushClient val first = new CelebornShufflePartitionPusher(client, 19, 3, encodedAttemptId, 12, 9, 604, 2416) @@ -775,7 +1006,8 @@ class CelebornShufflePartitionPusherSuite extends AnyFunSuite { client.beforePushBegins = () => () client.beforePushReturns = _ => () val state = recreated.get() - assert(state != null && (state ne original)) + assert(state eq original) + assert(client.cleanupCalls.get() == 1) assert(state.exception.get().getMessage == "Cleaned Up") assertNextMapWaitsForCompletion(client) { @@ -784,6 +1016,48 @@ class CelebornShufflePartitionPusherSuite extends AnyFunSuite { } } + test("abort interrupts a raw invocation blocked before publication") { + val client = new TransportRecordingCelebornPushClient + val first = + new CelebornShufflePartitionPusher(client, 19, 3, encodedAttemptId, 12, 9, 604, 2416) + val enteredRawPush = new CountDownLatch(1) + val blockRawPush = new CountDownLatch(1) + val pushFailure = new AtomicReference[Throwable]() + client.beforePushBegins = () => { + enteredRawPush.countDown() + try blockRawPush.await() + catch { + case interrupted: InterruptedException => + Thread.currentThread().interrupt() + val failure = + new IOException("Interrupted while waiting for Celeborn admission", interrupted) + client.currentState(19, 3, encodedAttemptId).exception.compareAndSet(null, failure) + throw failure + } + } + val bytes = frame(596) + val worker = new Thread(() => { + try first.pushPartitionData(0, bytes, bytes.length) + catch { case error: Throwable => pushFailure.set(error) } + }) + + try { + worker.start() + assert(enteredRawPush.await(5, TimeUnit.SECONDS)) + first.abort() + worker.join(5000) + assert(!worker.isAlive, "abort must interrupt a raw push waiting on Celeborn admission") + assert(pushFailure.get().isInstanceOf[IOException]) + assert(client.dataClientFactory.handler.outstandingPushes.isEmpty) + awaitCleanup(client) + } finally { + blockRawPush.countDown() + worker.interrupt() + worker.join(5000) + client.beforePushBegins = () => () + } + } + test("a raw invocation that throws after publication retains its outstanding callback") { val client = new TransportRecordingCelebornPushClient val first = @@ -806,7 +1080,7 @@ class CelebornShufflePartitionPusherSuite extends AnyFunSuite { } } - test("queued retry completion after cleanup releases admission without a transport callback") { + test("cleanup discards a queued retry without waiting for its executor") { val client = new TransportRecordingCelebornPushClient client.retriesBeforeFailure = 1 val first = @@ -819,12 +1093,26 @@ class CelebornShufflePartitionPusherSuite extends AnyFunSuite { 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) - } + val next = + new CelebornShufflePartitionPusher(client, 19, 3, encodedAttemptId + 1, 12, 9, 604, 2416) + val failure = new AtomicReference[Throwable]() + val worker = new Thread(() => { + try next.pushPartitionData(1, bytes, bytes.length) + catch { case error: Throwable => failure.set(error) } + }) + worker.start() + worker.join(5000) + assert(!worker.isAlive, "queued retries must not retain admission after cancellation") + assert(failure.get() == null) + + // The delegate executor may retain the cancelled wrapper, but it no longer retains or runs + // the stock retry and its payload. + client.retryExecutor.runNext() + assert(client.retryExecutor.pendingCount == 0) + assert(state.inFlightRequestTracker.totalInflightReqs.sum() == 1) + val nextState = client.currentState(19, 3, encodedAttemptId + 1) + client.complete(nextState) + assert(next.finish()(1) == bytes.length) } test("retry handoff to another transport keeps admission until its callback returns") { @@ -955,13 +1243,13 @@ class CelebornShufflePartitionPusherSuite extends AnyFunSuite { 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) + assert(cancelledState.inFlightRequestTracker.totalInflightReqs.sum() == 2) client.failTransport(cancelledState, new IOException("cancelled connection closed")) worker.join(5000) assert(!worker.isAlive) assert(failure.get() == null) - assert(cancelledState.inFlightRequestTracker.totalInflightReqs.sum() == 1) + assert(cancelledState.inFlightRequestTracker.totalInflightReqs.sum() == 2) val retryState = client.currentState(19, 3, encodedAttemptId + 1) client.complete(retryState) @@ -1084,6 +1372,22 @@ class CelebornShufflePartitionPusherSuite extends AnyFunSuite { CelebornTransportOwnershipTestHelper.assertCompletedRetriesForgetPayloads() } + test("Celeborn bootstrap failures do not escape shared client creation") { + CelebornTransportOwnershipTestHelper.assertBootstrapFailureDoesNotEscapeClientCreation() + } + + test("Celeborn bootstrap hooks remain only while their shared clients are owned") { + CelebornTransportOwnershipTestHelper.assertBootstrapHookFollowsClientLifetime() + } + + test("fatal Celeborn bootstrap errors do not leak shared hook registrations") { + CelebornTransportOwnershipTestHelper.assertFatalBootstrapErrorsDoNotLeakRegistrations() + } + + test("Celeborn hook release waits for an active bootstrap invocation") { + CelebornTransportOwnershipTestHelper.assertBootstrapReleaseWaitsForActiveInvocation() + } + test("failed push keeps native frame admission until the JNI caller retires its buffers") { Seq(false, true).foreach { publishedTransport => val client = new TransportRecordingCelebornPushClient @@ -1423,6 +1727,7 @@ final class TransportRecordingCelebornPushClient extends AsyncRecordingCelebornP def getDataClientFactory: RecordingCelebornTransportClientFactory = dataClientFactory var openConnectionBeforePush: Boolean = false + var openUninstrumentableConnectionBeforePush: Boolean = false var beforePushBegins: () => Unit = () => () var beforePushReturns: RecordingCelebornPushState => Unit = (_: RecordingCelebornPushState) => () @@ -1460,6 +1765,10 @@ final class TransportRecordingCelebornPushClient extends AsyncRecordingCelebornP openConnectionBeforePush = false dataClientFactory.openConnection() } + if (openUninstrumentableConnectionBeforePush) { + openUninstrumentableConnectionBeforePush = false + dataClientFactory.openUninstrumentableConnection() + } val state = currentState(shuffleId, mapId, attemptId) dataClientFactory.handler.add( state, @@ -1531,8 +1840,16 @@ final class RecordingCelebornTransportClientFactory { connectionPool.put("worker", new RecordingCelebornTransportClientPool(handler)) def openConnection(): Unit = { + openConnection(new io.netty.channel.embedded.EmbeddedChannel()) + } + + def openUninstrumentableConnection(): Unit = { + openConnection(null) + } + + private def openConnection(channel: io.netty.channel.Channel): Unit = { handler = new RecordingCelebornTransportResponseHandler - val pool = new RecordingCelebornTransportClientPool(handler) + val pool = new RecordingCelebornTransportClientPool(handler, channel) val bootstraps = clientBootstraps.iterator() while (bootstraps.hasNext) { bootstraps.next().doBootstrap(pool.clients(0)) @@ -1546,14 +1863,16 @@ trait RecordingCelebornTransportClientBootstrap { } final class RecordingCelebornTransportClientPool( - handler: RecordingCelebornTransportResponseHandler) { + handler: RecordingCelebornTransportResponseHandler, + channel: io.netty.channel.Channel = new io.netty.channel.embedded.EmbeddedChannel()) { val clients: Array[RecordingCelebornTransportClient] = - Array(new RecordingCelebornTransportClient(handler)) + Array(new RecordingCelebornTransportClient(handler, channel)) val locks: Array[Object] = Array(new Object) } -final class RecordingCelebornTransportClient(handler: RecordingCelebornTransportResponseHandler) { - val channel: io.netty.channel.Channel = new io.netty.channel.embedded.EmbeddedChannel() +final class RecordingCelebornTransportClient( + handler: RecordingCelebornTransportResponseHandler, + val channel: io.netty.channel.Channel) { def getChannel: io.netty.channel.Channel = channel def getHandler: RecordingCelebornTransportResponseHandler = handler } @@ -1614,7 +1933,11 @@ final class RecordingCelebornTransportCallback( retriesRemaining -= 1 submitRetry(this) } else { - pushState.exception.compareAndSet(null, new IOException(failure)) + val reportedFailure = failure match { + case io: IOException => io + case _ => new IOException(failure) + } + pushState.exception.compareAndSet(null, reportedFailure) } } }