This is an automated email from the ASF dual-hosted git repository.
SteNicholas pushed a commit to branch branch-0.7
in repository https://gitbox.apache.org/repos/asf/celeborn.git
The following commit(s) were added to refs/heads/branch-0.7 by this push:
new d63fbb1ce7 [CELEBORN-2392] Prevent false shuffle expiry during
concurrent registration
d63fbb1ce7 is described below
commit d63fbb1ce7106f2b895fe6d4d3a0842a463b7921
Author: Wayne Wang <[email protected]>
AuthorDate: Thu Oct 29 15:30:15 2026 +0800
[CELEBORN-2392] Prevent false shuffle expiry during concurrent registration
### What changes were proposed in this pull request?
Use `ConcurrentHashMap.newKeySet()` for the per-application sets in
`registeredAppAndShuffles` at all three creation sites:
- Normal registration in `updateRequestSlotsMeta()`.
- Lost-shuffle revision in `reviseLostShuffles()`.
- Snapshot restoration in `restoreMetaFromFile()`.
Add `RegisteredShuffleSetSuiteJ` with four deterministic tests covering the
concurrent-set invariant at each entry point, duplicate IDs, application
isolation, and snapshot round-trip preservation with subsequent updates. The
fixture releases its metrics source, rack resolver, and temporary files.
### Why are the changes needed?
The outer map is concurrent, but its values are plain `HashSet<Integer>`
instances. Metadata updates add shuffle IDs while worker-heartbeat processing
calls `contains()` on the same set. A concurrent outer map does not make these
nested operations safe.
During a `HashMap` resize, a heartbeat can observe a partially populated
new table and report an already registered shuffle as absent. The master then
returns that shuffle as expired, and the worker cleans up its committed data.
Subsequent fetches fail even though the shuffle was never unregistered.
This fixes concurrent access within the per-application set. It does not
make registration and heartbeat handling a single transaction or change the
expiry policy, RPCs, configuration, or snapshot schema.
### Does this PR resolve a correctness bug?
- [x] Yes
### Does this PR introduce _any_ user-facing change?
- [x] Yes
Active shuffle data is preserved when another shuffle is registered
concurrently.
### How was this patch tested?
Local validation on Java 17.0.14:
- **Regression red/green:** two pre-fix runs each had three expected
concurrent-set assertion failures and one passing snapshot control. The
identical four tests all pass after the fix.
- **Related suites:** 70 tests passed in total, including the four new
tests: `RegisteredShuffleSetSuiteJ` (4), `DefaultMetaSystemSuiteJ` (22),
`RatisMasterStatusSystemSuiteJ` (29), `MasterStateMachineSuiteJ` (6), and Scala
`MasterSuite` (9).
- **Formatting:** `spotless:check` for master and its dependencies, and
`git diff --check`, passed.
Additional local verification used external harnesses; these harnesses are
not part of the proposed patch or CI:
- Reverting each set factory independently makes only its corresponding
regression test fail.
- Ten normal suite runs plus injected test-body and post-allocation setup
failures released all 48 temporary directories and 96 executors, with no
fixture threads left running.
- A real Master/Worker mini-cluster exercised push, commit, heartbeat,
cleanup, and subsequent fetch for all three creation paths. A test-only Java
agent held the target `HashMap` at its actual resize window, immediately after
table publication and before migration. Metadata growth used the shared update
methods, and the native worker heartbeat was invoked explicitly to control
ordering.
- Before the fix, all three paths returned an active shuffle as expired,
physically deleted its files, and produced a subsequent
`FileNotFoundException`. The ID became visible again after registration
completed. After the fix, all three paths preserved the metadata and files and
returned the exact 55,296-byte payload. An unrelated application shuffle
remained readable in every scenario. The same harness and agent were used
before and after the fix.
Closes #3901 from wayne-x-wang/celeborn-2392.
Authored-by: Wayne Wang <[email protected]>
Signed-off-by: Nicholas Jiang <[email protected]>
(cherry picked from commit a7c07ece06edce67652fc9c147c4b7f164420a77)
Signed-off-by: Nicholas Jiang <[email protected]>
---
.../master/clustermeta/AbstractMetaManager.java | 8 +-
.../clustermeta/RegisteredShuffleSetSuiteJ.java | 160 +++++++++++++++++++++
2 files changed, 165 insertions(+), 3 deletions(-)
diff --git
a/master/src/main/java/org/apache/celeborn/service/deploy/master/clustermeta/AbstractMetaManager.java
b/master/src/main/java/org/apache/celeborn/service/deploy/master/clustermeta/AbstractMetaManager.java
index f726c28919..c0a47d0675 100644
---
a/master/src/main/java/org/apache/celeborn/service/deploy/master/clustermeta/AbstractMetaManager.java
+++
b/master/src/main/java/org/apache/celeborn/service/deploy/master/clustermeta/AbstractMetaManager.java
@@ -116,7 +116,7 @@ public abstract class AbstractMetaManager implements
IMetadataHandler {
String shuffleKey, String hostName, Map<String, Map<String, Integer>>
workerWithAllocations) {
Tuple2<String, Object> appIdShuffleId = Utils.splitShuffleKey(shuffleKey);
registeredAppAndShuffles
- .computeIfAbsent(appIdShuffleId._1(), v -> new HashSet<>())
+ .computeIfAbsent(appIdShuffleId._1(), v ->
ConcurrentHashMap.newKeySet())
.add((Integer) appIdShuffleId._2);
String appId = appIdShuffleId._1;
@@ -216,7 +216,9 @@ public abstract class AbstractMetaManager implements
IMetadataHandler {
}
public void reviseLostShuffles(String appId, List<Integer> lostShuffles) {
- registeredAppAndShuffles.computeIfAbsent(appId, v -> new
HashSet<>()).addAll(lostShuffles);
+ registeredAppAndShuffles
+ .computeIfAbsent(appId, v -> ConcurrentHashMap.newKeySet())
+ .addAll(lostShuffles);
}
public void deleteApp(String appId) {
@@ -437,7 +439,7 @@ public abstract class AbstractMetaManager implements
IMetadataHandler {
for (String shuffleKey : snapshotMetaInfo.getRegisteredShuffleList()) {
Tuple2<String, Object> appIdShuffleId =
Utils.splitShuffleKey(shuffleKey);
registeredAppAndShuffles
- .computeIfAbsent(appIdShuffleId._1, v -> new HashSet<>())
+ .computeIfAbsent(appIdShuffleId._1, v ->
ConcurrentHashMap.newKeySet())
.add((Integer) appIdShuffleId._2);
}
hostnameSet.addAll(snapshotMetaInfo.getHostnameSetList());
diff --git
a/master/src/test/java/org/apache/celeborn/service/deploy/master/clustermeta/RegisteredShuffleSetSuiteJ.java
b/master/src/test/java/org/apache/celeborn/service/deploy/master/clustermeta/RegisteredShuffleSetSuiteJ.java
new file mode 100644
index 0000000000..e20c69beca
--- /dev/null
+++
b/master/src/test/java/org/apache/celeborn/service/deploy/master/clustermeta/RegisteredShuffleSetSuiteJ.java
@@ -0,0 +1,160 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.celeborn.service.deploy.master.clustermeta;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertSame;
+import static org.junit.Assert.assertTrue;
+import static org.mockito.Mockito.mock;
+
+import java.io.File;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.Map;
+import java.util.Set;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.TimeUnit;
+
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.rules.TemporaryFolder;
+
+import org.apache.celeborn.common.CelebornConf;
+import org.apache.celeborn.common.network.CelebornRackResolver;
+import org.apache.celeborn.common.rpc.RpcEnv;
+import org.apache.celeborn.common.util.Utils;
+import org.apache.celeborn.service.deploy.master.ApplicationMetricsSource;
+
+public class RegisteredShuffleSetSuiteJ {
+ @Rule public final TemporaryFolder temporaryFolder = new TemporaryFolder();
+
+ private CelebornRackResolver rackResolver;
+ private ApplicationMetricsSource metricsSource;
+ private AbstractMetaManager metaManager;
+
+ @Before
+ public void setUp() {
+ CelebornConf conf = new CelebornConf();
+ rackResolver = new CelebornRackResolver(conf);
+ metricsSource = new ApplicationMetricsSource(conf);
+ metaManager =
+ new SingleMasterMetaManager(mock(RpcEnv.class), conf, rackResolver,
metricsSource);
+ }
+
+ @After
+ public void tearDown() throws InterruptedException {
+ try {
+ if (metricsSource != null) {
+ metricsSource.destroy();
+ assertTrue(metricsSource.metricsCleaner().awaitTermination(5,
TimeUnit.SECONDS));
+ }
+ } finally {
+ if (rackResolver != null) {
+ rackResolver.stop();
+ }
+ }
+ }
+
+ @Test
+ public void testRequestSlotsCreatesConcurrentShuffleSet() {
+ registerShuffle("app-a", 1);
+ Set<Integer> shuffles = metaManager.registeredAppAndShuffles.get("app-a");
+ registerShuffle("app-a", 2);
+ registerShuffle("app-a", 1);
+ registerShuffle("app-b", 3);
+
+ assertSame(shuffles, metaManager.registeredAppAndShuffles.get("app-a"));
+ assertEquals(new HashSet<>(Arrays.asList(1, 2)), shuffles);
+ assertEquals(Collections.singleton(3),
metaManager.registeredAppAndShuffles.get("app-b"));
+ assertConcurrentShuffleSet("app-a");
+ assertConcurrentShuffleSet("app-b");
+ }
+
+ @Test
+ public void testReviseLostShufflesCreatesConcurrentShuffleSet() {
+ metaManager.reviseLostShuffles("app-a", Arrays.asList(1, 2, 2));
+ Set<Integer> shuffles = metaManager.registeredAppAndShuffles.get("app-a");
+ metaManager.reviseLostShuffles("app-a", Arrays.asList(2, 3));
+ metaManager.reviseLostShuffles("app-b", Collections.singletonList(4));
+
+ assertSame(shuffles, metaManager.registeredAppAndShuffles.get("app-a"));
+ assertEquals(new HashSet<>(Arrays.asList(1, 2, 3)), shuffles);
+ assertEquals(Collections.singleton(4),
metaManager.registeredAppAndShuffles.get("app-b"));
+ assertConcurrentShuffleSet("app-a");
+ assertConcurrentShuffleSet("app-b");
+ }
+
+ @Test
+ public void testSnapshotRestoresConcurrentShuffleSets() throws Exception {
+ registerShuffle("app-a", 1);
+ metaManager.reviseLostShuffles("app-b", Arrays.asList(2, 3));
+ File snapshot = temporaryFolder.newFile("metadata.snapshot");
+ metaManager.writeMetaInfoToFile(snapshot);
+ metaManager.restoreMetaFromFile(snapshot);
+
+ assertEquals(Collections.singleton(1),
metaManager.registeredAppAndShuffles.get("app-a"));
+ assertEquals(
+ new HashSet<>(Arrays.asList(2, 3)),
metaManager.registeredAppAndShuffles.get("app-b"));
+ assertConcurrentShuffleSet("app-a");
+ assertConcurrentShuffleSet("app-b");
+ }
+
+ @Test
+ public void
testSnapshotPreservesRegisteredShufflesAndAcceptsFurtherUpdates() throws
Exception {
+ registerShuffle("app-a", 1);
+ registerShuffle("app-a", 2);
+ metaManager.reviseLostShuffles("app-b", Arrays.asList(3, 4));
+ Map<String, Set<Integer>> expected = new HashMap<>();
+ expected.put("app-a", new HashSet<>(Arrays.asList(1, 2)));
+ expected.put("app-b", new HashSet<>(Arrays.asList(3, 4)));
+ File snapshot = temporaryFolder.newFile("metadata.snapshot");
+ metaManager.writeMetaInfoToFile(snapshot);
+
+ registerShuffle("app-a", 99);
+ metaManager.deleteApp("app-b");
+ registerShuffle("app-c", 100);
+ metaManager.restoreMetaFromFile(snapshot);
+ assertEquals(expected, metaManager.registeredAppAndShuffles);
+
+ registerShuffle("app-a", 5);
+ metaManager.reviseLostShuffles("app-a", Arrays.asList(2, 6));
+ registerShuffle("app-b", 7);
+ expected.get("app-a").addAll(Arrays.asList(5, 6));
+ expected.get("app-b").add(7);
+ assertEquals(expected, metaManager.registeredAppAndShuffles);
+ }
+
+ private void registerShuffle(String appId, int shuffleId) {
+ metaManager.updateRequestSlotsMeta(
+ Utils.makeShuffleKey(appId, shuffleId), "localhost",
Collections.emptyMap());
+ }
+
+ private void assertConcurrentShuffleSet(String appId) {
+ Set<Integer> shuffles = metaManager.registeredAppAndShuffles.get(appId);
+ assertNotNull(shuffles);
+ // Check the concurrent-set invariant deterministically instead of relying
on a resize race.
+ assertTrue(
+ "Registered shuffles must use a concurrent key set: " + appId,
+ shuffles instanceof ConcurrentHashMap.KeySetView);
+ }
+}
