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 adadbd9248 [CELEBORN-2442] Optimize meta memory for very large
partitions
adadbd9248 is described below
commit adadbd924823416147680fae5131390a8f3833fd
Author: fengmingxiao <[email protected]>
AuthorDate: Tue Sep 15 14:18:26 2026 +0800
[CELEBORN-2442] Optimize meta memory for very large partitions
### What changes were proposed in this pull request?
This PR reduces the memory footprint of `PartitionLocation` metadata
while preserving the existing wire protocol and public getter behavior.
The main changes are:
- Introduce an immutable, weakly interned `WorkerEndpoint` to share
worker host and port information across partition locations.
- Move the cached `host:pushPort` and `host:fetchPort` strings into the
shared endpoint.
- Lazily initialize `StorageInfo` and `RoaringBitmap` using thread-safe
initialization.
- Keep `getMapIdBitMap()` non-null for compatibility, and add a
non-materializing accessor for serialization, HTTP API, and read-side
inspection.
- Avoid allocating empty bitmaps when locations are created by
`LifecycleManager` and `SlotsAllocator`.
- Preserve the current `storageTypes.head` selection behavior in
`LifecycleManager`.
- Keep string keys in `WorkerPartitionLocationInfo` to avoid parsing and
boxing overhead on the Worker hot path.
- Add allocation-free helpers for converting partition location IDs into
packed `long` values.
- Avoid materializing optional metadata from logging, protobuf
serialization, HTTP rendering, and range-read paths.
- Re-intern `WorkerEndpoint` instances after same-version Java
deserialization.
- Fix peer bitmap serialization so that the peer protobuf contains the
peer's bitmap instead of the primary location's bitmap.
- Make range reads fall back to the peer bitmap and fail open when bitmap
metadata is unavailable.
- Add JOL-based manual memory benchmarks and focused unit tests.
No protobuf schema is changed by this PR.
### Why are the changes needed?
A shuffle may contain millions of `PartitionLocation` instances.
Previously, every instance independently retained:
- Worker host and port fields.
- Cached host/port strings.
- An empty `StorageInfo`.
- An empty `RoaringBitmap`.
Most of this metadata is identical across locations or unused for the
majority of their lifetime. The eager allocations therefore contribute
significant heap usage and GC pressure on the Master, Worker, and client.
The JOL benchmark includes both the live location graph and the weak
interner overhead:
| Scenario | Before | After, including interner | Reduction |
|---|---:|---:|---:|
| Allocator, 10,000 pairs | 4,383,504 bytes | 2,148,056 bytes | 51.00% |
| Allocator, 1,000,000 pairs | 424,143,504 bytes | 184,308,056 bytes |
56.55% |
| Packed decoded pair shape, 10,000 pairs | 3,752,304 bytes | 2,148,056
bytes | 42.75% |
| Packed decoded pair shape, 1,000,000 pairs | 375,200,304 bytes |
184,308,056 bytes | 50.88% |
The additional serialization and range-read fixes are needed to ensure
that lazy bitmap allocation does not introduce data loss, incorrect peer
metadata, or compatibility regressions.
### Does this PR resolve a correctness bug?
- [x] Yes
It fixes peer bitmap serialization and ensures that missing bitmap
metadata does not cause an NPE or silently exclude a partition location during
range reads.
### Does this PR introduce _any_ user-facing change?
- [ ] Yes
There are no configuration, wire-protocol, or intended behavioral changes
for users.
### How was this patch tested?
The following verification was completed:
- Ran the `common`, `client`, and `master` module regression suites:
- 366 JUnit tests, with 1 ignored.
- 325 ScalaTest tests.
- 691 tests in total, with no failures or errors.
- Ran the final `PartitionLocationSuiteJ`:
- 20 tests passed.
- Covers lazy initialization, endpoint interning, endpoint setters, ID
parsing, peer references, compatibility getters, and same-version Java
serialization.
- Ran Worker partition and storage regression suites:
- 27 Java tests passed.
- 9 Scala tests passed.
- Ran the Spark 3.5 `LifecycleManagerReserveSlotsSuite` integration test:
- 1 test passed.
- All 11 reactor modules completed successfully.
- Compiled the affected modules through SBT:
- `celeborn-common/Test/compile`
- `celeborn-service/Compile/compile`
- Ran `spotless:check` and `git diff --check`.
- Manually ran `PartitionLocationMemorySuiteJ` with JOL for both the
10,000-pair and 1,000,000-pair scenarios.
Closes #3825 from FMX/optimize-meta-memory.
Authored-by: fengmingxiao <[email protected]>
Signed-off-by: zhengtao <[email protected]>
AI-Contributed/Feature: 0/406
AI-Contributed/UT: 0/547
---
.../celeborn/client/read/CelebornInputStream.java | 46 +++--
.../apache/celeborn/client/LifecycleManager.scala | 5 +-
.../read/CelebornInputStreamPeerFailoverTest.java | 53 +++++
common/pom.xml | 5 +
.../common/protocol/PartitionLocation.java | 206 +++++++++++---------
.../celeborn/common/protocol/WorkerEndpoint.java | 119 ++++++++++++
.../common/meta/WorkerPartitionLocationInfo.scala | 2 +-
.../apache/celeborn/common/util/PbSerDeUtils.scala | 6 +-
.../protocol/PartitionLocationMemorySuiteJ.java | 214 +++++++++++++++++++++
.../common/protocol/PartitionLocationOld.java | 68 +++++++
.../common/protocol/PartitionLocationSuiteJ.java | 165 +++++++++++++++-
.../celeborn/common/util/PbSerDeUtilsTest.scala | 47 +++++
.../deploy/master/slotsalloc/SlotsAllocator.java | 3 +-
pom.xml | 7 +
project/CelebornBuild.scala | 3 +
.../server/common/http/api/v1/ApiUtils.scala | 4 +-
16 files changed, 838 insertions(+), 115 deletions(-)
diff --git
a/client/src/main/java/org/apache/celeborn/client/read/CelebornInputStream.java
b/client/src/main/java/org/apache/celeborn/client/read/CelebornInputStream.java
index e9dc4b9625..23d7342e21 100644
---
a/client/src/main/java/org/apache/celeborn/client/read/CelebornInputStream.java
+++
b/client/src/main/java/org/apache/celeborn/client/read/CelebornInputStream.java
@@ -172,6 +172,34 @@ public abstract class CelebornInputStream extends
InputStream {
public abstract int partitionsRead();
+ static boolean shouldSkipLocation(
+ boolean rangeReadFilter, int startMapIndex, int endMapIndex,
PartitionLocation location) {
+ if (!rangeReadFilter || endMapIndex == Integer.MAX_VALUE) {
+ return false;
+ }
+ RoaringBitmap bitmap = location.getMapIdBitMapIfPresent();
+ if (bitmap == null && location.hasPeer()) {
+ bitmap = location.getPeer().getMapIdBitMapIfPresent();
+ }
+ if (bitmap == null) {
+ // Missing filter metadata cannot prove that this location is irrelevant.
+ return false;
+ }
+ if (startMapIndex >= endMapIndex) {
+ return true;
+ }
+ if (startMapIndex >= 0) {
+ return !bitmap.intersects((long) startMapIndex, (long) endMapIndex);
+ }
+ // Preserve the previous unsigned-int behavior for unexpected negative map
indexes.
+ for (int i = startMapIndex; i < endMapIndex; i++) {
+ if (bitmap.contains(i)) {
+ return false;
+ }
+ }
+ return true;
+ }
+
private static final class CelebornInputStreamImpl extends
CelebornInputStream {
private static final Random RAND = new Random();
@@ -362,22 +390,8 @@ public abstract class CelebornInputStream extends
InputStream {
}
private boolean skipLocation(int startMapIndex, int endMapIndex,
PartitionLocation location) {
- if (!rangeReadFilter) {
- return false;
- }
- if (endMapIndex == Integer.MAX_VALUE) {
- return false;
- }
- RoaringBitmap bitmap = location.getMapIdBitMap();
- if (bitmap == null && location.hasPeer()) {
- bitmap = location.getPeer().getMapIdBitMap();
- }
- for (int i = startMapIndex; i < endMapIndex; i++) {
- if (bitmap.contains(i)) {
- return false;
- }
- }
- return true;
+ return CelebornInputStream.shouldSkipLocation(
+ rangeReadFilter, startMapIndex, endMapIndex, location);
}
private Tuple2<PartitionLocation, PbStreamHandler> nextReadableLocation() {
diff --git
a/client/src/main/scala/org/apache/celeborn/client/LifecycleManager.scala
b/client/src/main/scala/org/apache/celeborn/client/LifecycleManager.scala
index 88b778395f..8e90ddd9ad 100644
--- a/client/src/main/scala/org/apache/celeborn/client/LifecycleManager.scala
+++ b/client/src/main/scala/org/apache/celeborn/client/LifecycleManager.scala
@@ -36,7 +36,6 @@ import scala.util.Random
import com.google.common.annotations.VisibleForTesting
import com.google.common.cache.{Cache, CacheBuilder}
-import org.roaringbitmap.RoaringBitmap
import org.apache.celeborn.client.LifecycleManager.{ShuffleAllocatedWorkers,
ShuffleFailedWorkers}
import org.apache.celeborn.client.listener.WorkerStatusListener
@@ -1660,7 +1659,7 @@ class LifecycleManager(val appUniqueId: String, val conf:
CelebornConf) extends
PartitionLocation.Mode.PRIMARY,
null,
new StorageInfo("", storageTypes.head, availableStorageTypes),
- new RoaringBitmap())
+ null)
if (pushReplicateEnabled) {
var replicaIndex = (primaryIndex + 1) % candidates.size
while (pushRackAwareEnabled && isOnSameRack(primaryIndex, replicaIndex)
@@ -1682,7 +1681,7 @@ class LifecycleManager(val appUniqueId: String, val conf:
CelebornConf) extends
PartitionLocation.Mode.REPLICA,
primaryLocation,
new StorageInfo("", storageTypes.head, availableStorageTypes),
- new RoaringBitmap())
+ null)
primaryLocation.setPeer(replicaLocation)
val primaryAndReplicaPairs =
slots.computeIfAbsent(candidates(replicaIndex), newLocationFunc)
primaryAndReplicaPairs._2.add(replicaLocation)
diff --git
a/client/src/test/java/org/apache/celeborn/client/read/CelebornInputStreamPeerFailoverTest.java
b/client/src/test/java/org/apache/celeborn/client/read/CelebornInputStreamPeerFailoverTest.java
index ec1003b7d0..7f25da1692 100644
---
a/client/src/test/java/org/apache/celeborn/client/read/CelebornInputStreamPeerFailoverTest.java
+++
b/client/src/test/java/org/apache/celeborn/client/read/CelebornInputStreamPeerFailoverTest.java
@@ -18,6 +18,7 @@
package org.apache.celeborn.client.read;
import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
import static org.junit.Assert.assertSame;
import static org.junit.Assert.assertThrows;
import static org.junit.Assert.assertTrue;
@@ -45,6 +46,7 @@ import java.util.concurrent.atomic.AtomicInteger;
import org.junit.Before;
import org.junit.Test;
+import org.roaringbitmap.RoaringBitmap;
import org.apache.celeborn.client.ShuffleClient;
import org.apache.celeborn.client.security.CryptoHandler;
@@ -285,6 +287,57 @@ public class CelebornInputStreamPeerFailoverTest {
Optional.<CryptoHandler>empty());
}
+ @Test
+ public void testRangeReadFilterReadsLocationWhenBothBitmapsAreAbsent() {
+ PartitionLocation primary = createPartitionLocation(PRIMARY_HOST);
+ PartitionLocation replica = createPartitionLocation(REPLICA_HOST);
+ primary.setPeer(replica);
+ replica.setPeer(primary);
+
+ assertFalse(CelebornInputStream.shouldSkipLocation(true, 0, 10, primary));
+ assertFalse(CelebornInputStream.shouldSkipLocation(false, 0, 10, primary));
+ }
+
+ @Test
+ public void testRangeReadFilterUsesPeerBitmapWhenPrimaryBitmapIsAbsent() {
+ PartitionLocation primary = createPartitionLocation(PRIMARY_HOST);
+ PartitionLocation replica = createPartitionLocation(REPLICA_HOST);
+ RoaringBitmap peerBitmap = new RoaringBitmap();
+ peerBitmap.add(5);
+ replica.setMapIdBitMap(peerBitmap);
+ primary.setPeer(replica);
+
+ assertFalse(CelebornInputStream.shouldSkipLocation(true, 0, 10, primary));
+ assertTrue(CelebornInputStream.shouldSkipLocation(true, 6, 10, primary));
+ }
+
+ @Test
+ public void testRangeReadFilterUsesHalfOpenRange() {
+ PartitionLocation location = createPartitionLocation(PRIMARY_HOST);
+ RoaringBitmap bitmap = new RoaringBitmap();
+ bitmap.add(5);
+ bitmap.add(10);
+ location.setMapIdBitMap(bitmap);
+
+ assertFalse(CelebornInputStream.shouldSkipLocation(true, 5, 10, location));
+ assertTrue(CelebornInputStream.shouldSkipLocation(true, 6, 10, location));
+ }
+
+ @Test
+ public void testRangeReadFilterHandlesEmptyAndSentinelRanges() {
+ PartitionLocation location = createPartitionLocation(PRIMARY_HOST);
+ RoaringBitmap bitmap = new RoaringBitmap();
+ bitmap.add(5);
+ location.setMapIdBitMap(bitmap);
+
+ assertTrue(CelebornInputStream.shouldSkipLocation(true, 5, 5, location));
+ assertTrue(CelebornInputStream.shouldSkipLocation(true, 10, 5, location));
+ assertTrue(CelebornInputStream.shouldSkipLocation(true, -1, -1, location));
+
+ bitmap.add(-1);
+ assertFalse(CelebornInputStream.shouldSkipLocation(true, -1, 0, location));
+ }
+
private CelebornInputStream createInputStream(String primaryHost, String
replicaHost)
throws IOException {
return createInputStream(primaryHost, replicaHost, null);
diff --git a/common/pom.xml b/common/pom.xml
index 9ff068ece1..7f8d59b897 100644
--- a/common/pom.xml
+++ b/common/pom.xml
@@ -158,6 +158,11 @@
<artifactId>mockito-core</artifactId>
<scope>test</scope>
</dependency>
+ <dependency>
+ <groupId>org.openjdk.jol</groupId>
+ <artifactId>jol-core</artifactId>
+ <scope>test</scope>
+ </dependency>
<dependency>
<groupId>org.apache.logging.log4j</groupId>
<artifactId>log4j-slf4j-impl</artifactId>
diff --git
a/common/src/main/java/org/apache/celeborn/common/protocol/PartitionLocation.java
b/common/src/main/java/org/apache/celeborn/common/protocol/PartitionLocation.java
index 46316d8bd2..68b99347f6 100644
---
a/common/src/main/java/org/apache/celeborn/common/protocol/PartitionLocation.java
+++
b/common/src/main/java/org/apache/celeborn/common/protocol/PartitionLocation.java
@@ -23,7 +23,13 @@ import org.roaringbitmap.RoaringBitmap;
import org.apache.celeborn.common.meta.WorkerInfo;
+/**
+ * Describes a partition location. Java serialization is retained for
same-version internal use; its
+ * serialized form is not a cross-version compatibility contract.
+ */
public class PartitionLocation implements Serializable {
+ private static final RoaringBitmap EMPTY_MAP_ID_BITMAP = new RoaringBitmap();
+
public enum Mode {
PRIMARY(0),
REPLICA(1);
@@ -54,26 +60,16 @@ public class PartitionLocation implements Serializable {
private int id;
private int epoch;
- private String host;
- private int rpcPort;
- private int pushPort;
- private int fetchPort;
- private int replicatePort;
+ private volatile WorkerEndpoint endpoint;
private Mode mode;
private PartitionLocation peer;
- private StorageInfo storageInfo;
- private RoaringBitmap mapIdBitMap;
- private transient String _hostPushPort;
- private transient String _hostFetchPort;
+ private volatile StorageInfo storageInfo;
+ private volatile RoaringBitmap mapIdBitMap;
public PartitionLocation(PartitionLocation loc) {
this.id = loc.id;
this.epoch = loc.epoch;
- this.host = loc.host;
- this.rpcPort = loc.rpcPort;
- this.pushPort = loc.pushPort;
- this.fetchPort = loc.fetchPort;
- this.replicatePort = loc.replicatePort;
+ this.endpoint = loc.endpoint;
this.mode = loc.mode;
this.peer = loc.peer;
this.storageInfo = loc.storageInfo;
@@ -89,18 +85,7 @@ public class PartitionLocation implements Serializable {
int fetchPort,
int replicatePort,
Mode mode) {
- this(
- id,
- epoch,
- host,
- rpcPort,
- pushPort,
- fetchPort,
- replicatePort,
- mode,
- null,
- new StorageInfo(),
- new RoaringBitmap());
+ this(id, epoch, host, rpcPort, pushPort, fetchPort, replicatePort, mode,
null, null, null);
}
public PartitionLocation(
@@ -113,18 +98,7 @@ public class PartitionLocation implements Serializable {
int replicatePort,
Mode mode,
PartitionLocation peer) {
- this(
- id,
- epoch,
- host,
- rpcPort,
- pushPort,
- fetchPort,
- replicatePort,
- mode,
- peer,
- new StorageInfo(),
- new RoaringBitmap());
+ this(id, epoch, host, rpcPort, pushPort, fetchPort, replicatePort, mode,
peer, null, null);
}
public PartitionLocation(
@@ -141,13 +115,9 @@ public class PartitionLocation implements Serializable {
RoaringBitmap mapIdBitMap) {
this.id = id;
this.epoch = epoch;
- this.host = host;
- this.rpcPort = rpcPort;
- this.pushPort = pushPort;
- this.fetchPort = fetchPort;
- this.replicatePort = replicatePort;
+ this.endpoint = WorkerEndpoint.apply(host, rpcPort, pushPort, fetchPort,
replicatePort);
this.mode = mode;
- this.peer = peer;
+ setPeer(peer);
this.storageInfo = hint;
this.mapIdBitMap = mapIdBitMap;
}
@@ -169,54 +139,69 @@ public class PartitionLocation implements Serializable {
}
public String getHost() {
- return host;
+ return endpoint.host();
}
- public void setHost(String host) {
- this.host = host;
+ public synchronized void setHost(String host) {
+ WorkerEndpoint current = endpoint;
+ this.endpoint =
+ WorkerEndpoint.apply(
+ host,
+ current.rpcPort(),
+ current.pushPort(),
+ current.fetchPort(),
+ current.replicatePort());
}
public int getPushPort() {
- return pushPort;
+ return endpoint.pushPort();
}
- public void setPushPort(int pushPort) {
- this.pushPort = pushPort;
+ public synchronized void setPushPort(int pushPort) {
+ WorkerEndpoint current = endpoint;
+ this.endpoint =
+ WorkerEndpoint.apply(
+ current.host(),
+ current.rpcPort(),
+ pushPort,
+ current.fetchPort(),
+ current.replicatePort());
}
public int getFetchPort() {
- return fetchPort;
+ return endpoint.fetchPort();
}
- public void setFetchPort(int fetchPort) {
- this.fetchPort = fetchPort;
+ public synchronized void setFetchPort(int fetchPort) {
+ WorkerEndpoint current = endpoint;
+ this.endpoint =
+ WorkerEndpoint.apply(
+ current.host(),
+ current.rpcPort(),
+ current.pushPort(),
+ fetchPort,
+ current.replicatePort());
}
public String hostAndPorts() {
return "host-rpcPort-pushPort-fetchPort-replicatePort:"
- + host
+ + getHost()
+ "-"
- + rpcPort
+ + getRpcPort()
+ "-"
- + pushPort
+ + getPushPort()
+ "-"
- + fetchPort
+ + getFetchPort()
+ "-"
- + replicatePort;
+ + getReplicatePort();
}
public String hostAndFetchPort() {
- if (_hostFetchPort == null) {
- _hostFetchPort = host + ":" + fetchPort;
- }
- return _hostFetchPort;
+ return endpoint.hostAndFetchPort();
}
public String hostAndPushPort() {
- if (_hostPushPort == null) {
- _hostPushPort = host + ":" + pushPort;
- }
- return _hostPushPort;
+ return endpoint.hostAndPushPort();
}
public Mode getMode() {
@@ -249,26 +234,54 @@ public class PartitionLocation implements Serializable {
}
public int getRpcPort() {
- return rpcPort;
+ return endpoint.rpcPort();
}
- public void setRpcPort(int rpcPort) {
- this.rpcPort = rpcPort;
+ public synchronized void setRpcPort(int rpcPort) {
+ WorkerEndpoint current = endpoint;
+ this.endpoint =
+ WorkerEndpoint.apply(
+ current.host(),
+ rpcPort,
+ current.pushPort(),
+ current.fetchPort(),
+ current.replicatePort());
}
public int getReplicatePort() {
- return replicatePort;
+ return endpoint.replicatePort();
}
- public void setReplicatePort(int replicatePort) {
- this.replicatePort = replicatePort;
+ public synchronized void setReplicatePort(int replicatePort) {
+ WorkerEndpoint current = endpoint;
+ this.endpoint =
+ WorkerEndpoint.apply(
+ current.host(),
+ current.rpcPort(),
+ current.pushPort(),
+ current.fetchPort(),
+ replicatePort);
}
public StorageInfo getStorageInfo() {
- return storageInfo;
+ return getStorageInfoOrCreate();
+ }
+
+ public StorageInfo getStorageInfoOrCreate() {
+ StorageInfo current = storageInfo;
+ if (current == null) {
+ synchronized (this) {
+ current = storageInfo;
+ if (current == null) {
+ current = new StorageInfo();
+ storageInfo = current;
+ }
+ }
+ }
+ return current;
}
- public void setStorageInfo(StorageInfo storageInfo) {
+ public synchronized void setStorageInfo(StorageInfo storageInfo) {
this.storageInfo = storageInfo;
}
@@ -280,15 +293,15 @@ public class PartitionLocation implements Serializable {
PartitionLocation o = (PartitionLocation) other;
return id == o.id
&& epoch == o.epoch
- && host.equals(o.host)
- && rpcPort == o.rpcPort
- && pushPort == o.pushPort
- && fetchPort == o.fetchPort;
+ && getHost().equals(o.getHost())
+ && getRpcPort() == o.getRpcPort()
+ && getPushPort() == o.getPushPort()
+ && getFetchPort() == o.getFetchPort();
}
@Override
public int hashCode() {
- return (id + epoch + host + rpcPort + pushPort + fetchPort).hashCode();
+ return (id + epoch + getHost() + getRpcPort() + getPushPort() +
getFetchPort()).hashCode();
}
@Override
@@ -303,15 +316,15 @@ public class PartitionLocation implements Serializable {
+ "-"
+ epoch
+ "\n host-rpcPort-pushPort-fetchPort-replicatePort:"
- + host
+ + getHost()
+ "-"
- + rpcPort
+ + getRpcPort()
+ "-"
- + pushPort
+ + getPushPort()
+ "-"
- + fetchPort
+ + getFetchPort()
+ "-"
- + replicatePort
+ + getReplicatePort()
+ "\n mode:"
+ mode
+ "\n peer:("
@@ -319,19 +332,38 @@ public class PartitionLocation implements Serializable {
+ ")\n storage hint:"
+ storageInfo
+ "\n mapIdBitMap:"
- + mapIdBitMap
+ + (mapIdBitMap == null ? EMPTY_MAP_ID_BITMAP : mapIdBitMap)
+ "]";
}
public WorkerInfo getWorker() {
- return new WorkerInfo(host, rpcPort, pushPort, fetchPort, replicatePort);
+ return new WorkerInfo(
+ getHost(), getRpcPort(), getPushPort(), getFetchPort(),
getReplicatePort());
}
public RoaringBitmap getMapIdBitMap() {
+ return getMapIdBitMapOrCreate();
+ }
+
+ public RoaringBitmap getMapIdBitMapIfPresent() {
return mapIdBitMap;
}
- public void setMapIdBitMap(RoaringBitmap mapIdBitMap) {
+ public RoaringBitmap getMapIdBitMapOrCreate() {
+ RoaringBitmap current = mapIdBitMap;
+ if (current == null) {
+ synchronized (this) {
+ current = mapIdBitMap;
+ if (current == null) {
+ current = new RoaringBitmap();
+ mapIdBitMap = current;
+ }
+ }
+ }
+ return current;
+ }
+
+ public synchronized void setMapIdBitMap(RoaringBitmap mapIdBitMap) {
this.mapIdBitMap = mapIdBitMap;
}
}
diff --git
a/common/src/main/java/org/apache/celeborn/common/protocol/WorkerEndpoint.java
b/common/src/main/java/org/apache/celeborn/common/protocol/WorkerEndpoint.java
new file mode 100644
index 0000000000..f96662d830
--- /dev/null
+++
b/common/src/main/java/org/apache/celeborn/common/protocol/WorkerEndpoint.java
@@ -0,0 +1,119 @@
+/*
+ * 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.common.protocol;
+
+import java.io.Serializable;
+
+import com.google.common.collect.Interner;
+import com.google.common.collect.Interners;
+
+/**
+ * Immutable and weakly interned worker endpoint shared by partition
locations. Java serialization
+ * is retained for same-version internal use; its serialized form is not a
cross-version
+ * compatibility contract.
+ */
+public final class WorkerEndpoint implements Serializable {
+ private static final Interner<WorkerEndpoint> INTERNED_ENDPOINTS =
Interners.newWeakInterner();
+
+ private final String host;
+ private final int rpcPort;
+ private final int pushPort;
+ private final int fetchPort;
+ private final int replicatePort;
+ private transient volatile String hostPushPort;
+ private transient volatile String hostFetchPort;
+
+ private WorkerEndpoint(String host, int rpcPort, int pushPort, int
fetchPort, int replicatePort) {
+ this.host = host;
+ this.rpcPort = rpcPort;
+ this.pushPort = pushPort;
+ this.fetchPort = fetchPort;
+ this.replicatePort = replicatePort;
+ }
+
+ public static WorkerEndpoint apply(
+ String host, int rpcPort, int pushPort, int fetchPort, int
replicatePort) {
+ WorkerEndpoint endpoint = new WorkerEndpoint(host, rpcPort, pushPort,
fetchPort, replicatePort);
+ return INTERNED_ENDPOINTS.intern(endpoint);
+ }
+
+ public String host() {
+ return host;
+ }
+
+ public int rpcPort() {
+ return rpcPort;
+ }
+
+ public int pushPort() {
+ return pushPort;
+ }
+
+ public int fetchPort() {
+ return fetchPort;
+ }
+
+ public int replicatePort() {
+ return replicatePort;
+ }
+
+ public String hostAndPushPort() {
+ String current = hostPushPort;
+ if (current == null) {
+ current = host + ":" + pushPort;
+ hostPushPort = current;
+ }
+ return current;
+ }
+
+ public String hostAndFetchPort() {
+ String current = hostFetchPort;
+ if (current == null) {
+ current = host + ":" + fetchPort;
+ hostFetchPort = current;
+ }
+ return current;
+ }
+
+ private Object readResolve() {
+ return apply(host, rpcPort, pushPort, fetchPort, replicatePort);
+ }
+
+ @Override
+ public boolean equals(Object other) {
+ if (!(other instanceof WorkerEndpoint)) {
+ return false;
+ }
+ WorkerEndpoint that = (WorkerEndpoint) other;
+ return rpcPort == that.rpcPort
+ && pushPort == that.pushPort
+ && fetchPort == that.fetchPort
+ && replicatePort == that.replicatePort
+ && host.equals(that.host);
+ }
+
+ @Override
+ public int hashCode() {
+ int result = host.hashCode();
+ result = 31 * result + rpcPort;
+ result = 31 * result + pushPort;
+ result = 31 * result + fetchPort;
+ result = 31 * result + replicatePort;
+ return result;
+ }
+}
diff --git
a/common/src/main/scala/org/apache/celeborn/common/meta/WorkerPartitionLocationInfo.scala
b/common/src/main/scala/org/apache/celeborn/common/meta/WorkerPartitionLocationInfo.scala
index 373f365658..c9930dbc5e 100644
---
a/common/src/main/scala/org/apache/celeborn/common/meta/WorkerPartitionLocationInfo.scala
+++
b/common/src/main/scala/org/apache/celeborn/common/meta/WorkerPartitionLocationInfo.scala
@@ -29,7 +29,7 @@ import org.apache.celeborn.common.util.{CollectionUtils,
JavaUtils}
class WorkerPartitionLocationInfo extends Logging {
- // key: ShuffleKey, values: (uniqueId -> PartitionLocation))
+ // key: ShuffleKey, values: (uniqueId -> PartitionLocation)
type PartitionInfo = ConcurrentHashMap[String, ConcurrentHashMap[String,
PartitionLocation]]
private[celeborn] val primaryPartitionLocations = new PartitionInfo
private[celeborn] val replicaPartitionLocations = new PartitionInfo
diff --git
a/common/src/main/scala/org/apache/celeborn/common/util/PbSerDeUtils.scala
b/common/src/main/scala/org/apache/celeborn/common/util/PbSerDeUtils.scala
index e9c407ce80..4392dacbe2 100644
--- a/common/src/main/scala/org/apache/celeborn/common/util/PbSerDeUtils.scala
+++ b/common/src/main/scala/org/apache/celeborn/common/util/PbSerDeUtils.scala
@@ -382,7 +382,7 @@ object PbSerDeUtils {
.setFetchPort(location.getFetchPort)
.setReplicatePort(location.getReplicatePort)
.setStorageInfo(StorageInfo.toPb(location.getStorageInfo))
- .setMapIdBitmap(Utils.roaringBitmapToByteString(location.getMapIdBitMap))
+
.setMapIdBitmap(Utils.roaringBitmapToByteString(location.getMapIdBitMapIfPresent))
if (location.hasPeer) {
val peerBuilder = PbPartitionLocation.newBuilder
if (location.getPeer.getMode eq Mode.PRIMARY) {
@@ -399,7 +399,7 @@ object PbSerDeUtils {
.setFetchPort(location.getPeer.getFetchPort)
.setReplicatePort(location.getPeer.getReplicatePort)
.setStorageInfo(StorageInfo.toPb(location.getPeer.getStorageInfo))
-
.setMapIdBitmap(Utils.roaringBitmapToByteString(location.getMapIdBitMap))
+
.setMapIdBitmap(Utils.roaringBitmapToByteString(location.getPeer.getMapIdBitMapIfPresent))
builder.setPeer(peerBuilder.build)
}
builder.build
@@ -562,7 +562,7 @@ object PbSerDeUtils {
pbPackedLocationsBuilder.addEpoches(location.getEpoch)
pbPackedLocationsBuilder.addWorkerIds(workerIdIndex(location.getWorker.toUniqueId))
pbPackedLocationsBuilder.addMapIdBitMap(
- Utils.roaringBitmapToByteString(location.getMapIdBitMap))
+ Utils.roaringBitmapToByteString(location.getMapIdBitMapIfPresent))
pbPackedLocationsBuilder.addTypes(location.getStorageInfo.getType.getValue)
pbPackedLocationsBuilder.addMountPoints(
mountPointsIndex(location.getStorageInfo.getMountPoint))
diff --git
a/common/src/test/java/org/apache/celeborn/common/protocol/PartitionLocationMemorySuiteJ.java
b/common/src/test/java/org/apache/celeborn/common/protocol/PartitionLocationMemorySuiteJ.java
new file mode 100644
index 0000000000..20fadb496b
--- /dev/null
+++
b/common/src/test/java/org/apache/celeborn/common/protocol/PartitionLocationMemorySuiteJ.java
@@ -0,0 +1,214 @@
+/*
+ * 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.common.protocol;
+
+import java.lang.reflect.Field;
+
+import org.junit.Ignore;
+import org.junit.Test;
+import org.openjdk.jol.info.GraphLayout;
+import org.roaringbitmap.RoaringBitmap;
+
+@Ignore("Manual JOL benchmark; run main to print retained-size comparisons.")
+public class PartitionLocationMemorySuiteJ {
+
+ private static final int ENDPOINT_COUNT = 2000;
+ private static final int DEFAULT_PAIR_COUNT = 10_000;
+ private static final String ALLOCATOR_SCENARIO = "allocator";
+ private static final String PACKED_DESERIALIZED_PAIR_SCENARIO =
"packed-deserialized-pair";
+
+ public static void main(String[] args) throws Exception {
+ String scenario = args.length > 0 ? args[0] : ALLOCATOR_SCENARIO;
+ int pairCount = args.length > 1 ? Integer.parseInt(args[1]) :
DEFAULT_PAIR_COUNT;
+ if (pairCount <= 0) {
+ throw new IllegalArgumentException("pairCount must be positive: " +
pairCount);
+ }
+ if (!ALLOCATOR_SCENARIO.equals(scenario)
+ && !PACKED_DESERIALIZED_PAIR_SCENARIO.equals(scenario)) {
+ throw new IllegalArgumentException(
+ "Unknown scenario: " + scenario + ". Expected allocator or
packed-deserialized-pair.");
+ }
+ new PartitionLocationMemorySuiteJ().printPeerPairFootprint(scenario,
pairCount, ENDPOINT_COUNT);
+ }
+
+ @Test
+ public void printPartitionLocationFootprint() throws Exception {
+ printPeerPairFootprint(ALLOCATOR_SCENARIO, DEFAULT_PAIR_COUNT,
ENDPOINT_COUNT);
+ }
+
+ private void printPeerPairFootprint(String scenario, int pairCount, int
endpointCount)
+ throws Exception {
+ String[] endpointHosts = endpointHosts(endpointCount);
+ boolean allocatorShaped = ALLOCATOR_SCENARIO.equals(scenario);
+ compare(
+ scenario + "-shaped, " + pairCount + " peer pairs across " +
endpointCount + " endpoints",
+ newOldLocationPairArray(pairCount, endpointHosts, allocatorShaped),
+ newLocationPairArray(pairCount, endpointHosts, allocatorShaped));
+ }
+
+ private String[] endpointHosts(int endpointCount) {
+ String[] hosts = new String[endpointCount];
+ for (int i = 0; i < endpointCount; i++) {
+ hosts[i] = "localhost-" + i;
+ }
+ return hosts;
+ }
+
+ private PartitionLocation[] newLocationPairArray(
+ int size, String[] endpointHosts, boolean allocatorShaped) {
+ PartitionLocation[] locations = new PartitionLocation[size * 2];
+ for (int i = 0; i < size; i++) {
+ int primaryEndpointIndex = i % endpointHosts.length;
+ int replicaEndpointIndex = (primaryEndpointIndex + 1) %
endpointHosts.length;
+ PartitionLocation primary =
+ newLocation(
+ i,
+ primaryEndpointIndex,
+ endpointHosts,
+ allocatorShaped,
+ PartitionLocation.Mode.PRIMARY);
+ PartitionLocation replica =
+ newLocation(
+ i,
+ replicaEndpointIndex,
+ endpointHosts,
+ allocatorShaped,
+ PartitionLocation.Mode.REPLICA);
+ primary.setPeer(replica);
+ replica.setPeer(primary);
+ locations[i * 2] = primary;
+ locations[i * 2 + 1] = replica;
+ }
+ return locations;
+ }
+
+ private PartitionLocation newLocation(
+ int id,
+ int endpointIndex,
+ String[] endpointHosts,
+ boolean allocatorShaped,
+ PartitionLocation.Mode mode) {
+ return new PartitionLocation(
+ id,
+ 0,
+ locationHost(endpointHosts[endpointIndex], allocatorShaped),
+ 1001 + endpointIndex,
+ 1002 + endpointIndex,
+ 1003 + endpointIndex,
+ 1004 + endpointIndex,
+ mode,
+ null,
+ newStorageInfo(),
+ null);
+ }
+
+ private PartitionLocationOld[] newOldLocationPairArray(
+ int size, String[] endpointHosts, boolean allocatorShaped) {
+ PartitionLocationOld[] locations = new PartitionLocationOld[size * 2];
+ for (int i = 0; i < size; i++) {
+ int primaryEndpointIndex = i % endpointHosts.length;
+ int replicaEndpointIndex = (primaryEndpointIndex + 1) %
endpointHosts.length;
+ PartitionLocationOld primary =
+ newOldLocation(
+ i,
+ primaryEndpointIndex,
+ endpointHosts,
+ allocatorShaped,
+ PartitionLocation.Mode.PRIMARY);
+ PartitionLocationOld replica =
+ newOldLocation(
+ i,
+ replicaEndpointIndex,
+ endpointHosts,
+ allocatorShaped,
+ PartitionLocation.Mode.REPLICA);
+ primary.setPeer(replica);
+ replica.setPeer(primary);
+ locations[i * 2] = primary;
+ locations[i * 2 + 1] = replica;
+ }
+ return locations;
+ }
+
+ private PartitionLocationOld newOldLocation(
+ int id,
+ int endpointIndex,
+ String[] endpointHosts,
+ boolean allocatorShaped,
+ PartitionLocation.Mode mode) {
+ return new PartitionLocationOld(
+ id,
+ 0,
+ locationHost(endpointHosts[endpointIndex], allocatorShaped),
+ 1001 + endpointIndex,
+ 1002 + endpointIndex,
+ 1003 + endpointIndex,
+ 1004 + endpointIndex,
+ mode,
+ null,
+ newStorageInfo(),
+ allocatorShaped ? new RoaringBitmap() : null);
+ }
+
+ private String locationHost(String endpointHost, boolean allocatorShaped) {
+ // Allocators reuse WorkerInfo.host. Packed protobuf decoding creates a
new host string per
+ // location while splitting the encoded worker ID. This scenario measures
one decoded pair's
+ // retained object shape, not an entire WorkerResource response.
+ return allocatorShaped ? endpointHost : new
String(endpointHost.toCharArray());
+ }
+
+ private StorageInfo newStorageInfo() {
+ return new StorageInfo("", StorageInfo.Type.MEMORY,
StorageInfo.ALL_TYPES_AVAILABLE_MASK);
+ }
+
+ private void compare(String label, Object oldValue, Object newValue) throws
Exception {
+ long oldSize = GraphLayout.parseInstance(oldValue).totalSize();
+ long newLocationSize = GraphLayout.parseInstance(newValue).totalSize();
+ long newSizeIncludingInterner =
+ GraphLayout.parseInstance(newValue, endpointInterner()).totalSize();
+ long internerOverhead = newSizeIncludingInterner - newLocationSize;
+ long saved = oldSize - newSizeIncludingInterner;
+ double savedPercentage = oldSize == 0 ? 0 : saved * 100.0 / oldSize;
+ int locationCount = java.lang.reflect.Array.getLength(oldValue);
+ System.out.printf(
+ "PartitionLocation footprint [%s]: old=%d bytes, "
+ + "newLocations=%d bytes, weakInternerOverhead=%d bytes, "
+ + "newIncludingLiveInterner=%d bytes, oldBytesPerLocation=%.2f, "
+ + "newBytesPerLocation=%.2f, savedIncludingLiveInterner=%d bytes
(%.2f%%)%n",
+ label,
+ oldSize,
+ newLocationSize,
+ internerOverhead,
+ newSizeIncludingInterner,
+ oldSize / (double) locationCount,
+ newSizeIncludingInterner / (double) locationCount,
+ saved,
+ savedPercentage);
+ if (newSizeIncludingInterner >= oldSize) {
+ throw new AssertionError(
+ "Optimized PartitionLocation retained size including its weak
interner must be smaller for "
+ + label);
+ }
+ }
+
+ private Object endpointInterner() throws Exception {
+ Field field = WorkerEndpoint.class.getDeclaredField("INTERNED_ENDPOINTS");
+ field.setAccessible(true);
+ return field.get(null);
+ }
+}
diff --git
a/common/src/test/java/org/apache/celeborn/common/protocol/PartitionLocationOld.java
b/common/src/test/java/org/apache/celeborn/common/protocol/PartitionLocationOld.java
new file mode 100644
index 0000000000..e45ef4bc52
--- /dev/null
+++
b/common/src/test/java/org/apache/celeborn/common/protocol/PartitionLocationOld.java
@@ -0,0 +1,68 @@
+/*
+ * 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.common.protocol;
+
+import java.io.Serializable;
+
+import org.roaringbitmap.RoaringBitmap;
+
+/** Snapshot of the pre-optimization layout, used only as the JOL comparison
baseline. */
+final class PartitionLocationOld implements Serializable {
+ private int id;
+ private int epoch;
+ private String host;
+ private int rpcPort;
+ private int pushPort;
+ private int fetchPort;
+ private int replicatePort;
+ private PartitionLocation.Mode mode;
+ private PartitionLocationOld peer;
+ private StorageInfo storageInfo;
+ private RoaringBitmap mapIdBitMap;
+ private transient String _hostPushPort;
+ private transient String _hostFetchPort;
+
+ PartitionLocationOld(
+ int id,
+ int epoch,
+ String host,
+ int rpcPort,
+ int pushPort,
+ int fetchPort,
+ int replicatePort,
+ PartitionLocation.Mode mode,
+ PartitionLocationOld peer,
+ StorageInfo storageInfo,
+ RoaringBitmap mapIdBitMap) {
+ this.id = id;
+ this.epoch = epoch;
+ this.host = host;
+ this.rpcPort = rpcPort;
+ this.pushPort = pushPort;
+ this.fetchPort = fetchPort;
+ this.replicatePort = replicatePort;
+ this.mode = mode;
+ this.peer = peer;
+ this.storageInfo = storageInfo;
+ this.mapIdBitMap = mapIdBitMap;
+ }
+
+ void setPeer(PartitionLocationOld peer) {
+ this.peer = peer;
+ }
+}
diff --git
a/common/src/test/java/org/apache/celeborn/common/protocol/PartitionLocationSuiteJ.java
b/common/src/test/java/org/apache/celeborn/common/protocol/PartitionLocationSuiteJ.java
index 0d0613a33e..210f2ad16a 100644
---
a/common/src/test/java/org/apache/celeborn/common/protocol/PartitionLocationSuiteJ.java
+++
b/common/src/test/java/org/apache/celeborn/common/protocol/PartitionLocationSuiteJ.java
@@ -18,6 +18,15 @@
package org.apache.celeborn.common.protocol;
import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertNotSame;
+import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertSame;
+
+import java.io.ByteArrayInputStream;
+import java.io.ByteArrayOutputStream;
+import java.io.ObjectInputStream;
+import java.io.ObjectOutputStream;
import org.junit.Test;
import org.roaringbitmap.RoaringBitmap;
@@ -172,6 +181,158 @@ public class PartitionLocationSuiteJ {
checkEqual(location1, location2, true);
}
+ @Test
+ public void testSetPeerMaintainsPeerReference() {
+ PartitionLocation primary =
+ new PartitionLocation(
+ partitionId, epoch, host, rpcPort, pushPort, fetchPort,
replicatePort, mode);
+ PartitionLocation replica =
+ new PartitionLocation(
+ partitionId,
+ epoch,
+ host,
+ rpcPort,
+ pushPort,
+ fetchPort,
+ replicatePort,
+ PartitionLocation.Mode.REPLICA);
+
+ primary.setPeer(replica);
+
+ assertEquals(true, primary.hasPeer());
+ assertSame(replica, primary.getPeer());
+
+ primary.setPeer(null);
+ assertEquals(false, primary.hasPeer());
+ }
+
+ @Test
+ public void testCopyPartitionLocationKeepsSharedFields() {
+ PartitionLocation primary =
+ new PartitionLocation(
+ partitionId, epoch, host, rpcPort, pushPort, fetchPort,
replicatePort, mode);
+ PartitionLocation replica =
+ new PartitionLocation(
+ partitionId,
+ epoch,
+ host,
+ rpcPort,
+ pushPort,
+ fetchPort,
+ replicatePort,
+ PartitionLocation.Mode.REPLICA);
+ primary.setPeer(replica);
+ StorageInfo storageInfo = primary.getStorageInfo();
+ RoaringBitmap bitmap = primary.getMapIdBitMapOrCreate();
+
+ PartitionLocation copy = new PartitionLocation(primary);
+
+ assertSame(replica, copy.getPeer());
+ assertSame(storageInfo, copy.getStorageInfo());
+ assertSame(bitmap, copy.getMapIdBitMap());
+ }
+
+ @Test
+ public void testLazyDefaultStorageInfoIsNotShared() {
+ PartitionLocation first =
+ new PartitionLocation(
+ partitionId, epoch, host, rpcPort, pushPort, fetchPort,
replicatePort, mode);
+ PartitionLocation second =
+ new PartitionLocation(
+ partitionId + 1, epoch, host, rpcPort, pushPort, fetchPort,
replicatePort, mode);
+
+ StorageInfo storageInfo = first.getStorageInfo();
+ storageInfo.availableStorageTypes = StorageInfo.HDFS_MASK;
+ storageInfo.setMountPoint("/mnt/disk1");
+
+ assertSame(storageInfo, first.getStorageInfoOrCreate());
+ assertNotSame(first.getStorageInfo(), second.getStorageInfo());
+ assertEquals(StorageInfo.HDFS_MASK,
first.getStorageInfo().availableStorageTypes);
+ assertEquals("/mnt/disk1", first.getStorageInfo().getMountPoint());
+ assertEquals(
+ StorageInfo.ALL_TYPES_AVAILABLE_MASK,
second.getStorageInfo().availableStorageTypes);
+ assertEquals("", second.getStorageInfo().getMountPoint());
+ }
+
+ @Test
+ public void testLazyMapIdBitmapKeepsCompatibleGetter() {
+ PartitionLocation location =
+ new PartitionLocation(
+ partitionId, epoch, host, rpcPort, pushPort, fetchPort,
replicatePort, mode);
+
+ assertNull(location.getMapIdBitMapIfPresent());
+ RoaringBitmap bitmap = location.getMapIdBitMapOrCreate();
+ bitmap.add(1);
+
+ assertSame(bitmap, location.getMapIdBitMap());
+ assertEquals(1, location.getMapIdBitMap().getCardinality());
+ }
+
+ @Test
+ public void testWorkerEndpointIsInterned() {
+ WorkerEndpoint first = WorkerEndpoint.apply(host, rpcPort, pushPort,
fetchPort, replicatePort);
+ WorkerEndpoint second = WorkerEndpoint.apply(host, rpcPort, pushPort,
fetchPort, replicatePort);
+ WorkerEndpoint different =
+ WorkerEndpoint.apply(host, rpcPort, pushPort + 1, fetchPort,
replicatePort);
+
+ assertSame(first, second);
+ assertNotSame(first, different);
+ }
+
+ @Test
+ public void testSameVersionJavaSerializationReinternsWorkerEndpoint() throws
Exception {
+ WorkerEndpoint endpoint =
+ WorkerEndpoint.apply(host, rpcPort, pushPort, fetchPort,
replicatePort);
+ ByteArrayOutputStream bytes = new ByteArrayOutputStream();
+ try (ObjectOutputStream output = new ObjectOutputStream(bytes)) {
+ output.writeObject(endpoint);
+ }
+
+ WorkerEndpoint deserialized;
+ try (ObjectInputStream input =
+ new ObjectInputStream(new ByteArrayInputStream(bytes.toByteArray()))) {
+ deserialized = (WorkerEndpoint) input.readObject();
+ }
+
+ assertSame(endpoint, deserialized);
+ assertEquals("localhost:1", deserialized.hostAndPushPort());
+ assertEquals("localhost:2", deserialized.hostAndFetchPort());
+ }
+
+ @Test
+ public void testHostPortCachesAreInvalidatedBySetters() {
+ PartitionLocation location =
+ new PartitionLocation(
+ partitionId, epoch, host, rpcPort, pushPort, fetchPort,
replicatePort, mode);
+
+ assertEquals("localhost:1", location.hostAndPushPort());
+ assertEquals("localhost:2", location.hostAndFetchPort());
+
+ location.setHost("remoteHost");
+ location.setRpcPort(13);
+ location.setPushPort(11);
+ location.setFetchPort(12);
+ location.setReplicatePort(14);
+
+ assertEquals("remoteHost:11", location.hostAndPushPort());
+ assertEquals("remoteHost:12", location.hostAndFetchPort());
+ assertEquals(13, location.getRpcPort());
+ assertEquals(14, location.getReplicatePort());
+ }
+
+ @Test
+ public void testWorkerEndpointKeepsWorkerInfoCompatible() {
+ PartitionLocation location1 =
+ new PartitionLocation(
+ partitionId, epoch, host, rpcPort, pushPort, fetchPort,
replicatePort, mode);
+ PartitionLocation location2 =
+ new PartitionLocation(
+ partitionId + 1, epoch, host, rpcPort, pushPort, fetchPort,
replicatePort, mode);
+
+ assertEquals(location1.getWorker(), location2.getWorker());
+ assertNotNull(location1.hostAndPorts());
+ }
+
@Test
public void testToStringOutput() {
PartitionLocation location1 =
@@ -209,7 +370,7 @@ public class PartitionLocationSuiteJ {
+ "
host-rpcPort-pushPort-fetchPort-replicatePort:localhost-3-1-2-4\n"
+ " mode:PRIMARY\n"
+ " peer:(empty)\n"
- + " storage hint:StorageInfo{type=MEMORY, mountPoint='',
finalResult=false, filePath=null, fileSize=0, chunkOffsets=null}\n"
+ + " storage hint:null\n"
+ " mapIdBitMap:{}]";
String exp2 =
"PartitionLocation[\n"
@@ -217,7 +378,7 @@ public class PartitionLocationSuiteJ {
+ "
host-rpcPort-pushPort-fetchPort-replicatePort:localhost-3-1-2-4\n"
+ " mode:PRIMARY\n"
+ "
peer:(host-rpcPort-pushPort-fetchPort-replicatePort:localhost-3-1-2-4)\n"
- + " storage hint:StorageInfo{type=MEMORY, mountPoint='',
finalResult=false, filePath=null, fileSize=0, chunkOffsets=null}\n"
+ + " storage hint:null\n"
+ " mapIdBitMap:{}]";
String exp3 =
"PartitionLocation[\n"
diff --git
a/common/src/test/scala/org/apache/celeborn/common/util/PbSerDeUtilsTest.scala
b/common/src/test/scala/org/apache/celeborn/common/util/PbSerDeUtilsTest.scala
index 15c4df0fc7..44028ad5ce 100644
---
a/common/src/test/scala/org/apache/celeborn/common/util/PbSerDeUtilsTest.scala
+++
b/common/src/test/scala/org/apache/celeborn/common/util/PbSerDeUtilsTest.scala
@@ -27,6 +27,7 @@ import scala.util.Random
import com.google.common.collect.Lists
import org.apache.hadoop.shaded.org.apache.commons.lang3.RandomStringUtils
+import org.roaringbitmap.RoaringBitmap
import org.apache.celeborn.CelebornFunSuite
import org.apache.celeborn.common.identity.UserIdentifier
@@ -461,6 +462,45 @@ class PbSerDeUtilsTest extends CelebornFunSuite {
assert(restoredPartitionLocation.equals(partitionLocation1))
}
+ test("fromAndToPbPartitionLocation preserves lazy and peer bitmaps") {
+ val primary =
+ new PartitionLocation(10, 0, "host1", 10, 11, 12, 13,
PartitionLocation.Mode.PRIMARY)
+ val replica =
+ new PartitionLocation(10, 0, "host2", 20, 21, 22, 23,
PartitionLocation.Mode.REPLICA)
+ val primaryBitmap = new RoaringBitmap()
+ primaryBitmap.add(1)
+ val replicaBitmap = new RoaringBitmap()
+ replicaBitmap.add(2)
+ primary.setMapIdBitMap(primaryBitmap)
+ replica.setMapIdBitMap(replicaBitmap)
+ primary.setPeer(replica)
+ replica.setPeer(primary)
+
+ val restored = PbSerDeUtils.fromPbPartitionLocation(
+ PbSerDeUtils.toPbPartitionLocation(primary))
+
+ assert(restored.hasPeer)
+ assert(restored.getMapIdBitMap.contains(1))
+ assert(!restored.getMapIdBitMap.contains(2))
+ assert(restored.getPeer.getMapIdBitMap.contains(2))
+ assert(!restored.getPeer.getMapIdBitMap.contains(1))
+
+ val (packedPrimaries, _) = fromPbPackedPartitionLocationsPair(
+ toPbPackedPartitionLocationsPair(List(primary)))
+ val packedRestored = packedPrimaries.get(0)
+ assert(packedRestored.hasPeer)
+ assert(packedRestored.getMapIdBitMap.contains(1))
+ assert(!packedRestored.getMapIdBitMap.contains(2))
+ assert(packedRestored.getPeer.getMapIdBitMap.contains(2))
+ assert(!packedRestored.getPeer.getMapIdBitMap.contains(1))
+
+ val withoutBitmap =
+ new PartitionLocation(11, 0, "host3", 30, 31, 32, 33,
PartitionLocation.Mode.PRIMARY)
+ val restoredWithoutBitmap = PbSerDeUtils.fromPbPartitionLocation(
+ PbSerDeUtils.toPbPartitionLocation(withoutBitmap))
+ assert(restoredWithoutBitmap.getMapIdBitMapIfPresent == null)
+ }
+
test("fromAndToPbWorkerResource") {
val pbWorkerResource = PbSerDeUtils.toPbWorkerResource(workerResource)
val restoredWorkerResource =
PbSerDeUtils.fromPbWorkerResource(pbWorkerResource)
@@ -694,6 +734,9 @@ class PbSerDeUtilsTest extends CelebornFunSuite {
assert(primaryLocations.zip(locs1.asScala).count(x => x._1 != x._2) == 0)
assert(replicaLocations.zip(locs2.asScala).count(x => x._1 != x._2) == 0)
+ (locs1.asScala ++ locs2.asScala).foreach { loc =>
+ assert(loc.hasPeer)
+ }
assert(packedWorkerResourceSize < workerResourceSize)
log.info(s"Packed size : ${packedWorkerResourceSize} unpacked size
:${workerResourceSize}")
@@ -705,6 +748,10 @@ class PbSerDeUtilsTest extends CelebornFunSuite {
testSerializationPerformance(100)
}
+ test("packed partition location serde remains compact at scale") {
+ testSerializationPerformance(1000)
+ }
+
test("GetReduceFileGroup with primary and replica locations") {
val shuffleMap: util.Map[Integer, util.Set[PartitionLocation]] =
JavaUtils.newConcurrentHashMap()
diff --git
a/master/src/main/java/org/apache/celeborn/service/deploy/master/slotsalloc/SlotsAllocator.java
b/master/src/main/java/org/apache/celeborn/service/deploy/master/slotsalloc/SlotsAllocator.java
index 8f1cf5d4ab..05e2bee264 100644
---
a/master/src/main/java/org/apache/celeborn/service/deploy/master/slotsalloc/SlotsAllocator.java
+++
b/master/src/main/java/org/apache/celeborn/service/deploy/master/slotsalloc/SlotsAllocator.java
@@ -26,7 +26,6 @@ import java.util.stream.Collectors;
import scala.Tuple2;
import scala.Tuple3;
-import org.roaringbitmap.RoaringBitmap;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
@@ -548,7 +547,7 @@ public class SlotsAllocator {
mode,
peer,
storageInfo,
- new RoaringBitmap());
+ null);
}
private static void addLocation(
diff --git a/pom.xml b/pom.xml
index 5b51809361..0bab4860e4 100644
--- a/pom.xml
+++ b/pom.xml
@@ -88,6 +88,7 @@
<google.jsr305.version>1.3.9</google.jsr305.version>
<grpc.version>1.44.0</grpc.version>
<guava.version>33.1.0-jre</guava.version>
+ <jol.version>0.17</jol.version>
<junit.version>4.13.2</junit.version>
<leveldb.version>1.8</leveldb.version>
<log4j2.version>2.25.4</log4j2.version>
@@ -795,6 +796,12 @@
<version>${swagger-ui.version}</version>
</dependency>
+ <dependency>
+ <groupId>org.openjdk.jol</groupId>
+ <artifactId>jol-core</artifactId>
+ <version>${jol.version}</version>
+ <scope>test</scope>
+ </dependency>
<dependency>
<groupId>junit</groupId>
<artifactId>junit</artifactId>
diff --git a/project/CelebornBuild.scala b/project/CelebornBuild.scala
index 5c5218701c..6c9abcd735 100644
--- a/project/CelebornBuild.scala
+++ b/project/CelebornBuild.scala
@@ -57,6 +57,7 @@ object Dependencies {
val junitInterfaceVersion = "0.13.3"
// don't forget update `junitInterfaceVersion` when we upgrade junit
val junitVersion = "4.13.2"
+ val jolVersion = "0.17"
val leveldbJniVersion = "1.8"
val log4j2Version = "2.25.4"
val disruptorVersion = "3.4.4"
@@ -249,6 +250,7 @@ object Dependencies {
// https://www.scala-sbt.org/1.x/docs/Testing.html
val junitInterface = "com.github.sbt" % "junit-interface" %
junitInterfaceVersion
val junit = "junit" % "junit" % junitVersion
+ val jolCore = "org.openjdk.jol" % "jol-core" % jolVersion
val mockitoCore = "org.mockito" % "mockito-core" % mockitoVersion
val mockitoInline = "org.mockito" % "mockito-inline" % mockitoVersion
val scalatestMockito = "org.mockito" %% "mockito-scala-scalatest" %
scalatestMockitoVersion
@@ -706,6 +708,7 @@ object CelebornCommon {
Dependencies.jacksonCore,
Dependencies.jacksonDatabind,
Dependencies.jacksonAnnotations,
+ Dependencies.jolCore % "test",
Dependencies.log4jSlf4jImpl % "test",
Dependencies.log4j12Api % "test",
// SSL support
diff --git
a/service/src/main/scala/org/apache/celeborn/server/common/http/api/v1/ApiUtils.scala
b/service/src/main/scala/org/apache/celeborn/server/common/http/api/v1/ApiUtils.scala
index 9fac3025df..4de6566fa8 100644
---
a/service/src/main/scala/org/apache/celeborn/server/common/http/api/v1/ApiUtils.scala
+++
b/service/src/main/scala/org/apache/celeborn/server/common/http/api/v1/ApiUtils.scala
@@ -159,7 +159,9 @@ object ApiUtils {
case StorageInfo.Type.S3 =>
locationData.storage(StorageEnum.S3)
}
-
Option(partitionLocation.getMapIdBitMap).map(_.toString).foreach(locationData.mapIdBitMap)
+ Option(partitionLocation.getMapIdBitMapIfPresent)
+ .map(_.toString)
+ .foreach(locationData.mapIdBitMap)
locationData
}
}
