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..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,6 +602,26 @@ 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: InterruptedException => + Thread.currentThread().interrupt() + throw ex + case ex: Exception => + onClientCreateFailure(ex) + None + } + } + @VisibleForTesting private[celeborn] def createClientsInParallel( locationsByHostPort: Seq[(String, Seq[PartitionLocation])], @@ -611,19 +633,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) } } } 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..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 @@ -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 @@ -164,39 +164,82 @@ 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(failureCount === 1) - assert(clients(failedLocation.hostAndFetchPort) eq client) + assert(clients.isEmpty) + assert(attempts.get() === 1) + assert(failures.get() === 1) } finally { streamCreatorPool.shutdownNow() } } + test("propagate client creation interruption without reporting a worker failure") { + val location = newLocation(0, "worker-0", 19098) + val failureReported = new AtomicBoolean(false) + + 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 { + 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 +248,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 +283,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 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 08c0d277b38..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 @@ -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,15 @@ public TransportClient retryCreateClient( return null; } + private static InterruptedException findInterruptedException(Throwable throwable) { + for (Throwable cause : Throwables.getCausalChain(throwable)) { + if (cause instanceof InterruptedException) { + return (InterruptedException) cause; + } + } + return null; + } + public TransportClient createClient( String remoteHost, int remotePort, int partitionId, ChannelInboundHandlerAdapter decoder) throws IOException, InterruptedException { @@ -315,14 +326,21 @@ 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()) { + 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 +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(); + } + } +}