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 d1d69cb0ae [CELEBORN-2406] Avoid blocking GetReducerFileGroup RPC
under ConcurrentHashMap bin lock in updateFileGroup
d1d69cb0ae is described below
commit d1d69cb0aead510f88b145f5a042e0841a9f89ce
Author: Venkata Krishnan Sowrirajan <[email protected]>
AuthorDate: Mon Aug 10 15:07:42 2026 +0800
[CELEBORN-2406] Avoid blocking GetReducerFileGroup RPC under
ConcurrentHashMap bin lock in updateFileGroup
### What changes were proposed in this pull request?
`ShuffleClientImpl.updateFileGroup` loaded the reducer file group via
`reduceFileGroupsMap.compute(shuffleId, ...)`. `ConcurrentHashMap.compute`
holds the per-key bin lock for the entire remapping function, so the blocking
`GetReducerFileGroup` RPC to the `LifecycleManager` ran while holding that lock.
At high reduce parallelism the first task of a shuffle takes the bin lock
and issues the RPC; every other reduce task for the same `shuffleId` then
blocks on that bin lock, even cache hits. On a large shuffle the stage sits
idle for minutes with reduce threads BLOCKED in `updateFileGroup` and no fetch
activity, i.e. a load convoy.
This PR moves the RPC out from under the bin lock:
- Lock-free `get()` fast path, so cache hits never take a lock.
- Cold loads serialize on a dedicated per-shuffle monitor
(`fileGroupLoadLocks`, cleared in `cleanupShuffle`) with a double-check, so
only one RPC is issued per shuffle.
- Caching semantics unchanged: a cached tuple whose `_1()` is null still
reloads.
`FlinkShuffleClientImpl` overrides `updateFileGroup` and is unaffected; the
Spark plugin uses the base `ShuffleClientImpl` and was exposed.
### Why are the changes needed?
Observed in production (tens of thousands of reduce partitions, ~900
executors): a stage stalled over 20 minutes with reduce threads blocked in
`updateFileGroup` -> `ConcurrentHashMap.compute`. A slow RPC under the bin lock
turns a transient lookup into a multi-minute wedge.
### Does this PR introduce any user-facing change?
No.
### How was this patch tested?
New
`ShuffleClientSuiteJ#testUpdateReducerFileGroupConcurrentLoadIssuesSingleRpc`:
16 threads concurrently load the same shuffle against a mock RPC that blocks
500ms; asserts all callers finish and exactly one RPC is issued.
Closes #3776 from venkata91/celeborn-updatefilegroup-convoy-fix.
Lead-authored-by: Venkata Krishnan Sowrirajan <[email protected]>
Co-authored-by: Venkata krishnan Sowrirajan <[email protected]>
Signed-off-by: zhengtao <[email protected]>
AI-Contributed/Feature: 0/28
AI-Contributed/UT: 0/54
---
.../apache/celeborn/client/ShuffleClientImpl.java | 28 +++++++----
.../celeborn/client/ShuffleClientSuiteJ.java | 54 ++++++++++++++++++++++
2 files changed, 72 insertions(+), 10 deletions(-)
diff --git
a/client/src/main/java/org/apache/celeborn/client/ShuffleClientImpl.java
b/client/src/main/java/org/apache/celeborn/client/ShuffleClientImpl.java
index 210151241f..5e84f84ee4 100644
--- a/client/src/main/java/org/apache/celeborn/client/ShuffleClientImpl.java
+++ b/client/src/main/java/org/apache/celeborn/client/ShuffleClientImpl.java
@@ -188,6 +188,10 @@ public class ShuffleClientImpl extends ShuffleClient {
protected final Map<Integer, Tuple3<ReduceFileGroups, String, Exception>>
reduceFileGroupsMap =
JavaUtils.newConcurrentHashMap();
+ // Per-shuffle monitor serializing first-time file group loads, so the
blocking
+ // GetReducerFileGroup RPC runs outside reduceFileGroupsMap's bin lock (see
updateFileGroup).
+ private final Map<Integer, Object> fileGroupLoadLocks =
JavaUtils.newConcurrentHashMap();
+
private final TransportMessagesHelper messagesHelper = new
TransportMessagesHelper();
public ShuffleClientImpl(String appUniqueId, CelebornConf conf,
UserIdentifier userIdentifier) {
@@ -1865,6 +1869,7 @@ public class ShuffleClientImpl extends ShuffleClient {
// clear status
reducePartitionMap.remove(shuffleId);
reduceFileGroupsMap.remove(shuffleId);
+ fileGroupLoadLocks.remove(shuffleId);
mapperEndMap.remove(shuffleId);
stageEndShuffleSet.remove(shuffleId);
splitting.remove(shuffleId);
@@ -1975,16 +1980,19 @@ public class ShuffleClientImpl extends ShuffleClient {
public ReduceFileGroups updateFileGroup(
int shuffleId, int partitionId, boolean isSegmentGranularityVisible)
throws CelebornIOException {
- Tuple3<ReduceFileGroups, String, Exception> fileGroupTuple =
- reduceFileGroupsMap.compute(
- shuffleId,
- (id, existsTuple) -> {
- if (existsTuple == null || existsTuple._1() == null) {
- return loadFileGroupInternal(shuffleId,
isSegmentGranularityVisible);
- } else {
- return existsTuple;
- }
- });
+ // Cache hits take no lock. compute() would hold the bin lock across the
blocking RPC and
+ // convoy every reduce task of the shuffle, so cold loads serialize on a
per-shuffle monitor
+ // (double-checked) instead: one RPC in flight, and the RPC runs off the
map's bin lock.
+ Tuple3<ReduceFileGroups, String, Exception> fileGroupTuple =
reduceFileGroupsMap.get(shuffleId);
+ if (fileGroupTuple == null || fileGroupTuple._1() == null) {
+ synchronized (fileGroupLoadLocks.computeIfAbsent(shuffleId, id -> new
Object())) {
+ fileGroupTuple = reduceFileGroupsMap.get(shuffleId);
+ if (fileGroupTuple == null || fileGroupTuple._1() == null) {
+ fileGroupTuple = loadFileGroupInternal(shuffleId,
isSegmentGranularityVisible);
+ reduceFileGroupsMap.put(shuffleId, fileGroupTuple);
+ }
+ }
+ }
if (fileGroupTuple._1() == null) {
throw new CelebornIOException(
loadFileGroupException(shuffleId, partitionId,
(fileGroupTuple._2())),
diff --git
a/client/src/test/java/org/apache/celeborn/client/ShuffleClientSuiteJ.java
b/client/src/test/java/org/apache/celeborn/client/ShuffleClientSuiteJ.java
index 74fe6379c1..46d9a4f5b7 100644
--- a/client/src/test/java/org/apache/celeborn/client/ShuffleClientSuiteJ.java
+++ b/client/src/test/java/org/apache/celeborn/client/ShuffleClientSuiteJ.java
@@ -29,8 +29,10 @@ import java.util.Collections;
import java.util.HashMap;
import java.util.Map;
import java.util.Set;
+import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;
+import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.atomic.AtomicReference;
import scala.reflect.ClassTag;
@@ -631,6 +633,58 @@ public class ShuffleClientSuiteJ {
Assert.assertTrue(exception.getCause() instanceof TimeoutException);
}
+ @Test
+ public void testUpdateReducerFileGroupConcurrentLoadIssuesSingleRpc()
+ throws InterruptedException {
+ // Concurrent first-time loads must dedup to a single GetReducerFileGroup
RPC (the RPC once
+ // ran under reduceFileGroupsMap's bin lock, convoying every reduce task
of the shuffle).
+ CelebornConf conf = new CelebornConf();
+ AtomicInteger rpcCount = new AtomicInteger(0);
+ when(endpointRef.askSync(any(), any(), any(Integer.class),
any(Long.class), any()))
+ .thenAnswer(
+ t -> {
+ rpcCount.incrementAndGet();
+ Thread.sleep(500);
+ return GetReducerFileGroupResponse$.MODULE$.apply(
+ StatusCode.SUCCESS,
+ new HashMap<>(),
+ new int[0],
+ Collections.emptySet(),
+ Collections.emptyMap(),
+ new byte[0],
+ SerdeVersion.V1);
+ });
+
+ shuffleClient =
+ new ShuffleClientImpl(TEST_APPLICATION_ID, conf, new
UserIdentifier("mock", "mock"));
+ shuffleClient.setupLifecycleManagerRef(endpointRef);
+
+ int threads = 16;
+ CountDownLatch done = new CountDownLatch(threads);
+ AtomicReference<Exception> failure = new AtomicReference<>();
+ Thread[] workers = new Thread[threads];
+ for (int i = 0; i < threads; i++) {
+ Thread t =
+ new Thread(
+ () -> {
+ try {
+ shuffleClient.updateFileGroup(0, 0);
+ } catch (Exception e) {
+ failure.compareAndSet(null, e);
+ } finally {
+ done.countDown();
+ }
+ },
+ "test-updateFileGroup-" + i);
+ t.setDaemon(true);
+ workers[i] = t;
+ t.start();
+ }
+ Assert.assertTrue("all callers should finish", done.await(30,
TimeUnit.SECONDS));
+ Assert.assertNull(failure.get());
+ Assert.assertEquals("first-time load must issue exactly one RPC", 1,
rpcCount.get());
+ }
+
@Test
public void testSuccessfulReadReducePartitionEnd() throws IOException {
CelebornConf conf = new CelebornConf();
