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 2 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
33 changes: 27 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) {
Thread.currentThread().interrupt();
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Could we maintain only the thread's interrupt state while still invoking callback.onFailure? And I've noticed some inconsistencies in how InterruptedException is handled in the Celeborn client. In some cases it throws an exception, while in other cases it is simply ignored.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Are you suggesting that we can invoke callback.onFailure(InterruptedException e) to interrupt the current thread? That seems reasonable. I believe that ignoring InterruptedException is unsafe, especially when we catch it without any specific handling.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Contributor Author

yes, better keep that behavior If there is no particular reason.

} 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 @@ -876,6 +883,9 @@ Map<Integer, Integer> reviveBatch(

return results;
} catch (Exception e) {
if (e instanceof InterruptedException) {
Thread.currentThread().interrupt();
}
StringBuilder partitionIds = new StringBuilder();
StringBuilder epochs = new StringBuilder();
requests.forEach(
Expand Down Expand Up @@ -1232,8 +1242,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) {
Thread.currentThread().interrupt();
} else {
wrappedCallback.onFailure(
new CelebornIOException(StatusCode.PUSH_DATA_CREATE_CONNECTION_FAIL_PRIMARY, e));
}
}
} else {
// add batch data
Expand Down Expand Up @@ -1586,8 +1600,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) {
Thread.currentThread().interrupt();
} else {
wrappedCallback.onFailure(
new CelebornIOException(StatusCode.PUSH_DATA_CREATE_CONNECTION_FAIL_PRIMARY, e));
}
}
}

Expand Down Expand Up @@ -1700,6 +1718,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