[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 small, bounded executor instead of on fetch event-loop threads, so index I/O does not block those threads and several sorted-file resolutions can run concurrently. 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 or new configuration settings are introduced.

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. Pending sort-completion futures are scoped to the exact shuffle and file and are coordinated with sorting timeouts, shuffle expiration, and worker shutdown.

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 affected-suite run passed 34 Java tests and 3 Scala tests. After the final concurrency refinements, the focused timeout, retry, and fetch regression subset was rerun and passed 10 Java tests and the same 3 Scala tests.

The regression coverage exercises nonblocking individual and batched stream opens, ordered batch responses, asynchronous sort failures, shared waiters, shuffle-key isolation, targeted cleanup, timeout retries, independent index resolution, failed-sort retry ownership, and requests arriving after sorter shutdown.

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