[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]

The retry-bounding fix (and the headOption change in createClientsInParallel) only covers the parallel client-creation path. The sequential makeOpenStreamList still attempts creation per partition-location: when tryCreateClient returns None, onClientCreateFailure only excludes + warns and never populates workerRequestMap, so the next location sharing the same hostPort re-enters if (!workerRequestMap.containsKey(hostPort)) (line 217) and calls createClient against the same dead endpoint again — the N × maxRetries blow-up this PR sets out to remove.

So with celeborn.client.spark.batch.openStream.parallelClientCreation.enabled=false, the stated goal ("Bound … client creation retries") isn't achieved. This isn't a regression (it matches the old behavior), but since the parallel path is now bounded it's an inconsistency worth either closing (e.g. record the failed hostPort so subsequent same-host locations short-circuit) or calling out as a known limitation.

[SteNicholas]

isInterruption returns true when either (a) the current thread's interrupt flag is set, or (b) any exception in throwable's cause chain is an InterruptedException. Both are broader than "this failure is a task cancellation."

Concern: a genuine, retryable fetch failure whose cause chain happens to wrap an InterruptedException (e.g. a transient Netty blocking-await interruption that got wrapped, without the task actually being cancelled) is now treated as an interruption in getNextChunk / createReaderWithRetry — the code closes the reader and throws without attempting the peer/replica, whereas the pre-PR path would have failed over. It also flips the live thread's interrupt flag via interruptedIOException. In the common cancellation case that's exactly right; the risk is the false-positive where an InterruptedException in the chain doesn't represent a real cancellation, which would strand a recoverable read on the primary and turn it into a task failure. Worth confirming that scenario can't arise from a non-cancellation source.

[SteNicholas]

Consistency: this catches only a bare InterruptedException. A wrapped one (IOException(InterruptedException)) falls through to case ex: Exception below and is reported as a worker failure via excludeFailedFetchLocation — the exact misclassification that TransportClientFactory.findInterruptedException and CelebornInputStream.isInterruption (both added in this PR) walk the cause chain to prevent. It's masked today only because the real createClient path goes through retryCreateClient, which already unwraps to a bare InterruptedException — a fragile coupling.

Note the three interrupt-detection helpers this PR introduces already disagree: findInterruptedException ignores the thread flag, isInterruption checks it first, and tryCreateClient checks neither the chain nor the flag mid-createClient. Consider extracting one shared helper (e.g. in common ExceptionUtils) and using it in all three sites so they can't drift.

Related latent gap in createClientsInParallel below: a pool task that rethrows InterruptedException surfaces at futures.foreach(_.get()) as a java.util.concurrent.ExecutionException, which the catch { case ex: InterruptedException } there won't match — so the re-interrupt is skipped and a raw ExecutionException escapes read(). It's hard to hit today (the waiting thread's own .get() usually throws InterruptedException first, and pool-thread flags are cleared between tasks), but unwrapping ExecutionException(InterruptedException) there would close it.