From 8225f4f945defe10b2149e5fe331488817503ce6 Mon Sep 17 00:00:00 2001 From: Chao Sun Date: Wed, 24 Jun 2026 22:45:51 -0700 Subject: [PATCH 1/2] [CELEBORN-2371] Bound Spark batch-open client creation retries and stop them on interruption --- .../celeborn/CelebornShuffleReader.scala | 8 +- .../celeborn/CelebornShuffleReaderSuite.scala | 89 ++++++++++++++++++- .../client/TransportClientFactory.java | 49 ++++++++-- .../network/TransportClientFactorySuiteJ.java | 23 +++++ ...TransportClientFactoryInterruptSuiteJ.java | 57 ++++++++++++ 5 files changed, 218 insertions(+), 8 deletions(-) create mode 100644 common/src/test/java/org/apache/celeborn/common/network/client/TransportClientFactoryInterruptSuiteJ.java diff --git a/client-spark/spark-3/src/main/scala/org/apache/spark/shuffle/celeborn/CelebornShuffleReader.scala b/client-spark/spark-3/src/main/scala/org/apache/spark/shuffle/celeborn/CelebornShuffleReader.scala index 155fc088616..88ece4a36bd 100644 --- a/client-spark/spark-3/src/main/scala/org/apache/spark/shuffle/celeborn/CelebornShuffleReader.scala +++ b/client-spark/spark-3/src/main/scala/org/apache/spark/shuffle/celeborn/CelebornShuffleReader.scala @@ -600,6 +600,9 @@ class CelebornShuffleReader[K, C]( object CelebornShuffleReader { var streamCreatorPool: ThreadPoolExecutor = null + // TransportClientFactory already retries each attempt; allow one extra pooled-client attempt. + private val MAX_CLIENT_CREATION_ATTEMPTS_PER_HOST = 2 + @VisibleForTesting private[celeborn] def createClientsInParallel( locationsByHostPort: Seq[(String, Seq[PartitionLocation])], @@ -611,9 +614,12 @@ object CelebornShuffleReader { val futures = locationsByHostPort.map { case (hostPort, locations) => streamCreatorPool.submit(new Runnable { override def run(): Unit = { - val locationsIterator = locations.iterator + val locationsIterator = locations.iterator.take(MAX_CLIENT_CREATION_ATTEMPTS_PER_HOST) var clientCreated = false while (!clientCreated && locationsIterator.hasNext) { + if (Thread.currentThread().isInterrupted) { + throw new InterruptedException("Client creation interrupted") + } val location = locationsIterator.next() try { clientsByHostPort.put(hostPort, createClient(location)) diff --git a/client-spark/spark-3/src/test/scala/org/apache/spark/shuffle/celeborn/CelebornShuffleReaderSuite.scala b/client-spark/spark-3/src/test/scala/org/apache/spark/shuffle/celeborn/CelebornShuffleReaderSuite.scala index d2cec3abf88..eaf81bc4c28 100644 --- a/client-spark/spark-3/src/test/scala/org/apache/spark/shuffle/celeborn/CelebornShuffleReaderSuite.scala +++ b/client-spark/spark-3/src/test/scala/org/apache/spark/shuffle/celeborn/CelebornShuffleReaderSuite.scala @@ -20,7 +20,7 @@ package org.apache.spark.shuffle.celeborn import java.io.IOException import java.nio.file.Files import java.util.concurrent.{CountDownLatch, TimeoutException, TimeUnit} -import java.util.concurrent.atomic.{AtomicBoolean, AtomicReference} +import java.util.concurrent.atomic.{AtomicBoolean, AtomicInteger, AtomicReference} import org.apache.spark.{Dependency, ShuffleDependency, TaskContext} import org.apache.spark.shuffle.ShuffleReadMetricsReporter @@ -188,6 +188,30 @@ class CelebornShuffleReaderSuite extends AnyFunSuite { } } + test("bound failed batch open stream client creation attempts per worker") { + val locations = (0 until 100).map(id => newLocation(id, "worker-0", 19098)) + val streamCreatorPool = ThreadUtils.newDaemonCachedThreadPool("test-create-client", 1, 60) + val attempts = new AtomicInteger() + val failures = new AtomicInteger() + + try { + val clients = CelebornShuffleReader.createClientsInParallel( + Seq(locations.head.hostAndFetchPort -> locations), + streamCreatorPool, + _ => { + attempts.incrementAndGet() + throw new IOException("boom") + }, + (_, _, _) => failures.incrementAndGet()) + + assert(clients.isEmpty) + assert(attempts.get() === 2) + assert(failures.get() === 2) + } finally { + streamCreatorPool.shutdownNow() + } + } + test("cancel batch open stream client creation when waiting thread is interrupted") { val blockedLocation = newLocation(0, "worker-0", 19098) val retryLocation = newLocation(1, "worker-0", 19098) @@ -255,6 +279,69 @@ class CelebornShuffleReaderSuite extends AnyFunSuite { } } + test("do not retry client creation after cancellation during failure callback") { + val failedLocation = newLocation(0, "worker-0", 19098) + val retryLocation = newLocation(1, "worker-0", 19098) + val retryClient = Mockito.mock(classOf[TransportClient]) + val streamCreatorPool = ThreadUtils.newDaemonCachedThreadPool("test-create-client", 1, 60) + val failureCallbackStarted = new CountDownLatch(1) + val releaseFailureCallback = new CountDownLatch(1) + val failureCallbackInterrupted = new CountDownLatch(1) + val retried = new AtomicBoolean(false) + val callerInterrupted = new AtomicBoolean(false) + val callerFailure = new AtomicReference[Throwable]() + + val caller = new Thread(new Runnable { + override def run(): Unit = { + try { + CelebornShuffleReader.createClientsInParallel( + Seq(failedLocation.hostAndFetchPort -> Seq(failedLocation, retryLocation)), + streamCreatorPool, + location => { + if (location eq failedLocation) throw new IOException("boom") + retried.set(true) + retryClient + }, + (_, _, _) => { + failureCallbackStarted.countDown() + try { + releaseFailureCallback.await() + } catch { + case _: InterruptedException => + Thread.currentThread().interrupt() + failureCallbackInterrupted.countDown() + } + }) + callerFailure.set(new AssertionError("Expected waiting thread to be interrupted")) + } catch { + case _: InterruptedException => + callerInterrupted.set(Thread.currentThread().isInterrupted) + case ex: Throwable => + callerFailure.set(ex) + } + } + }) + + try { + caller.start() + assert(failureCallbackStarted.await(5, TimeUnit.SECONDS)) + caller.interrupt() + caller.join(TimeUnit.SECONDS.toMillis(5)) + assert(!caller.isAlive) + assert(failureCallbackInterrupted.await(5, TimeUnit.SECONDS)) + streamCreatorPool.shutdown() + assert(streamCreatorPool.awaitTermination(5, TimeUnit.SECONDS)) + assert(callerFailure.get() == null) + assert(callerInterrupted.get()) + assert(!retried.get()) + } finally { + releaseFailureCallback.countDown() + streamCreatorPool.shutdownNow() + caller.interrupt() + caller.join(TimeUnit.SECONDS.toMillis(5)) + } + } + private def newLocation(id: Int, host: String, fetchPort: Int): PartitionLocation = new PartitionLocation(id, 0, host, 0, 0, fetchPort, 0, PartitionLocation.Mode.PRIMARY) } diff --git a/common/src/main/java/org/apache/celeborn/common/network/client/TransportClientFactory.java b/common/src/main/java/org/apache/celeborn/common/network/client/TransportClientFactory.java index 08c0d277b38..34a46cd63ab 100644 --- a/common/src/main/java/org/apache/celeborn/common/network/client/TransportClientFactory.java +++ b/common/src/main/java/org/apache/celeborn/common/network/client/TransportClientFactory.java @@ -28,6 +28,7 @@ import java.util.concurrent.atomic.AtomicReference; import java.util.function.Supplier; +import com.google.common.annotations.VisibleForTesting; import com.google.common.base.Preconditions; import com.google.common.base.Throwables; import com.google.common.collect.Lists; @@ -160,9 +161,10 @@ public TransportClient retryCreateClient( try { return createClient(remoteHost, remotePort, partitionId, supplier.get()); } catch (Exception e) { - if (e instanceof InterruptedException) { + InterruptedException interruptedException = findInterruptedException(e); + if (interruptedException != null) { Thread.currentThread().interrupt(); - throw e; + throw interruptedException; } numTries++; logger.warn( @@ -182,6 +184,17 @@ public TransportClient retryCreateClient( return null; } + private static InterruptedException findInterruptedException(Throwable throwable) { + Throwable current = throwable; + while (current != null) { + if (current instanceof InterruptedException) { + return (InterruptedException) current; + } + current = current.getCause(); + } + return null; + } + public TransportClient createClient( String remoteHost, int remotePort, int partitionId, ChannelInboundHandlerAdapter decoder) throws IOException, InterruptedException { @@ -315,14 +328,19 @@ public void initChannel(SocketChannel ch) { long preConnect = System.nanoTime(); ChannelFuture cf = bootstrap.connect(address); if (connectTimeoutMs <= 0) { - cf.await(); + awaitWithChannelCleanup( + () -> { + cf.await(); + return true; + }, + cf); assert cf.isDone(); if (cf.isCancelled()) { throw new IOException(String.format("Connecting to %s cancelled", address)); } else if (!cf.isSuccess()) { throw new IOException(String.format("Failed to connect to %s", address), cf.cause()); } - } else if (!cf.await(connectTimeoutMs)) { + } else if (!awaitWithChannelCleanup(() -> cf.await(connectTimeoutMs), cf)) { closeChannel(cf); throw new CelebornIOException( String.format("Connecting to %s timed out (%s ms)", address, connectTimeoutMs)); @@ -351,7 +369,7 @@ public void operationComplete(final Future handshakeFuture) { } } }); - if (!future.await(connectionTimeoutMs)) { + if (!awaitWithChannelCleanup(() -> future.await(connectionTimeoutMs), cf)) { closeChannel(cf); throw new IOException( String.format("Failed to connect to %s within connection timeout", address)); @@ -403,6 +421,25 @@ public void operationComplete(final Future handshakeFuture) { return client; } + @FunctionalInterface + @VisibleForTesting + interface InterruptibleAwait { + boolean await() throws InterruptedException; + } + + @VisibleForTesting + static boolean awaitWithChannelCleanup( + InterruptibleAwait interruptibleAwait, ChannelFuture channelFuture) + throws InterruptedException { + try { + return interruptibleAwait.await(); + } catch (InterruptedException e) { + closeChannel(channelFuture); + Thread.currentThread().interrupt(); + throw e; + } + } + /** Close all connections in the connection pool, and shutdown the worker thread pool. */ @Override public void close() { @@ -428,7 +465,7 @@ public TransportContext getContext() { return context; } - private void closeChannel(ChannelFuture channelFuture) { + private static void closeChannel(ChannelFuture channelFuture) { try { channelFuture.channel().close(); } catch (Exception e) { diff --git a/common/src/test/java/org/apache/celeborn/common/network/TransportClientFactorySuiteJ.java b/common/src/test/java/org/apache/celeborn/common/network/TransportClientFactorySuiteJ.java index 7ec7190aa06..3c51a175d3f 100644 --- a/common/src/test/java/org/apache/celeborn/common/network/TransportClientFactorySuiteJ.java +++ b/common/src/test/java/org/apache/celeborn/common/network/TransportClientFactorySuiteJ.java @@ -260,4 +260,27 @@ public void testRetryCreateClient() throws IOException, InterruptedException { factory.retryCreateClient("xxx", 10, 1, TransportFrameDecoder::new); Assert.assertEquals(transportClient, client); } + + @Test + public void doNotRetryCreateClientWhenInterruptedExceptionIsWrapped() throws Exception { + TransportClientFactory factory = Mockito.spy(context.createClientFactory()); + InterruptedException interruptedException = new InterruptedException("test"); + Mockito.doThrow(new IOException("wrapped", interruptedException)) + .when(factory) + .createClient(anyString(), anyInt(), anyInt(), any()); + + try { + InterruptedException thrown = + assertThrows( + InterruptedException.class, + () -> factory.retryCreateClient("xxx", 10, 1, TransportFrameDecoder::new)); + + assertSame(interruptedException, thrown); + assertTrue(Thread.currentThread().isInterrupted()); + Mockito.verify(factory, Mockito.times(1)) + .createClient(anyString(), anyInt(), anyInt(), any()); + } finally { + Thread.interrupted(); + } + } } diff --git a/common/src/test/java/org/apache/celeborn/common/network/client/TransportClientFactoryInterruptSuiteJ.java b/common/src/test/java/org/apache/celeborn/common/network/client/TransportClientFactoryInterruptSuiteJ.java new file mode 100644 index 00000000000..e6735b27fa7 --- /dev/null +++ b/common/src/test/java/org/apache/celeborn/common/network/client/TransportClientFactoryInterruptSuiteJ.java @@ -0,0 +1,57 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.celeborn.common.network.client; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertThrows; +import static org.junit.Assert.assertTrue; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import io.netty.channel.Channel; +import io.netty.channel.ChannelFuture; +import org.junit.Test; + +public class TransportClientFactoryInterruptSuiteJ { + + @Test + public void closeChannelWhenClientCreationIsInterrupted() throws Exception { + ChannelFuture channelFuture = mock(ChannelFuture.class); + Channel channel = mock(Channel.class); + when(channelFuture.channel()).thenReturn(channel); + + try { + InterruptedException exception = + assertThrows( + InterruptedException.class, + () -> + TransportClientFactory.awaitWithChannelCleanup( + () -> { + throw new InterruptedException("test"); + }, + channelFuture)); + + assertEquals("test", exception.getMessage()); + assertTrue(Thread.currentThread().isInterrupted()); + verify(channel).close(); + } finally { + Thread.interrupted(); + } + } +} From 311bc8c1cb80483a03b46ded061d705bb1336e2e Mon Sep 17 00:00:00 2001 From: Chao Sun Date: Tue, 7 Jul 2026 17:52:05 -0700 Subject: [PATCH 2/2] [CELEBORN-2371] Address batch-open interruption review feedback --- .../celeborn/CelebornShuffleReader.scala | 65 ++++--- .../celeborn/CelebornShuffleReaderSuite.scala | 168 ++++++------------ .../client/read/CelebornInputStream.java | 40 +++++ .../CelebornInputStreamPeerFailoverTest.java | 119 ++++++++++++- .../client/TransportClientFactory.java | 10 +- 5 files changed, 256 insertions(+), 146 deletions(-) diff --git a/client-spark/spark-3/src/main/scala/org/apache/spark/shuffle/celeborn/CelebornShuffleReader.scala b/client-spark/spark-3/src/main/scala/org/apache/spark/shuffle/celeborn/CelebornShuffleReader.scala index 88ece4a36bd..c4db67e3d24 100644 --- a/client-spark/spark-3/src/main/scala/org/apache/spark/shuffle/celeborn/CelebornShuffleReader.scala +++ b/client-spark/spark-3/src/main/scala/org/apache/spark/shuffle/celeborn/CelebornShuffleReader.scala @@ -215,21 +215,23 @@ class CelebornShuffleReader[K, C]( partCnt += 1 val hostPort = location.hostAndFetchPort if (!workerRequestMap.containsKey(hostPort)) { - try { - val client = shuffleClient.getDataClientFactory().createClient( - location.getHost, - location.getFetchPort) + CelebornShuffleReader.tryCreateClient( + location, + loc => + shuffleClient.getDataClientFactory().createClient( + loc.getHost, + loc.getFetchPort), + ex => { + shuffleClient.excludeFailedFetchLocation(hostPort, ex) + logWarning( + s"Failed to create client for $shuffleKey-${location.getId} from host: ${hostPort}. " + + s"Shuffle reader will try its replica if exists.") + }).foreach { client => val pbOpenStreamList = PbOpenStreamList.newBuilder() pbOpenStreamList.setShuffleKey(shuffleKey) workerRequestMap.put( hostPort, (client, new JArrayList[PartitionLocation], pbOpenStreamList)) - } catch { - case ex: Exception => - shuffleClient.excludeFailedFetchLocation(hostPort, ex) - logWarning( - s"Failed to create client for $shuffleKey-${location.getId} from host: ${hostPort}. " + - s"Shuffle reader will try its replica if exists.") } } workerRequestMap.get(hostPort) match { @@ -600,8 +602,25 @@ class CelebornShuffleReader[K, C]( object CelebornShuffleReader { var streamCreatorPool: ThreadPoolExecutor = null - // TransportClientFactory already retries each attempt; allow one extra pooled-client attempt. - private val MAX_CLIENT_CREATION_ATTEMPTS_PER_HOST = 2 + @VisibleForTesting + private[celeborn] def tryCreateClient( + location: PartitionLocation, + createClient: PartitionLocation => TransportClient, + onClientCreateFailure: Exception => Unit): Option[TransportClient] = { + if (Thread.currentThread().isInterrupted) { + throw new InterruptedException("Client creation interrupted") + } + try { + Some(createClient(location)) + } catch { + case ex: InterruptedException => + Thread.currentThread().interrupt() + throw ex + case ex: Exception => + onClientCreateFailure(ex) + None + } + } @VisibleForTesting private[celeborn] def createClientsInParallel( @@ -614,22 +633,12 @@ object CelebornShuffleReader { val futures = locationsByHostPort.map { case (hostPort, locations) => streamCreatorPool.submit(new Runnable { override def run(): Unit = { - val locationsIterator = locations.iterator.take(MAX_CLIENT_CREATION_ATTEMPTS_PER_HOST) - var clientCreated = false - while (!clientCreated && locationsIterator.hasNext) { - if (Thread.currentThread().isInterrupted) { - throw new InterruptedException("Client creation interrupted") - } - val location = locationsIterator.next() - try { - clientsByHostPort.put(hostPort, createClient(location)) - clientCreated = true - } catch { - case ex: InterruptedException => - Thread.currentThread().interrupt() - throw ex - case ex: Exception => - onClientCreateFailure(hostPort, location, ex) + locations.headOption.foreach { location => + tryCreateClient( + location, + createClient, + ex => onClientCreateFailure(hostPort, location, ex)).foreach { client => + clientsByHostPort.put(hostPort, client) } } } diff --git a/client-spark/spark-3/src/test/scala/org/apache/spark/shuffle/celeborn/CelebornShuffleReaderSuite.scala b/client-spark/spark-3/src/test/scala/org/apache/spark/shuffle/celeborn/CelebornShuffleReaderSuite.scala index eaf81bc4c28..cb362491744 100644 --- a/client-spark/spark-3/src/test/scala/org/apache/spark/shuffle/celeborn/CelebornShuffleReaderSuite.scala +++ b/client-spark/spark-3/src/test/scala/org/apache/spark/shuffle/celeborn/CelebornShuffleReaderSuite.scala @@ -164,31 +164,7 @@ class CelebornShuffleReaderSuite extends AnyFunSuite { } } - test("retry failed batch open stream client creation for the same worker") { - val failedLocation = newLocation(0, "worker-0", 19098) - val retryLocation = newLocation(1, "worker-0", 19098) - val client = Mockito.mock(classOf[TransportClient]) - val streamCreatorPool = ThreadUtils.newDaemonCachedThreadPool("test-create-client", 1, 60) - var failureCount = 0 - - try { - val clients = CelebornShuffleReader.createClientsInParallel( - Seq(failedLocation.hostAndFetchPort -> Seq(failedLocation, retryLocation)), - streamCreatorPool, - location => { - if (location eq failedLocation) throw new IOException("boom") - client - }, - (_, _, _) => failureCount += 1) - - assert(failureCount === 1) - assert(clients(failedLocation.hostAndFetchPort) eq client) - } finally { - streamCreatorPool.shutdownNow() - } - } - - test("bound failed batch open stream client creation attempts per worker") { + test("attempt batch open stream client creation once per worker") { val locations = (0 until 100).map(id => newLocation(id, "worker-0", 19098)) val streamCreatorPool = ThreadUtils.newDaemonCachedThreadPool("test-create-client", 1, 60) val attempts = new AtomicInteger() @@ -205,89 +181,66 @@ class CelebornShuffleReaderSuite extends AnyFunSuite { (_, _, _) => failures.incrementAndGet()) assert(clients.isEmpty) - assert(attempts.get() === 2) - assert(failures.get() === 2) + assert(attempts.get() === 1) + assert(failures.get() === 1) } finally { streamCreatorPool.shutdownNow() } } - test("cancel batch open stream client creation when waiting thread is interrupted") { - val blockedLocation = newLocation(0, "worker-0", 19098) - val retryLocation = newLocation(1, "worker-0", 19098) - val retryClient = Mockito.mock(classOf[TransportClient]) - val streamCreatorPool = ThreadUtils.newDaemonCachedThreadPool("test-create-client", 1, 60) - val clientStarted = new CountDownLatch(1) - val releaseClient = new CountDownLatch(1) - val clientInterrupted = new CountDownLatch(1) - val retried = new AtomicBoolean(false) + test("propagate client creation interruption without reporting a worker failure") { + val location = newLocation(0, "worker-0", 19098) val failureReported = new AtomicBoolean(false) - val callerInterrupted = new AtomicBoolean(false) - val callerFailure = new AtomicReference[Throwable]() - val caller = new Thread(new Runnable { - override def run(): Unit = { - try { - CelebornShuffleReader.createClientsInParallel( - Seq(blockedLocation.hostAndFetchPort -> Seq(blockedLocation, retryLocation)), - streamCreatorPool, - location => { - if (location eq blockedLocation) { - clientStarted.countDown() - try { - releaseClient.await() - retryClient - } catch { - case ex: InterruptedException => - clientInterrupted.countDown() - throw ex - } - } else { - retried.set(true) - retryClient - } - }, - (_, _, _) => failureReported.set(true)) - callerFailure.set(new AssertionError("Expected waiting thread to be interrupted")) - } catch { - case _: InterruptedException => - callerInterrupted.set(Thread.currentThread().isInterrupted) - case ex: Throwable => - callerFailure.set(ex) - } + try { + val exception = intercept[InterruptedException] { + CelebornShuffleReader.tryCreateClient( + location, + _ => throw new InterruptedException("test"), + _ => failureReported.set(true)) } - }) + + assert(exception.getMessage === "test") + assert(Thread.currentThread().isInterrupted) + assert(!failureReported.get()) + } finally { + Thread.interrupted() + } + } + + test("skip client creation when the thread is already interrupted") { + val location = newLocation(0, "worker-0", 19098) + val clientCreationAttempted = new AtomicBoolean(false) + val failureReported = new AtomicBoolean(false) try { - caller.start() - assert(clientStarted.await(5, TimeUnit.SECONDS)) - caller.interrupt() - caller.join(TimeUnit.SECONDS.toMillis(5)) - assert(!caller.isAlive) - assert(clientInterrupted.await(5, TimeUnit.SECONDS)) - streamCreatorPool.shutdown() - assert(streamCreatorPool.awaitTermination(5, TimeUnit.SECONDS)) - assert(callerFailure.get() == null) - assert(callerInterrupted.get()) - assert(!retried.get()) + Thread.currentThread().interrupt() + intercept[InterruptedException] { + CelebornShuffleReader.tryCreateClient( + location, + _ => { + clientCreationAttempted.set(true) + Mockito.mock(classOf[TransportClient]) + }, + _ => failureReported.set(true)) + } + + assert(Thread.currentThread().isInterrupted) + assert(!clientCreationAttempted.get()) assert(!failureReported.get()) } finally { - releaseClient.countDown() - streamCreatorPool.shutdownNow() - caller.interrupt() - caller.join(TimeUnit.SECONDS.toMillis(5)) + Thread.interrupted() } } - test("do not retry client creation after cancellation during failure callback") { - val failedLocation = newLocation(0, "worker-0", 19098) - val retryLocation = newLocation(1, "worker-0", 19098) - val retryClient = Mockito.mock(classOf[TransportClient]) + test("cancel batch open stream client creation when waiting thread is interrupted") { + val blockedLocation = newLocation(0, "worker-0", 19098) + val client = Mockito.mock(classOf[TransportClient]) val streamCreatorPool = ThreadUtils.newDaemonCachedThreadPool("test-create-client", 1, 60) - val failureCallbackStarted = new CountDownLatch(1) - val releaseFailureCallback = new CountDownLatch(1) - val failureCallbackInterrupted = new CountDownLatch(1) - val retried = new AtomicBoolean(false) + val clientStarted = new CountDownLatch(1) + val releaseClient = new CountDownLatch(1) + val clientInterrupted = new CountDownLatch(1) + val failureReported = new AtomicBoolean(false) val callerInterrupted = new AtomicBoolean(false) val callerFailure = new AtomicReference[Throwable]() @@ -295,23 +248,20 @@ class CelebornShuffleReaderSuite extends AnyFunSuite { override def run(): Unit = { try { CelebornShuffleReader.createClientsInParallel( - Seq(failedLocation.hostAndFetchPort -> Seq(failedLocation, retryLocation)), + Seq(blockedLocation.hostAndFetchPort -> Seq(blockedLocation)), streamCreatorPool, - location => { - if (location eq failedLocation) throw new IOException("boom") - retried.set(true) - retryClient - }, - (_, _, _) => { - failureCallbackStarted.countDown() + _ => { + clientStarted.countDown() try { - releaseFailureCallback.await() + releaseClient.await() + client } catch { - case _: InterruptedException => - Thread.currentThread().interrupt() - failureCallbackInterrupted.countDown() + case ex: InterruptedException => + clientInterrupted.countDown() + throw ex } - }) + }, + (_, _, _) => failureReported.set(true)) callerFailure.set(new AssertionError("Expected waiting thread to be interrupted")) } catch { case _: InterruptedException => @@ -324,18 +274,18 @@ class CelebornShuffleReaderSuite extends AnyFunSuite { try { caller.start() - assert(failureCallbackStarted.await(5, TimeUnit.SECONDS)) + assert(clientStarted.await(5, TimeUnit.SECONDS)) caller.interrupt() caller.join(TimeUnit.SECONDS.toMillis(5)) assert(!caller.isAlive) - assert(failureCallbackInterrupted.await(5, TimeUnit.SECONDS)) + assert(clientInterrupted.await(5, TimeUnit.SECONDS)) streamCreatorPool.shutdown() assert(streamCreatorPool.awaitTermination(5, TimeUnit.SECONDS)) assert(callerFailure.get() == null) assert(callerInterrupted.get()) - assert(!retried.get()) + assert(!failureReported.get()) } finally { - releaseFailureCallback.countDown() + releaseClient.countDown() streamCreatorPool.shutdownNow() caller.interrupt() caller.join(TimeUnit.SECONDS.toMillis(5)) diff --git a/client/src/main/java/org/apache/celeborn/client/read/CelebornInputStream.java b/client/src/main/java/org/apache/celeborn/client/read/CelebornInputStream.java index 37e0be3e375..dd9419d4faf 100644 --- a/client/src/main/java/org/apache/celeborn/client/read/CelebornInputStream.java +++ b/client/src/main/java/org/apache/celeborn/client/read/CelebornInputStream.java @@ -27,6 +27,7 @@ import scala.Tuple2; import com.github.luben.zstd.ZstdException; +import com.google.common.base.Throwables; import com.google.common.util.concurrent.Uninterruptibles; import io.netty.buffer.ByteBuf; import net.jpountz.lz4.LZ4Exception; @@ -48,6 +49,7 @@ import org.apache.celeborn.common.protocol.*; import org.apache.celeborn.common.unsafe.Platform; import org.apache.celeborn.common.util.ExceptionMaker; +import org.apache.celeborn.common.util.ExceptionUtils; import org.apache.celeborn.common.util.Utils; import org.apache.celeborn.common.write.LocationPushFailedBatches; @@ -438,6 +440,23 @@ private boolean isExcluded(PartitionLocation location) { } } + private static boolean isInterruption(Throwable throwable) { + if (Thread.currentThread().isInterrupted()) { + return true; + } + for (Throwable cause : Throwables.getCausalChain(throwable)) { + if (cause instanceof InterruptedException) { + return true; + } + } + return false; + } + + private static IOException interruptedIOException(Exception exception) { + Thread.currentThread().interrupt(); + return ExceptionUtils.wrapThrowableToIOException(exception); + } + private PartitionReader createReaderWithRetry( PartitionLocation location, PbStreamHandler pbStreamHandler) throws IOException { return createReaderWithRetry(location, pbStreamHandler, Optional.empty()); @@ -464,6 +483,9 @@ private PartitionReader createReaderWithRetry( checkpointMetadata); return reader; } catch (Exception e) { + if (isInterruption(e)) { + throw interruptedIOException(e); + } lastException = e; shuffleClient.excludeFailedFetchLocation(location.hostAndFetchPort(), e); fetchChunkRetryCnt++; @@ -493,6 +515,9 @@ private PartitionReader createReaderWithRetry( .toByteArray()); client.sendRpc(bufferStreamEnd.toByteBuffer()); } catch (InterruptedException | IOException | RuntimeException ex) { + if (isInterruption(ex)) { + throw interruptedIOException(ex); + } logger.warn( "Close {} stream {} failed", location.hostAndFetchPort(), @@ -528,6 +553,15 @@ private ByteBuf getNextChunk() throws IOException { } return currentReader.next(); } catch (Exception e) { + if (isInterruption(e)) { + IOException interrupted = interruptedIOException(e); + try { + currentReader.close(); + } catch (RuntimeException closeException) { + interrupted.addSuppressed(closeException); + } + throw interrupted; + } shuffleClient.excludeFailedFetchLocation( currentReader.getLocation().hostAndFetchPort(), e); fetchChunkRetryCnt++; @@ -890,6 +924,9 @@ private boolean fillBuffer() throws IOException { } return hasData; } catch (LZ4Exception | ZstdException | IOException e) { + if (isInterruption(e)) { + throw interruptedIOException(e); + } logger.error( "Failed to fill buffer from chunk. AppShuffleId {}, shuffleId {}, partitionId {}, location {}", appShuffleId, @@ -918,6 +955,9 @@ private boolean fillBuffer() throws IOException { } throw ioe; } catch (Exception e) { + if (isInterruption(e)) { + throw interruptedIOException(e); + } logger.error( "Failed to fill buffer from chunk. AppShuffleId {}, shuffleId {}, partitionId {}, location {}", appShuffleId, diff --git a/client/src/test/java/org/apache/celeborn/client/read/CelebornInputStreamPeerFailoverTest.java b/client/src/test/java/org/apache/celeborn/client/read/CelebornInputStreamPeerFailoverTest.java index a7dd7db3178..84e2bbb6b78 100644 --- a/client/src/test/java/org/apache/celeborn/client/read/CelebornInputStreamPeerFailoverTest.java +++ b/client/src/test/java/org/apache/celeborn/client/read/CelebornInputStreamPeerFailoverTest.java @@ -17,6 +17,10 @@ package org.apache.celeborn.client.read; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertSame; +import static org.junit.Assert.assertThrows; +import static org.junit.Assert.assertTrue; import static org.junit.Assert.fail; import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.anyInt; @@ -24,7 +28,10 @@ import static org.mockito.ArgumentMatchers.anyString; import static org.mockito.Mockito.atLeast; import static org.mockito.Mockito.doNothing; +import static org.mockito.Mockito.doThrow; import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.times; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; @@ -48,8 +55,9 @@ import org.apache.celeborn.common.protocol.PartitionLocation; import org.apache.celeborn.common.protocol.PbStreamHandler; import org.apache.celeborn.common.protocol.StorageInfo; +import org.apache.celeborn.common.util.ExceptionMaker; -/** Tests for CelebornInputStream peer failover when stream cleanup throws RuntimeException. */ +/** Tests for CelebornInputStream peer failover and interruption handling. */ public class CelebornInputStreamPeerFailoverTest { private static final String SHUFFLE_KEY = "appid-1-1"; @@ -142,6 +150,103 @@ public void testPeerFailoverWithIOExceptionDuringCleanup() throws Exception { } } + @Test + public void testInterruptedClientCreationDoesNotRetryOrFailOver() throws Exception { + InterruptedException interruptedException = new InterruptedException("cancelled"); + IOException wrappedException = new IOException("wrapped", interruptedException); + when(clientFactory.createClient(anyString(), anyInt())).thenThrow(wrappedException); + + try { + IOException thrown = + assertThrows(IOException.class, () -> createInputStream(PRIMARY_HOST, REPLICA_HOST)); + + assertSame(wrappedException, thrown); + assertTrue(Thread.currentThread().isInterrupted()); + verify(clientFactory, times(1)).createClient(anyString(), anyInt()); + verify(shuffleClient, never()).excludeFailedFetchLocation(anyString(), any()); + } finally { + Thread.interrupted(); + } + } + + @Test + public void testInterruptedCleanupDoesNotFailOver() throws Exception { + AtomicInteger primaryAttempts = new AtomicInteger(); + AtomicInteger replicaAttempts = new AtomicInteger(); + InterruptedException interruptedException = new InterruptedException("cancelled"); + IOException cleanupException = new IOException("cleanup interrupted", interruptedException); + + when(clientFactory.createClient(anyString(), anyInt())) + .thenAnswer( + invocation -> { + String host = invocation.getArgument(0); + if (PRIMARY_HOST.equals(host)) { + if (primaryAttempts.incrementAndGet() == 1) { + throw new IOException("Worker Not Registered!"); + } + throw cleanupException; + } + replicaAttempts.incrementAndGet(); + return mockReplicaClient(); + }); + + try { + IOException thrown = + assertThrows(IOException.class, () -> createInputStream(PRIMARY_HOST, REPLICA_HOST)); + + assertSame(cleanupException, thrown); + assertTrue(Thread.currentThread().isInterrupted()); + assertEquals(2, primaryAttempts.get()); + assertEquals(0, replicaAttempts.get()); + } finally { + Thread.interrupted(); + } + } + + @Test + public void testInterruptedReconnectDoesNotReportFetchFailureOrFailOver() throws Exception { + AtomicInteger primaryAttempts = new AtomicInteger(); + AtomicInteger replicaAttempts = new AtomicInteger(); + TransportClient inactiveClient = mock(TransportClient.class); + when(inactiveClient.isActive()).thenReturn(false, true); + InterruptedException interruptedException = new InterruptedException("cancelled"); + IOException reconnectException = new IOException("reconnect interrupted", interruptedException); + RuntimeException closeException = new RuntimeException("close failed"); + doThrow(closeException).when(inactiveClient).sendRpc(any(ByteBuffer.class)); + + when(clientFactory.createClient(anyString(), anyInt())) + .thenAnswer( + invocation -> { + String host = invocation.getArgument(0); + if (PRIMARY_HOST.equals(host)) { + if (primaryAttempts.incrementAndGet() == 1) { + return inactiveClient; + } + throw reconnectException; + } + replicaAttempts.incrementAndGet(); + return mockReplicaClient(); + }); + ExceptionMaker exceptionMaker = mock(ExceptionMaker.class); + + try { + CelebornInputStream inputStream = + createInputStream(PRIMARY_HOST, REPLICA_HOST, exceptionMaker); + IOException thrown = assertThrows(IOException.class, inputStream::read); + + assertSame(reconnectException, thrown.getCause()); + assertEquals(1, thrown.getSuppressed().length); + assertSame(closeException, thrown.getSuppressed()[0]); + assertTrue(Thread.currentThread().isInterrupted()); + assertEquals(2, primaryAttempts.get()); + assertEquals(0, replicaAttempts.get()); + verify(shuffleClient, never()).excludeFailedFetchLocation(anyString(), any()); + verify(shuffleClient, never()).reportShuffleFetchFailure(anyInt(), anyInt(), anyLong()); + } finally { + Thread.interrupted(); + } + } + /** Tests that all retries are exhausted and an exception is thrown when there is no peer. */ @Test(expected = CelebornIOException.class) public void testFailureWithoutPeer() throws Exception { @@ -177,7 +282,13 @@ public void testFailureWithoutPeer() throws Exception { false); } - private void createInputStream(String primaryHost, String replicaHost) throws IOException { + private CelebornInputStream createInputStream(String primaryHost, String replicaHost) + throws IOException { + return createInputStream(primaryHost, replicaHost, null); + } + + private CelebornInputStream createInputStream( + String primaryHost, String replicaHost, ExceptionMaker exceptionMaker) throws IOException { PartitionLocation primary = createPartitionLocation(primaryHost); PartitionLocation replica = createPartitionLocation(replicaHost); primary.setPeer(replica); @@ -189,7 +300,7 @@ private void createInputStream(String primaryHost, String replicaHost) throws IO ArrayList streamHandlers = new ArrayList<>(); streamHandlers.add(PbStreamHandler.newBuilder().setStreamId(123L).setNumChunks(10).build()); - CelebornInputStream.create( + return CelebornInputStream.create( conf, clientFactory, SHUFFLE_KEY, @@ -207,7 +318,7 @@ private void createInputStream(String primaryHost, String replicaHost) throws IO 1, 1, 0, - null, + exceptionMaker, new TestMetricsCallback(), false); } diff --git a/common/src/main/java/org/apache/celeborn/common/network/client/TransportClientFactory.java b/common/src/main/java/org/apache/celeborn/common/network/client/TransportClientFactory.java index 34a46cd63ab..5fd21fd3ffe 100644 --- a/common/src/main/java/org/apache/celeborn/common/network/client/TransportClientFactory.java +++ b/common/src/main/java/org/apache/celeborn/common/network/client/TransportClientFactory.java @@ -185,12 +185,10 @@ public TransportClient retryCreateClient( } private static InterruptedException findInterruptedException(Throwable throwable) { - Throwable current = throwable; - while (current != null) { - if (current instanceof InterruptedException) { - return (InterruptedException) current; + for (Throwable cause : Throwables.getCausalChain(throwable)) { + if (cause instanceof InterruptedException) { + return (InterruptedException) cause; } - current = current.getCause(); } return null; } @@ -336,8 +334,10 @@ public void initChannel(SocketChannel ch) { cf); assert cf.isDone(); if (cf.isCancelled()) { + closeChannel(cf); throw new IOException(String.format("Connecting to %s cancelled", address)); } else if (!cf.isSuccess()) { + closeChannel(cf); throw new IOException(String.format("Failed to connect to %s", address), cf.cause()); } } else if (!awaitWithChannelCleanup(() -> cf.await(connectTimeoutMs), cf)) {