gortiz commented on code in PR #19786: URL: https://github.com/apache/pinot/pull/19786#discussion_r4280903328
########## pinot-common/src/test/java/org/apache/pinot/common/utils/helix/CompactZNRecordSerializerTest.java: ########## @@ -0,0 +1,411 @@ +/** + * 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.pinot.common.utils.helix; + +import java.io.IOException; +import java.nio.charset.StandardCharsets; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import java.util.Random; +import java.util.SortedMap; +import java.util.TreeMap; +import org.apache.helix.HelixProperty; +import org.apache.helix.model.ExternalView; +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.util.GZipCompressionUtil; +import org.apache.pinot.spi.utils.CommonConstants; +import org.apache.pinot.spi.utils.CommonConstants.Helix.StateModel.SegmentStateModel; +import org.testng.annotations.DataProvider; +import org.testng.annotations.Test; + +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertFalse; +import static org.testng.Assert.assertNotNull; +import static org.testng.Assert.assertNotSame; +import static org.testng.Assert.assertNull; +import static org.testng.Assert.assertSame; +import static org.testng.Assert.assertThrows; +import static org.testng.Assert.assertTrue; + + +/// Checks that [CompactZNRecordSerializer] returns records equal to those of [ZNRecordSerializer], with interned +/// names, shared instance-state maps and immutable map fields. +public class CompactZNRecordSerializerTest { + private static final ZNRecordSerializer DEFAULT_SERIALIZER = new ZNRecordSerializer(); + private static final CompactZNRecordSerializer COMPACT_SERIALIZER = new CompactZNRecordSerializer(); + + @DataProvider + public Object[][] compression() { + return new Object[][]{{false}, {true}}; + } + + private static byte[] bytes(ZNRecord record, boolean compressed) + throws IOException { + byte[] bytes = DEFAULT_SERIALIZER.serialize(record); + if (compressed && !GZipCompressionUtil.isCompressed(bytes)) { + bytes = GZipCompressionUtil.compress(bytes); + } + assertEquals(GZipCompressionUtil.isCompressed(bytes), compressed); + return bytes; + } + + private static ZNRecord assertSameAsDefault(byte[] bytes) + throws IOException { + ZNRecord expected = (ZNRecord) DEFAULT_SERIALIZER.deserialize(bytes); + // Call the compact parser directly so that a parse failure cannot hide behind the fallback + ZNRecord actual = CompactZNRecordSerializer.deserializeCompact(bytes); + assertEquals(COMPACT_SERIALIZER.deserialize(bytes), actual); + assertNotNull(expected); + assertNotNull(actual); + assertEquals(actual.getId(), expected.getId()); + assertEquals(actual, expected); + assertEquals(actual.getSimpleFields(), expected.getSimpleFields()); + assertEquals(actual.getListFields(), expected.getListFields()); + assertEquals(actual.getMapFields(), expected.getMapFields()); + assertEquals(actual.getRawPayload(), expected.getRawPayload()); + // Same iteration order for simple and list fields; map fields are sorted + assertEquals(new ArrayList<>(actual.getSimpleFields().keySet()), + new ArrayList<>(expected.getSimpleFields().keySet())); + assertEquals(new ArrayList<>(actual.getListFields().keySet()), new ArrayList<>(expected.getListFields().keySet())); + assertEquals(new ArrayList<>(actual.getMapFields().keySet()), + new ArrayList<>(new TreeMap<>(expected.getMapFields()).keySet())); + return actual; + } + + private static ZNRecord sampleIdealState() { + ZNRecord record = new ZNRecord("myTable_REALTIME"); + record.setSimpleField("IDEAL_STATE_MODE", "CUSTOMIZED"); + record.setSimpleField("REBALANCE_MODE", "CUSTOMIZED"); + record.setSimpleField("HELIX_ENABLED", "true"); + record.setSimpleField(CommonConstants.IdealState.HYBRID_TABLE_TIME_BOUNDARY, "1700000000000"); + record.setListField("emptyList", new ArrayList<>()); + record.setListField("myTable__0__1__20240101T0000Z", new ArrayList<>(List.of("Server_b_1", "Server_a_1"))); + for (int i = 0; i < 50; i++) { + Map<String, String> stateMap = new LinkedHashMap<>(); + // Unsorted on purpose + stateMap.put("Server_" + (i % 3 + 5) + "_8098", SegmentStateModel.ONLINE); + stateMap.put("Server_" + (i % 3 + 1) + "_8098", i % 7 == 0 ? SegmentStateModel.CONSUMING : "ONLINE"); + record.setMapField("myTable__" + (49 - i) + "__0__20240101T0000Z", stateMap); + } + record.setMapField("myTable__unknown", Map.of("Server_1_8098", "SOME_NEW_STATE", "Server_2_8098", "ERROR")); + record.setMapField("myTable__empty", Map.of()); + return record; + } + + @Test(dataProvider = "compression") + public void testRoundTripIdealState(boolean compressed) + throws IOException { + ZNRecord actual = assertSameAsDefault(bytes(sampleIdealState(), compressed)); + + // IdealState and ExternalView wrappers work on the compact record + IdealState idealState = new IdealState(actual); + assertTrue(idealState.isEnabled()); + assertEquals(idealState.getRecord().getSimpleField(CommonConstants.IdealState.HYBRID_TABLE_TIME_BOUNDARY), + "1700000000000"); + assertEquals(idealState.getPartitionSet().size(), 52); + assertEquals(idealState.getInstanceStateMap("myTable__unknown").get("Server_1_8098"), "SOME_NEW_STATE"); + ExternalView externalView = new ExternalView(actual); + assertEquals(externalView.getStateMap("myTable__empty"), Map.of()); + assertTrue(externalView.getStateMap("myTable__0__0__20240101T0000Z") instanceof SortedMap); + } + + @Test(dataProvider = "compression") + public void testRoundTripEmptyAndMissingFields(boolean compressed) + throws IOException { + assertSameAsDefault(bytes(new ZNRecord("empty"), compressed)); + ZNRecord withPayload = new ZNRecord("payload"); + withPayload.setRawPayload(new byte[]{1, 2, 3, 4, 5}); + assertSameAsDefault(bytes(withPayload, compressed)); + + // Hand-written JSON: missing fields, unknown properties, null and non-string scalar values, duplicate keys + String json = "{\"id\":\"t_OFFLINE\",\"unknown\":{\"a\":[1,{\"b\":null}]}," + + "\"simpleFields\":{\"n\":12,\"b\":true,\"nul\":null,\"dup\":\"1\",\"dup\":\"2\"}," + + "\"mapFields\":{\"s2\":{\"z\":\"ONLINE\",\"a\":\"OFFLINE\",\"a\":\"ONLINE\"},\"s1\":{\"x\":null,\"y\":7}," + + "\"s0\":null,\"s2\":{\"q\":\"DROPPED\"}},\"listFields\":{\"l\":[\"x\",null,\"y\"],\"nl\":null}}"; + byte[] bytes = json.getBytes(StandardCharsets.UTF_8); + assertSameAsDefault(bytes); + assertSameAsDefault(GZipCompressionUtil.compress(bytes)); + // Only the id + assertSameAsDefault("{\"id\":\"only\"}".getBytes(StandardCharsets.UTF_8)); + } + + @Test + public void testEscapesAndTrailingContent() + throws IOException { + // JSON escapes in field names, map keys and values (quote, backslash, control character, unicode) + String json = "{\"id\":\"t\\\"1\",\"simpleFields\":{\"k\\\\\":\"v\\n\\u00e9\"}," + + "\"mapFields\":{\"seg\\u0041\":{\"Server\\\"1\":\"ONLINE\",\"Server_2\":\"ON\\u004cINE\"}}," + + "\"listFields\":{\"l\":[\"Server\\\"1\",\"\\u2603\"]}}"; + byte[] bytes = json.getBytes(StandardCharsets.UTF_8); + ZNRecord record = assertSameAsDefault(bytes); + assertSameAsDefault(GZipCompressionUtil.compress(bytes)); + assertEquals(record.getId(), "t\"1"); + assertEquals(record.getMapField("segA"), Map.of("Server\"1", "ONLINE", "Server_2", "ONLINE")); + // An escaped known state still maps to the shared constant + assertSame(record.getMapField("segA").get("Server_2"), SegmentStateModel.ONLINE); + + // Content after the record is ignored, as ObjectMapper.readValue does + byte[] trailing = "{\"id\":\"x\",\"mapFields\":{\"s\":{\"a\":\"ONLINE\"}}} {\"id\":\"y\"}" + .getBytes(StandardCharsets.UTF_8); + assertEquals(assertSameAsDefault(trailing).getId(), "x"); + assertEquals(assertSameAsDefault(GZipCompressionUtil.compress(trailing)).getId(), "x"); + } + + @Test + public void testRoundTripRandomized() + throws IOException { + Random random = new Random(42); + for (int iteration = 0; iteration < 20; iteration++) { + ZNRecord record = new ZNRecord("table_" + iteration); + int numSegments = random.nextInt(200); + for (int i = 0; i < numSegments; i++) { + Map<String, String> stateMap = new LinkedHashMap<>(); + int numReplicas = random.nextInt(5); + for (int r = 0; r < numReplicas; r++) { + String[] states = {"ONLINE", "OFFLINE", "CONSUMING", "ERROR", "OTHER"}; + stateMap.put("Server_" + random.nextInt(10), states[random.nextInt(states.length)]); + } + record.setMapField("segment_" + random.nextInt(1000), stateMap); + } + assertSameAsDefault(bytes(record, random.nextBoolean())); + } + } + + @Test + public void testInterningAndSharing() + throws IOException { + ZNRecord actual = (ZNRecord) COMPACT_SERIALIZER.deserialize(bytes(sampleIdealState(), true)); + Map<String, Map<String, String>> mapFields = actual.getMapFields(); + assertTrue(mapFields instanceof ImmutableSortedArrayMap); + for (Map.Entry<String, Map<String, String>> entry : mapFields.entrySet()) { + assertSame(entry.getKey(), entry.getKey().intern()); + Map<String, String> stateMap = entry.getValue(); + assertTrue(stateMap instanceof ImmutableSortedArrayMap); + for (Map.Entry<String, String> stateEntry : stateMap.entrySet()) { + assertSame(stateEntry.getKey(), stateEntry.getKey().intern()); + assertSame(stateEntry.getValue(), stateEntry.getValue().intern()); + } + } + // Segment 49 - i: i = 48 is ONLINE on Server_5, i = 42 is CONSUMING on Server_1 + assertSame(mapFields.get("myTable__1__0__20240101T0000Z").get("Server_5_8098"), SegmentStateModel.ONLINE); + assertSame(mapFields.get("myTable__7__0__20240101T0000Z").get("Server_1_8098"), SegmentStateModel.CONSUMING); + for (String instance : actual.getListField("myTable__0__1__20240101T0000Z")) { + assertSame(instance, instance.intern()); + } + + // Equal instance-state maps share one instance; different ones do not + // Segment 49 - i: i = 1 and i = 4 have the same servers and both ONLINE + Map<String, String> stateMap48 = mapFields.get("myTable__48__0__20240101T0000Z"); + Map<String, String> stateMap45 = mapFields.get("myTable__45__0__20240101T0000Z"); + assertEquals(stateMap48, stateMap45); + assertSame(stateMap48, stateMap45); + // i = 0 is CONSUMING on Server_1, i = 3 is ONLINE on Server_1 + assertNotSame(mapFields.get("myTable__49__0__20240101T0000Z"), mapFields.get("myTable__46__0__20240101T0000Z")); + + // Map fields are immutable, scalar fields are not + assertThrows(UnsupportedOperationException.class, () -> actual.setMapField("x", Map.of())); + assertThrows(UnsupportedOperationException.class, () -> stateMap48.put("x", "ONLINE")); + actual.setVersion(7); + actual.setCreationTime(1L); + actual.setModifiedTime(2L); + actual.setEphemeralOwner(3L); + actual.setSimpleField("new", "value"); + assertEquals(actual.getVersion(), 7); + + // HelixProperty copies the outer map into a TreeMap and keeps the inner maps + IdealState idealState = new IdealState(actual); + assertSame(idealState.getInstanceStateMap("myTable__48__0__20240101T0000Z"), stateMap48); + assertEquals(idealState.getRecord().getMapFields(), mapFields); + } + + @Test + public void testSortInPlace() { + String[] keys = {"c", "a", "b", "a", "d", null}; + Object[] values = {"1", "2", "3", "4", "5", null}; + int size = CompactZNRecordSerializer.sortInPlace(keys, values, 5); + assertEquals(size, 4); + assertEquals(Arrays.copyOf(keys, size), new String[]{"a", "b", "c", "d"}); + assertEquals(Arrays.copyOf(values, size), new Object[]{"4", "3", "1", "5"}); + } + + @Test + public void testMalformedInputFallsBack() + throws IOException { + // Both deserializers return null for malformed JSON + byte[] malformed = "{\"id\":\"x\",\"mapFields\":{\"s\":[1]}}".getBytes(StandardCharsets.UTF_8); + assertNull(DEFAULT_SERIALIZER.deserialize(malformed)); + assertNull(COMPACT_SERIALIZER.deserialize(malformed)); + // Gzip-compressed: the compact parser fails and falls back to the default deserializer + byte[] compressedMalformed = GZipCompressionUtil.compress(malformed); + assertThrows(IOException.class, () -> CompactZNRecordSerializer.deserializeCompact(compressedMalformed)); + assertNull(DEFAULT_SERIALIZER.deserialize(compressedMalformed)); + assertNull(COMPACT_SERIALIZER.deserialize(compressedMalformed)); Review Comment: Added the test. A gzip record with `"rawPayload":[1,2,3]` fails the direct compact parse, and `deserialize()` falls back to the same record that `ZNRecordSerializer` returns. ########## pinot-common/src/test/java/org/apache/pinot/common/utils/helix/CompactZNRecordSerializerTest.java: ########## @@ -0,0 +1,411 @@ +/** + * 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.pinot.common.utils.helix; + +import java.io.IOException; +import java.nio.charset.StandardCharsets; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import java.util.Random; +import java.util.SortedMap; +import java.util.TreeMap; +import org.apache.helix.HelixProperty; +import org.apache.helix.model.ExternalView; +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.util.GZipCompressionUtil; +import org.apache.pinot.spi.utils.CommonConstants; +import org.apache.pinot.spi.utils.CommonConstants.Helix.StateModel.SegmentStateModel; +import org.testng.annotations.DataProvider; +import org.testng.annotations.Test; + +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertFalse; +import static org.testng.Assert.assertNotNull; +import static org.testng.Assert.assertNotSame; +import static org.testng.Assert.assertNull; +import static org.testng.Assert.assertSame; +import static org.testng.Assert.assertThrows; +import static org.testng.Assert.assertTrue; + + +/// Checks that [CompactZNRecordSerializer] returns records equal to those of [ZNRecordSerializer], with interned +/// names, shared instance-state maps and immutable map fields. +public class CompactZNRecordSerializerTest { + private static final ZNRecordSerializer DEFAULT_SERIALIZER = new ZNRecordSerializer(); + private static final CompactZNRecordSerializer COMPACT_SERIALIZER = new CompactZNRecordSerializer(); + + @DataProvider + public Object[][] compression() { + return new Object[][]{{false}, {true}}; + } + + private static byte[] bytes(ZNRecord record, boolean compressed) + throws IOException { + byte[] bytes = DEFAULT_SERIALIZER.serialize(record); + if (compressed && !GZipCompressionUtil.isCompressed(bytes)) { + bytes = GZipCompressionUtil.compress(bytes); + } + assertEquals(GZipCompressionUtil.isCompressed(bytes), compressed); + return bytes; + } + + private static ZNRecord assertSameAsDefault(byte[] bytes) + throws IOException { + ZNRecord expected = (ZNRecord) DEFAULT_SERIALIZER.deserialize(bytes); + // Call the compact parser directly so that a parse failure cannot hide behind the fallback + ZNRecord actual = CompactZNRecordSerializer.deserializeCompact(bytes); + assertEquals(COMPACT_SERIALIZER.deserialize(bytes), actual); + assertNotNull(expected); + assertNotNull(actual); + assertEquals(actual.getId(), expected.getId()); + assertEquals(actual, expected); + assertEquals(actual.getSimpleFields(), expected.getSimpleFields()); + assertEquals(actual.getListFields(), expected.getListFields()); + assertEquals(actual.getMapFields(), expected.getMapFields()); + assertEquals(actual.getRawPayload(), expected.getRawPayload()); + // Same iteration order for simple and list fields; map fields are sorted + assertEquals(new ArrayList<>(actual.getSimpleFields().keySet()), + new ArrayList<>(expected.getSimpleFields().keySet())); + assertEquals(new ArrayList<>(actual.getListFields().keySet()), new ArrayList<>(expected.getListFields().keySet())); + assertEquals(new ArrayList<>(actual.getMapFields().keySet()), + new ArrayList<>(new TreeMap<>(expected.getMapFields()).keySet())); + return actual; + } + + private static ZNRecord sampleIdealState() { + ZNRecord record = new ZNRecord("myTable_REALTIME"); + record.setSimpleField("IDEAL_STATE_MODE", "CUSTOMIZED"); + record.setSimpleField("REBALANCE_MODE", "CUSTOMIZED"); + record.setSimpleField("HELIX_ENABLED", "true"); + record.setSimpleField(CommonConstants.IdealState.HYBRID_TABLE_TIME_BOUNDARY, "1700000000000"); + record.setListField("emptyList", new ArrayList<>()); + record.setListField("myTable__0__1__20240101T0000Z", new ArrayList<>(List.of("Server_b_1", "Server_a_1"))); + for (int i = 0; i < 50; i++) { + Map<String, String> stateMap = new LinkedHashMap<>(); + // Unsorted on purpose + stateMap.put("Server_" + (i % 3 + 5) + "_8098", SegmentStateModel.ONLINE); + stateMap.put("Server_" + (i % 3 + 1) + "_8098", i % 7 == 0 ? SegmentStateModel.CONSUMING : "ONLINE"); + record.setMapField("myTable__" + (49 - i) + "__0__20240101T0000Z", stateMap); + } + record.setMapField("myTable__unknown", Map.of("Server_1_8098", "SOME_NEW_STATE", "Server_2_8098", "ERROR")); + record.setMapField("myTable__empty", Map.of()); + return record; + } + + @Test(dataProvider = "compression") + public void testRoundTripIdealState(boolean compressed) + throws IOException { + ZNRecord actual = assertSameAsDefault(bytes(sampleIdealState(), compressed)); + + // IdealState and ExternalView wrappers work on the compact record + IdealState idealState = new IdealState(actual); + assertTrue(idealState.isEnabled()); + assertEquals(idealState.getRecord().getSimpleField(CommonConstants.IdealState.HYBRID_TABLE_TIME_BOUNDARY), + "1700000000000"); + assertEquals(idealState.getPartitionSet().size(), 52); + assertEquals(idealState.getInstanceStateMap("myTable__unknown").get("Server_1_8098"), "SOME_NEW_STATE"); + ExternalView externalView = new ExternalView(actual); + assertEquals(externalView.getStateMap("myTable__empty"), Map.of()); + assertTrue(externalView.getStateMap("myTable__0__0__20240101T0000Z") instanceof SortedMap); + } + + @Test(dataProvider = "compression") + public void testRoundTripEmptyAndMissingFields(boolean compressed) + throws IOException { + assertSameAsDefault(bytes(new ZNRecord("empty"), compressed)); + ZNRecord withPayload = new ZNRecord("payload"); + withPayload.setRawPayload(new byte[]{1, 2, 3, 4, 5}); + assertSameAsDefault(bytes(withPayload, compressed)); + + // Hand-written JSON: missing fields, unknown properties, null and non-string scalar values, duplicate keys + String json = "{\"id\":\"t_OFFLINE\",\"unknown\":{\"a\":[1,{\"b\":null}]}," + + "\"simpleFields\":{\"n\":12,\"b\":true,\"nul\":null,\"dup\":\"1\",\"dup\":\"2\"}," + + "\"mapFields\":{\"s2\":{\"z\":\"ONLINE\",\"a\":\"OFFLINE\",\"a\":\"ONLINE\"},\"s1\":{\"x\":null,\"y\":7}," + + "\"s0\":null,\"s2\":{\"q\":\"DROPPED\"}},\"listFields\":{\"l\":[\"x\",null,\"y\"],\"nl\":null}}"; + byte[] bytes = json.getBytes(StandardCharsets.UTF_8); + assertSameAsDefault(bytes); + assertSameAsDefault(GZipCompressionUtil.compress(bytes)); + // Only the id + assertSameAsDefault("{\"id\":\"only\"}".getBytes(StandardCharsets.UTF_8)); + } + + @Test + public void testEscapesAndTrailingContent() + throws IOException { + // JSON escapes in field names, map keys and values (quote, backslash, control character, unicode) + String json = "{\"id\":\"t\\\"1\",\"simpleFields\":{\"k\\\\\":\"v\\n\\u00e9\"}," + + "\"mapFields\":{\"seg\\u0041\":{\"Server\\\"1\":\"ONLINE\",\"Server_2\":\"ON\\u004cINE\"}}," + + "\"listFields\":{\"l\":[\"Server\\\"1\",\"\\u2603\"]}}"; + byte[] bytes = json.getBytes(StandardCharsets.UTF_8); + ZNRecord record = assertSameAsDefault(bytes); + assertSameAsDefault(GZipCompressionUtil.compress(bytes)); + assertEquals(record.getId(), "t\"1"); + assertEquals(record.getMapField("segA"), Map.of("Server\"1", "ONLINE", "Server_2", "ONLINE")); + // An escaped known state still maps to the shared constant + assertSame(record.getMapField("segA").get("Server_2"), SegmentStateModel.ONLINE); + + // Content after the record is ignored, as ObjectMapper.readValue does + byte[] trailing = "{\"id\":\"x\",\"mapFields\":{\"s\":{\"a\":\"ONLINE\"}}} {\"id\":\"y\"}" + .getBytes(StandardCharsets.UTF_8); + assertEquals(assertSameAsDefault(trailing).getId(), "x"); + assertEquals(assertSameAsDefault(GZipCompressionUtil.compress(trailing)).getId(), "x"); + } + + @Test + public void testRoundTripRandomized() + throws IOException { + Random random = new Random(42); + for (int iteration = 0; iteration < 20; iteration++) { + ZNRecord record = new ZNRecord("table_" + iteration); + int numSegments = random.nextInt(200); + for (int i = 0; i < numSegments; i++) { + Map<String, String> stateMap = new LinkedHashMap<>(); + int numReplicas = random.nextInt(5); + for (int r = 0; r < numReplicas; r++) { + String[] states = {"ONLINE", "OFFLINE", "CONSUMING", "ERROR", "OTHER"}; + stateMap.put("Server_" + random.nextInt(10), states[random.nextInt(states.length)]); + } + record.setMapField("segment_" + random.nextInt(1000), stateMap); + } + assertSameAsDefault(bytes(record, random.nextBoolean())); + } + } + + @Test + public void testInterningAndSharing() + throws IOException { + ZNRecord actual = (ZNRecord) COMPACT_SERIALIZER.deserialize(bytes(sampleIdealState(), true)); + Map<String, Map<String, String>> mapFields = actual.getMapFields(); + assertTrue(mapFields instanceof ImmutableSortedArrayMap); + for (Map.Entry<String, Map<String, String>> entry : mapFields.entrySet()) { + assertSame(entry.getKey(), entry.getKey().intern()); + Map<String, String> stateMap = entry.getValue(); + assertTrue(stateMap instanceof ImmutableSortedArrayMap); + for (Map.Entry<String, String> stateEntry : stateMap.entrySet()) { + assertSame(stateEntry.getKey(), stateEntry.getKey().intern()); + assertSame(stateEntry.getValue(), stateEntry.getValue().intern()); + } + } + // Segment 49 - i: i = 48 is ONLINE on Server_5, i = 42 is CONSUMING on Server_1 + assertSame(mapFields.get("myTable__1__0__20240101T0000Z").get("Server_5_8098"), SegmentStateModel.ONLINE); + assertSame(mapFields.get("myTable__7__0__20240101T0000Z").get("Server_1_8098"), SegmentStateModel.CONSUMING); + for (String instance : actual.getListField("myTable__0__1__20240101T0000Z")) { + assertSame(instance, instance.intern()); + } + + // Equal instance-state maps share one instance; different ones do not + // Segment 49 - i: i = 1 and i = 4 have the same servers and both ONLINE + Map<String, String> stateMap48 = mapFields.get("myTable__48__0__20240101T0000Z"); + Map<String, String> stateMap45 = mapFields.get("myTable__45__0__20240101T0000Z"); + assertEquals(stateMap48, stateMap45); + assertSame(stateMap48, stateMap45); + // i = 0 is CONSUMING on Server_1, i = 3 is ONLINE on Server_1 + assertNotSame(mapFields.get("myTable__49__0__20240101T0000Z"), mapFields.get("myTable__46__0__20240101T0000Z")); + + // Map fields are immutable, scalar fields are not + assertThrows(UnsupportedOperationException.class, () -> actual.setMapField("x", Map.of())); + assertThrows(UnsupportedOperationException.class, () -> stateMap48.put("x", "ONLINE")); + actual.setVersion(7); + actual.setCreationTime(1L); + actual.setModifiedTime(2L); + actual.setEphemeralOwner(3L); + actual.setSimpleField("new", "value"); + assertEquals(actual.getVersion(), 7); + + // HelixProperty copies the outer map into a TreeMap and keeps the inner maps + IdealState idealState = new IdealState(actual); + assertSame(idealState.getInstanceStateMap("myTable__48__0__20240101T0000Z"), stateMap48); + assertEquals(idealState.getRecord().getMapFields(), mapFields); + } + + @Test + public void testSortInPlace() { + String[] keys = {"c", "a", "b", "a", "d", null}; + Object[] values = {"1", "2", "3", "4", "5", null}; + int size = CompactZNRecordSerializer.sortInPlace(keys, values, 5); + assertEquals(size, 4); + assertEquals(Arrays.copyOf(keys, size), new String[]{"a", "b", "c", "d"}); + assertEquals(Arrays.copyOf(values, size), new Object[]{"4", "3", "1", "5"}); + } + + @Test + public void testMalformedInputFallsBack() + throws IOException { + // Both deserializers return null for malformed JSON + byte[] malformed = "{\"id\":\"x\",\"mapFields\":{\"s\":[1]}}".getBytes(StandardCharsets.UTF_8); + assertNull(DEFAULT_SERIALIZER.deserialize(malformed)); + assertNull(COMPACT_SERIALIZER.deserialize(malformed)); + // Gzip-compressed: the compact parser fails and falls back to the default deserializer + byte[] compressedMalformed = GZipCompressionUtil.compress(malformed); + assertThrows(IOException.class, () -> CompactZNRecordSerializer.deserializeCompact(compressedMalformed)); + assertNull(DEFAULT_SERIALIZER.deserialize(compressedMalformed)); + assertNull(COMPACT_SERIALIZER.deserialize(compressedMalformed)); + assertNull(COMPACT_SERIALIZER.deserialize("not json".getBytes(StandardCharsets.UTF_8))); + assertNull(COMPACT_SERIALIZER.deserialize(new byte[0])); + assertNull(COMPACT_SERIALIZER.deserialize(null)); + } + + @Test + public void testCorruptGzipTrailerFallsBack() + throws IOException { + byte[] bytes = bytes(sampleIdealState(), true); + assertNotNull(CompactZNRecordSerializer.deserializeCompact(bytes)); + // The gzip trailer is the CRC-32 (4 bytes) and the uncompressed size (4 bytes). The deflate data stays valid. + byte[] badCrc = bytes.clone(); + badCrc[badCrc.length - 8] ^= 0x01; + assertThrows(IOException.class, () -> CompactZNRecordSerializer.deserializeCompact(badCrc)); + assertNull(DEFAULT_SERIALIZER.deserialize(badCrc)); + assertNull(COMPACT_SERIALIZER.deserialize(badCrc)); + byte[] badSize = bytes.clone(); + badSize[badSize.length - 1] ^= 0x01; + assertThrows(IOException.class, () -> CompactZNRecordSerializer.deserializeCompact(badSize)); + assertNull(COMPACT_SERIALIZER.deserialize(badSize)); + } + + @Test(dataProvider = "compression") + public void testOnlyGzipUsesCompactParser(boolean compressed) + throws IOException { + byte[] bytes = bytes(sampleIdealState(), compressed); + ZNRecord expected = (ZNRecord) DEFAULT_SERIALIZER.deserialize(bytes); + ZNRecord actual = (ZNRecord) COMPACT_SERIALIZER.deserialize(bytes); + assertNotNull(actual); + assertEquals(actual, expected); + assertEquals(actual.getMapFields(), expected.getMapFields()); + assertEquals(HelixHelper.isCompact(actual), compressed); + Map<String, String> stateMap = actual.getMapField("myTable__0__0__20240101T0000Z"); + if (compressed) { + assertTrue(actual.getMapFields() instanceof ImmutableSortedArrayMap); + assertTrue(stateMap instanceof ImmutableSortedArrayMap); + } else { + // A plain znode goes to the default deserializer: same map types, mutable maps + assertEquals(actual.getMapFields().getClass(), expected.getMapFields().getClass()); + assertEquals(stateMap.getClass(), expected.getMapField("myTable__0__0__20240101T0000Z").getClass()); + assertFalse(stateMap instanceof ImmutableSortedArrayMap); + stateMap.put("Server_new_8098", SegmentStateModel.ONLINE); + actual.setMapField("newSegment", new TreeMap<>()); + assertEquals(actual.getMapFields().size(), expected.getMapFields().size() + 1); + } + } + + @Test + public void testSerializeDelegates() + throws IOException { + ZNRecord record = sampleIdealState(); + byte[] expected = DEFAULT_SERIALIZER.serialize(record); + assertEquals(COMPACT_SERIALIZER.serialize(record), expected); + // A compact record serializes to the same bytes as an equal record from the default deserializer + ZNRecord compact = CompactZNRecordSerializer.deserializeCompact(expected); + ZNRecord standard = (ZNRecord) DEFAULT_SERIALIZER.deserialize(expected); + assertTrue(HelixHelper.isCompact(compact)); + assertEquals(DEFAULT_SERIALIZER.deserialize(COMPACT_SERIALIZER.serialize(compact)), standard); + } + + @Test(dataProvider = "compression") + public void testToIdealStateAndExternalViewShareCompactRecord(boolean compressed) Review Comment: Moved both tests into `HelixHelperTest`. They still use the compact serializer for fixtures. ########## pinot-common/src/test/java/org/apache/pinot/common/utils/helix/ImmutableSortedArrayMapTest.java: ########## @@ -0,0 +1,207 @@ +/** + * 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.pinot.common.utils.helix; + +import java.util.AbstractMap; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.HashSet; +import java.util.Iterator; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import java.util.NoSuchElementException; +import java.util.Set; +import java.util.SortedMap; +import java.util.TreeMap; +import org.testng.annotations.Test; + +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertFalse; +import static org.testng.Assert.assertNull; +import static org.testng.Assert.assertSame; +import static org.testng.Assert.assertThrows; +import static org.testng.Assert.assertTrue; + + +/// Checks that [ImmutableSortedArrayMap] behaves like a [TreeMap] for reads and rejects every mutation. +public class ImmutableSortedArrayMapTest { + + private static Map<String, String> unsortedSample() { + Map<String, String> map = new LinkedHashMap<>(); + map.put("Server_c_8098", "ONLINE"); + map.put("Server_a_8098", "CONSUMING"); + map.put("Server_b_8098", "OFFLINE"); + map.put("Server_d_8098", null); + return map; + } + + @Test + public void testReadsMatchTreeMap() { + Map<String, String> source = unsortedSample(); + TreeMap<String, String> expected = new TreeMap<>(source); + ImmutableSortedArrayMap<String> map = ImmutableSortedArrayMap.copyOf(source); + + assertEquals(map.size(), expected.size()); + assertFalse(map.isEmpty()); + assertNull(map.comparator()); + assertEquals(new ArrayList<>(map.keySet()), new ArrayList<>(expected.keySet())); + assertEquals(new ArrayList<>(map.values()), new ArrayList<>(expected.values())); + assertEquals(new ArrayList<>(map.entrySet()), new ArrayList<>(expected.entrySet())); + assertEquals(map.toString(), expected.toString()); + assertEquals(map.firstKey(), expected.firstKey()); + assertEquals(map.lastKey(), expected.lastKey()); + for (String key : List.of("Server_a_8098", "Server_d_8098", "Server_e_8098", "")) { + assertEquals(map.get(key), expected.get(key), key); + assertEquals(map.containsKey(key), expected.containsKey(key), key); + assertEquals(map.getOrDefault(key, "x"), expected.getOrDefault(key, "x"), key); + } + // Lookups with a non-interned equal key and with keys of another type + assertEquals(map.get(new String("Server_b_8098")), "OFFLINE"); + assertNull(map.get(42)); + assertFalse(map.containsKey(null)); + assertNull(map.get(null)); + assertTrue(map.containsValue("ONLINE")); + assertTrue(map.containsValue(null)); + assertFalse(map.containsValue("ERROR")); + assertFalse(map.containsValue(42)); + assertTrue(map.entrySet().contains(new AbstractMap.SimpleEntry<>("Server_a_8098", "CONSUMING"))); + assertFalse(map.entrySet().contains(new AbstractMap.SimpleEntry<>("Server_a_8098", "ONLINE"))); + assertTrue(map.values().contains(null)); + } + + @Test + public void testEqualsAndHashCodeParity() { + Map<String, String> source = unsortedSample(); + TreeMap<String, String> treeMap = new TreeMap<>(source); + HashMap<String, String> hashMap = new HashMap<>(source); + ImmutableSortedArrayMap<String> map = ImmutableSortedArrayMap.copyOf(source); + + assertEquals(map, treeMap); + assertEquals(treeMap, map); + assertEquals(map, hashMap); + assertEquals(hashMap, map); + assertEquals(map.hashCode(), treeMap.hashCode()); + assertEquals(map.hashCode(), hashMap.hashCode()); + assertEquals(map, ImmutableSortedArrayMap.copyOf(hashMap)); + + Set<String> hashSet = new HashSet<>(source.keySet()); + assertEquals(map.keySet(), hashSet); + assertEquals(hashSet, map.keySet()); + assertEquals(map.keySet(), treeMap.keySet()); + assertEquals(map.keySet().hashCode(), hashSet.hashCode()); + assertEquals(map.entrySet(), treeMap.entrySet()); + assertEquals(treeMap.entrySet(), map.entrySet()); + assertEquals(map.entrySet().hashCode(), treeMap.entrySet().hashCode()); + + // A HashMap keyed by key sets works across implementations, as in ReplicaGroupInstanceSelector + Map<Set<String>, String> byKeySet = new HashMap<>(); + byKeySet.put(treeMap.keySet(), "value"); + assertEquals(byKeySet.get(map.keySet()), "value"); + byKeySet.clear(); + byKeySet.put(map.keySet(), "value"); + assertEquals(byKeySet.get(hashSet), "value"); + + Map<String, String> different = new TreeMap<>(source); + different.put("Server_a_8098", "ONLINE"); + assertFalse(map.equals(different)); + assertFalse(different.equals(map)); + assertFalse(map.equals(ImmutableSortedArrayMap.copyOf(different))); + } + + @Test + public void testSortedViews() { + TreeMap<String, String> expected = new TreeMap<>(unsortedSample()); + ImmutableSortedArrayMap<String> map = ImmutableSortedArrayMap.copyOf(expected); + for (String from : List.of("", "Server_a_8098", "Server_b", "Server_c_8098", "Z")) { + assertEquals(map.tailMap(from), expected.tailMap(from), from); + assertEquals(map.headMap(from), expected.headMap(from), from); + for (String to : List.of("Server_b", "Server_c_8098", "Z")) { + if (from.compareTo(to) <= 0) { + SortedMap<String, String> subMap = map.subMap(from, to); + assertEquals(subMap, expected.subMap(from, to), from + "-" + to); + assertEquals(new ArrayList<>(subMap.keySet()), new ArrayList<>(expected.subMap(from, to).keySet())); + } + } + } + assertThrows(IllegalArgumentException.class, () -> map.subMap("b", "a")); + assertThrows(NullPointerException.class, () -> map.headMap(null)); + assertThrows(NoSuchElementException.class, () -> ImmutableSortedArrayMap.empty().firstKey()); + assertThrows(NoSuchElementException.class, () -> ImmutableSortedArrayMap.empty().lastKey()); + } + + @Test + public void testDuplicateKeysLastWins() { + String[] keys = {"b", "a", "b", "c", "a"}; + Object[] values = {"1", "2", "3", "4", "5"}; + ImmutableSortedArrayMap<String> map = ImmutableSortedArrayMap.sortAndWrap(keys, values, keys.length); + assertEquals(map, Map.of("a", "5", "b", "3", "c", "4")); + assertEquals(new ArrayList<>(map.keySet()), List.of("a", "b", "c")); + + // Sorted input with a trailing capacity is trimmed + String[] sortedKeys = {"a", "b", null, null}; + Object[] sortedValues = {"1", "2", null, null}; + assertEquals(ImmutableSortedArrayMap.sortAndWrap(sortedKeys, sortedValues, 2), Map.of("a", "1", "b", "2")); + assertSame(ImmutableSortedArrayMap.sortAndWrap(new String[0], new Object[0], 0), ImmutableSortedArrayMap.empty()); + } + + @Test + public void testMutatorsThrow() { + ImmutableSortedArrayMap<String> map = ImmutableSortedArrayMap.copyOf(unsortedSample()); + assertThrows(UnsupportedOperationException.class, () -> map.put("x", "y")); + assertThrows(UnsupportedOperationException.class, () -> map.remove("Server_a_8098")); + assertThrows(UnsupportedOperationException.class, () -> map.remove("Server_a_8098", "CONSUMING")); + assertThrows(UnsupportedOperationException.class, () -> map.putAll(Map.of("x", "y"))); + assertThrows(UnsupportedOperationException.class, map::clear); + assertThrows(UnsupportedOperationException.class, () -> map.putIfAbsent("x", "y")); + assertThrows(UnsupportedOperationException.class, () -> map.replace("Server_a_8098", "y")); + assertThrows(UnsupportedOperationException.class, () -> map.replace("Server_a_8098", "CONSUMING", "y")); + assertThrows(UnsupportedOperationException.class, () -> map.replaceAll((k, v) -> v)); + assertThrows(UnsupportedOperationException.class, () -> map.computeIfAbsent("x", k -> "y")); + assertThrows(UnsupportedOperationException.class, () -> map.computeIfPresent("Server_a_8098", (k, v) -> v)); + assertThrows(UnsupportedOperationException.class, () -> map.compute("x", (k, v) -> v)); + assertThrows(UnsupportedOperationException.class, () -> map.merge("x", "y", (a, b) -> a)); + assertThrows(UnsupportedOperationException.class, () -> map.keySet().remove("Server_a_8098")); + assertThrows(UnsupportedOperationException.class, () -> map.keySet().add("x")); + assertThrows(UnsupportedOperationException.class, () -> map.keySet().clear()); + assertThrows(UnsupportedOperationException.class, () -> map.keySet().retainAll(Set.of())); + assertThrows(UnsupportedOperationException.class, () -> map.values().remove("ONLINE")); + assertThrows(UnsupportedOperationException.class, () -> map.values().clear()); + assertThrows(UnsupportedOperationException.class, () -> map.entrySet().clear()); + assertThrows(UnsupportedOperationException.class, + () -> map.entrySet().remove(new AbstractMap.SimpleEntry<>("Server_a_8098", "CONSUMING"))); + assertThrows(UnsupportedOperationException.class, () -> map.entrySet().iterator().next().setValue("x")); + Iterator<String> iterator = map.keySet().iterator(); + iterator.next(); + assertThrows(UnsupportedOperationException.class, iterator::remove); + assertEquals(map, new TreeMap<>(unsortedSample())); + } + + @Test + public void testTreeMapCopyIsLinear() { + // TreeMap takes the buildFromSorted path for a SortedMap with the same (natural) comparator, so the copy made by + // HelixProperty does not compare keys. This test checks the result, the comparator() == null check enables the + // fast path. + ImmutableSortedArrayMap<String> map = ImmutableSortedArrayMap.copyOf(unsortedSample()); + TreeMap<String, String> copy = new TreeMap<>(); + copy.putAll(map); + assertEquals(copy, map); + assertEquals(new ArrayList<>(copy.keySet()), new ArrayList<>(map.keySet())); Review Comment: Renamed to `testTreeMapCopyCompatibility`. ########## pinot-broker/src/test/java/org/apache/pinot/broker/routing/manager/CompactAssignmentRoutingManagerTest.java: ########## @@ -0,0 +1,409 @@ +/** + * 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.pinot.broker.routing.manager; + +import java.util.ArrayList; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.TreeMap; +import java.util.TreeSet; +import java.util.concurrent.TimeUnit; +import org.apache.helix.AccessOption; +import org.apache.helix.HelixAdmin; +import org.apache.helix.HelixConstants.ChangeType; +import org.apache.helix.HelixManager; +import org.apache.helix.HelixManagerFactory; +import org.apache.helix.InstanceType; +import org.apache.helix.PropertyKey; +import org.apache.helix.manager.zk.ZKHelixAdmin; +import org.apache.helix.model.InstanceConfig; +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.helix.zookeeper.impl.client.ZkClient; +import org.apache.pinot.broker.broker.helix.CompactAssignmentZkHelixManager; +import org.apache.pinot.common.metadata.ZKMetadataProvider; +import org.apache.pinot.common.metadata.segment.SegmentPartitionMetadata; +import org.apache.pinot.common.metadata.segment.SegmentZKMetadata; +import org.apache.pinot.common.metrics.BrokerMetrics; +import org.apache.pinot.common.request.BrokerRequest; +import org.apache.pinot.common.utils.ZkStarter; +import org.apache.pinot.common.utils.helix.HelixHelper; +import org.apache.pinot.core.routing.RoutingTable; +import org.apache.pinot.core.routing.SegmentsToQuery; +import org.apache.pinot.core.routing.TablePartitionInfo; +import org.apache.pinot.core.routing.TablePartitionReplicatedServersInfo; +import org.apache.pinot.core.routing.timeboundary.TimeBoundaryInfo; +import org.apache.pinot.core.transport.ServerInstance; +import org.apache.pinot.core.transport.server.routing.stats.ServerRoutingStatsManager; +import org.apache.pinot.segment.spi.partition.metadata.ColumnPartitionMetadata; +import org.apache.pinot.spi.config.table.ColumnPartitionConfig; +import org.apache.pinot.spi.config.table.RoutingConfig; +import org.apache.pinot.spi.config.table.SegmentPartitionConfig; +import org.apache.pinot.spi.config.table.TableConfig; +import org.apache.pinot.spi.config.table.TableType; +import org.apache.pinot.spi.data.FieldSpec; +import org.apache.pinot.spi.data.Schema; +import org.apache.pinot.spi.env.PinotConfiguration; +import org.apache.pinot.spi.utils.CommonConstants; +import org.apache.pinot.spi.utils.CommonConstants.Helix.StateModel.SegmentStateModel; +import org.apache.pinot.spi.utils.builder.TableConfigBuilder; +import org.apache.pinot.spi.utils.builder.TableNameBuilder; +import org.apache.pinot.sql.parsers.CalciteSqlCompiler; +import org.testng.annotations.AfterClass; +import org.testng.annotations.BeforeClass; +import org.testng.annotations.Test; + +import static org.mockito.Mockito.mock; +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertFalse; +import static org.testng.Assert.assertNotNull; +import static org.testng.Assert.assertTrue; + + +/// Checks that a [BrokerRoutingManager] on a [CompactAssignmentZkHelixManager] routes exactly like one on the default +/// Helix manager. The routing manager wraps the compact records without copying them (see +/// [HelixHelper#toIdealState]), so every routing component (instance selector, segment pruners, time +/// boundary manager, segment partition metadata manager) reads the immutable outer map of the parsed record. +/// +/// The test writes a hybrid table into a local ZooKeeper, builds the routing in both managers, compares what they +/// return, then changes the assignment and compares again after the assignment change. +public class CompactAssignmentRoutingManagerTest { + private static final String CLUSTER_NAME = "CompactAssignmentRoutingManagerTest"; + private static final String RAW_TABLE_NAME = "myTable"; + private static final String OFFLINE_TABLE_NAME = TableNameBuilder.OFFLINE.tableNameWithType(RAW_TABLE_NAME); + private static final String REALTIME_TABLE_NAME = TableNameBuilder.REALTIME.tableNameWithType(RAW_TABLE_NAME); + private static final String TIME_COLUMN = "ts"; + private static final String PARTITION_COLUMN = "memberId"; + private static final int NUM_PARTITIONS = 4; + private static final int NUM_SERVERS = 4; + private static final long BASE_TIME_MS = 1_700_000_000_000L; + private static final long HOUR_MS = TimeUnit.HOURS.toMillis(1); + private static final List<String> QUERIES = List.of("SELECT * FROM myTable", + "SELECT * FROM myTable WHERE memberId = 2", + "SELECT * FROM myTable WHERE ts > " + (BASE_TIME_MS + 10 * HOUR_MS), + "SELECT * FROM myTable WHERE memberId IN (1, 3) AND ts < " + (BASE_TIME_MS + 6 * HOUR_MS)); + + private ZkStarter.ZookeeperInstance _zookeeperInstance; + private ZkClient _zkClient; + private HelixManager _defaultHelixManager; + private HelixManager _compactHelixManager; + private ZkHelixPropertyStore<ZNRecord> _propertyStore; + private BrokerRoutingManager _defaultRoutingManager; + private BrokerRoutingManager _compactRoutingManager; + + @BeforeClass + public void setUp() + throws Exception { + _zookeeperInstance = ZkStarter.startLocalZkServer(); + String zkUrl = _zookeeperInstance.getZkUrl(); + HelixAdmin helixAdmin = new ZKHelixAdmin.Builder().setZkAddress(zkUrl).build(); + try { + helixAdmin.addCluster(CLUSTER_NAME); + for (int i = 0; i < NUM_SERVERS; i++) { + InstanceConfig instanceConfig = new InstanceConfig(server(i)); + instanceConfig.setHostName("host" + i); + instanceConfig.setPort(Integer.toString(8000 + i)); + helixAdmin.addInstance(CLUSTER_NAME, instanceConfig); + } + } finally { + helixAdmin.close(); + } + _zkClient = new ZkClient.Builder().setZkServer(zkUrl).setZkSerializer(new ZNRecordSerializer()).build(); + + _defaultHelixManager = + HelixManagerFactory.getZKHelixManager(CLUSTER_NAME, "Broker_default_8099", InstanceType.SPECTATOR, zkUrl); + _defaultHelixManager.connect(); + _compactHelixManager = + new CompactAssignmentZkHelixManager(CLUSTER_NAME, "Broker_compact_8099", InstanceType.SPECTATOR, zkUrl); + _compactHelixManager.connect(); + _propertyStore = _defaultHelixManager.getHelixPropertyStore(); + + Schema schema = new Schema.SchemaBuilder().setSchemaName(RAW_TABLE_NAME) + .addSingleValueDimension(PARTITION_COLUMN, FieldSpec.DataType.INT) + .addDateTime(TIME_COLUMN, FieldSpec.DataType.LONG, "1:MILLISECONDS:EPOCH", "1:MILLISECONDS") + .build(); + ZKMetadataProvider.setSchema(_propertyStore, schema); + ZKMetadataProvider.setTableConfig(_propertyStore, tableConfig(TableType.OFFLINE)); + ZKMetadataProvider.setTableConfig(_propertyStore, tableConfig(TableType.REALTIME)); + + // Offline: 20 segments on 2 replicas each; one replica in ERROR and one segment missing from the external view + Map<String, Map<String, String>> offlineIdealState = new TreeMap<>(); + Map<String, Map<String, String>> offlineExternalView = new TreeMap<>(); + for (int i = 0; i < 20; i++) { + String segment = offlineSegment(i); + setSegmentZkMetadata(OFFLINE_TABLE_NAME, segment, i % NUM_PARTITIONS, BASE_TIME_MS + i * HOUR_MS, false); + Map<String, String> instanceStateMap = replicas(i, SegmentStateModel.ONLINE); + offlineIdealState.put(segment, instanceStateMap); + if (i == 7) { + continue; + } + Map<String, String> externalViewStateMap = new TreeMap<>(instanceStateMap); + if (i == 5) { + externalViewStateMap.put(server(i % NUM_SERVERS), SegmentStateModel.ERROR); + } + offlineExternalView.put(segment, externalViewStateMap); + } + writeAssignment(OFFLINE_TABLE_NAME, offlineIdealState, offlineExternalView, true); + + // Realtime: 3 sequences per partition, the last one CONSUMING + Map<String, Map<String, String>> realtimeAssignment = new TreeMap<>(); + for (int partition = 0; partition < NUM_PARTITIONS; partition++) { + for (int sequence = 0; sequence < 3; sequence++) { + String segment = RAW_TABLE_NAME + "__" + partition + "__" + sequence + "__20240101T0000Z"; + boolean consuming = sequence == 2; + setSegmentZkMetadata(REALTIME_TABLE_NAME, segment, partition, + BASE_TIME_MS + (15 + sequence) * HOUR_MS, consuming); + realtimeAssignment.put(segment, + replicas(partition, consuming ? SegmentStateModel.CONSUMING : SegmentStateModel.ONLINE)); + } + } + writeAssignment(REALTIME_TABLE_NAME, realtimeAssignment, realtimeAssignment, true); + + _defaultRoutingManager = newRoutingManager(_defaultHelixManager); + _compactRoutingManager = newRoutingManager(_compactHelixManager); + } + + @AfterClass + public void tearDown() { + if (_defaultRoutingManager != null) { + _defaultRoutingManager.stop(); + } + if (_compactRoutingManager != null) { + _compactRoutingManager.stop(); + } + if (_defaultHelixManager != null) { + _defaultHelixManager.disconnect(); + } + if (_compactHelixManager != null) { + _compactHelixManager.disconnect(); + } + if (_zkClient != null) { + _zkClient.close(); + } + ZkStarter.stopLocalZkServer(_zookeeperInstance); + } + + private static String server(int i) { + return "Server_host" + i + "_" + (8000 + i); + } + + private static String offlineSegment(int i) { + return OFFLINE_TABLE_NAME + "_" + i; + } + + private static Map<String, String> replicas(int i, String state) { + Map<String, String> instanceStateMap = new TreeMap<>(); + instanceStateMap.put(server(i % NUM_SERVERS), state); + instanceStateMap.put(server((i + 1) % NUM_SERVERS), state); + return instanceStateMap; + } + + private static TableConfig tableConfig(TableType tableType) { + return new TableConfigBuilder(tableType).setTableName(RAW_TABLE_NAME) + .setTimeColumnName(TIME_COLUMN) + .setSegmentPartitionConfig(new SegmentPartitionConfig( + Map.of(PARTITION_COLUMN, new ColumnPartitionConfig("Modulo", NUM_PARTITIONS)))) + .setRoutingConfig(new RoutingConfig(null, + List.of(RoutingConfig.PARTITION_SEGMENT_PRUNER_TYPE, RoutingConfig.TIME_SEGMENT_PRUNER_TYPE), null, false)) + .build(); + } + + private void setSegmentZkMetadata(String tableNameWithType, String segment, int partition, long endTimeMs, + boolean consuming) { + SegmentZKMetadata segmentZKMetadata = new SegmentZKMetadata(segment); + // Old segments, so that the instance selectors do not treat them as new + segmentZKMetadata.setCreationTime(BASE_TIME_MS); + segmentZKMetadata.setPushTime(BASE_TIME_MS); + segmentZKMetadata.setPartitionMetadata(new SegmentPartitionMetadata(Map.of(PARTITION_COLUMN, + new ColumnPartitionMetadata("Modulo", NUM_PARTITIONS, Set.of(partition), null)))); + if (consuming) { + segmentZKMetadata.setStatus(CommonConstants.Segment.Realtime.Status.IN_PROGRESS); + } else { + segmentZKMetadata.setStartTime(endTimeMs - HOUR_MS); + segmentZKMetadata.setEndTime(endTimeMs); + segmentZKMetadata.setTimeUnit(TimeUnit.MILLISECONDS); + segmentZKMetadata.setTotalDocs(1000); + } + ZKMetadataProvider.setSegmentZKMetadata(_propertyStore, tableNameWithType, segmentZKMetadata); + } + + private void writeAssignment(String tableNameWithType, Map<String, Map<String, String>> idealState, + Map<String, Map<String, String>> externalView, boolean create) { + PropertyKey.Builder keyBuilder = new PropertyKey.Builder(CLUSTER_NAME); + ZNRecord idealStateRecord = new ZNRecord(tableNameWithType); + idealStateRecord.setSimpleField("IDEAL_STATE_MODE", "CUSTOMIZED"); + idealStateRecord.setSimpleField("REBALANCE_MODE", "CUSTOMIZED"); + idealStateRecord.setSimpleField("STATE_MODEL_DEF_REF", "SegmentOnlineOfflineStateModel"); + idealStateRecord.setSimpleField("NUM_PARTITIONS", Integer.toString(idealState.size())); + idealStateRecord.setSimpleField("HELIX_ENABLED", "true"); + idealStateRecord.setMapFields(idealState); + ZNRecord externalViewRecord = new ZNRecord(tableNameWithType); + externalViewRecord.setMapFields(externalView); + // Compressed, because the compact serializer only parses gzip-compressed znodes itself + idealStateRecord.setBooleanField(ZNRecord.ENABLE_COMPRESSION_BOOLEAN_FIELD, true); + externalViewRecord.setBooleanField(ZNRecord.ENABLE_COMPRESSION_BOOLEAN_FIELD, true); + String idealStatePath = keyBuilder.idealStates(tableNameWithType).getPath(); + String externalViewPath = keyBuilder.externalView(tableNameWithType).getPath(); + if (create) { + _zkClient.createPersistent(idealStatePath, idealStateRecord); + _zkClient.createPersistent(externalViewPath, externalViewRecord); + } else { + _zkClient.writeData(idealStatePath, idealStateRecord); + _zkClient.writeData(externalViewPath, externalViewRecord); + } + } + + private BrokerRoutingManager newRoutingManager(HelixManager helixManager) { + PinotConfiguration config = new PinotConfiguration(); + config.setProperty(CommonConstants.Broker.CONFIG_OF_ENABLE_PARTITION_METADATA_MANAGER, true); + BrokerRoutingManager routingManager = + new BrokerRoutingManager(mock(BrokerMetrics.class), mock(ServerRoutingStatsManager.class), config); + routingManager.init(helixManager); + routingManager.processClusterChange(ChangeType.INSTANCE_CONFIG); + routingManager.buildRouting(OFFLINE_TABLE_NAME); + routingManager.buildRouting(REALTIME_TABLE_NAME); + return routingManager; + } + + @Test + public void testSameRoutingAsDefaultReader() { + assertReadsAreCompact(); + Map<String, Object> expected = snapshot(_defaultRoutingManager); + assertEquals(snapshot(_compactRoutingManager), expected); + + // Sanity checks so that the comparison is not vacuous + assertEquals(expected.get(OFFLINE_TABLE_NAME + ".timeBoundary"), + TIME_COLUMN + "=" + (BASE_TIME_MS + 19 * HOUR_MS - TimeUnit.DAYS.toMillis(1))); + assertEquals(expected.get(OFFLINE_TABLE_NAME + ".disabled"), false); + assertNotNull(expected.get(OFFLINE_TABLE_NAME + ".partitions")); + assertNotNull(expected.get(REALTIME_TABLE_NAME + ".replicatedServers")); + @SuppressWarnings("unchecked") + Map<String, Object> filtered = (Map<String, Object>) expected.get(OFFLINE_TABLE_NAME + "." + QUERIES.get(1) + ".0"); + assertEquals(filtered.get("numPruned"), 15); + @SuppressWarnings("unchecked") + Map<String, Object> all = (Map<String, Object>) expected.get(OFFLINE_TABLE_NAME + "." + QUERIES.get(0) + ".0"); + assertEquals(all.get("unavailable"), List.of(offlineSegment(7))); + + // Change the offline assignment: add a segment, recover the ERROR replica, drop a segment + Map<String, Map<String, String>> idealState = new TreeMap<>(); + Map<String, Map<String, String>> externalView = new TreeMap<>(); + for (int i = 1; i < 21; i++) { + String segment = offlineSegment(i); + if (i == 20) { + setSegmentZkMetadata(OFFLINE_TABLE_NAME, segment, i % NUM_PARTITIONS, BASE_TIME_MS + i * HOUR_MS, false); + } + idealState.put(segment, replicas(i, SegmentStateModel.ONLINE)); + externalView.put(segment, replicas(i, SegmentStateModel.ONLINE)); + } + writeAssignment(OFFLINE_TABLE_NAME, idealState, externalView, false); + _defaultRoutingManager.processClusterChange(ChangeType.IDEAL_STATE); + _defaultRoutingManager.processClusterChange(ChangeType.EXTERNAL_VIEW); + _compactRoutingManager.processClusterChange(ChangeType.IDEAL_STATE); + _compactRoutingManager.processClusterChange(ChangeType.EXTERNAL_VIEW); + + assertReadsAreCompact(); + Map<String, Object> expectedAfterChange = snapshot(_defaultRoutingManager); + assertEquals(snapshot(_compactRoutingManager), expectedAfterChange); + assertEquals(expectedAfterChange.get(OFFLINE_TABLE_NAME + ".timeBoundary"), + TIME_COLUMN + "=" + (BASE_TIME_MS + 20 * HOUR_MS - TimeUnit.DAYS.toMillis(1))); + @SuppressWarnings("unchecked") + Map<String, Object> allAfterChange = + (Map<String, Object>) expectedAfterChange.get(OFFLINE_TABLE_NAME + "." + QUERIES.get(0) + ".0"); + assertEquals(allAfterChange.get("unavailable"), List.of()); + } + + /// The compact manager reads compact records and the routing manager wraps them without a copy; the default + /// manager reads regular records. + private void assertReadsAreCompact() { Review Comment: The routing manager builds those objects in private methods and does not keep them anywhere a test can reach without a lot of plumbing. So I narrowed the comment to the raw reads it checks, and pointed to the `HelixHelperTest` case that checks the no-copy wrap. ########## pinot-common/src/main/java/org/apache/pinot/common/utils/helix/CompactZNRecordSerializer.java: ########## @@ -0,0 +1,472 @@ +/** + * 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.pinot.common.utils.helix; + +import com.fasterxml.jackson.core.JsonFactory; +import com.fasterxml.jackson.core.JsonParseException; +import com.fasterxml.jackson.core.JsonParser; +import com.fasterxml.jackson.core.JsonToken; +import java.io.ByteArrayInputStream; +import java.io.IOException; +import java.io.OutputStream; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.HashMap; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import java.util.zip.GZIPInputStream; +import javax.annotation.Nullable; +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.helix.zookeeper.zkclient.serialize.ZkSerializer; +import org.apache.pinot.spi.utils.CommonConstants.Helix.StateModel.SegmentStateModel; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + + +/// A [ZkSerializer] that reads gzip-compressed ideal state and external view znodes with less allocation than +/// [ZNRecordSerializer], and returns an equal [ZNRecord]. +/// +/// Only gzip-compressed znodes use the compact parser. A plain (uncompressed) znode goes to [ZNRecordSerializer] +/// unchanged. Pinot compresses large table ideal states, so plain znodes are the small ones, and there the compact +/// parser does not pay off: in a size sweep it allocated about 1x to 3.5x as much and used 1x to 1.5x the CPU of +/// [ZNRecordSerializer] up to about 200 segments. On gzip-compressed znodes of any size it allocated about 0.06x to +/// 0.3x as much and used about 0.25x to 0.75x the CPU. +/// +/// Compared with [ZNRecordSerializer#deserialize], for a gzip-compressed znode: +/// - It streams a gzip-compressed znode through a [GZIPInputStream] straight into the JSON parser. It does not +/// inflate the whole payload into a growing byte array first. +/// - Its [JsonFactory] does not canonicalize or intern field names, so the parser does not build a symbol table with +/// one entry per segment. +/// - Instance names (the keys of the instance-state maps and the values of the list fields) are interned once per +/// parse. Routing keys its server maps by interned instance ids, so lookups keep their identity fast path. +/// - State values come from a fixed table of the known segment states. An unknown state is interned. +/// - Segment names (the keys of the map fields) are interned, so they keep their identity with the names that the +/// routing components already hold. +/// - Equal instance-state maps within one znode share one [ImmutableSortedArrayMap]. The outer segment map is an +/// [ImmutableSortedArrayMap] too. +/// +/// The map fields of a record from the compact parser are immutable: any mutation throws +/// [UnsupportedOperationException]. Simple fields and list fields are mutable [LinkedHashMap]s, as with +/// [ZNRecordSerializer]. Scalar setters of the record (version, creation time, etc.) work as usual. +/// +/// When the compact parser fails, the serializer logs a warning and falls back to [ZNRecordSerializer]. [#serialize] +/// always delegates to [ZNRecordSerializer]. +/// +/// [HelixHelper#toIdealState] and [HelixHelper#toExternalView] wrap a record from the compact parser without the copy +/// of the outer segment map that `new IdealState(record)` and `new ExternalView(record)` make. Those constructors +/// share the inner instance-state maps, so an [org.apache.helix.model.IdealState] or +/// [org.apache.helix.model.ExternalView] built from a compact record has immutable instance-state maps either way. +/// A record from [ZNRecordSerializer] (a plain znode or the fallback) still goes through those constructors. +/// +/// Thread-safety: the serializer is stateless apart from the shared [JsonFactory], which is thread-safe. Each call +/// uses its own parser and dictionaries, so concurrent calls are safe. +public class CompactZNRecordSerializer implements ZkSerializer { + private static final Logger LOGGER = LoggerFactory.getLogger(CompactZNRecordSerializer.class); + + private static final JsonFactory JSON_FACTORY = JsonFactory.builder() + .disable(JsonFactory.Feature.CANONICALIZE_FIELD_NAMES) + .disable(JsonFactory.Feature.INTERN_FIELD_NAMES) + .build(); + + /// Jackson has its own input buffer, so the inflater only needs a small one. + private static final int GZIP_BUFFER_SIZE = 2 * 1024; + /// Small, so that the fixed cost per parse stays low; the arrays grow geometrically. + private static final int INITIAL_SEGMENT_CAPACITY = 16; + private static final int INITIAL_STATE_MAP_TABLE_CAPACITY = 16; + private static final int INITIAL_INSTANCE_CAPACITY = 8; + + /// The states a segment replica can have. The constants are string literals, so they are already interned. + private static final String[] KNOWN_STATES = { + SegmentStateModel.ONLINE, SegmentStateModel.CONSUMING, SegmentStateModel.OFFLINE, SegmentStateModel.ERROR, + "DROPPED" Review Comment: Now uses `HelixDefinedState.DROPPED.name()`. ########## pinot-common/src/main/java/org/apache/pinot/common/utils/helix/CompactZNRecordSerializer.java: ########## @@ -0,0 +1,472 @@ +/** + * 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.pinot.common.utils.helix; + +import com.fasterxml.jackson.core.JsonFactory; +import com.fasterxml.jackson.core.JsonParseException; +import com.fasterxml.jackson.core.JsonParser; +import com.fasterxml.jackson.core.JsonToken; +import java.io.ByteArrayInputStream; +import java.io.IOException; +import java.io.OutputStream; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.HashMap; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import java.util.zip.GZIPInputStream; +import javax.annotation.Nullable; +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.helix.zookeeper.zkclient.serialize.ZkSerializer; +import org.apache.pinot.spi.utils.CommonConstants.Helix.StateModel.SegmentStateModel; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + + +/// A [ZkSerializer] that reads gzip-compressed ideal state and external view znodes with less allocation than +/// [ZNRecordSerializer], and returns an equal [ZNRecord]. +/// +/// Only gzip-compressed znodes use the compact parser. A plain (uncompressed) znode goes to [ZNRecordSerializer] +/// unchanged. Pinot compresses large table ideal states, so plain znodes are the small ones, and there the compact +/// parser does not pay off: in a size sweep it allocated about 1x to 3.5x as much and used 1x to 1.5x the CPU of +/// [ZNRecordSerializer] up to about 200 segments. On gzip-compressed znodes of any size it allocated about 0.06x to +/// 0.3x as much and used about 0.25x to 0.75x the CPU. +/// +/// Compared with [ZNRecordSerializer#deserialize], for a gzip-compressed znode: +/// - It streams a gzip-compressed znode through a [GZIPInputStream] straight into the JSON parser. It does not +/// inflate the whole payload into a growing byte array first. +/// - Its [JsonFactory] does not canonicalize or intern field names, so the parser does not build a symbol table with +/// one entry per segment. +/// - Instance names (the keys of the instance-state maps and the values of the list fields) are interned once per +/// parse. Routing keys its server maps by interned instance ids, so lookups keep their identity fast path. +/// - State values come from a fixed table of the known segment states. An unknown state is interned. +/// - Segment names (the keys of the map fields) are interned, so they keep their identity with the names that the +/// routing components already hold. +/// - Equal instance-state maps within one znode share one [ImmutableSortedArrayMap]. The outer segment map is an +/// [ImmutableSortedArrayMap] too. +/// +/// The map fields of a record from the compact parser are immutable: any mutation throws +/// [UnsupportedOperationException]. Simple fields and list fields are mutable [LinkedHashMap]s, as with +/// [ZNRecordSerializer]. Scalar setters of the record (version, creation time, etc.) work as usual. +/// +/// When the compact parser fails, the serializer logs a warning and falls back to [ZNRecordSerializer]. [#serialize] +/// always delegates to [ZNRecordSerializer]. +/// +/// [HelixHelper#toIdealState] and [HelixHelper#toExternalView] wrap a record from the compact parser without the copy +/// of the outer segment map that `new IdealState(record)` and `new ExternalView(record)` make. Those constructors +/// share the inner instance-state maps, so an [org.apache.helix.model.IdealState] or +/// [org.apache.helix.model.ExternalView] built from a compact record has immutable instance-state maps either way. +/// A record from [ZNRecordSerializer] (a plain znode or the fallback) still goes through those constructors. +/// +/// Thread-safety: the serializer is stateless apart from the shared [JsonFactory], which is thread-safe. Each call +/// uses its own parser and dictionaries, so concurrent calls are safe. +public class CompactZNRecordSerializer implements ZkSerializer { + private static final Logger LOGGER = LoggerFactory.getLogger(CompactZNRecordSerializer.class); + + private static final JsonFactory JSON_FACTORY = JsonFactory.builder() + .disable(JsonFactory.Feature.CANONICALIZE_FIELD_NAMES) + .disable(JsonFactory.Feature.INTERN_FIELD_NAMES) + .build(); + + /// Jackson has its own input buffer, so the inflater only needs a small one. + private static final int GZIP_BUFFER_SIZE = 2 * 1024; + /// Small, so that the fixed cost per parse stays low; the arrays grow geometrically. + private static final int INITIAL_SEGMENT_CAPACITY = 16; + private static final int INITIAL_STATE_MAP_TABLE_CAPACITY = 16; + private static final int INITIAL_INSTANCE_CAPACITY = 8; + + /// The states a segment replica can have. The constants are string literals, so they are already interned. + private static final String[] KNOWN_STATES = { + SegmentStateModel.ONLINE, SegmentStateModel.CONSUMING, SegmentStateModel.OFFLINE, SegmentStateModel.ERROR, + "DROPPED" + }; + + private final ZNRecordSerializer _delegate; + + public CompactZNRecordSerializer() { + this(new ZNRecordSerializer()); + } + + public CompactZNRecordSerializer(ZNRecordSerializer delegate) { + _delegate = delegate; + } + + @Override + public byte[] serialize(Object data) { + return _delegate.serialize(data); + } + + @Nullable + @Override + public Object deserialize(byte[] bytes) { + if (bytes == null || bytes.length == 0) { + // Reading a parent or empty node, same as ZNRecordSerializer + return null; + } + if (!GZipCompressionUtil.isCompressed(bytes)) { + // Plain znodes are small; the default deserializer is cheaper there (see the class documentation) + return _delegate.deserialize(bytes); + } + try { + return deserializeCompact(bytes); + } catch (Exception e) { + LOGGER.warn("Caught exception while deserializing ZNRecord with the compact reader, falling back to the " + + "default deserializer", e); + return _delegate.deserialize(bytes); + } + } + + /// Parses the given znode content, gzip-compressed or plain, with the compact parser. Throws on malformed input; + /// [#deserialize] then falls back. [#deserialize] only passes gzip-compressed content; tests and benchmarks also + /// call it with plain content. + static ZNRecord deserializeCompact(byte[] bytes) + throws IOException { + if (!GZipCompressionUtil.isCompressed(bytes)) { + try (JsonParser parser = JSON_FACTORY.createParser(bytes)) { + return new Reader(parser).readRecord(); + } + } + GZIPInputStream gzipInputStream = new GZIPInputStream(new ByteArrayInputStream(bytes), GZIP_BUFFER_SIZE); + try (JsonParser parser = JSON_FACTORY.createParser(gzipInputStream)) { + ZNRecord record = new Reader(parser).readRecord(); + // The parser stops at the end of the record. GZIPInputStream verifies the CRC-32 and the size in the gzip + // trailer only when it reaches the end of the stream, so read the rest of the stream (as + // GZipCompressionUtil.uncompress does). A bad trailer throws, and deserialize falls back. + gzipInputStream.transferTo(OutputStream.nullOutputStream()); + return record; + } + } + + /// The state of one parse: the parser, the dictionaries, and scratch buffers. Not thread-safe; one per call. + private static final class Reader { + private final JsonParser _parser; + private final HashMap<String, String> _instanceNames = new HashMap<>(); + private final StateMapTable _stateMaps = new StateMapTable(); + private String[] _instanceKeys = new String[INITIAL_INSTANCE_CAPACITY]; + private Object[] _instanceStates = new Object[INITIAL_INSTANCE_CAPACITY]; + + Reader(JsonParser parser) { + _parser = parser; + } + + ZNRecord readRecord() + throws IOException { + expect(_parser.nextToken(), JsonToken.START_OBJECT); + String id = null; + Map<String, String> simpleFields = null; + Map<String, List<String>> listFields = null; + Map<String, Map<String, String>> mapFields = null; + byte[] rawPayload = null; + boolean hasSimpleFields = false; + boolean hasListFields = false; + boolean hasMapFields = false; + JsonToken token; + while ((token = _parser.nextToken()) == JsonToken.FIELD_NAME) { + String name = _parser.currentName(); + token = _parser.nextToken(); + switch (name) { + case "id": + id = token == JsonToken.VALUE_NULL ? null : scalarText(token); + break; + case "simpleFields": + hasSimpleFields = true; + simpleFields = token == JsonToken.VALUE_NULL ? null : readSimpleFields(token); + break; + case "listFields": + hasListFields = true; + listFields = token == JsonToken.VALUE_NULL ? null : readListFields(token); + break; + case "mapFields": + hasMapFields = true; + mapFields = token == JsonToken.VALUE_NULL ? null : readMapFields(token); + break; + case "rawPayload": + rawPayload = token == JsonToken.VALUE_NULL ? null : _parser.getBinaryValue(); + break; + default: + // ZNRecord ignores unknown properties + _parser.skipChildren(); + break; + } + } + // Like ObjectMapper.readValue (FAIL_ON_TRAILING_TOKENS is off), content after the record is ignored + expect(token, JsonToken.END_OBJECT); + + ZNRecord record = new ZNRecord(id); + if (hasSimpleFields) { + record.setSimpleFields(simpleFields); + } + if (hasListFields) { + record.setListFields(listFields); + } + if (hasMapFields) { + record.setMapFields(mapFields); + } + if (rawPayload != null) { Review Comment: Added a round-trip test for an empty rawPayload, plain and gzip. Both readers return a non-null empty array. -- 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]
