[CELEBORN-2371] Bound Spark batch-open client creation retries and stop them on interruption

[sunchao]

Why are the changes needed?

CELEBORN-2371 follows up on the parallel Spark batch-open client creation added by #3692.

Batch-open locations are grouped by host:fetchPort, but each worker task previously walked multiple PartitionLocation entries. Since every createClient invocation already receives TransportClientFactory's full retry budget, that outer loop could multiply connection latency without targeting a different endpoint.

Cancellation also needs to terminate consistently. A direct or wrapped InterruptedException must stop reader setup and fetch retry flow before it is recorded as a worker failure, added to shared exclusion state, retried against a peer, or reported as a shuffle fetch failure. The unlimited-timeout transport branch also needed explicit cleanup on cancelled or failed connection futures.

What changes were proposed in this PR?

Keep one retry budget per worker endpoint

Each grouped host:fetchPort now uses one representative location for client creation. TransportClientFactory remains the owner of connection retries, controlled by the existing celeborn.data.io.maxRetries; there is no second outer retry setting or multiplicative retry budget.

Propagate cancellation through reader paths

CelebornShuffleReader now uses a shared interrupt-aware client-creation helper in both sequential and parallel batch-open paths. It checks a pre-existing interrupt flag, restores the flag when createClient throws InterruptedException, and exits without invoking the ordinary failure callback.

CelebornInputStream now detects interruption throughout reader creation, failed-stream cleanup, reconnect, and buffer-fill paths. It exits before exclusion, retry, peer failover, or shuffle-fetch-failure reporting. If reader cleanup itself fails while propagating cancellation, the cleanup failure is retained as a suppressed exception without masking the original interruption.

Preserve interruption and cleanup in transport bootstrap

TransportClientFactory uses Guava's Throwables.getCausalChain() to detect wrapped InterruptedException, restores the interrupt flag, and stops retrying immediately. TCP-connect and TLS-handshake waits close the in-progress channel before propagating interruption. The unlimited-timeout connection path now also closes cancelled and failed channel futures explicitly.

How was this PR tested?

Formatting was applied with:

./build/mvn --no-transfer-progress -DskipTests spotless:apply

The focused transport tests passed (11 tests):

./build/mvn --no-transfer-progress -pl common -am \
  -Dtest=TransportClientFactorySuiteJ,TransportClientFactoryInterruptSuiteJ \
  -DwildcardSuites=none clean test

The focused input-stream peer-failover and interruption tests passed (6 tests):

./build/mvn --no-transfer-progress -pl client -am \
  -Dtest=CelebornInputStreamPeerFailoverTest \
  -DwildcardSuites=none clean test

The Spark 3.5 reader suite passed (7 tests):

./build/mvn --no-transfer-progress -Pspark-3.5 -pl client-spark/spark-3 -am \
  -Dtest=none \
  -DwildcardSuites=org.apache.spark.shuffle.celeborn.CelebornShuffleReaderSuite \
  clean test

Spark 4.0 / Scala 2.13 production and test compilation also passed:

./build/mvn --no-transfer-progress -Pspark-4.0 -pl client-spark/spark-3 -am \
  -DskipTests clean test

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

[SteNicholas]

@sunchao, thanks — the direction is right and the tests are focused. The bound on the outer loop, propagating interruption through bootstrap, and closing the in-progress channel on interrupt are all good improvements. A few points worth addressing before merge (none are hard blockers, but the first two are worth a response):

1. The second per-host attempt is redundant. All locations grouped under a hostPort share the same host:fetchPort (grouping keys by hostAndFetchPort, and the createClient lambda uses only getHost/getFetchPort), so attempt #2 re-targets the same worker. See the inline note. MAX=1 would give the same result; if you keep 2, the "preserves same-worker fallback" comment is misleading.

2. Interrupt handling is now inconsistent across createClient callers. findInterruptedException makes retryCreateClient set the interrupt flag and throw a bare InterruptedException on a wrapped interrupt, but only createClientsInParallel has matching case ex: InterruptedException handling. The other callers — CelebornInputStream.createReaderWithRetry (catches Exception, then excludeFailedFetchLocation + Uninterruptibles.sleepUninterruptibly which re-asserts the flag), the non-parallel makeOpenStreamList, and the push/replicate paths — swallow it generically, so a cancellation gets recorded as a fetch/push failure (polluting shared cross-task exclusion state) and, because Netty's await() throws immediately when the flag is preset, cascades into fast-failing the remaining retries. Much of this cascade is pre-existing for direct interrupts; this PR widens it to wrapped ones. Worth either handling InterruptedException consistently at these call sites, or noting the trade-off.

3. Minor — worker-thread interrupt surfaces as ExecutionException. When a worker run() throws InterruptedException (the new line-620 check, or a wrapped interrupt on the worker thread rethrown at line 628), futures.foreach(_.get()) raises ExecutionException, which is not matched by the case ex: InterruptedException at ~642 — so it escapes without restoring the caller's interrupt status. Low impact, but the new code makes "InterruptedException out of run()" more reachable.

4. Nit — reuse. findInterruptedException duplicates the cause-chain walk in MasterClient.findMasterNotLeaderException; Guava Throwables.getCausalChain() (already imported here) or commons-lang3 ExceptionUtils.throwableOfType / a shared ExceptionUtils helper would avoid a third hand-rolled copy (and could carry a cycle guard).

[sunchao]

Follow-up on the two review-summary-only notes:

• The cause-chain lookup now uses Throwables.getCausalChain().
• I intentionally left worker-side ExecutionException behavior unchanged. Cancellation of the waiting caller makes Future.get() throw InterruptedException; that path restores the caller-thread interrupt flag and cancels remaining workers in finally. If only a worker is interrupted, the caller receives ExecutionException; setting the caller-thread interrupt flag in that case would transfer a thread-local signal across threads. The worker helper restores its own flag before rethrowing.

The full follow-up is in 311bc8c1c, and the PR description now reflects the final one-attempt design and focused validation.

[SteNicholas]

@sunchao, thanks for the interruption/cancellation propagation changes. The core approach is sound: each interrupt check is placed at the top of its catch, before the excludeFailedFetchLocation / reportShuffleFetchFailure side effects; retryCreateClient now unwraps a wrapped InterruptedException via the cause chain; awaitWithChannelCleanup closes the channel and restores the interrupt flag before rethrowing; and the added channel cleanup on the unlimited-timeout (connectTimeoutMs <= 0) path fixes a real leak. The headOption change in the parallel path is safe — every location in a hostAndFetchPort group targets the identical getHost/getFetchPort endpoint, so walking the rest was only re-retrying the same host. The touched suites pass locally (CelebornInputStreamPeerFailoverTest 6, TransportClientFactorySuiteJ 9, TransportClientFactoryInterruptSuiteJ 1).

A few non-blocking notes below: one incompleteness (the sequential path isn't bounded), one design consideration (breadth of the interrupt classification), and a consistency/consolidation point on the three new interrupt-detection helpers.

[SteNicholas]

@sunchao, could you please address above comments?

[sunchao]

@SteNicholas sorry lost track of this. Updated the PR and addressed the comments. Please take another look!

[Copilot]

Pull request overview

Copilot reviewed 8 out of 8 changed files in this pull request and generated no new comments.

[SteNicholas]

LGTM. Thanks for updates.