[CELEBORN-2065] Avoid blocking fetch threads while partition files sort

[sunchao]

Why are the changes needed?

Some reduce readers need only a particular range of map outputs rather than an entire partition. When that happens, the worker may first need to sort the partition file and build an index before it can determine which portion of the file to return. This can occur with skew handling and other reads that split a reduce partition into multiple map ranges.

Today, waiting for that sort ties up the wrong resource. FetchHandler calls PartitionFilesSorter.getSortedFileInfo synchronously on a fetch-server Netty event-loop thread. If the sorted file is not ready, that thread polls every 50 ms until sorting finishes or celeborn.worker.sortPartition.timeout expires, which defaults to 220 seconds. The thread is not doing useful work during that time, but it cannot handle another fetch request either.

For example, consider a worker with 32 fetch event-loop threads:

1. Thirty-two concurrent readers request map ranges from partition files that are still being sorted. Several readers might even be waiting for different ranges of the same partition.
2. Each request blocks one fetch thread, although requests for the same partition are all waiting for the same underlying sort.
3. A 33rd reader requests an ordinary, already-available partition that does not need sorting. It still cannot be served because there is no fetch thread left to process its request.

Consequently, a slow sort can delay unrelated shuffle reads and make a worker appear unable to serve data that is already available. Larger fetch thread pools only raise the number of concurrent requests needed to reproduce the problem; they do not address the blocking dependency.

CELEBORN-2065 tracks this issue. The same fetch-thread starvation was discussed in #3593 and #3652, but neither proposal was merged.

What changes were proposed in this PR?

This PR changes disk-backed reduce-stream opening from a synchronous wait into an asynchronous continuation whenever a sorted view of the partition is needed. Instead of occupying a fetch event-loop thread until a partition is sorted, the worker registers interest in the result and returns the thread to the fetch server without waiting for the sort to finish. The existing sorter continues doing the actual work in the background; once the sorted file is ready, the requested map range is resolved and the stream response is sent.

Before:
  fetch event-loop thread -> wait for sorting -> resolve map range -> reply

After:
  fetch event-loop thread -> register completion -> return
  background sorter       -> shared completion
  bounded resolver        -> resolve map range -> reply

The completion is shared by shuffle and partition file, so concurrent readers of different map ranges wait for one sort rather than each holding a fetch thread. Sorted-index resolution runs on a configurable executor (celeborn.worker.sortPartition.resolve.threads, default 4) instead of on fetch event-loop threads. Filesystem preparation, including local and DFS existence checks and removal of leftover output files, also runs on the sorter executor. In the example above, the 32 range readers register their continuations and release their fetch threads, allowing the already-available 33rd request to be handled while sorting continues.

The asynchronous path is integrated into both individual and batched reduce-stream opens. A batch still produces one response in its original request order, with ordinary sort and stream-open failures represented per requested stream; legacy requests retain their existing response format. The existing synchronous sorter API remains available to other callers, and no RPC protocol changes are introduced. The SortedFileWaiters gauge exposes pending asynchronous readers.

Because completion and retries can now occur concurrently, the sorter also coordinates the lifecycle of its output and pending readers more carefully. Successful results are published only after the sorted output has been finalized. When sorting fails, output handles are closed before a retry can take ownership, preventing overlapping writers and stale attempts from disrupting their replacements. The underlying sort completion is shared by shuffle and file, but each reader keeps its own timeout and cancellation. A reader joining a sort in progress therefore receives the full configured wait budget; one reader timing out or cancelling does not fail another reader or cancel the sort.

Shuffle cleanup also covers readers whose sorting has already finished but whose index resolution is queued or running. Index loading happens outside short lifecycle guards; cache publication and stream registration are checked against the request's original lifecycle. Cleanup rejects stale admissions before retiring sorter state, so an in-flight request cannot recreate cache entries or register a stream after its cleanup pass. Failures from an expired sort cannot affect a replacement sort with the same key.

How was this PR tested?

Applied repository-wide formatting:

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

Ran the affected worker Java suites and relevant Scala storage suite:

build/mvn --no-transfer-progress -pl worker -am \
  -Dtest=FetchHandlerSuiteJ,DiskPartitionFilesSorterSuiteJ,DiskReducePartitionDataWriterSuiteJ,MemoryReducePartitionDataWriterSuiteJ \
  -DwildcardSuites=org.apache.celeborn.service.deploy.worker.storage.PartitionMetaHandlerSuite \
  -DfailIfNoTests=false -Dsurefire.failIfNoSpecifiedTests=false test

The complete affected-suite run passed 42 Java tests and 3 Scala tests on Java 11 / Scala 2.12.18. The focused development run also passed 9 Java tests and the same 3 Scala tests before the final retry-ownership and shared-index error regressions were added. Repository-wide Spotless formatting completed successfully.

Regression coverage includes independent waiter deadlines and cancellation, shared sorting and index loading, nonblocking DFS preparation, cleanup during raw-file lookup and queued/running index resolution, stale-sort failure isolation, ordered batch responses, and the existing local/memory writer and fetch paths.

[Copilot]

Pull request overview

This PR refactors the worker fetch path so that reduce-stream opens that depend on partition-file sorting no longer synchronously wait on Netty fetch event-loop threads. Instead, stream opening continues asynchronously once the partition has been sorted and the map-range index has been resolved on a bounded executor, reducing fetch-thread starvation under concurrent range reads.

Changes:

• Introduces an async PartitionFilesSorter.getSortedFileInfoAsync(...) API with shared completion per (shuffleKey, fileId) and timeout coordination.
• Refactors FetchHandler reduce open-stream handling (single + batched) to register async continuations rather than blocking while sorting/index resolution is pending.
• Expands/updates unit tests to cover non-blocking behavior, shared waiters, timeout retries, shuffle-key isolation, and failure propagation.

Reviewed changes

Copilot reviewed 6 out of 6 changed files in this pull request and generated 1 comment.

Show a summary per file

File	Description
worker/src/main/scala/org/apache/celeborn/service/deploy/worker/FetchHandler.scala	Converts reduce open-stream handling (single and batch) to an async continuation model to avoid blocking fetch threads.
worker/src/main/java/org/apache/celeborn/service/deploy/worker/storage/PartitionFilesSorter.java	Adds async sorted-file completion tracking, bounded resolve executor, timeout scheduling, and shutdown/cleanup coordination.
worker/src/test/java/org/apache/celeborn/service/deploy/worker/FetchHandlerSuiteJ.java	Adds tests asserting open-stream calls do not block while sorted-file info is pending and preserves batch ordering.
worker/src/test/java/org/apache/celeborn/service/deploy/worker/storage/local/DiskPartitionFilesSorterSuiteJ.java	Adds concurrency and lifecycle tests for async sorting waiters, retries, and cleanup behavior.
worker/src/test/java/org/apache/celeborn/service/deploy/worker/storage/local/DiskReducePartitionDataWriterSuiteJ.java	Updates mocks to use the new async sorter API.
worker/src/test/java/org/apache/celeborn/service/deploy/worker/storage/memory/MemoryReducePartitionDataWriterSuiteJ.java	Updates mocks to use the new async sorter API.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[sunchao]

Fixed in 33ea2be. FetchHandler now validates that fileName, startIndex, endIndex, and readLocalShuffle contain the same number of entries before starting OPEN_STREAM_TIME. Malformed batches increment OPEN_STREAM_FAIL_COUNT and return an RPC failure without starting the timer. FetchHandlerSuiteJ.testBatchOpenStreamRejectsMismatchedRequestFieldLengths covers both missing and extra entries.

[SteNicholas]

Could we add config option fot number of sorted file resolve threads?

[sunchao]

Done

[SteNicholas]

Could this add log for MemoryFileInfo?

[sunchao]

Done

[SteNicholas]

Does this need to add metric or log for sorted file waiter count to verify?

[sunchao]

Added a new metric SortedFileWaiters

[SteNicholas]

@sunchao, thanks for contribution of this greate improvement. I have left some comments for this pull request. PTAL.

[sunchao]

Hi @SteNicholas , appreciate if you can take another look on this PR, thanks!

[SteNicholas]

sortCompletionFuture is removed from sortCompletionFutures as soon as sorting completes, while the future returned by thenApplyAsync(resolve) is not tracked. If shuffle cleanup runs after sorting but while resolve is queued or running, cleanup cannot fail this waiter; the continuation can still succeed, repopulate the index cache, and let FetchHandler register a stream after its one-time stream cleanup. I reproduced this by blocking resolve, invoking cleanup, and observing the waiter succeed after release. Please track or cancel the derived request future, or validate a per-shuffle generation before returning and registering the result.

[sunchao]

Fixed in 6f8522f. Pending requests are now tracked through index resolution, and cache publication and stream registration are guarded by the request's original shuffle lifecycle. Worker cleanup expires fetch admissions before retiring sorter state, so delayed file lookups cannot admit work after cleanup.

Added regression coverage for cleanup while index resolution is queued or paused after parsing, and while the initial file lookup is blocked. After releasing and draining the work, the requests fail without recreating cache entries or streams.

[SteNicholas]

This method is called by the asynchronous open-stream path, but the constructor synchronously performs local exists and delete calls and, for DFS storage, hadoopFs.exists and hadoopFs.delete. Remote metadata operations can therefore block the fetch event loop before work reaches the sorter executor, undermining the goal of this PR. Please move construction and file-system preparation onto a worker executor.

[sunchao]

Fixed in 6f8522f. FileSorter construction now only records metadata; local/DFS existence checks and deletion of leftover output files run in prepareFiles() on the sorter executor.

testDfsSorterPreparationDoesNotBlockReceiveCaller uses the real fetch handler and sorter, blocks DFS exists, and verifies that receive and subsequent work on the request executor finish before the metadata operation is released.

[SteNicholas]

All waiters reuse this future and therefore share one timeout that starts with the first waiter. A later request inherits only the remaining time, potentially almost zero, while a request arriving just after this future times out gets a fresh full timeout. The previous synchronous path applied the timeout per request. Please share the raw sort-completion future but apply a timeout to each returned waiter, or document and test the new global sort-attempt timeout semantics.

[sunchao]

Fixed in 6f8522f by restoring per-request timeouts. Only the underlying sort completion is shared; each waiter has its own deadline/cancellation signal. A late reader gets its full wait budget, and another reader timing out does not cancel its wait or the sort.

testAsyncSortedFileWaitersHaveIndependentTimeouts explicitly fires the first waiter's timer, verifies the later waiter and retry remain pending, and then lets both succeed from one physical sort. It also covers cancellation isolation and repeated abandoned readers without retained sort-future callbacks.

[Copilot]

Pull request overview

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

[sunchao]

gentle ping @SteNicholas . Please take another look, thanks!