xiangfu0 commented on code in PR #19738:
URL: https://github.com/apache/pinot/pull/19738#discussion_r4278295591
##########
pinot-common/src/test/java/org/apache/pinot/common/utils/helix/HelixHelperTest.java:
##########
@@ -88,4 +106,105 @@ public void testRemoveDisabledPartitions() {
assertTrue(HelixHelper.removeDisabledPartitions(instanceConfig));
assertTrue(instanceConfig.getDisabledPartitionsMap().isEmpty());
}
+
+ @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 {
Review Comment:
This guard test is the right call — pinning both the reflected fields and
the serializer-written properties means a Helix upgrade that changes what a
ZNRecord stores will fail loudly instead of silently dropping a field from the
clone. Reaching the `protected static mapper` via a subclass is a bit coupled
to Helix internals, but since the whole point is to break on a Helix change,
that's acceptable.
##########
pinot-common/src/main/java/org/apache/pinot/common/utils/helix/HelixHelper.java:
##########
@@ -78,9 +78,31 @@ private HelixHelper() {
public static final String BROKER_RESOURCE =
CommonConstants.Helix.BROKER_RESOURCE_INSTANCE;
+ /// 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()));
Review Comment:
Nice — the ordering equivalence here is subtle and you got it right: outer
maps end up as a sorted `TreeMap` (matching the deep-copy the
`IdealState(ZNRecord)` constructor does on the deserialized record) and inner
maps as a source-order `LinkedHashMap`. I verified this against Helix 2.0.1 and
it matches the round trip exactly.
One tiny edge: `new TreeMap<>(...)` would NPE on a null outer key, whereas
the JSON round trip couldn't carry one anyway. Helix segment/field keys are
never null, so this is just a note, not a concern.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
