This is an automated email from the ASF dual-hosted git repository.
Jackie-Jiang pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/pinot.git
The following commit(s) were added to refs/heads/master by this push:
new 0d624f598f2 Clone IdealState without a serialization round trip
(#19738)
0d624f598f2 is described below
commit 0d624f598f236994c784737dfa6d29f609f21e85
Author: Suvodeep Pyne <[email protected]>
AuthorDate: Fri Oct 16 13:07:53 2026 -0700
Clone IdealState without a serialization round trip (#19738)
---
.../pinot/common/utils/helix/HelixHelper.java | 35 +++++--
.../common/utils/helix/IdealStateSingleCommit.java | 9 +-
.../pinot/common/utils/helix/HelixHelperTest.java | 113 +++++++++++++++++++++
3 files changed, 143 insertions(+), 14 deletions(-)
diff --git
a/pinot-common/src/main/java/org/apache/pinot/common/utils/helix/HelixHelper.java
b/pinot-common/src/main/java/org/apache/pinot/common/utils/helix/HelixHelper.java
index ac102349d0d..6517ffe002b 100644
---
a/pinot-common/src/main/java/org/apache/pinot/common/utils/helix/HelixHelper.java
+++
b/pinot-common/src/main/java/org/apache/pinot/common/utils/helix/HelixHelper.java
@@ -21,10 +21,12 @@ package org.apache.pinot.common.utils.helix;
import com.google.common.base.Preconditions;
import java.util.ArrayList;
import java.util.HashSet;
+import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.Random;
import java.util.Set;
+import java.util.TreeMap;
import java.util.function.Function;
import java.util.function.Supplier;
import java.util.stream.Collectors;
@@ -45,7 +47,6 @@ import org.apache.helix.model.InstanceConfig;
import org.apache.helix.model.builder.HelixConfigScopeBuilder;
import org.apache.helix.store.zk.ZkHelixPropertyStore;
import org.apache.helix.zookeeper.datamodel.ZNRecord;
-import org.apache.helix.zookeeper.datamodel.serializer.ZNRecordSerializer;
import org.apache.pinot.common.helix.ExtraInstanceConfig;
import org.apache.pinot.common.metadata.ZKMetadataProvider;
import org.apache.pinot.common.utils.config.TagNameUtils;
@@ -73,7 +74,6 @@ public class HelixHelper {
RetryPolicies.randomDelayRetryPolicy(20, 100L, 200L);
private static final Logger LOGGER =
LoggerFactory.getLogger(HelixHelper.class);
- private static final ZNRecordSerializer ZN_RECORD_SERIALIZER = new
ZNRecordSerializer();
private static final IdealStateGroupCommit IDEAL_STATE_GROUP_COMMIT = new
IdealStateGroupCommit();
private static final String ONLINE = "ONLINE";
@@ -90,14 +90,37 @@ public class HelixHelper {
return HelixManagerFactory.getZKHelixManager(clusterName, instanceName,
instanceType, zkAddress);
}
+ /// Returns a deep copy of the IdealState with the same content as cloning
it through a serialization round trip of
+ /// its ZNRecord: id, simple fields, map fields, list fields and raw
payload, with the same container types and
+ /// iteration order. Like the round trip, the ZK stat fields (e.g. version)
are not copied. Unlike the round trip, the
+ /// serializer's write policies (list field bound, size limit) are not
applied; they still apply when the IdealState
+ /// is written.
+ ///
+ /// Copying the fields directly avoids serializing (and compressing) the
whole ZNRecord and parsing it back, which
+ /// takes seconds for an IdealState with hundreds of thousands of segments
and widens the window for version
+ /// conflicts when updating it.
public static IdealState cloneIdealState(IdealState idealState) {
- return new IdealState(
- (ZNRecord)
ZN_RECORD_SERIALIZER.deserialize(ZN_RECORD_SERIALIZER.serialize(idealState.getRecord())));
+ ZNRecord record = idealState.getRecord();
+ IdealState copy = new IdealState(record.getId());
+ ZNRecord copyRecord = copy.getRecord();
+ copyRecord.setSimpleFields(new TreeMap<>(record.getSimpleFields()));
+ TreeMap<String, Map<String, String>> mapFields = new
TreeMap<>(record.getMapFields());
+ mapFields.replaceAll((key, value) -> value != null ? new
LinkedHashMap<>(value) : null);
+ copyRecord.setMapFields(mapFields);
+ TreeMap<String, List<String>> listFields = new
TreeMap<>(record.getListFields());
+ listFields.replaceAll((key, value) -> value != null ? new
ArrayList<>(value) : null);
+ copyRecord.setListFields(listFields);
+ byte[] rawPayload = record.getRawPayload();
+ if (rawPayload != null) {
+ copyRecord.setRawPayload(rawPayload.clone());
+ }
+ return copy;
}
/// Returns whether the given record came from the compact parser of
[CompactZNRecordSerializer], that is, its map
- /// fields are an [ImmutableSortedArrayMap]. A record from
[ZNRecordSerializer] returns false, including the records
- /// that the compact serializer delegates to it (plain znodes and the
fallback).
+ /// fields are an [ImmutableSortedArrayMap]. A record from
+ ///
[ZNRecordSerializer][org.apache.helix.zookeeper.datamodel.serializer.ZNRecordSerializer]
returns false, including
+ /// the records that the compact serializer delegates to it (plain znodes
and the fallback).
public static boolean isCompact(ZNRecord record) {
return record.getMapFields() instanceof ImmutableSortedArrayMap;
}
diff --git
a/pinot-common/src/main/java/org/apache/pinot/common/utils/helix/IdealStateSingleCommit.java
b/pinot-common/src/main/java/org/apache/pinot/common/utils/helix/IdealStateSingleCommit.java
index b6eced13267..e7c9b361bb1 100644
---
a/pinot-common/src/main/java/org/apache/pinot/common/utils/helix/IdealStateSingleCommit.java
+++
b/pinot-common/src/main/java/org/apache/pinot/common/utils/helix/IdealStateSingleCommit.java
@@ -27,7 +27,6 @@ import org.apache.helix.HelixManager;
import org.apache.helix.PropertyKey;
import org.apache.helix.model.IdealState;
import org.apache.helix.zookeeper.datamodel.ZNRecord;
-import org.apache.helix.zookeeper.datamodel.serializer.ZNRecordSerializer;
import org.apache.helix.zookeeper.zkclient.exception.ZkBadVersionException;
import org.apache.pinot.common.metrics.ControllerMeter;
import org.apache.pinot.common.metrics.ControllerMetrics;
@@ -42,7 +41,6 @@ public class IdealStateSingleCommit {
private static final Logger LOGGER =
LoggerFactory.getLogger(IdealStateSingleCommit.class);
private static final String ENABLE_COMPRESSIONS_KEY = "enableCompression";
private static final int NUM_PARTITIONS_THRESHOLD_TO_ENABLE_COMPRESSION =
1000;
- private static final ZNRecordSerializer ZN_RECORD_SERIALIZER = new
ZNRecordSerializer();
private IdealStateSingleCommit() {
}
@@ -64,7 +62,7 @@ public class IdealStateSingleCommit {
// Make a copy of the idealState above to pass it to the updater
// NOTE: new IdealState(idealState.getRecord()) does not work
because it's shallow copy for map fields and
// list fields
- IdealState idealStateCopy = cloneIdealState(idealState);
+ IdealState idealStateCopy = HelixHelper.cloneIdealState(idealState);
IdealState updatedIdealState;
try {
@@ -149,9 +147,4 @@ public class IdealStateSingleCommit {
private static class IdealStateWrapper {
IdealState _idealState;
}
-
- public static IdealState cloneIdealState(IdealState idealState) {
- return new IdealState(
- (ZNRecord)
ZN_RECORD_SERIALIZER.deserialize(ZN_RECORD_SERIALIZER.serialize(idealState.getRecord())));
- }
}
diff --git
a/pinot-common/src/test/java/org/apache/pinot/common/utils/helix/HelixHelperTest.java
b/pinot-common/src/test/java/org/apache/pinot/common/utils/helix/HelixHelperTest.java
index dd84cea5143..da248b03ece 100644
---
a/pinot-common/src/test/java/org/apache/pinot/common/utils/helix/HelixHelperTest.java
+++
b/pinot-common/src/test/java/org/apache/pinot/common/utils/helix/HelixHelperTest.java
@@ -18,10 +18,21 @@
*/
package org.apache.pinot.common.utils.helix;
+import com.fasterxml.jackson.databind.ObjectMapper;
+import com.fasterxml.jackson.databind.introspect.BeanPropertyDefinition;
+import java.io.ByteArrayInputStream;
import java.io.IOException;
+import java.lang.reflect.Field;
+import java.lang.reflect.Modifier;
+import java.util.ArrayList;
import java.util.Arrays;
+import java.util.HashSet;
+import java.util.LinkedHashMap;
import java.util.List;
+import java.util.Map;
+import java.util.Set;
import java.util.TreeMap;
+import java.util.stream.Collectors;
import org.apache.helix.HelixProperty;
import org.apache.helix.model.ExternalView;
import org.apache.helix.model.IdealState;
@@ -30,6 +41,7 @@ import org.apache.helix.zookeeper.datamodel.ZNRecord;
import org.apache.helix.zookeeper.datamodel.serializer.ZNRecordSerializer;
import org.apache.helix.zookeeper.util.GZipCompressionUtil;
import org.apache.pinot.spi.utils.CommonConstants;
+import org.apache.pinot.spi.utils.JsonUtils;
import org.testng.annotations.Test;
import static org.testng.Assert.assertEquals;
@@ -191,4 +203,105 @@ public class HelixHelperTest {
record.setModifiedTime(2000L);
record.setEphemeralOwner(3000L);
}
+
+ @Test
+ public void testCloneIdealState() {
+ IdealState idealState = new IdealState("myTable_REALTIME");
+ ZNRecord record = idealState.getRecord();
+ record.setSimpleField("REPLICAS", "2");
+ record.setBooleanField("enableCompression", true);
+ // Insertion order that differs from the sorted order, as in a
deserialized ZNRecord
+ Map<String, Map<String, String>> mapFields = new LinkedHashMap<>();
+ Map<String, String> instanceStateMap = new LinkedHashMap<>();
+ instanceStateMap.put("Server_2", "CONSUMING");
+ instanceStateMap.put("Server_1", "CONSUMING");
+ mapFields.put("segment_2", instanceStateMap);
+ mapFields.put("segment_1", new LinkedHashMap<>(Map.of("Server_1",
"ONLINE")));
+ record.setMapFields(mapFields);
+ record.setListField("segment_2", new ArrayList<>(List.of("Server_2",
"Server_1")));
+ record.setRawPayload(new byte[]{1, 2, 3});
+ record.setVersion(5);
+
+ // Expected result of cloning through a serialization round trip
+ ZNRecordSerializer serializer = new ZNRecordSerializer();
+ ZNRecord expected =
+ new IdealState((ZNRecord)
serializer.deserialize(serializer.serialize(record))).getRecord();
+ ZNRecord copy = HelixHelper.cloneIdealState(idealState).getRecord();
+
+ // Same content and iteration order
+ assertEquals(copy.getId(), expected.getId());
+ assertEquals(copy, expected);
+ assertEquals(new ArrayList<>(copy.getMapFields().keySet()), new
ArrayList<>(expected.getMapFields().keySet()));
+ assertEquals(new ArrayList<>(copy.getMapField("segment_2").keySet()),
+ new ArrayList<>(expected.getMapField("segment_2").keySet()));
+ assertEquals(copy.getRawPayload(), expected.getRawPayload());
+ assertEquals(copy.getVersion(), expected.getVersion());
+
+ // Modifying the copy does not modify the original
+ assertNotSame(copy.getRawPayload(), record.getRawPayload());
+ copy.getSimpleFields().put("REPLICAS", "3");
+ copy.getMapField("segment_2").put("Server_2", "ONLINE");
+ copy.getMapFields().remove("segment_1");
+ copy.getListField("segment_2").add("Server_3");
+ assertEquals(record.getSimpleField("REPLICAS"), "2");
+ assertEquals(record.getMapField("segment_2").get("Server_2"), "CONSUMING");
+ assertTrue(record.getMapFields().containsKey("segment_1"));
+ assertEquals(record.getListField("segment_2"), List.of("Server_2",
"Server_1"));
+ }
+
+ /// cloneIdealState copies a ZNRecord field by field, so it must know every
field an IdealState stores in ZK. If a
+ /// Helix upgrade fails this test, make cloneIdealState copy the new state
(or leave it out if it is not stored, like
+ /// the ZK stat fields), then update the expected values.
+ @Test
+ public void testCloneIdealStateKnowsAllStoredFields()
+ throws IOException {
+ // Instance fields, including inherited ones, sorted
+ assertEquals(getInstanceFieldNames(ZNRecord.class),
+ List.of("_creationTime", "_deltaList", "_ephemeralOwner",
"_modifiedTime", "_serializer", "_version", "id",
+ "listFields", "mapFields", "rawPayload", "simpleFields"));
+ assertEquals(getInstanceFieldNames(IdealState.class), List.of("_record",
"_stat"));
+
+ // Of these, only the following are stored in ZK: the properties the
serializer's mapper can write, whatever their
+ // values, and the properties it writes for a record with all of them set
+ Set<String> storedFields = Set.of("id", "simpleFields", "mapFields",
"listFields", "rawPayload");
+ ObjectMapper mapper = ZNRecordSerializerMapper.get();
+
assertEquals(mapper.getSerializationConfig().introspect(mapper.constructType(ZNRecord.class)).findProperties()
+ .stream()
+ .filter(BeanPropertyDefinition::couldSerialize)
+ .map(BeanPropertyDefinition::getName)
+ .collect(Collectors.toSet()), storedFields);
+ ZNRecord record = new ZNRecord("myTable_REALTIME");
+ record.setSimpleField("REPLICAS", "2");
+ record.setMapField("segment_1", Map.of("Server_1", "ONLINE"));
+ record.setListField("segment_1", List.of("Server_1"));
+ record.setRawPayload(new byte[]{1});
+ byte[] bytes = new ZNRecordSerializer().serialize(record);
+ if (GZipCompressionUtil.isCompressed(bytes)) {
+ bytes = GZipCompressionUtil.uncompress(new ByteArrayInputStream(bytes));
+ }
+ Set<String> writtenFields = new HashSet<>();
+
JsonUtils.bytesToJsonNode(bytes).fieldNames().forEachRemaining(writtenFields::add);
+ assertEquals(writtenFields, storedFields);
+ }
+
+ /// Keeps duplicates, so that a field hiding an inherited one is caught.
+ private static List<String> getInstanceFieldNames(Class<?> clazz) {
+ List<String> fieldNames = new ArrayList<>();
+ for (Class<?> c = clazz; c != Object.class; c = c.getSuperclass()) {
+ for (Field field : c.getDeclaredFields()) {
+ if (!Modifier.isStatic(field.getModifiers()) && !field.isSynthetic()) {
+ fieldNames.add(field.getName());
+ }
+ }
+ }
+ fieldNames.sort(null);
+ return fieldNames;
+ }
+
+ /// Exposes the ObjectMapper that ZNRecordSerializer writes ZNRecords with.
+ private static class ZNRecordSerializerMapper extends ZNRecordSerializer {
+ static ObjectMapper get() {
+ return mapper;
+ }
+ }
}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
