This is an automated email from the ASF dual-hosted git repository.
zaynt4606 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 fca9da2368 [CELEBORN-2482] Report pushes failed by push-timeout
channel close as push timeout instead of connection exception
fca9da2368 is described below
commit fca9da23687c4d484f449abce9f87ed66eaf2ae1
Author: afterincomparableyum
<[email protected]>
AuthorDate: Tue Oct 20 11:32:35 2026 +0800
[CELEBORN-2482] Report pushes failed by push-timeout channel close as push
timeout instead of connection exception
### What changes were proposed in this pull request?
When a push request times out and its timeout handler closes the channel
(CELEBORN-2458), `TransportResponseHandler` now records the timeout status
(`PUSH_DATA_TIMEOUT_PRIMARY`/`REPLICA`) before the close. The other pushes on
that channel that fail because of the close are reported with this status
instead of as connection exceptions. That covers both `channelInactive` and
write failures on the closed channel.
### Why are the changes needed?
Since CELEBORN-2458, a push timeout closes the channel. The other pushes
still outstanding on it are then failed with `Connection from ... closed` (or
`Failed to send request ...`), which `ExceptionUtils.connectFail` maps to
`PUSH_DATA_CONNECTION_EXCEPTION_PRIMARY`/`REPLICA`. These collateral failures
overwrite the worker's `PUSH_DATA_TIMEOUT_*` entry in the excluded workers, so
a worker that stopped responding gets misreported as a connection failure. This
also makes `PushDataTimeout [...]
```
- celeborn spark integration test - pushdata timeout will add to
pushExcludedWorkers *** FAILED ***
PUSH_DATA_CONNECTION_EXCEPTION_PRIMARY did not equal
PUSH_DATA_TIMEOUT_PRIMARY, and PUSH_DATA_CONNECTION_EXCEPTION_PRIMARY did not
equal PUSH_DATA_TIMEOUT_REPLICA (PushDataTimeoutTest.scala:157)
```
In the [failed
run](https://github.com/apache/celeborn/actions/runs/38396707096/job/115293643837),
a push to a worker timed out (`PUSH_DATA_TIMEOUT_PRIMARY`). In the same
millisecond, `Still have 1 requests outstanding when connection from ... is
closed` failed a second push on that channel, and the worker was then reported
as `PUSH_DATA_CONNECTION_EXCEPTION_PRIMARY`.
### Does this PR resolve a correctness bug?
- [ ] Yes
### Does this PR introduce _any_ user-facing change?
- [ ] Yes
### How was this patch tested?
- New tests in `TransportResponseHandlerSuiteJ`. The new push-timeout test
fails without this change.
- A push timeout that closes the channel fails the other outstanding
pushes, and later write failures, with the timeout status.
- A channel close without a push timeout still fails pushes with the
connection exception.
- `TransportResponseHandlerSuiteJ` (10 tests) and
`TransportClientFactorySuiteJ` (11 tests): passed.
- `PushDataTimeoutTest` with Spark 3.5: 5 tests passed.
Closes #3876 from afterincomparableyum/push-timeout-close-status.
Authored-by: afterincomparableyum
<[email protected]>
Signed-off-by: zhengtao <[email protected]>
AI-Contributed/Feature: 0/42
AI-Contributed/UT: 0/100
---
.../network/client/TransportResponseHandler.java | 38 +++++---
.../celeborn/common/write/PushRequestInfo.java | 4 +
.../network/TransportResponseHandlerSuiteJ.java | 100 +++++++++++++++++++++
3 files changed, 132 insertions(+), 10 deletions(-)
diff --git
a/common/src/main/java/org/apache/celeborn/common/network/client/TransportResponseHandler.java
b/common/src/main/java/org/apache/celeborn/common/network/client/TransportResponseHandler.java
index 5074a5cc2e..92f953fb0a 100644
---
a/common/src/main/java/org/apache/celeborn/common/network/client/TransportResponseHandler.java
+++
b/common/src/main/java/org/apache/celeborn/common/network/client/TransportResponseHandler.java
@@ -59,6 +59,8 @@ public class TransportResponseHandler extends
MessageHandler<ResponseMessage> {
private final ConcurrentHashMap<Long, RpcResponseCallback> outstandingRpcs;
private final ConcurrentHashMap<Long, PushRequestInfo> outstandingPushes;
+ private volatile StatusCode pushTimeoutCloseStatus;
+
/** Records the time (in system nanoseconds) that the last fetch or RPC
request was sent. */
private final AtomicLong timeOfLastRequestNs;
@@ -129,6 +131,7 @@ public class TransportResponseHandler extends
MessageHandler<ResponseMessage> {
}
public void failExpiredPushRequest() {
+ StatusCode timeoutStatus = pushTimeoutStatus();
long currentTime = System.currentTimeMillis();
Iterator<Map.Entry<Long, PushRequestInfo>> iter =
outstandingPushes.entrySet().iterator();
while (iter.hasNext()) {
@@ -136,18 +139,15 @@ public class TransportResponseHandler extends
MessageHandler<ResponseMessage> {
if (entry.getValue().dueTime <= currentTime) {
PushRequestInfo info = removePushRequest(entry.getKey());
if (info != null) {
+ if (timeoutStatus != null && info.hasTimeoutHandler()) {
+ pushTimeoutCloseStatus = timeoutStatus;
+ }
if (info.channelFuture != null) {
info.channelFuture.cancel(true);
}
info.handleTimeout();
- String module = conf.getModuleName();
- // When module name equals to DATA_MODULE, mean shuffle client push
data, else means
- // do data replication.
- if (TransportModuleConstants.DATA_MODULE.equals(module)) {
- info.callback.onFailure(new
CelebornIOException(StatusCode.PUSH_DATA_TIMEOUT_PRIMARY));
- } else if (TransportModuleConstants.PUSH_MODULE.equals(module)
- || TransportModuleConstants.REPLICATE_MODULE.equals(module)) {
- info.callback.onFailure(new
CelebornIOException(StatusCode.PUSH_DATA_TIMEOUT_REPLICA));
+ if (timeoutStatus != null) {
+ info.callback.onFailure(new CelebornIOException(timeoutStatus));
}
info.channelFuture = null;
info.callback = null;
@@ -156,6 +156,24 @@ public class TransportResponseHandler extends
MessageHandler<ResponseMessage> {
}
}
+ private StatusCode pushTimeoutStatus() {
+ String module = conf.getModuleName();
+ // When module name equals to DATA_MODULE, mean shuffle client push data,
else means
+ // do data replication.
+ if (TransportModuleConstants.DATA_MODULE.equals(module)) {
+ return StatusCode.PUSH_DATA_TIMEOUT_PRIMARY;
+ } else if (TransportModuleConstants.PUSH_MODULE.equals(module)
+ || TransportModuleConstants.REPLICATE_MODULE.equals(module)) {
+ return StatusCode.PUSH_DATA_TIMEOUT_REPLICA;
+ }
+ return null;
+ }
+
+ private Throwable pushFailureCause(Throwable cause) {
+ StatusCode timeoutStatus = pushTimeoutCloseStatus;
+ return timeoutStatus == null ? cause : new
CelebornIOException(timeoutStatus, cause);
+ }
+
public void failExpiredFetchRequest() {
long currentTime = System.currentTimeMillis();
Iterator<Map.Entry<StreamChunkSlice, FetchRequestInfo>> iter =
@@ -256,7 +274,7 @@ public class TransportResponseHandler extends
MessageHandler<ResponseMessage> {
long requestId = pushRequestIter.next();
PushRequestInfo info = removePushRequest(requestId);
if (info != null) {
- info.callback.onFailure(cause);
+ info.callback.onFailure(pushFailureCause(cause));
}
} catch (Exception e) {
logger.warn("RpcResponseCallback.onFailure throws exception", e);
@@ -437,7 +455,7 @@ public class TransportResponseHandler extends
MessageHandler<ResponseMessage> {
if (info != null) {
RpcResponseCallback callback = info.callback;
if (callback != null) {
- callback.onFailure(new CelebornIOException(errorMsg, cause));
+ callback.onFailure(pushFailureCause(new CelebornIOException(errorMsg,
cause)));
} else {
logger.warn(
"PushRequestInfo {} callback is null when handle push request
failure", pushRequestId);
diff --git
a/common/src/main/java/org/apache/celeborn/common/write/PushRequestInfo.java
b/common/src/main/java/org/apache/celeborn/common/write/PushRequestInfo.java
index 5d543aa05a..d2384dce50 100644
--- a/common/src/main/java/org/apache/celeborn/common/write/PushRequestInfo.java
+++ b/common/src/main/java/org/apache/celeborn/common/write/PushRequestInfo.java
@@ -41,6 +41,10 @@ public class PushRequestInfo {
this.channelFuture = future;
}
+ public boolean hasTimeoutHandler() {
+ return timeoutHandler != null;
+ }
+
public void handleTimeout() {
if (timeoutHandler != null) {
timeoutHandler.run();
diff --git
a/common/src/test/java/org/apache/celeborn/common/network/TransportResponseHandlerSuiteJ.java
b/common/src/test/java/org/apache/celeborn/common/network/TransportResponseHandlerSuiteJ.java
index 8b6adece82..242016a0d3 100644
---
a/common/src/test/java/org/apache/celeborn/common/network/TransportResponseHandlerSuiteJ.java
+++
b/common/src/test/java/org/apache/celeborn/common/network/TransportResponseHandlerSuiteJ.java
@@ -18,12 +18,17 @@
package org.apache.celeborn.common.network;
import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
import static org.mockito.Mockito.*;
import java.nio.ByteBuffer;
+import java.nio.channels.ClosedChannelException;
import io.netty.channel.ChannelFuture;
+import io.netty.channel.ChannelPromise;
+import io.netty.channel.DefaultChannelPromise;
import io.netty.channel.local.LocalChannel;
+import io.netty.util.concurrent.ImmediateEventExecutor;
import org.junit.Test;
import org.mockito.InOrder;
import org.mockito.Mockito;
@@ -207,4 +212,99 @@ public class TransportResponseHandlerSuiteJ {
error.getMessage().startsWith(StatusCode.PUSH_DATA_TIMEOUT_REPLICA.name())));
assertEquals(0, handler.numOutstandingRequests());
}
+
+ @Test
+ public void failOtherPushesWithTimeoutStatusWhenPushTimeoutClosesChannel() {
+ TransportResponseHandler handler =
+ new TransportResponseHandler(
+ Utils.fromCelebornConf(new CelebornConf(),
TransportModuleConstants.DATA_MODULE, 8),
+ new LocalChannel());
+ RpcResponseCallback expiredCallback = mock(RpcResponseCallback.class);
+ RpcResponseCallback pendingCallback = mock(RpcResponseCallback.class);
+ RpcResponseCallback writeFailedCallback = mock(RpcResponseCallback.class);
+ handler.addPushRequest(
+ 1,
+ new PushRequestInfo(
+ System.currentTimeMillis() - 1, expiredCallback,
handler::channelInactive));
+ handler.addPushRequest(
+ 2,
+ new PushRequestInfo(
+ System.currentTimeMillis() + 30000L, pendingCallback,
mock(Runnable.class)));
+
+ handler.failExpiredPushRequest();
+
+
verify(expiredCallback).onFailure(failureStartingWith(StatusCode.PUSH_DATA_TIMEOUT_PRIMARY));
+
verify(pendingCallback).onFailure(failureStartingWith(StatusCode.PUSH_DATA_TIMEOUT_PRIMARY));
+ assertEquals(0, handler.numOutstandingRequests());
+
+ handler.addPushRequest(
+ 3,
+ new PushRequestInfo(
+ System.currentTimeMillis() + 30000L, writeFailedCallback,
mock(Runnable.class)));
+ handler.handlePushFailure(
+ 3, "Failed to send request PUSH 3 to <unknown remote>", new
ClosedChannelException());
+ verify(writeFailedCallback)
+ .onFailure(failureStartingWith(StatusCode.PUSH_DATA_TIMEOUT_PRIMARY));
+ }
+
+ @Test
+ public void
failOtherPushesWithTimeoutStatusWhenCancelledWriteClosesChannel() {
+ TransportResponseHandler handler =
+ new TransportResponseHandler(
+ Utils.fromCelebornConf(new CelebornConf(),
TransportModuleConstants.DATA_MODULE, 8),
+ new LocalChannel());
+ RpcResponseCallback expiredCallback = mock(RpcResponseCallback.class);
+ RpcResponseCallback pendingCallback = mock(RpcResponseCallback.class);
+ PushRequestInfo expiredInfo =
+ new PushRequestInfo(System.currentTimeMillis() - 1, expiredCallback,
mock(Runnable.class));
+ // A still-queued write: cancelling it notifies the listener inline, which
closes the channel
+ // like PushChannelListener does on write failure.
+ ChannelPromise pushFuture =
+ new DefaultChannelPromise(new LocalChannel(),
ImmediateEventExecutor.INSTANCE);
+ pushFuture.addListener(
+ future -> {
+ if (!future.isSuccess()) {
+ handler.channelInactive();
+ handler.handlePushFailure(1, "Failed to send request PUSH 1",
future.cause());
+ }
+ });
+ expiredInfo.setChannelFuture(pushFuture);
+ handler.addPushRequest(1, expiredInfo);
+ handler.addPushRequest(
+ 2,
+ new PushRequestInfo(
+ System.currentTimeMillis() + 30000L, pendingCallback,
mock(Runnable.class)));
+
+ handler.failExpiredPushRequest();
+
+ assertTrue(pushFuture.isCancelled());
+
verify(expiredCallback).onFailure(failureStartingWith(StatusCode.PUSH_DATA_TIMEOUT_PRIMARY));
+
verify(pendingCallback).onFailure(failureStartingWith(StatusCode.PUSH_DATA_TIMEOUT_PRIMARY));
+ assertEquals(0, handler.numOutstandingRequests());
+ }
+
+ @Test
+ public void
failPushesWithConnectionExceptionWhenChannelClosesWithoutPushTimeout() {
+ TransportResponseHandler handler =
+ new TransportResponseHandler(
+ Utils.fromCelebornConf(new CelebornConf(),
TransportModuleConstants.DATA_MODULE, 8),
+ new LocalChannel());
+ RpcResponseCallback callback = mock(RpcResponseCallback.class);
+ handler.addPushRequest(
+ 1,
+ new PushRequestInfo(System.currentTimeMillis() + 30000L, callback,
mock(Runnable.class)));
+
+ handler.channelInactive();
+
+ verify(callback)
+ .onFailure(
+ Mockito.<Throwable>argThat(
+ error ->
+ error.getMessage().startsWith("Connection from ")
+ && error.getMessage().endsWith(" closed")));
+ }
+
+ private static Throwable failureStartingWith(StatusCode statusCode) {
+ return Mockito.<Throwable>argThat(error ->
error.getMessage().startsWith(statusCode.name()));
+ }
}
