Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

[CELEBORN-1670] Avoid swallowing InterruptedException in ShuffleClientImpl #2849

Open
wants to merge 4 commits into
base: main
Choose a base branch
from
+80 −10
Open

[CELEBORN-1670] Avoid swallowing InterruptedException in ShuffleClientImpl #2849

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
3a47623
[CELEBORN-1670] Avoid swallowing InterruptedException in ShuffleClien…
jiang13021 Oct 24, 2024
7f1c080
trigger CI
jiang13021 Oct 25, 2024
1b78a3c
fix
jiang13021 Nov 5, 2024
089bd28
trigger CI
jiang13021 Oct 25, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
43 changes: 37 additions & 6 deletions client/src/main/java/org/apache/celeborn/client/ShuffleClientImpl.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -354,8 +354,12 @@ private void submitRetryPushData(
batchId,
newLoc,
e);
pushDataRpcResponseCallback.onFailure(
new CelebornIOException(StatusCode.PUSH_DATA_CREATE_CONNECTION_FAIL_PRIMARY, e));
if (e instanceof InterruptedException) {
pushDataRpcResponseCallback.onFailure(e);
} else {
pushDataRpcResponseCallback.onFailure(
new CelebornIOException(StatusCode.PUSH_DATA_CREATE_CONNECTION_FAIL_PRIMARY, e));
}
}
}
}
Expand Down Expand Up @@ -691,6 +695,9 @@ private ConcurrentHashMap<Integer, PartitionLocation> registerShuffleInternal(
numRetries - 1);
}
} catch (Exception e) {
if (e instanceof InterruptedException) {
Thread.currentThread().interrupt();
}
logger.error(
"Exception raised while registering shuffle {} with {} mapper and {} partitions.",
shuffleId,
Expand Down Expand Up @@ -889,6 +896,9 @@ Map<Integer, Integer> reviveBatch(
partitionIds,
epochs,
e);
if (e instanceof InterruptedException) {
Thread.currentThread().interrupt();
}
return null;
}
}
Expand Down Expand Up @@ -1151,6 +1161,11 @@ public void onFailure(Throwable e) {
if (pushState.exception.get() != null) {
return;
}
if (e instanceof InterruptedException) {
Thread.currentThread().interrupt();
callback.onFailure(e);
return;
}
StatusCode cause = getPushDataFailCause(e.getMessage());
if (remainReviveTimes <= 0) {
if (e instanceof CelebornIOException) {
Expand Down Expand Up @@ -1232,8 +1247,12 @@ public void onFailure(Throwable e) {
nextBatchId,
loc,
e);
wrappedCallback.onFailure(
new CelebornIOException(StatusCode.PUSH_DATA_CREATE_CONNECTION_FAIL_PRIMARY, e));
if (e instanceof InterruptedException) {
wrappedCallback.onFailure(e);
} else {
wrappedCallback.onFailure(
new CelebornIOException(StatusCode.PUSH_DATA_CREATE_CONNECTION_FAIL_PRIMARY, e));
}
}
} else {
// add batch data
Expand Down Expand Up @@ -1506,6 +1525,11 @@ public void onFailure(Throwable e) {
if (pushState.exception.get() != null) {
return;
}
if (e instanceof InterruptedException) {
Thread.currentThread().interrupt();
callback.onFailure(e);
return;
}
StatusCode cause = getPushDataFailCause(e.getMessage());
if (remainReviveTimes <= 0) {
if (e instanceof CelebornIOException) {
Expand Down Expand Up @@ -1586,8 +1610,12 @@ public void onFailure(Throwable e) {
Arrays.toString(batchIds),
addressPair,
e);
wrappedCallback.onFailure(
new CelebornIOException(StatusCode.PUSH_DATA_CREATE_CONNECTION_FAIL_PRIMARY, e));
if (e instanceof InterruptedException) {
wrappedCallback.onFailure(e);
} else {
wrappedCallback.onFailure(
new CelebornIOException(StatusCode.PUSH_DATA_CREATE_CONNECTION_FAIL_PRIMARY, e));
}
}
}

Expand Down Expand Up @@ -1700,6 +1728,9 @@ protected Tuple2<ReduceFileGroups, String> loadFileGroupInternal(
}
}
} catch (Exception e) {
if (e instanceof InterruptedException) {
Thread.currentThread().interrupt();
}
logger.error("Exception raised while call GetReducerFileGroup for {}.", shuffleId, e);
exceptionMsg = e.getMessage();
}
Expand Down
47 changes: 43 additions & 4 deletions client/src/test/java/org/apache/celeborn/client/ShuffleClientSuiteJ.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,11 +17,10 @@

package org.apache.celeborn.client;

import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertThrows;
import static org.junit.Assert.assertTrue;
import static org.junit.Assert.*;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.anyLong;
import static org.mockito.BDDMockito.willAnswer;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.when;

Expand Down Expand Up @@ -120,6 +119,31 @@ public void testPushData() throws IOException, InterruptedException {
}
}

@Test
public void testPushDataAndInterrupted() throws IOException, InterruptedException {
CelebornConf conf = setupEnv(CompressionCodec.NONE, StatusCode.SUCCESS, true);
try {
shuffleClient.pushData(
TEST_SHUFFLE_ID,
TEST_ATTEMPT_ID,
TEST_ATTEMPT_ID,
TEST_REDUCRE_ID,
TEST_BUF1,
0,
TEST_BUF1.length,
1,
1);
Thread.sleep(10 * 1000); // waiting for interrupt
fail();
} catch (Exception e) {
if (e instanceof InterruptedException) {
assertTrue(true);
} else {
fail();
}
}
}

@Test
public void testMergeData() throws IOException, InterruptedException {
for (CompressionCodec codec : CompressionCodec.values()) {
Expand Down Expand Up @@ -205,6 +229,12 @@ private CelebornConf setupEnv(CompressionCodec codec) throws IOException, Interr

private CelebornConf setupEnv(CompressionCodec codec, StatusCode statusCode)
throws IOException, InterruptedException {
return setupEnv(codec, statusCode, false);
}

private CelebornConf setupEnv(
CompressionCodec codec, StatusCode statusCode, boolean interruptWhenPushData)
throws IOException, InterruptedException {
CelebornConf conf = new CelebornConf();
conf.set(CelebornConf.SHUFFLE_COMPRESSION_CODEC().key(), codec.name());
conf.set(CelebornConf.CLIENT_PUSH_RETRY_THREADS().key(), "1");
Expand Down Expand Up @@ -346,7 +376,16 @@ public Void get(long timeout, TimeUnit unit) {
}
};

when(client.pushData(any(), anyLong(), any())).thenAnswer(t -> mockedFuture);
if (interruptWhenPushData) {
willAnswer(
invocation -> {
throw new InterruptedException("test");
})
.given(client)
.pushData(any(), anyLong(), any());
} else {
when(client.pushData(any(), anyLong(), any())).thenAnswer(t -> mockedFuture);
}
when(clientFactory.createClient(
primaryLocation.getHost(), primaryLocation.getPushPort(), TEST_REDUCRE_ID))
.thenAnswer(t -> client);
Expand Down
Loading