[CELEBORN-2458] Invalidate transport client after push timeout

[venkata91]

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

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.

[Copilot]

🟡 Changes recommended

新增的超时处理 hook 需要做异常隔离（避免 timeoutHandler 抛错中断超时扫描/回调），且新增测试里过短的 Mockito.after(200) 可能导致用例不稳定。

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Pull request overview

该 PR 针对推送超时后仍可能复用旧连接（尤其是副本推送）的问题，通过在超时路径中使 TransportClient 失效并关闭底层 Netty channel，确保后续请求能建立新连接，同时将超时副本 peer 临时加入 unavailablePeers 以复用现有 fast-fail 语义。

Changes:

• 在 push 请求超时处理时触发超时回调（用于使 TransportClient 失效并关闭连接），并新增单测验证回调顺序与连接替换行为。
• Worker 侧将副本推送超时的 peer 放入 unavailablePeers，并抽取/补充过期清理逻辑及对应测试。
• Spark IT 用例增加对超时后 unavailablePeers 被填充的断言，并调整 fast-fail 配置以稳定验证。

File summaries

File	Description
worker/src/main/scala/org/apache/celeborn/service/deploy/worker/Worker.scala	抽取 removeExpiredUnavailablePeers 方法，用于统一清理过期的 fast-fail peer。
worker/src/test/scala/org/apache/celeborn/service/deploy/worker/WorkerSuite.scala	增加 removeExpiredUnavailablePeers 的过期清理测试覆盖。
worker/src/main/scala/org/apache/celeborn/service/deploy/worker/PushDataHandler.scala	在副本推送超时（PUSH_DATA_TIMEOUT_REPLICA）场景将 peer 写入 unavailablePeers。
tests/spark-it/src/test/scala/org/apache/celeborn/tests/spark/PushDataTimeoutTest.scala	增加 IT 断言：副本超时应导致某些 worker 的 unavailablePeers 非空，并清理测试状态。
common/src/main/java/org/apache/celeborn/common/write/PushRequestInfo.java	为 push 请求引入 timeoutHandler 并提供 handleTimeout() 入口。
common/src/main/java/org/apache/celeborn/common/network/client/TransportResponseHandler.java	push 超时扫描中调用 info.handleTimeout()，为连接失效等超时动作提供挂钩。
common/src/main/java/org/apache/celeborn/common/network/client/TransportClient.java	push/merged push 构造 PushRequestInfo 时注入超时 handler（失效并关闭 channel）。
common/src/test/java/org/apache/celeborn/common/network/TransportResponseHandlerSuiteJ.java	新增单测验证：先执行 timeoutHandler，再触发 push failure callback。
common/src/test/java/org/apache/celeborn/common/network/TransportClientFactorySuiteJ.java	新增 pooled client 场景单测：push 超时后 client 被替换、后续新连接推送成功。

Review details

Suppressed comments (1)

common/src/main/java/org/apache/celeborn/common/network/client/TransportResponseHandler.java:146

• 在 failExpiredPushRequest() 中新增的 info.handleTimeout() 如果抛出异常，会直接中断超时扫描逻辑，导致当前超时请求未回调 onFailure、并可能遗留 outstandingPushes（后续也不会被正确清理）。建议对 timeoutHandler 做保护性 try/catch，并记录日志，避免影响主超时处理流程。

          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)) {

• Files reviewed: 9/9 changed files
• Comments generated: 1
• Review effort level: Lite

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[Copilot]

Warning

Copilot couldn't run its full agentic review because it didn't start before the timeout. Make sure your repository has a runner available, or add a copilot-code-review.yml file specifying one with the runs-on attribute. See the docs for more details.

Pull request overview

Copilot reviewed 9 out of 9 changed files in this pull request and generated 2 comments.

[venkata91]

Addressed. The first request now uses a 1 ms timeout and the test waits 10 ms before calling failExpiredPushRequest(), ensuring dueTime is deterministically in the past even with a strict comparison.

[venkata91]

Addressed. invalidate() now delegates timedOut state management to the existing timeOut() method before closing the channel, avoiding duplicated state-setting logic.

[venkata91]

Gentle ping for reviews!

[venkata91]

Production evidence for this, from a ~500-worker 0.6.2 fleet.

We see sustained PUSH_DATA_TIMEOUT_REPLICA bursts that match this failure mode exactly. Over 7 days in one cluster:

ReplicateDataTimeoutCount              8,448,683   95.8%
ReplicateDataConnectionExceptionCount    310,022
ReplicateDataCreateConnectionFailCount    57,049
ReplicateDataFailCount                         0
ReplicateDataWriteFailCount                    0

ReplicateDataFailCount and ReplicateDataWriteFailCount are exactly zero, so no replica ever rejected a write. It is all requests expiring in failExpiredPushRequest.

A representative burst:

02:12  worker-447 restarts on node A, re-registers
02:33  worker-447 restarts again, on node B (new pod IP), re-registers
02:33+ replication timeout burst begins
02:49  PrimaryPushDataTime p99 pins at 305,285 ms (client push timeout is 300s),
       sustained ~15 min; baseline is 0.02-10 ms
03:46  same primary/peer pairs timing out repeatedly, same second, many partitions

Two things corroborate the stale-pooled-client diagnosis:

• Zero "Still have N requests outstanding when connection ... is closed" log lines during the burst. channelInactive never fired, so the channel did report active, exactly as described here.
• PbWorkerLost was 0 throughout. These are graceful restarts with stable worker identity, so the master re-registers rather than declaring the worker lost. Nothing upstream of the transport layer notices.

On scale: without invalidation the poisoned client stays eligible for reuse, so every subsequent push to that peer expires at 300s as well. One 20-minute window recorded 1,974,971 increments, 34% of a three-day total. The repeat rate looks unbounded. Adding unavailablePeers on the timeout branch bounds it to roughly one generation per peer per replicate.fastFail.duration.

Client-visible cost is real: PrimaryPushDataTime is the time a Spark executor's push is held, so a pinned p99 is shuffleWriteTime directly.

Would be glad to see this merged, and we would need a branch-0.6 backport to pick it up.

[github-actions]

This PR is stale because it has been open 20 days with no activity. Remove stale label or comment or this will be closed in 10 days.