sunchao opened a new pull request, #3746: URL: https://github.com/apache/celeborn/pull/3746
## Why are the changes needed? [CELEBORN-2371](https://issues.apache.org/jira/browse/CELEBORN-2371) follows up on the parallel Spark batch-open client creation added by [#3692](https://github.com/apache/celeborn/pull/3692). The reader now creates clients for different workers in parallel, but each worker task still walks every `PartitionLocation` until one `createClient` call succeeds. A `createClient` call already has `TransportClientFactory`'s complete retry budget, so the outer location loop can multiply that budget. For example, suppose one reducer has 100 partition locations on `worker-a` and that worker is unavailable. The task for `worker-a` can make 100 outer `createClient` calls, and every call can run the configured transport retries. Because the reader waits for every worker task, the unavailable worker can keep the whole batch-open setup pending long after healthy workers have finished. Cancellation also needs to terminate client creation consistently. If cancellation interrupts a failure callback, that callback can restore the interrupt flag and return; without checking the flag, the worker task starts the next location attempt. At the transport layer, synchronous bootstrap can wrap an `InterruptedException` in another exception, so the retry loop currently treats it as retryable. An interrupt while waiting for TCP connect or TLS setup also leaves cleanup of the in-progress channel implicit. ## What changes were proposed in this PR? ### Bound the per-worker outer retry loop Each worker task now tries at most two partition locations. The first attempt keeps the normal path, and the second preserves same-worker fallback. Each attempt still receives the existing internal `TransportClientFactory` retry budget; this change only prevents that complete budget from being repeated once per partition location. ### Stop the Spark worker task after cancellation Before starting another same-worker attempt, `CelebornShuffleReader` checks the thread's interrupt flag and exits with `InterruptedException`. This complements the existing future cancellation path and prevents a callback that observed cancellation from falling through to another client creation. ### Preserve interruption through transport bootstrap `TransportClientFactory` now searches an exception's cause chain for `InterruptedException`, restores the thread's interrupt flag, and stops retrying immediately. The TCP-connect and TLS-handshake waits also share an interrupt-aware helper that closes the in-progress channel before propagating the interruption. The tests cover the two-attempt bound, cancellation during the failure callback, a wrapped `InterruptedException`, and channel cleanup during an interrupted wait. ## How was this PR tested? Formatting was applied with: ```text ./build/mvn --no-transfer-progress -DskipTests spotless:apply ``` The focused transport tests passed (11 tests): ```text ./build/mvn --no-transfer-progress -pl common -am \ -Dtest=TransportClientFactorySuiteJ,TransportClientFactoryInterruptSuiteJ \ -DwildcardSuites=none clean test ``` The Spark 3.5 reader suite passed (7 tests): ```text ./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: ```text ./build/mvn --no-transfer-progress -Pspark-4.0 -pl client-spark/spark-3 -am \ -DskipTests clean test ``` -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
