This is an automated email from the ASF dual-hosted git repository.
SteNicholas pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/celeborn.git
The following commit(s) were added to refs/heads/main by this push:
new f0c8139621 [CELEBORN-2371] Bound Spark batch-open client creation
retries and stop them on interruption
f0c8139621 is described below
commit f0c8139621f3da73ccf30566fb344ffb2ff47cf9
Author: Chao Sun <[email protected]>
AuthorDate: Fri Jul 24 10:14:00 2026 +0800
[CELEBORN-2371] Bound Spark batch-open client creation retries and stop
them on interruption
## Why are the changes needed?
[CELEBORN-2371](https://issues.apache.org/jira/browse/CELEBORN-2371)
follows up on the parallel Spark batch-open client creation added by
[#3692](https://github.com/apache/celeborn/pull/3692).
Batch-open locations are grouped by `host:fetchPort`, but each worker task
previously walked multiple `PartitionLocation` entries. Since every
`createClient` invocation already receives `TransportClientFactory`'s full
retry budget, that outer loop could multiply connection latency without
targeting a different endpoint.
Cancellation also needs to terminate consistently. A direct or wrapped
`InterruptedException` must stop reader setup and fetch retry flow before it is
recorded as a worker failure, added to shared exclusion state, retried against
a peer, or reported as a shuffle fetch failure. The unlimited-timeout transport
branch also needed explicit cleanup on cancelled or failed connection futures.
## What changes were proposed in this PR?
### Keep one retry budget per worker endpoint
Each grouped `host:fetchPort` now uses one representative location for
client creation. `TransportClientFactory` remains the owner of connection
retries, controlled by the existing `celeborn.data.io.maxRetries`; there is no
second outer retry setting or multiplicative retry budget.
### Propagate cancellation through reader paths
`CelebornShuffleReader` now uses a shared interrupt-aware client-creation
helper in both sequential and parallel batch-open paths. It checks a
pre-existing interrupt flag, restores the flag when `createClient` throws
`InterruptedException`, and exits without invoking the ordinary failure
callback.
`CelebornInputStream` now detects interruption throughout reader creation,
failed-stream cleanup, reconnect, and buffer-fill paths. It exits before
exclusion, retry, peer failover, or shuffle-fetch-failure reporting. If reader
cleanup itself fails while propagating cancellation, the cleanup failure is
retained as a suppressed exception without masking the original interruption.
### Preserve interruption and cleanup in transport bootstrap
`TransportClientFactory` uses Guava's `Throwables.getCausalChain()` to
detect wrapped `InterruptedException`, restores the interrupt flag, and stops
retrying immediately. TCP-connect and TLS-handshake waits close the in-progress
channel before propagating interruption. The unlimited-timeout connection path
now also closes cancelled and failed channel futures explicitly.
## How was this PR tested?
Formatting was applied with:
```text
./build/mvn --no-transfer-progress -DskipTests spotless:apply
```
The focused transport tests passed (11 tests):
```text
./build/mvn --no-transfer-progress -pl common -am \
-Dtest=TransportClientFactorySuiteJ,TransportClientFactoryInterruptSuiteJ
\
-DwildcardSuites=none clean test
```
The focused input-stream peer-failover and interruption tests passed (6
tests):
```text
./build/mvn --no-transfer-progress -pl client -am \
-Dtest=CelebornInputStreamPeerFailoverTest \
-DwildcardSuites=none clean test
```
The Spark 3.5 reader suite passed (7 tests):
```text
./build/mvn --no-transfer-progress -Pspark-3.5 -pl client-spark/spark-3 -am
\
-Dtest=none \
-DwildcardSuites=org.apache.spark.shuffle.celeborn.CelebornShuffleReaderSuite \
clean test
```
Spark 4.0 / Scala 2.13 production and test compilation also passed:
```text
./build/mvn --no-transfer-progress -Pspark-4.0 -pl client-spark/spark-3 -am
\
-DskipTests clean test
```
Closes #3746 from
sunchao/dev/chao/codex/bound-batch-open-client-retries-oss.
Authored-by: Chao Sun <[email protected]>
Signed-off-by: Nicholas Jiang <[email protected]>
---
.../shuffle/celeborn/CelebornShuffleReader.scala | 92 +++++++++----
.../celeborn/CelebornShuffleReaderSuite.scala | 151 ++++++++++++++++-----
.../celeborn/client/read/CelebornInputStream.java | 32 +++++
.../read/CelebornInputStreamPeerFailoverTest.java | 119 +++++++++++++++-
.../network/client/TransportClientFactory.java | 41 +++++-
.../celeborn/common/util/ExceptionUtils.java | 11 ++
.../network/TransportClientFactorySuiteJ.java | 23 ++++
.../TransportClientFactoryInterruptSuiteJ.java | 57 ++++++++
8 files changed, 458 insertions(+), 68 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 32863c6a54..87bf9ffd6e 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
@@ -18,8 +18,8 @@
package org.apache.spark.shuffle.celeborn
import java.io.IOException
-import java.util.{ArrayList => JArrayList, HashMap => JHashMap, Map => JMap,
Optional, Set => JSet}
-import java.util.concurrent.{ConcurrentHashMap, ThreadPoolExecutor,
TimeoutException, TimeUnit}
+import java.util.{ArrayList => JArrayList, HashMap => JHashMap, HashSet =>
JHashSet, Map => JMap, Optional, Set => JSet}
+import java.util.concurrent.{ConcurrentHashMap, ExecutionException,
ThreadPoolExecutor, TimeoutException, TimeUnit}
import java.util.concurrent.atomic.AtomicReference
import java.util.function.BiFunction
@@ -48,7 +48,7 @@ import
org.apache.celeborn.common.network.protocol.TransportMessage
import org.apache.celeborn.common.protocol._
import org.apache.celeborn.common.protocol.message.{ControlMessages,
StatusCode}
import
org.apache.celeborn.common.protocol.message.ControlMessages.GetReducerFileGroupResponse
-import org.apache.celeborn.common.util.{JavaUtils, ThreadUtils, Utils}
+import org.apache.celeborn.common.util.{ExceptionUtils, JavaUtils,
ThreadUtils, Utils}
class CelebornShuffleReader[K, C](
handle: CelebornShuffleHandle[K, _, C],
@@ -236,27 +236,31 @@ class CelebornShuffleReader[K, C](
val parallelClientCreationEnabled =
conf.batchOpenStreamParallelClientCreationEnabled
val locationsByHostPort =
new mutable.LinkedHashMap[String, JArrayList[PartitionLocation]]()
+ val attemptedClientHostPorts = new JHashSet[String]()
def makeOpenStreamList(locations: JSet[PartitionLocation]): Unit = {
locations.asScala.foreach { location =>
partCnt += 1
val hostPort = location.hostAndFetchPort
if (!workerRequestMap.containsKey(hostPort)) {
- try {
- val client = shuffleClient.getDataClientFactory().createClient(
- location.getHost,
- location.getFetchPort)
+ CelebornShuffleReader.tryCreateClientOncePerEndpoint(
+ location,
+ attemptedClientHostPorts,
+ 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 {
@@ -627,6 +631,42 @@ class CelebornShuffleReader[K, C](
object CelebornShuffleReader {
var streamCreatorPool: ThreadPoolExecutor = null
+ @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: Exception =>
+ val interruptedException = ExceptionUtils.findInterruptedException(ex)
+ if (interruptedException != null) {
+ Thread.currentThread().interrupt()
+ throw interruptedException
+ } else {
+ onClientCreateFailure(ex)
+ None
+ }
+ }
+ }
+
+ @VisibleForTesting
+ private[celeborn] def tryCreateClientOncePerEndpoint(
+ location: PartitionLocation,
+ attemptedClientHostPorts: JSet[String],
+ createClient: PartitionLocation => TransportClient,
+ onClientCreateFailure: Exception => Unit): Option[TransportClient] = {
+ if (attemptedClientHostPorts.add(location.hostAndFetchPort)) {
+ tryCreateClient(location, createClient, onClientCreateFailure)
+ } else {
+ None
+ }
+ }
+
@VisibleForTesting
private[celeborn] def createClientsInParallel(
locationsByHostPort: Seq[(String, Seq[PartitionLocation])],
@@ -638,19 +678,12 @@ object CelebornShuffleReader {
val futures = locationsByHostPort.map { case (hostPort, locations) =>
streamCreatorPool.submit(new Runnable {
override def run(): Unit = {
- val locationsIterator = locations.iterator
- var clientCreated = false
- while (!clientCreated && locationsIterator.hasNext) {
- 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)
}
}
}
@@ -664,6 +697,13 @@ object CelebornShuffleReader {
case ex: InterruptedException =>
Thread.currentThread().interrupt()
throw ex
+ case ex: ExecutionException =>
+ val interruptedException = ExceptionUtils.findInterruptedException(ex)
+ if (interruptedException != null) {
+ Thread.currentThread().interrupt()
+ throw interruptedException
+ }
+ throw ex
} finally {
if (!waitCompleted) {
futures.foreach(_.cancel(true))
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 901d7797b5..acd68bfff7 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
@@ -19,8 +19,9 @@ package org.apache.spark.shuffle.celeborn
import java.io.IOException
import java.nio.file.Files
+import java.util.HashSet
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
@@ -164,39 +165,131 @@ 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])
+ 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)
- var failureCount = 0
+ val attempts = new AtomicInteger()
+ val failures = new AtomicInteger()
try {
val clients = CelebornShuffleReader.createClientsInParallel(
- Seq(failedLocation.hostAndFetchPort -> Seq(failedLocation,
retryLocation)),
+ Seq(locations.head.hostAndFetchPort -> locations),
streamCreatorPool,
- location => {
- if (location eq failedLocation) throw new IOException("boom")
- client
+ _ => {
+ attempts.incrementAndGet()
+ throw new IOException("boom")
},
- (_, _, _) => failureCount += 1)
+ (_, _, _) => failures.incrementAndGet())
+
+ assert(clients.isEmpty)
+ assert(attempts.get() === 1)
+ assert(failures.get() === 1)
+ } finally {
+ streamCreatorPool.shutdownNow()
+ }
+ }
+
+ test("attempt sequential batch open stream client creation once per worker
endpoint") {
+ val firstLocation = newLocation(0, "worker-0", 19098)
+ val duplicateEndpoint = newLocation(1, "worker-0", 19098)
+ val differentEndpoint = newLocation(2, "worker-0", 19099)
+ val attemptedClientHostPorts = new HashSet[String]()
+ val attempts = new AtomicInteger()
+ val failures = new AtomicInteger()
+
+ Seq(firstLocation, duplicateEndpoint, differentEndpoint).foreach {
location =>
+ val client = CelebornShuffleReader.tryCreateClientOncePerEndpoint(
+ location,
+ attemptedClientHostPorts,
+ _ => {
+ attempts.incrementAndGet()
+ throw new IOException("boom")
+ },
+ _ => failures.incrementAndGet())
+ assert(client.isEmpty)
+ }
+
+ assert(attempts.get() === 2)
+ assert(failures.get() === 2)
+ }
+
+ test("propagate wrapped client creation interruption without reporting a
worker failure") {
+ val location = newLocation(0, "worker-0", 19098)
+ val failureReported = new AtomicBoolean(false)
+ val interruptedException = new InterruptedException("test")
+
+ try {
+ val exception = intercept[InterruptedException] {
+ CelebornShuffleReader.tryCreateClient(
+ location,
+ _ => throw new IOException("wrapped", interruptedException),
+ _ => failureReported.set(true))
+ }
+
+ assert(exception eq interruptedException)
+ assert(Thread.currentThread().isInterrupted)
+ assert(!failureReported.get())
+ } finally {
+ Thread.interrupted()
+ }
+ }
+
+ test("propagate interruption across parallel client creation future
boundary") {
+ val location = newLocation(0, "worker-0", 19098)
+ val interruptedException = new InterruptedException("test")
+ val failureReported = new AtomicBoolean(false)
+ val streamCreatorPool =
ThreadUtils.newDaemonCachedThreadPool("test-create-client", 1, 60)
- assert(failureCount === 1)
- assert(clients(failedLocation.hostAndFetchPort) eq client)
+ try {
+ val exception = intercept[InterruptedException] {
+ CelebornShuffleReader.createClientsInParallel(
+ Seq(location.hostAndFetchPort -> Seq(location)),
+ streamCreatorPool,
+ _ => throw new IOException("wrapped", interruptedException),
+ (_, _, _) => failureReported.set(true))
+ }
+
+ assert(exception eq interruptedException)
+ assert(Thread.currentThread().isInterrupted)
+ assert(!failureReported.get())
} finally {
+ Thread.interrupted()
streamCreatorPool.shutdownNow()
}
}
+ 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 {
+ 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 {
+ Thread.interrupted()
+ }
+ }
+
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 client = 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)
val failureReported = new AtomicBoolean(false)
val callerInterrupted = new AtomicBoolean(false)
val callerFailure = new AtomicReference[Throwable]()
@@ -205,22 +298,17 @@ class CelebornShuffleReaderSuite extends AnyFunSuite {
override def run(): Unit = {
try {
CelebornShuffleReader.createClientsInParallel(
- Seq(blockedLocation.hostAndFetchPort -> Seq(blockedLocation,
retryLocation)),
+ Seq(blockedLocation.hostAndFetchPort -> Seq(blockedLocation)),
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
+ _ => {
+ clientStarted.countDown()
+ try {
+ releaseClient.await()
+ client
+ } catch {
+ case ex: InterruptedException =>
+ clientInterrupted.countDown()
+ throw ex
}
},
(_, _, _) => failureReported.set(true))
@@ -245,7 +333,6 @@ class CelebornShuffleReaderSuite extends AnyFunSuite {
assert(streamCreatorPool.awaitTermination(5, TimeUnit.SECONDS))
assert(callerFailure.get() == null)
assert(callerInterrupted.get())
- assert(!retried.get())
assert(!failureReported.get())
} finally {
releaseClient.countDown()
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 d76b1eb9d3..614ce5926f 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
@@ -49,6 +49,7 @@ import
org.apache.celeborn.common.network.protocol.TransportMessage;
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;
@@ -448,6 +449,16 @@ public abstract class CelebornInputStream extends
InputStream {
}
}
+ private static boolean isInterruption(Throwable throwable) {
+ return Thread.currentThread().isInterrupted()
+ || ExceptionUtils.findInterruptedException(throwable) != null;
+ }
+
+ 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());
@@ -474,6 +485,9 @@ public abstract class CelebornInputStream extends
InputStream {
checkpointMetadata);
return reader;
} catch (Exception e) {
+ if (isInterruption(e)) {
+ throw interruptedIOException(e);
+ }
lastException = e;
shuffleClient.excludeFailedFetchLocation(location.hostAndFetchPort(), e);
fetchChunkRetryCnt++;
@@ -503,6 +517,9 @@ public abstract class CelebornInputStream extends
InputStream {
.toByteArray());
client.sendRpc(bufferStreamEnd.toByteBuffer());
} catch (InterruptedException | IOException | RuntimeException
ex) {
+ if (isInterruption(ex)) {
+ throw interruptedIOException(ex);
+ }
logger.warn(
"Close {} stream {} failed",
location.hostAndFetchPort(),
@@ -538,6 +555,15 @@ public abstract class CelebornInputStream extends
InputStream {
}
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++;
@@ -946,6 +972,9 @@ public abstract class CelebornInputStream extends
InputStream {
}
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,
@@ -974,6 +1003,9 @@ public abstract class CelebornInputStream extends
InputStream {
}
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 456486b9e8..ec1003b7d0 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.anyLong;
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;
@@ -50,8 +57,9 @@ import org.apache.celeborn.common.protocol.MessageType;
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";
@@ -144,6 +152,103 @@ public class CelebornInputStreamPeerFailoverTest {
}
}
+ @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 {
@@ -180,7 +285,13 @@ public class CelebornInputStreamPeerFailoverTest {
Optional.<CryptoHandler>empty());
}
- 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);
@@ -192,7 +303,7 @@ public class CelebornInputStreamPeerFailoverTest {
ArrayList<PbStreamHandler> streamHandlers = new ArrayList<>();
streamHandlers.add(PbStreamHandler.newBuilder().setStreamId(123L).setNumChunks(10).build());
- CelebornInputStream.create(
+ return CelebornInputStream.create(
conf,
clientFactory,
SHUFFLE_KEY,
@@ -210,7 +321,7 @@ public class CelebornInputStreamPeerFailoverTest {
1,
1,
0,
- null,
+ exceptionMaker,
new TestMetricsCallback(),
false,
Optional.<CryptoHandler>empty());
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 08c0d277b3..b79383277a 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.TimeUnit;
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;
@@ -47,6 +48,7 @@ import org.apache.celeborn.common.network.TransportContext;
import
org.apache.celeborn.common.network.sasl.registration.RegistrationClientBootstrap;
import org.apache.celeborn.common.network.server.TransportChannelHandler;
import org.apache.celeborn.common.network.util.*;
+import org.apache.celeborn.common.util.ExceptionUtils;
import org.apache.celeborn.common.util.JavaUtils;
import org.apache.celeborn.common.util.Utils;
@@ -160,9 +162,10 @@ public class TransportClientFactory implements Closeable {
try {
return createClient(remoteHost, remotePort, partitionId,
supplier.get());
} catch (Exception e) {
- if (e instanceof InterruptedException) {
+ InterruptedException interruptedException =
ExceptionUtils.findInterruptedException(e);
+ if (interruptedException != null) {
Thread.currentThread().interrupt();
- throw e;
+ throw interruptedException;
}
numTries++;
logger.warn(
@@ -315,14 +318,21 @@ public class TransportClientFactory implements Closeable {
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()) {
+ 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 (!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 +361,7 @@ public class TransportClientFactory implements Closeable {
}
}
});
- 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 +413,25 @@ public class TransportClientFactory implements Closeable {
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 +457,7 @@ public class TransportClientFactory implements Closeable {
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/main/java/org/apache/celeborn/common/util/ExceptionUtils.java
b/common/src/main/java/org/apache/celeborn/common/util/ExceptionUtils.java
index 8f00251771..39cf14d617 100644
--- a/common/src/main/java/org/apache/celeborn/common/util/ExceptionUtils.java
+++ b/common/src/main/java/org/apache/celeborn/common/util/ExceptionUtils.java
@@ -21,6 +21,8 @@ import java.io.IOException;
import java.io.PrintWriter;
import java.io.StringWriter;
+import com.google.common.base.Throwables;
+
import org.apache.celeborn.common.exception.CelebornIOException;
import org.apache.celeborn.common.exception.PartitionUnRetryAbleException;
@@ -52,6 +54,15 @@ public class ExceptionUtils {
}
}
+ public static InterruptedException findInterruptedException(Throwable
throwable) {
+ for (Throwable cause : Throwables.getCausalChain(throwable)) {
+ if (cause instanceof InterruptedException) {
+ return (InterruptedException) cause;
+ }
+ }
+ return null;
+ }
+
public static String stringifyException(Throwable exception) {
if (exception == null) {
return "(null)";
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 7ec7190aa0..3c51a175d3 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 class TransportClientFactorySuiteJ {
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 0000000000..e6735b27fa
--- /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();
+ }
+ }
+}
