(celeborn) branch main updated: [CELEBORN-2434] Clamp endMapIndex in sliceSortedBufferByMapRange

[za...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

zaynt4606 pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/celeborn.git


The following commit(s) were added to refs/heads/main by this push:
     new c090945514 [CELEBORN-2434] Clamp endMapIndex in sliceSortedBufferByMapRange
c090945514 is described below

commit c090945514d39f888ac27a706e4807df7dd815b8
Author: yew1eb <ye...@gmail.com>
AuthorDate: Sat Oct 10 17:56:56 2026 +0800

    [CELEBORN-2434] Clamp endMapIndex in sliceSortedBufferByMapRange
    
    ### What changes were proposed in this pull request?
    
    Add the missing `endMapIndex == Integer.MAX_VALUE` clamp in `ShuffleBlockInfoUtils.sliceSortedBufferByMapRange`, mirroring the existing guard in the sibling method `getChunkOffsetsFromShuffleBlockInfos`. Also add a unit test covering the full re-read (`Integer.MAX_VALUE`) case.
    
    ### Why are the changes needed?
    
    When a memory-stored shuffle partition is read once with a map range (e.g. AQE skew join) and then re-read in full (`endIndex = Int.MaxValue`, e.g. Spark ReusedExchange), `FetchHandler.handleReduceOpenStreamInternal` passes the raw `Int.MaxValue` into `getSortedFileInfo` because `FileInfo.addStream` returns false for the already-sorted file. The memory branch then loops `startMapIndex..Integer.MAX_VALUE` in `sliceSortedBufferByMapRange` — billions of `TreeMap.get` calls on the Netty f [...]
    
    The disk path is not affected (it goes through the already-clamped `getChunkOffsetsFromShuffleBlockInfos`).
    
    Symptoms / evidence:
    
    • MemoryStorageReusedExchangeSuite (a single tiny test with 1000 rows) takes 28–54 minutes instead of ~1 minute, both in CI and locally.
    • Client side repeatedly logs BatchOpenStream for N cost 240006ms (exactly the 240s fetch timeout), followed by retries that stall again.
    • Client logs Connection ... has been quiet for 4.0 m while there are outstanding requests against both workers' fetch ports.
    • A thread dump taken while hung shows the fetch event loop spinning:
    
    ```
      "fetch-server-34-2" RUNNABLE
        at java.util.TreeMap.getEntry(TreeMap.java:359)
        at java.util.TreeMap.get(TreeMap.java:278)
        at ShuffleBlockInfoUtils.sliceSortedBufferByMapRange(ShuffleBlockInfoUtils.java:140)
        at PartitionFilesSorter.getSortedFileInfo(PartitionFilesSorter.java:224)
        at FetchHandler.handleReduceOpenStreamInternal(FetchHandler.scala:277)
    ```
    
    In production this means: memory storage + a range read followed by a full re-read of the same partition can wedge a worker's fetch event loop for an
    unbounded time (the loop count scales with Int.MaxValue, multiplied by the number of files in the batch request).
    
    ### Does this PR resolve a correctness bug?
    
    - [ ] Yes
    
    ### Does this PR introduce _any_ user-facing change?
    
    - [ ] Yes
    
    ### How was this patch tested?
    
    - Added `testSliceSortedBufferByMapRangeWithMaxEndIndex` to `ShuffleBlockInfoUtilsTest`, asserting a full re-read with `Integer.MAX_VALUE` produces exactly the same result as an explicit full-range read.
    - `ShuffleBlockInfoUtilsTest`: all 6 tests pass.
    
    Closes #3815 from yew1eb/CELEBORN-2434.
    
    Authored-by: yew1eb <ye...@gmail.com>
    Signed-off-by: zhengtao <sh...@alibaba-inc.com>
    
    AI-Contributed/Feature: 0/7
    AI-Contributed/UT: 0/60
---
 .../common/util/ShuffleBlockInfoUtils.java         |  7 ++-
 .../common/util/ShuffleBlockInfoUtilsTest.java     | 60 ++++++++++++++++++++++
 2 files changed, 66 insertions(+), 1 deletion(-)

diff --git a/common/src/main/java/org/apache/celeborn/common/util/ShuffleBlockInfoUtils.java b/common/src/main/java/org/apache/celeborn/common/util/ShuffleBlockInfoUtils.java
index d62bf95d39..9c784d3e76 100644
--- a/common/src/main/java/org/apache/celeborn/common/util/ShuffleBlockInfoUtils.java
+++ b/common/src/main/java/org/apache/celeborn/common/util/ShuffleBlockInfoUtils.java
@@ -136,7 +136,12 @@ public class ShuffleBlockInfoUtils {
     int offset = 0;
     int length = 0;
     boolean blockBoundary = true;
-    for (int i = startMapIndex; i < endMapIndex; i++) {
+    int maxMapIndex = endMapIndex;
+    if (endMapIndex == Integer.MAX_VALUE) {
+      // not a range read
+      maxMapIndex = indexMap.keySet().stream().max(Integer::compareTo).get() + 1;
+    }
+    for (int i = startMapIndex; i < maxMapIndex; i++) {
       List<ShuffleBlockInfo> blockInfos = indexMap.get(i);
       if (blockInfos != null) {
         for (ShuffleBlockInfo blockInfo : blockInfos) {
diff --git a/common/src/test/java/org/apache/celeborn/common/util/ShuffleBlockInfoUtilsTest.java b/common/src/test/java/org/apache/celeborn/common/util/ShuffleBlockInfoUtilsTest.java
index 72d7e6af54..dea55d3a6c 100644
--- a/common/src/test/java/org/apache/celeborn/common/util/ShuffleBlockInfoUtilsTest.java
+++ b/common/src/test/java/org/apache/celeborn/common/util/ShuffleBlockInfoUtilsTest.java
@@ -143,4 +143,64 @@ public class ShuffleBlockInfoUtilsTest {
         "Target buffer should remain empty with an empty indexMap",
         targetByteBuf.numComponents() == 0);
   }
+
+  @Test
+  public void testSliceSortedBufferByMapRangeWithMaxEndIndex() {
+    Map<Integer, List<ShuffleBlockInfo>> indexMap = new HashMap<>();
+
+    indexMap.put(0, Arrays.asList(new ShuffleBlockInfo(0, 50), new ShuffleBlockInfo(50, 30)));
+    indexMap.put(1, Arrays.asList(new ShuffleBlockInfo(0, 20), new ShuffleBlockInfo(20, 30)));
+
+    for (ShuffleBlockInfo blockInfo :
+        indexMap.values().stream().flatMap(List::stream).toArray(ShuffleBlockInfo[]::new)) {
+      byte[] data = new byte[(int) blockInfo.length];
+      Arrays.fill(data, (byte) blockInfo.offset); // Fill data with offset value for simplicity
+      sortedByteBuf.addComponents(Unpooled.wrappedBuffer(data));
+    }
+
+    // A full re-read passes endMapIndex == Integer.MAX_VALUE. It must be clamped to
+    // maxMapIndex + 1 instead of looping up to Integer.MAX_VALUE, and must produce
+    // exactly the same result as an explicit full-range read.
+    CompositeByteBuf targetByteBufWithMaxEndIndex = Unpooled.compositeBuffer();
+    ShuffleBlockInfoUtils.sliceSortedBufferByMapRange(
+        0,
+        Integer.MAX_VALUE,
+        indexMap,
+        sortedByteBuf,
+        targetByteBufWithMaxEndIndex,
+        shuffleChunkSize);
+
+    ShuffleBlockInfoUtils.sliceSortedBufferByMapRange(
+        0, 2, indexMap, sortedByteBuf, targetByteBuf, shuffleChunkSize);
+
+    Assert.assertEquals(
+        "Unexpected number of components in target buffer",
+        targetByteBuf.numComponents(),
+        targetByteBufWithMaxEndIndex.numComponents());
+    // addComponent does not advance the target's writerIndex, so readableBytes() on the
+    // composite buffers is always 0. Compare the per-component bytes instead.
+    long expectedBytes = 50 + 30 + 20 + 30;
+    Assert.assertEquals(
+        "Unexpected sliced bytes for full re-read",
+        expectedBytes,
+        totalComponentBytes(targetByteBufWithMaxEndIndex));
+    Assert.assertEquals(
+        "Full re-read must produce the same bytes as an explicit full-range read",
+        totalComponentBytes(targetByteBuf),
+        totalComponentBytes(targetByteBufWithMaxEndIndex));
+    for (int i = 0; i < targetByteBuf.numComponents(); i++) {
+      Assert.assertEquals(
+          "Component " + i + " content mismatch",
+          targetByteBuf.component(i),
+          targetByteBufWithMaxEndIndex.component(i));
+    }
+  }
+
+  private static long totalComponentBytes(CompositeByteBuf buf) {
+    long total = 0;
+    for (int i = 0; i < buf.numComponents(); i++) {
+      total += buf.component(i).readableBytes();
+    }
+    return total;
+  }
 }