You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@celeborn.apache.org by za...@apache.org on 2026/10/13 06:54:07 UTC
(celeborn) branch main updated: [CELEBORN-2458] Invalidate transport client after push timeout
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 626a07cfc0 [CELEBORN-2458] Invalidate transport client after push timeout
626a07cfc0 is described below
commit 626a07cfc055927f6436cba1d97bd4436e779c22
Author: Venkata Krishnan Sowrirajan <vs...@microsoft.com>
AuthorDate: Tue Oct 13 14:47:54 2026 +0800
[CELEBORN-2458] Invalidate transport client after push timeout
## Problem
A push request timeout leaves the pooled `TransportClient` eligible for reuse while its Netty channel still reports active. After a Worker restart or Pod IP replacement, long-running applications can therefore repeatedly send replica pushes over a stale connection and receive `PUSH_DATA_TIMEOUT_REPLICA`.
JIRA: [CELEBORN-2458](https://issues.apache.org/jira/browse/CELEBORN-2458)
## Fix
- Invalidate and asynchronously close the transport client when a push request expires.
- Ensure `TransportClientFactory` creates a fresh connection for the next request.
- Temporarily add a timed-out replica peer to `unavailablePeers`, reusing the existing replica fast-fail duration.
- Preserve higher-level retry/revive semantics rather than replaying the payload in the transport layer.
- Keep connection invalidation private to the transport implementation.
## Verification
- Added a pooled-client test with multiple outstanding pushes, exactly-once callbacks, connection replacement, and a successful push over the new channel.
- Added Spark integration assertions for normal and merged replica timeouts populating `unavailablePeers`.
- Added Worker coverage for fast-fail expiration and subsequent peer recovery.
- `TransportResponseHandlerSuiteJ` and `TransportClientFactorySuiteJ`: passed.
- `WorkerSuite`: 7 passed.
- `PushDataTimeoutTest` with Spark 3.5: 5 passed.
- Spark 3.5 `common`, `worker`, and `spark-it` package build: passed.
Closes #3842 from venkata91/vsowrirajan/invalidate-push-timeout-client.
Lead-authored-by: Venkata Krishnan Sowrirajan <vs...@microsoft.com>
Co-authored-by: Venkata krishnan Sowrirajan <vs...@linkedin.com>
Signed-off-by: zhengtao <sh...@alibaba-inc.com>
AI-Contributed/Feature: 0/38
AI-Contributed/UT: 0/147
(cherry picked from commit c32ccce1f8a2640b8b12002b2ff749a1973df623)
Signed-off-by: zhengtao <sh...@alibaba-inc.com>
---
.../common/network/client/TransportClient.java | 9 ++-
.../network/client/TransportResponseHandler.java | 1 +
.../celeborn/common/write/PushRequestInfo.java | 12 +++
.../network/TransportClientFactorySuiteJ.java | 86 ++++++++++++++++++++++
.../network/TransportResponseHandlerSuiteJ.java | 32 ++++++++
.../celeborn/tests/spark/PushDataTimeoutTest.scala | 6 +-
.../service/deploy/worker/PushDataHandler.scala | 2 +
.../celeborn/service/deploy/worker/Worker.scala | 14 ++--
.../service/deploy/worker/WorkerSuite.scala | 23 ++++++
9 files changed, 177 insertions(+), 8 deletions(-)
diff --git a/common/src/main/java/org/apache/celeborn/common/network/client/TransportClient.java b/common/src/main/java/org/apache/celeborn/common/network/client/TransportClient.java
index b0a61d054f..7aed2de605 100644
--- a/common/src/main/java/org/apache/celeborn/common/network/client/TransportClient.java
+++ b/common/src/main/java/org/apache/celeborn/common/network/client/TransportClient.java
@@ -220,7 +220,7 @@ public class TransportClient implements Closeable {
long requestId = requestId();
long dueTime = System.currentTimeMillis() + pushDataTimeout;
- PushRequestInfo info = new PushRequestInfo(dueTime, callback);
+ PushRequestInfo info = new PushRequestInfo(dueTime, callback, this::invalidate);
handler.addPushRequest(requestId, info);
pushData.requestId = requestId;
PushChannelListener listener = new PushChannelListener(requestId, rpcSendoutCallback);
@@ -237,7 +237,7 @@ public class TransportClient implements Closeable {
long requestId = requestId();
long dueTime = System.currentTimeMillis() + pushDataTimeout;
- PushRequestInfo info = new PushRequestInfo(dueTime, callback);
+ PushRequestInfo info = new PushRequestInfo(dueTime, callback, this::invalidate);
handler.addPushRequest(requestId, info);
pushMergedData.requestId = requestId;
@@ -350,6 +350,11 @@ public class TransportClient implements Closeable {
this.timedOut = true;
}
+ private void invalidate() {
+ this.timedOut = true;
+ channel.close();
+ }
+
@VisibleForTesting
public TransportResponseHandler getHandler() {
return handler;
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 4ad8373c02..5074a5cc2e 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
@@ -139,6 +139,7 @@ public class TransportResponseHandler extends MessageHandler<ResponseMessage> {
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.
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 691b433512..5d543aa05a 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
@@ -25,13 +25,25 @@ public class PushRequestInfo {
public ChannelFuture channelFuture;
public long dueTime;
public RpcResponseCallback callback;
+ private final Runnable timeoutHandler;
public PushRequestInfo(long dueTime, RpcResponseCallback callback) {
+ this(dueTime, callback, null);
+ }
+
+ public PushRequestInfo(long dueTime, RpcResponseCallback callback, Runnable timeoutHandler) {
this.dueTime = dueTime;
this.callback = callback;
+ this.timeoutHandler = timeoutHandler;
}
public void setChannelFuture(ChannelFuture future) {
this.channelFuture = future;
}
+
+ public void handleTimeout() {
+ if (timeoutHandler != null) {
+ timeoutHandler.run();
+ }
+ }
}
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 3c51a175d3..989cddc164 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
@@ -25,7 +25,10 @@ import static org.mockito.ArgumentMatchers.anyString;
import static org.mockito.Mockito.mock;
import java.io.IOException;
+import java.nio.ByteBuffer;
import java.util.*;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicInteger;
import org.junit.After;
@@ -35,12 +38,19 @@ import org.junit.Test;
import org.mockito.Mockito;
import org.apache.celeborn.common.CelebornConf;
+import org.apache.celeborn.common.network.buffer.NioManagedBuffer;
+import org.apache.celeborn.common.network.client.RpcResponseCallback;
import org.apache.celeborn.common.network.client.TransportClient;
import org.apache.celeborn.common.network.client.TransportClientFactory;
+import org.apache.celeborn.common.network.protocol.PushData;
+import org.apache.celeborn.common.network.protocol.RequestMessage;
+import org.apache.celeborn.common.network.protocol.RpcResponse;
import org.apache.celeborn.common.network.server.BaseMessageHandler;
import org.apache.celeborn.common.network.server.TransportServer;
import org.apache.celeborn.common.network.util.TransportConf;
import org.apache.celeborn.common.network.util.TransportFrameDecoder;
+import org.apache.celeborn.common.protocol.TransportModuleConstants;
+import org.apache.celeborn.common.protocol.message.StatusCode;
import org.apache.celeborn.common.util.JavaUtils;
import org.apache.celeborn.common.util.ThreadUtils;
@@ -184,6 +194,82 @@ public class TransportClientFactorySuiteJ {
factory.close();
}
+ @Test
+ public void replaceClientAfterPushTimeout() throws Exception {
+ AtomicBoolean respond = new AtomicBoolean(false);
+ BaseMessageHandler handler =
+ new BaseMessageHandler() {
+ @Override
+ public void receive(TransportClient client, RequestMessage message) {
+ PushData pushData = (PushData) message;
+ if (respond.get()) {
+ client
+ .getChannel()
+ .writeAndFlush(
+ new RpcResponse(
+ pushData.requestId, new NioManagedBuffer(ByteBuffer.allocate(0))));
+ }
+ }
+
+ @Override
+ public boolean checkRegistered() {
+ return true;
+ }
+ };
+ TransportConf conf =
+ new TransportConf(TransportModuleConstants.REPLICATE_MODULE, new CelebornConf());
+
+ try (TransportContext timeoutContext = new TransportContext(conf, handler);
+ TransportServer timeoutServer = timeoutContext.createServer();
+ TransportClientFactory factory = timeoutContext.createClientFactory()) {
+ TransportClient timedOutClient =
+ factory.createClient(getLocalHost(), timeoutServer.getPort());
+ RpcResponseCallback timedOutCallback = mock(RpcResponseCallback.class);
+ RpcResponseCallback outstandingCallback = mock(RpcResponseCallback.class);
+ timedOutClient.pushData(
+ new PushData(
+ (byte) 0,
+ "shuffleKey",
+ "partitionId-1",
+ new NioManagedBuffer(ByteBuffer.allocate(1))),
+ 0,
+ timedOutCallback);
+ timedOutClient.pushData(
+ new PushData(
+ (byte) 0,
+ "shuffleKey",
+ "partitionId-2",
+ new NioManagedBuffer(ByteBuffer.allocate(1))),
+ TimeUnit.MINUTES.toMillis(1),
+ outstandingCallback);
+
+ timedOutClient.getHandler().failExpiredPushRequest();
+
+ Mockito.verify(timedOutCallback, Mockito.timeout(5000).times(1))
+ .onFailure(
+ Mockito.argThat(
+ error ->
+ error.getMessage().startsWith(StatusCode.PUSH_DATA_TIMEOUT_REPLICA.name())));
+ Mockito.verify(outstandingCallback, Mockito.timeout(5000).times(1)).onFailure(Mockito.any());
+
+ respond.set(true);
+ TransportClient replacementClient =
+ factory.createClient(getLocalHost(), timeoutServer.getPort());
+ assertNotSame(timedOutClient, replacementClient);
+ RpcResponseCallback successCallback = mock(RpcResponseCallback.class);
+ replacementClient.pushData(
+ new PushData(
+ (byte) 0,
+ "shuffleKey",
+ "partitionId-3",
+ new NioManagedBuffer(ByteBuffer.allocate(1))),
+ TimeUnit.SECONDS.toMillis(5),
+ successCallback);
+ Mockito.verify(successCallback, Mockito.timeout(5000)).onSuccess(Mockito.any());
+ Mockito.verify(successCallback, Mockito.never()).onFailure(Mockito.any());
+ }
+ }
+
@Test
public void closeBlockClientsWithFactory() throws IOException, InterruptedException {
TransportClientFactory factory = context.createClientFactory();
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 3dd32334fb..8b6adece82 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
@@ -25,6 +25,8 @@ import java.nio.ByteBuffer;
import io.netty.channel.ChannelFuture;
import io.netty.channel.local.LocalChannel;
import org.junit.Test;
+import org.mockito.InOrder;
+import org.mockito.Mockito;
import org.apache.celeborn.common.CelebornConf;
import org.apache.celeborn.common.network.buffer.NioManagedBuffer;
@@ -33,6 +35,7 @@ import org.apache.celeborn.common.network.client.RpcResponseCallback;
import org.apache.celeborn.common.network.client.TransportResponseHandler;
import org.apache.celeborn.common.network.protocol.*;
import org.apache.celeborn.common.protocol.TransportModuleConstants;
+import org.apache.celeborn.common.protocol.message.StatusCode;
import org.apache.celeborn.common.read.FetchRequestInfo;
import org.apache.celeborn.common.util.Utils;
import org.apache.celeborn.common.write.PushRequestInfo;
@@ -175,4 +178,33 @@ public class TransportResponseHandlerSuiteJ {
verify(callback, times(1)).onFailure(any());
assertEquals(0, handler.numOutstandingRequests());
}
+
+ @Test
+ public void invokeTimeoutHandlerBeforePushFailureCallback() {
+ TransportResponseHandler handler =
+ new TransportResponseHandler(
+ Utils.fromCelebornConf(
+ new CelebornConf(), TransportModuleConstants.REPLICATE_MODULE, 8),
+ new LocalChannel());
+ Runnable timeoutHandler = mock(Runnable.class);
+ RpcResponseCallback callback = mock(RpcResponseCallback.class);
+ PushRequestInfo info =
+ new PushRequestInfo(System.currentTimeMillis() - 1, callback, timeoutHandler);
+ ChannelFuture pushFuture = mock(ChannelFuture.class);
+ info.setChannelFuture(pushFuture);
+ handler.addPushRequest(12345, info);
+
+ handler.failExpiredPushRequest();
+
+ verify(pushFuture).cancel(true);
+ InOrder timeoutOrder = inOrder(timeoutHandler, callback);
+ timeoutOrder.verify(timeoutHandler).run();
+ timeoutOrder
+ .verify(callback)
+ .onFailure(
+ Mockito.<Throwable>argThat(
+ error ->
+ error.getMessage().startsWith(StatusCode.PUSH_DATA_TIMEOUT_REPLICA.name())));
+ assertEquals(0, handler.numOutstandingRequests());
+ }
}
diff --git a/tests/spark-it/src/test/scala/org/apache/celeborn/tests/spark/PushDataTimeoutTest.scala b/tests/spark-it/src/test/scala/org/apache/celeborn/tests/spark/PushDataTimeoutTest.scala
index ea398968cd..0163343073 100644
--- a/tests/spark-it/src/test/scala/org/apache/celeborn/tests/spark/PushDataTimeoutTest.scala
+++ b/tests/spark-it/src/test/scala/org/apache/celeborn/tests/spark/PushDataTimeoutTest.scala
@@ -39,7 +39,8 @@ class PushDataTimeoutTest extends AnyFunSuite
logInfo("test initialized, setup celeborn mini cluster")
val workerConf = Map(
CelebornConf.TEST_CLIENT_PUSH_PRIMARY_DATA_TIMEOUT.key -> "true",
- CelebornConf.TEST_WORKER_PUSH_REPLICA_DATA_TIMEOUT.key -> "true")
+ CelebornConf.TEST_WORKER_PUSH_REPLICA_DATA_TIMEOUT.key -> "true",
+ CelebornConf.WORKER_REPLICATE_FAST_FAIL_DURATION.key -> "10m")
// required at least 4 workers, the reason behind this requirement is that when replication is
// enabled, there is a possibility that two workers might be added to the excluded list due to
// primary/replica timeout issues, then there are not enough workers to do replication if
@@ -53,6 +54,7 @@ class PushDataTimeoutTest extends AnyFunSuite
PushDataHandler.pushReplicaDataTimeoutTested.set(false)
PushDataHandler.pushPrimaryMergeDataTimeoutTested.set(false)
PushDataHandler.pushReplicaMergeDataTimeoutTested.set(false)
+ workerInfos.keys.foreach(_.unavailablePeers.clear())
}
override def afterEach(): Unit = {
@@ -87,6 +89,7 @@ class PushDataTimeoutTest extends AnyFunSuite
assert(PushDataHandler.pushPrimaryDataTimeoutTested.get())
if (enabled) {
assert(PushDataHandler.pushReplicaDataTimeoutTested.get())
+ assert(workerInfos.keys.exists(!_.unavailablePeers.isEmpty))
}
}
}
@@ -116,6 +119,7 @@ class PushDataTimeoutTest extends AnyFunSuite
assert(PushDataHandler.pushPrimaryMergeDataTimeoutTested.get())
if (enabled) {
assert(PushDataHandler.pushReplicaMergeDataTimeoutTested.get())
+ assert(workerInfos.keys.exists(!_.unavailablePeers.isEmpty))
}
}
}
diff --git a/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/PushDataHandler.scala b/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/PushDataHandler.scala
index 7c64593d25..53f201fc45 100644
--- a/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/PushDataHandler.scala
+++ b/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/PushDataHandler.scala
@@ -352,6 +352,7 @@ class PushDataHandler(val workerSource: WorkerSource) extends BaseMessageHandler
workerSource.incCounter(WorkerSource.REPLICATE_DATA_WRITE_FAIL_COUNT)
callbackWithTimer.onFailure(e)
} else if (e.getMessage.startsWith(StatusCode.PUSH_DATA_TIMEOUT_REPLICA.name())) {
+ unavailablePeers.put(peerWorker, System.currentTimeMillis())
workerSource.incCounter(WorkerSource.REPLICATE_DATA_TIMEOUT_COUNT)
callbackWithTimer.onFailure(e)
} else if (ExceptionUtils.connectFail(e.getMessage)) {
@@ -735,6 +736,7 @@ class PushDataHandler(val workerSource: WorkerSource) extends BaseMessageHandler
workerSource.incCounter(WorkerSource.REPLICATE_DATA_WRITE_FAIL_COUNT)
pushMergedDataCallback.onFailure(e)
} else if (e.getMessage.startsWith(StatusCode.PUSH_DATA_TIMEOUT_REPLICA.name())) {
+ unavailablePeers.put(peerWorker, System.currentTimeMillis())
workerSource.incCounter(WorkerSource.REPLICATE_DATA_TIMEOUT_COUNT)
pushMergedDataCallback.onFailure(e)
} else if (ExceptionUtils.connectFail(e.getMessage)) {
diff --git a/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/Worker.scala b/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/Worker.scala
index 8d713f2700..f1b7d21f86 100644
--- a/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/Worker.scala
+++ b/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/Worker.scala
@@ -587,11 +587,7 @@ private[celeborn] class Worker(
checkFastFailTask = forwardMessageScheduler.scheduleWithFixedDelay(
new Runnable {
override def run(): Unit = Utils.tryLogNonFatalError {
- unavailablePeers.entrySet().forEach { entry: JMap.Entry[WorkerInfo, Long] =>
- if (System.currentTimeMillis() - entry.getValue > replicaFastFailDuration) {
- unavailablePeers.remove(entry.getKey)
- }
- }
+ removeExpiredUnavailablePeers(System.currentTimeMillis())
}
},
0,
@@ -869,6 +865,14 @@ private[celeborn] class Worker(
})
}
+ private[worker] def removeExpiredUnavailablePeers(currentTime: Long): Unit = {
+ unavailablePeers.entrySet().forEach { entry: JMap.Entry[WorkerInfo, Long] =>
+ if (currentTime - entry.getValue > replicaFastFailDuration) {
+ unavailablePeers.remove(entry.getKey)
+ }
+ }
+ }
+
private def removeAppResourceConsumption(applicationIds: Iterable[String]): Unit = {
applicationIds.foreach { applicationId => removeAppResourceConsumption(applicationId) }
}
diff --git a/worker/src/test/scala/org/apache/celeborn/service/deploy/worker/WorkerSuite.scala b/worker/src/test/scala/org/apache/celeborn/service/deploy/worker/WorkerSuite.scala
index 09c4c78b51..cf6df0d76a 100644
--- a/worker/src/test/scala/org/apache/celeborn/service/deploy/worker/WorkerSuite.scala
+++ b/worker/src/test/scala/org/apache/celeborn/service/deploy/worker/WorkerSuite.scala
@@ -38,6 +38,7 @@ import org.scalatest.funsuite.AnyFunSuite
import org.apache.celeborn.common.CelebornConf
import org.apache.celeborn.common.client.MasterClient
import org.apache.celeborn.common.identity.UserIdentifier
+import org.apache.celeborn.common.meta.WorkerInfo
import org.apache.celeborn.common.protocol._
import org.apache.celeborn.common.protocol.message.ControlMessages.{CommitFilesResponse, HeartbeatFromWorkerResponse}
import org.apache.celeborn.common.protocol.message.StatusCode
@@ -135,6 +136,28 @@ class WorkerSuite extends AnyFunSuite with BeforeAndAfterEach with MiniClusterFe
Assert.assertEquals(1, allWriters.size())
}
+ test("remove expired unavailable peers") {
+ conf.set(CelebornConf.WORKER_STORAGE_DIRS.key, "/tmp")
+ conf.set(CelebornConf.WORKER_REPLICATE_FAST_FAIL_DURATION.key, "100ms")
+ worker = new Worker(conf, workerArgs)
+ val expiredPeer =
+ new WorkerInfo("expired", 1, 2, 3, 4, -1, new util.HashMap(), null)
+ val activePeer =
+ new WorkerInfo("active", 5, 6, 7, 8, -1, new util.HashMap(), null)
+ val currentTime = System.currentTimeMillis()
+ worker.unavailablePeers.put(expiredPeer, currentTime - 101)
+ worker.unavailablePeers.put(activePeer, currentTime)
+
+ worker.removeExpiredUnavailablePeers(currentTime)
+
+ assert(!worker.unavailablePeers.containsKey(expiredPeer))
+ assert(worker.unavailablePeers.containsKey(activePeer))
+
+ worker.removeExpiredUnavailablePeers(currentTime + 101)
+
+ assert(worker.unavailablePeers.isEmpty)
+ }
+
test("handle top resource consumption") {
conf.set(CelebornConf.WORKER_STORAGE_DIRS.key, "/tmp")
conf.set(CelebornConf.METRICS_WORKER_APP_TOP_RESOURCE_CONSUMPTION_COUNT, 5)