Jackie-Jiang commented on code in PR #19786: URL: https://github.com/apache/pinot/pull/19786#discussion_r4263910511
########## 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(); + } + } Review Comment: The production caller already filters out plain input before calling this method. Could we replace this branch with `assert GZipCompressionUtil.isCompressed(bytes)` and make the method contract gzip-only? Direct-parser tests can compress their inputs, while public `deserialize()` tests retain coverage for plain-input delegation. ########## pinot-common/src/main/java/org/apache/pinot/common/utils/helix/ImmutableSortedArrayMap.java: ########## @@ -0,0 +1,582 @@ +/** + * 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.AbstractCollection; +import java.util.AbstractMap; +import java.util.AbstractSet; +import java.util.Arrays; +import java.util.Collection; +import java.util.Comparator; +import java.util.Iterator; +import java.util.Map; +import java.util.NoSuchElementException; +import java.util.Objects; +import java.util.Set; +import java.util.SortedMap; +import java.util.function.BiFunction; +import java.util.function.Function; +import javax.annotation.Nullable; + + +/// An immutable [SortedMap] with [String] keys in natural order, backed by a sorted key array and a parallel value +/// array. +/// +/// [CompactZNRecordSerializer] uses it for the map fields of ideal states and external views: one instance per +/// distinct instance-state map, and one for the outer segment map. Lookups use binary search with an identity fast +/// path, so interned keys compare cheaply. +/// +/// The map follows the [Map] and [SortedMap] contracts: [#equals], [#hashCode] and [#toString] match those of a +/// [java.util.TreeMap] with the same content, and [#keySet], [#entrySet] and [#values] are [AbstractSet] and +/// [AbstractCollection] views with the standard `equals` and `hashCode`. [#comparator] returns `null`, so +/// `new TreeMap<>(map)` and `TreeMap.putAll` copy it in linear time. Keys are never `null`; values can be `null`. +/// Every mutator throws [UnsupportedOperationException]. Review Comment: The documented contract is stronger than the implementation. Range copies discard bounds, so `map.subMap("b", "d").headMap("z")` succeeds where the equivalent TreeMap range throws; the `range()` comment calling copies indistinguishable from views is therefore inaccurate. Also, inherited no-op view operations such as `keySet().addAll(List.of())` or `values().removeIf(v -> false)` return normally despite the 'every mutator throws' promise. These do not affect current routing use, but please document the limitations or implement the stronger guarantees. ########## pinot-common/src/main/java/org/apache/pinot/common/utils/helix/ImmutableSortedArrayMap.java: ########## @@ -0,0 +1,582 @@ +/** + * 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.AbstractCollection; +import java.util.AbstractMap; +import java.util.AbstractSet; +import java.util.Arrays; +import java.util.Collection; +import java.util.Comparator; +import java.util.Iterator; +import java.util.Map; +import java.util.NoSuchElementException; +import java.util.Objects; +import java.util.Set; +import java.util.SortedMap; +import java.util.function.BiFunction; +import java.util.function.Function; +import javax.annotation.Nullable; + + +/// An immutable [SortedMap] with [String] keys in natural order, backed by a sorted key array and a parallel value +/// array. +/// +/// [CompactZNRecordSerializer] uses it for the map fields of ideal states and external views: one instance per +/// distinct instance-state map, and one for the outer segment map. Lookups use binary search with an identity fast +/// path, so interned keys compare cheaply. +/// +/// The map follows the [Map] and [SortedMap] contracts: [#equals], [#hashCode] and [#toString] match those of a +/// [java.util.TreeMap] with the same content, and [#keySet], [#entrySet] and [#values] are [AbstractSet] and +/// [AbstractCollection] views with the standard `equals` and `hashCode`. [#comparator] returns `null`, so +/// `new TreeMap<>(map)` and `TreeMap.putAll` copy it in linear time. Keys are never `null`; values can be `null`. +/// Every mutator throws [UnsupportedOperationException]. +/// +/// Thread-safety: the map is immutable and safe to share between threads once published. The cached hash codes and +/// views are computed lazily with benign races: every thread computes the same value. +public final class ImmutableSortedArrayMap<V> extends AbstractMap<String, V> implements SortedMap<String, V> { + private static final ImmutableSortedArrayMap<?> EMPTY = new ImmutableSortedArrayMap<>(new String[0], new Object[0]); + + private final String[] _keys; + private final Object[] _values; + + // Lazily computed, benign races. 0 means "not computed yet" (a real hash code of 0 is recomputed each time). + private int _hashCode; + private int _keySetHashCode; + private KeySet _keySet; + private EntrySet _entrySet; + private Values _valuesView; + + /// Wraps the given arrays without copying them. The keys must be non-null, distinct and sorted in natural order, + /// and the caller must not modify either array afterwards. + ImmutableSortedArrayMap(String[] keys, Object[] values) { + assert keys.length == values.length; + _keys = keys; + _values = values; + } + + /// Returns the empty map. + @SuppressWarnings("unchecked") + public static <V> ImmutableSortedArrayMap<V> empty() { + return (ImmutableSortedArrayMap<V>) EMPTY; + } + + /// Returns an immutable sorted copy of the given map. Keys must be non-null. + public static <V> ImmutableSortedArrayMap<V> copyOf(Map<String, ? extends V> map) { + if (map instanceof ImmutableSortedArrayMap) { + @SuppressWarnings("unchecked") + ImmutableSortedArrayMap<V> immutableMap = (ImmutableSortedArrayMap<V>) map; + return immutableMap; + } + int size = map.size(); + if (size == 0) { + return empty(); + } + String[] keys = new String[size]; + Object[] values = new Object[size]; + int i = 0; + for (Map.Entry<String, ? extends V> entry : map.entrySet()) { + keys[i] = Objects.requireNonNull(entry.getKey(), "Null key"); + values[i] = entry.getValue(); + i++; + } + return sortAndWrap(keys, values, size); + } + + /// Sorts the first `size` entries of the given parallel arrays by key and wraps them. When a key appears more than + /// once, the last value wins, like repeated [Map#put] calls. The arrays may be modified and must not be used by the + /// caller afterwards. + static <V> ImmutableSortedArrayMap<V> sortAndWrap(String[] keys, Object[] values, int size) { + if (size == 0) { + return empty(); + } + boolean strictlySorted = true; + for (int i = 1; i < size; i++) { + if (keys[i - 1].compareTo(keys[i]) >= 0) { + strictlySorted = false; + break; + } + } + if (strictlySorted) { + if (keys.length == size) { + return new ImmutableSortedArrayMap<>(keys, values); + } + return new ImmutableSortedArrayMap<>(Arrays.copyOf(keys, size), Arrays.copyOf(values, size)); + } + // Stable sort, so the last occurrence of a duplicated key stays the last in its run. + Integer[] order = new Integer[size]; + for (int i = 0; i < size; i++) { + order[i] = i; + } + Arrays.sort(order, Comparator.comparing(i -> keys[i])); + String[] sortedKeys = new String[size]; + Object[] sortedValues = new Object[size]; + int numUnique = 0; + for (int i = 0; i < size; i++) { + int index = order[i]; + String key = keys[index]; + if (numUnique > 0 && sortedKeys[numUnique - 1].equals(key)) { + sortedValues[numUnique - 1] = values[index]; + } else { + sortedKeys[numUnique] = key; + sortedValues[numUnique] = values[index]; + numUnique++; + } + } + if (numUnique != size) { + sortedKeys = Arrays.copyOf(sortedKeys, numUnique); + sortedValues = Arrays.copyOf(sortedValues, numUnique); + } + return new ImmutableSortedArrayMap<>(sortedKeys, sortedValues); + } + + /// Returns the index of the key, or `-(insertionPoint + 1)` when absent. + private int indexOf(String key) { + int low = 0; + int high = _keys.length - 1; + while (low <= high) { + int mid = (low + high) >>> 1; + String midKey = _keys[mid]; + int cmp = midKey == key ? 0 : midKey.compareTo(key); + if (cmp < 0) { + low = mid + 1; + } else if (cmp > 0) { + high = mid - 1; + } else { + return mid; + } + } + return -(low + 1); + } + + /// Returns the index of the first key that is greater than or equal to the given key. + private int lowerBound(String key) { + int index = indexOf(key); + return index >= 0 ? index : -(index + 1); + } + + /// Returns the key at the given position in sorted order. + String keyAt(int index) { + return _keys[index]; + } + + /// Returns the value at the given position in sorted order. + @SuppressWarnings("unchecked") + V valueAt(int index) { + return (V) _values[index]; + } + + @Override + public int size() { + return _keys.length; + } + + @Override + public boolean isEmpty() { + return _keys.length == 0; + } + + @Override + public boolean containsKey(Object key) { + return key instanceof String && indexOf((String) key) >= 0; + } + + @Override + public boolean containsValue(Object value) { + for (Object v : _values) { + if (v == value || (v != null && v.equals(value))) { + return true; + } + } + return false; + } + + @Nullable + @Override + public V get(Object key) { + if (!(key instanceof String)) { + return null; + } + int index = indexOf((String) key); + return index >= 0 ? valueAt(index) : null; + } + + @Override + public V getOrDefault(Object key, V defaultValue) { + if (!(key instanceof String)) { + return defaultValue; + } + int index = indexOf((String) key); + return index >= 0 ? valueAt(index) : defaultValue; + } + + @Nullable + @Override + public Comparator<? super String> comparator() { + return null; + } + + @Override + public String firstKey() { + if (_keys.length == 0) { + throw new NoSuchElementException(); + } + return _keys[0]; + } + + @Override + public String lastKey() { + if (_keys.length == 0) { + throw new NoSuchElementException(); + } + return _keys[_keys.length - 1]; + } + + @Override + public SortedMap<String, V> subMap(String fromKey, String toKey) { + Objects.requireNonNull(fromKey); + Objects.requireNonNull(toKey); + if (fromKey.compareTo(toKey) > 0) { + throw new IllegalArgumentException("fromKey > toKey"); + } + return range(lowerBound(fromKey), lowerBound(toKey)); + } + + @Override + public SortedMap<String, V> headMap(String toKey) { + return range(0, lowerBound(Objects.requireNonNull(toKey))); + } + + @Override + public SortedMap<String, V> tailMap(String fromKey) { + return range(lowerBound(Objects.requireNonNull(fromKey)), _keys.length); + } + + /// Returns the entries in `[from, to)`. The result is a copy, which is indistinguishable from a view because + /// neither map can change. + private SortedMap<String, V> range(int from, int to) { + if (from == 0 && to == _keys.length) { + return this; + } + if (from >= to) { + return empty(); + } + return new ImmutableSortedArrayMap<>(Arrays.copyOfRange(_keys, from, to), Arrays.copyOfRange(_values, from, to)); + } + + @Override + public Set<String> keySet() { + KeySet keySet = _keySet; + if (keySet == null) { + keySet = new KeySet(); + _keySet = keySet; + } + return keySet; + } + + @Override + public Collection<V> values() { + Values values = _valuesView; + if (values == null) { + values = new Values(); + _valuesView = values; + } + return values; + } + + @Override + public Set<Map.Entry<String, V>> entrySet() { + EntrySet entrySet = _entrySet; + if (entrySet == null) { + entrySet = new EntrySet(); + _entrySet = entrySet; + } + return entrySet; + } + + @Override + public boolean equals(Object o) { + if (o == this) { + return true; + } + if (o instanceof ImmutableSortedArrayMap) { + ImmutableSortedArrayMap<?> that = (ImmutableSortedArrayMap<?>) o; + return Arrays.equals(_keys, that._keys) && Arrays.equals(_values, that._values); + } + return super.equals(o); + } + + @Override + public int hashCode() { + int hashCode = _hashCode; + if (hashCode == 0) { + for (int i = 0; i < _keys.length; i++) { + hashCode += _keys[i].hashCode() ^ Objects.hashCode(_values[i]); + } + _hashCode = hashCode; Review Comment: The cached hash assumes each value has stable equality/hash semantics, but public `copyOf()` accepts arbitrary mutable values and retains their references. For example, compute this map's hash with a mutable list value, then change that list: it still compares equal to an ordinary map with the updated value, but their hashes differ. `entrySet().hashCode()` inherits the same issue. Current parser values are immutable, so this does not affect routing. Please narrow/document the utility's immutable-value requirement, or remove value-dependent caching if it remains a general-purpose public map. `copyOf()` currently only has test callers. ########## 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) { + record.setRawPayload(rawPayload); + } + return record; + } + + private Map<String, String> readSimpleFields(JsonToken token) + throws IOException { + expect(token, JsonToken.START_OBJECT); + Map<String, String> simpleFields = new LinkedHashMap<>(); + while ((token = _parser.nextToken()) == JsonToken.FIELD_NAME) { + // ZNRecordSerializer interns all field names + String key = _parser.currentName().intern(); + token = _parser.nextToken(); + simpleFields.put(key, token == JsonToken.VALUE_NULL ? null : scalarText(token)); + } + expect(token, JsonToken.END_OBJECT); + return simpleFields; + } + + private Map<String, List<String>> readListFields(JsonToken token) + throws IOException { + expect(token, JsonToken.START_OBJECT); + Map<String, List<String>> listFields = new LinkedHashMap<>(); + while ((token = _parser.nextToken()) == JsonToken.FIELD_NAME) { + String key = _parser.currentName().intern(); + token = _parser.nextToken(); + if (token == JsonToken.VALUE_NULL) { + listFields.put(key, null); + continue; + } + expect(token, JsonToken.START_ARRAY); + List<String> list = new ArrayList<>(); + while ((token = _parser.nextToken()) != JsonToken.END_ARRAY) { + // List fields hold instance names (e.g. the preference lists of SEMI_AUTO resources) + list.add(token == JsonToken.VALUE_NULL ? null : instanceName(scalarText(token))); + } + listFields.put(key, list); + } + expect(token, JsonToken.END_OBJECT); + return listFields; + } + + private Map<String, Map<String, String>> readMapFields(JsonToken token) + throws IOException { + expect(token, JsonToken.START_OBJECT); + String[] segments = new String[INITIAL_SEGMENT_CAPACITY]; + Object[] stateMaps = new Object[INITIAL_SEGMENT_CAPACITY]; + int numSegments = 0; + while ((token = _parser.nextToken()) == JsonToken.FIELD_NAME) { + String segment = _parser.currentName().intern(); + token = _parser.nextToken(); + Map<String, String> stateMap = token == JsonToken.VALUE_NULL ? null : readStateMap(token); + if (numSegments == segments.length) { + int newCapacity = segments.length * 2; + segments = Arrays.copyOf(segments, newCapacity); + stateMaps = Arrays.copyOf(stateMaps, newCapacity); + } + segments[numSegments] = segment; + stateMaps[numSegments] = stateMap; + numSegments++; + } + expect(token, JsonToken.END_OBJECT); + return ImmutableSortedArrayMap.sortAndWrap(segments, stateMaps, numSegments); + } + + private Map<String, String> readStateMap(JsonToken token) + throws IOException { + expect(token, JsonToken.START_OBJECT); + int size = 0; + while ((token = _parser.nextToken()) == JsonToken.FIELD_NAME) { + String instance = instanceName(_parser.currentName()); + token = _parser.nextToken(); + String state = token == JsonToken.VALUE_NULL ? null : stateValue(token); + if (size == _instanceKeys.length) { + _instanceKeys = Arrays.copyOf(_instanceKeys, size * 2); + _instanceStates = Arrays.copyOf(_instanceStates, size * 2); + } + _instanceKeys[size] = instance; + _instanceStates[size] = state; + size++; + } + expect(token, JsonToken.END_OBJECT); + size = sortInPlace(_instanceKeys, _instanceStates, size); + return _stateMaps.intern(_instanceKeys, _instanceStates, size); + } + + /// Returns the canonical (interned) instance name equal to the given name. + private String instanceName(String name) { + String canonical = _instanceNames.get(name); + if (canonical == null) { + canonical = name.intern(); + _instanceNames.put(canonical, canonical); + } + return canonical; + } + + /// Returns the canonical state for the current value token without allocating for the known states. + private String stateValue(JsonToken token) + throws IOException { + if (token != JsonToken.VALUE_STRING) { + return scalarText(token).intern(); + } + char[] chars = _parser.getTextCharacters(); + int offset = _parser.getTextOffset(); + int length = _parser.getTextLength(); + for (String knownState : KNOWN_STATES) { + if (matches(knownState, chars, offset, length)) { + return knownState; + } + } + return new String(chars, offset, length).intern(); + } + + /// Returns the text of a scalar token. Like Jackson's `Map<String, String>` binding, numbers and booleans become + /// their text, and objects or arrays fail. + private String scalarText(JsonToken token) + throws IOException { + if (!token.isScalarValue()) { + throw new JsonParseException(_parser, "Expected a scalar value but got: " + token); + } + return _parser.getText(); + } + + private void expect(@Nullable JsonToken actual, JsonToken expected) + throws JsonParseException { + if (actual != expected) { + throw new JsonParseException(_parser, "Expected " + expected + " but got: " + actual); + } + } + } + + private static boolean matches(String s, char[] chars, int offset, int length) { + if (s.length() != length) { + return false; + } + for (int i = 0; i < length; i++) { + if (s.charAt(i) != chars[offset + i]) { + return false; + } + } + return true; + } + + /// Sorts the first `size` entries of the parallel arrays by key with an insertion sort (instance-state maps hold a + /// handful of replicas). When a key appears more than once, the last value wins. Returns the number of distinct + /// keys, which are compacted at the start of the arrays. + static int sortInPlace(String[] keys, Object[] values, int size) { + int numUnique = 0; + for (int i = 0; i < size; i++) { + String key = keys[i]; + Object value = values[i]; + int low = 0; + int high = numUnique - 1; + int found = -1; + while (low <= high) { + int mid = (low + high) >>> 1; + int cmp = keys[mid].compareTo(key); + if (cmp < 0) { + low = mid + 1; + } else if (cmp > 0) { + high = mid - 1; + } else { + found = mid; + break; + } + } + if (found >= 0) { + values[found] = value; + continue; + } + System.arraycopy(keys, low, keys, low + 1, numUnique - low); + System.arraycopy(values, low, values, low + 1, numUnique - low); + keys[low] = key; + values[low] = value; + numUnique++; + } + return numUnique; + } + + /// Deduplicates instance-state maps by content within one parse. Keys and values passed in are canonical (interned + /// or from the known state table), so content equality is reference equality element by element. + private static final class StateMapTable { + private ImmutableSortedArrayMap<?>[] _maps = new ImmutableSortedArrayMap<?>[INITIAL_STATE_MAP_TABLE_CAPACITY]; + private int[] _hashes = new int[INITIAL_STATE_MAP_TABLE_CAPACITY]; + private int _size; + + @SuppressWarnings("unchecked") + Map<String, String> intern(String[] keys, Object[] values, int size) { + int hash = 1; + for (int i = 0; i < size; i++) { + hash = 31 * hash + keys[i].hashCode(); + hash = 31 * hash + (values[i] == null ? 0 : values[i].hashCode()); + } + // Avoid 0, which marks an empty slot + hash = hash == 0 ? 1 : hash; Review Comment: This normalization and comment appear unnecessary: occupancy is determined by `_maps[slot] == null`, and `_hashes[slot]` is consulted only for occupied slots. Zero is therefore a valid stored hash; there is no need to remap it to one. ########## pinot-common/src/main/java/org/apache/pinot/common/utils/helix/ImmutableSortedArrayMap.java: ########## @@ -0,0 +1,582 @@ +/** + * 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.AbstractCollection; +import java.util.AbstractMap; +import java.util.AbstractSet; +import java.util.Arrays; +import java.util.Collection; +import java.util.Comparator; +import java.util.Iterator; +import java.util.Map; +import java.util.NoSuchElementException; +import java.util.Objects; +import java.util.Set; +import java.util.SortedMap; +import java.util.function.BiFunction; +import java.util.function.Function; +import javax.annotation.Nullable; + + +/// An immutable [SortedMap] with [String] keys in natural order, backed by a sorted key array and a parallel value +/// array. +/// +/// [CompactZNRecordSerializer] uses it for the map fields of ideal states and external views: one instance per +/// distinct instance-state map, and one for the outer segment map. Lookups use binary search with an identity fast +/// path, so interned keys compare cheaply. +/// +/// The map follows the [Map] and [SortedMap] contracts: [#equals], [#hashCode] and [#toString] match those of a +/// [java.util.TreeMap] with the same content, and [#keySet], [#entrySet] and [#values] are [AbstractSet] and +/// [AbstractCollection] views with the standard `equals` and `hashCode`. [#comparator] returns `null`, so +/// `new TreeMap<>(map)` and `TreeMap.putAll` copy it in linear time. Keys are never `null`; values can be `null`. +/// Every mutator throws [UnsupportedOperationException]. +/// +/// Thread-safety: the map is immutable and safe to share between threads once published. The cached hash codes and +/// views are computed lazily with benign races: every thread computes the same value. +public final class ImmutableSortedArrayMap<V> extends AbstractMap<String, V> implements SortedMap<String, V> { + private static final ImmutableSortedArrayMap<?> EMPTY = new ImmutableSortedArrayMap<>(new String[0], new Object[0]); + + private final String[] _keys; + private final Object[] _values; + + // Lazily computed, benign races. 0 means "not computed yet" (a real hash code of 0 is recomputed each time). + private int _hashCode; + private int _keySetHashCode; + private KeySet _keySet; + private EntrySet _entrySet; + private Values _valuesView; + + /// Wraps the given arrays without copying them. The keys must be non-null, distinct and sorted in natural order, + /// and the caller must not modify either array afterwards. + ImmutableSortedArrayMap(String[] keys, Object[] values) { + assert keys.length == values.length; + _keys = keys; + _values = values; + } + + /// Returns the empty map. + @SuppressWarnings("unchecked") + public static <V> ImmutableSortedArrayMap<V> empty() { + return (ImmutableSortedArrayMap<V>) EMPTY; + } + + /// Returns an immutable sorted copy of the given map. Keys must be non-null. + public static <V> ImmutableSortedArrayMap<V> copyOf(Map<String, ? extends V> map) { + if (map instanceof ImmutableSortedArrayMap) { + @SuppressWarnings("unchecked") + ImmutableSortedArrayMap<V> immutableMap = (ImmutableSortedArrayMap<V>) map; + return immutableMap; + } + int size = map.size(); + if (size == 0) { + return empty(); + } + String[] keys = new String[size]; + Object[] values = new Object[size]; + int i = 0; + for (Map.Entry<String, ? extends V> entry : map.entrySet()) { + keys[i] = Objects.requireNonNull(entry.getKey(), "Null key"); + values[i] = entry.getValue(); + i++; + } + return sortAndWrap(keys, values, size); + } + + /// Sorts the first `size` entries of the given parallel arrays by key and wraps them. When a key appears more than + /// once, the last value wins, like repeated [Map#put] calls. The arrays may be modified and must not be used by the + /// caller afterwards. + static <V> ImmutableSortedArrayMap<V> sortAndWrap(String[] keys, Object[] values, int size) { + if (size == 0) { + return empty(); + } + boolean strictlySorted = true; + for (int i = 1; i < size; i++) { + if (keys[i - 1].compareTo(keys[i]) >= 0) { + strictlySorted = false; + break; + } + } + if (strictlySorted) { + if (keys.length == size) { + return new ImmutableSortedArrayMap<>(keys, values); + } + return new ImmutableSortedArrayMap<>(Arrays.copyOf(keys, size), Arrays.copyOf(values, size)); + } + // Stable sort, so the last occurrence of a duplicated key stays the last in its run. + Integer[] order = new Integer[size]; + for (int i = 0; i < size; i++) { + order[i] = i; + } + Arrays.sort(order, Comparator.comparing(i -> keys[i])); + String[] sortedKeys = new String[size]; + Object[] sortedValues = new Object[size]; + int numUnique = 0; + for (int i = 0; i < size; i++) { + int index = order[i]; + String key = keys[index]; + if (numUnique > 0 && sortedKeys[numUnique - 1].equals(key)) { + sortedValues[numUnique - 1] = values[index]; + } else { + sortedKeys[numUnique] = key; + sortedValues[numUnique] = values[index]; + numUnique++; + } + } + if (numUnique != size) { + sortedKeys = Arrays.copyOf(sortedKeys, numUnique); + sortedValues = Arrays.copyOf(sortedValues, numUnique); + } + return new ImmutableSortedArrayMap<>(sortedKeys, sortedValues); + } + + /// Returns the index of the key, or `-(insertionPoint + 1)` when absent. + private int indexOf(String key) { + int low = 0; + int high = _keys.length - 1; + while (low <= high) { + int mid = (low + high) >>> 1; + String midKey = _keys[mid]; + int cmp = midKey == key ? 0 : midKey.compareTo(key); + if (cmp < 0) { + low = mid + 1; + } else if (cmp > 0) { + high = mid - 1; + } else { + return mid; + } + } + return -(low + 1); + } + + /// Returns the index of the first key that is greater than or equal to the given key. + private int lowerBound(String key) { + int index = indexOf(key); + return index >= 0 ? index : -(index + 1); + } + + /// Returns the key at the given position in sorted order. + String keyAt(int index) { + return _keys[index]; + } + + /// Returns the value at the given position in sorted order. + @SuppressWarnings("unchecked") + V valueAt(int index) { + return (V) _values[index]; + } + + @Override + public int size() { + return _keys.length; + } + + @Override + public boolean isEmpty() { + return _keys.length == 0; + } + + @Override + public boolean containsKey(Object key) { + return key instanceof String && indexOf((String) key) >= 0; + } + + @Override + public boolean containsValue(Object value) { + for (Object v : _values) { + if (v == value || (v != null && v.equals(value))) { + return true; + } + } + return false; + } + + @Nullable + @Override + public V get(Object key) { + if (!(key instanceof String)) { + return null; + } + int index = indexOf((String) key); + return index >= 0 ? valueAt(index) : null; + } + + @Override + public V getOrDefault(Object key, V defaultValue) { + if (!(key instanceof String)) { + return defaultValue; + } + int index = indexOf((String) key); + return index >= 0 ? valueAt(index) : defaultValue; + } + + @Nullable + @Override + public Comparator<? super String> comparator() { + return null; + } + + @Override + public String firstKey() { + if (_keys.length == 0) { + throw new NoSuchElementException(); + } + return _keys[0]; + } + + @Override + public String lastKey() { + if (_keys.length == 0) { + throw new NoSuchElementException(); + } + return _keys[_keys.length - 1]; + } + + @Override + public SortedMap<String, V> subMap(String fromKey, String toKey) { + Objects.requireNonNull(fromKey); + Objects.requireNonNull(toKey); + if (fromKey.compareTo(toKey) > 0) { + throw new IllegalArgumentException("fromKey > toKey"); + } + return range(lowerBound(fromKey), lowerBound(toKey)); + } + + @Override + public SortedMap<String, V> headMap(String toKey) { + return range(0, lowerBound(Objects.requireNonNull(toKey))); + } + + @Override + public SortedMap<String, V> tailMap(String fromKey) { + return range(lowerBound(Objects.requireNonNull(fromKey)), _keys.length); + } + + /// Returns the entries in `[from, to)`. The result is a copy, which is indistinguishable from a view because + /// neither map can change. + private SortedMap<String, V> range(int from, int to) { + if (from == 0 && to == _keys.length) { + return this; + } + if (from >= to) { + return empty(); + } + return new ImmutableSortedArrayMap<>(Arrays.copyOfRange(_keys, from, to), Arrays.copyOfRange(_values, from, to)); + } + + @Override + public Set<String> keySet() { + KeySet keySet = _keySet; + if (keySet == null) { + keySet = new KeySet(); + _keySet = keySet; + } + return keySet; + } + + @Override + public Collection<V> values() { + Values values = _valuesView; + if (values == null) { + values = new Values(); + _valuesView = values; + } + return values; + } + + @Override + public Set<Map.Entry<String, V>> entrySet() { + EntrySet entrySet = _entrySet; + if (entrySet == null) { + entrySet = new EntrySet(); + _entrySet = entrySet; + } + return entrySet; + } + + @Override + public boolean equals(Object o) { + if (o == this) { + return true; + } + if (o instanceof ImmutableSortedArrayMap) { + ImmutableSortedArrayMap<?> that = (ImmutableSortedArrayMap<?>) o; + return Arrays.equals(_keys, that._keys) && Arrays.equals(_values, that._values); + } + return super.equals(o); + } + + @Override + public int hashCode() { + int hashCode = _hashCode; + if (hashCode == 0) { + for (int i = 0; i < _keys.length; i++) { + hashCode += _keys[i].hashCode() ^ Objects.hashCode(_values[i]); + } + _hashCode = hashCode; + } + return hashCode; + } + + // Mutators. AbstractMap already throws for put; the remaining ones are overridden so that no default method + // reaches a mutable code path. + + @Override + public V put(String key, V value) { + throw new UnsupportedOperationException(); + } + + @Override + public V remove(Object key) { + throw new UnsupportedOperationException(); + } + + @Override + public void putAll(Map<? extends String, ? extends V> m) { + throw new UnsupportedOperationException(); + } + + @Override + public void clear() { + throw new UnsupportedOperationException(); + } + + @Override + public void replaceAll(BiFunction<? super String, ? super V, ? extends V> function) { + throw new UnsupportedOperationException(); + } + + @Override + public V putIfAbsent(String key, V value) { + throw new UnsupportedOperationException(); + } + + @Override + public boolean remove(Object key, Object value) { + throw new UnsupportedOperationException(); + } + + @Override + public boolean replace(String key, V oldValue, V newValue) { + throw new UnsupportedOperationException(); + } + + @Override + public V replace(String key, V value) { + throw new UnsupportedOperationException(); + } + + @Override + public V computeIfAbsent(String key, Function<? super String, ? extends V> mappingFunction) { + throw new UnsupportedOperationException(); + } + + @Override + public V computeIfPresent(String key, BiFunction<? super String, ? super V, ? extends V> remappingFunction) { + throw new UnsupportedOperationException(); + } + + @Override + public V compute(String key, BiFunction<? super String, ? super V, ? extends V> remappingFunction) { + throw new UnsupportedOperationException(); + } + + @Override + public V merge(String key, V value, BiFunction<? super V, ? super V, ? extends V> remappingFunction) { + throw new UnsupportedOperationException(); + } + + /// Iterator over the positions of the map. `remove` throws [UnsupportedOperationException] (the default). + private abstract class ArrayIterator<T> implements Iterator<T> { + private int _next; + + @Override + public boolean hasNext() { + return _next < _keys.length; + } + + @Override + public T next() { + if (_next >= _keys.length) { + throw new NoSuchElementException(); + } + return elementAt(_next++); + } + + abstract T elementAt(int index); + } + + private final class KeySet extends AbstractSet<String> { + @Override + public Iterator<String> iterator() { + return new ArrayIterator<String>() { + @Override + String elementAt(int index) { + return _keys[index]; + } + }; + } + + @Override + public int size() { + return _keys.length; + } + + @Override + public boolean contains(Object o) { + return containsKey(o); + } + + @Override + public boolean equals(Object o) { + if (o == this) { + return true; + } + if (o instanceof ImmutableSortedArrayMap<?>.KeySet) { + return Arrays.equals(_keys, ((ImmutableSortedArrayMap<?>.KeySet) o).keys()); + } + return super.equals(o); + } + + @Override + public int hashCode() { + int hashCode = _keySetHashCode; + if (hashCode == 0) { + for (String key : _keys) { + hashCode += key.hashCode(); + } + _keySetHashCode = hashCode; + } + return hashCode; + } + + private String[] keys() { + return _keys; + } + + @Override + public boolean add(String s) { + throw new UnsupportedOperationException(); + } + + @Override + public boolean remove(Object o) { + throw new UnsupportedOperationException(); + } + + @Override + public void clear() { + throw new UnsupportedOperationException(); + } + } + + private final class Values extends AbstractCollection<V> { + @Override + public Iterator<V> iterator() { + return new ArrayIterator<V>() { + @Override + V elementAt(int index) { + return valueAt(index); + } + }; + } + + @Override + public int size() { + return _keys.length; + } + + @Override + public boolean contains(Object o) { + return containsValue(o); + } + + @Override + public boolean add(V v) { + throw new UnsupportedOperationException(); + } + + @Override + public boolean remove(Object o) { + throw new UnsupportedOperationException(); + } + + @Override + public void clear() { + throw new UnsupportedOperationException(); + } + } + + private final class EntrySet extends AbstractSet<Map.Entry<String, V>> { + @Override + public Iterator<Map.Entry<String, V>> iterator() { + return new ArrayIterator<Map.Entry<String, V>>() { + @Override + Map.Entry<String, V> elementAt(int index) { + return new AbstractMap.SimpleImmutableEntry<>(_keys[index], valueAt(index)); + } + }; + } + + @Override + public int size() { + return _keys.length; + } + + @Override + public boolean contains(Object o) { + if (!(o instanceof Map.Entry)) { + return false; + } + Map.Entry<?, ?> entry = (Map.Entry<?, ?>) o; + Object key = entry.getKey(); + if (!(key instanceof String)) { + return false; + } + int index = indexOf((String) key); + return index >= 0 && Objects.equals(_values[index], entry.getValue()); + } + + @Override + public boolean equals(Object o) { + return super.equals(o); + } Review Comment: This override only delegates to `super.equals()` and can be removed. As smaller cleanup, the `put()` override and each collection view's `add()` override also repeat the same UnsupportedOperationException behavior already provided by AbstractMap/AbstractCollection. The other mutator overrides can remain where they intentionally reject no-op operations or prevent callback execution. ########## 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: This helper verifies compact raw accessor reads, but it does not verify that the routing manager wraps them without copying. Reverting the routing calls to ordinary IdealState/ExternalView constructors would leave these assertions passing. Please narrow the comment to what is checked, or add a focused assertion on the routing wrapper path if this is meant to guard the no-copy wiring. ########## 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: The two `testToIdealStateAndExternalView...` methods primarily test HelixHelper behavior. Could we move them into the existing HelixHelperTest to keep tests with their owning utility? They can still use the compact serializer for their fixtures. ########## pinot-broker/src/main/java/org/apache/pinot/broker/broker/helix/CompactAssignmentZkHelixManager.java: ########## @@ -0,0 +1,126 @@ +/** + * 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.broker.helix; + +import com.google.common.annotations.VisibleForTesting; +import com.google.common.base.Preconditions; +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.ZKHelixManager; +import org.apache.helix.zookeeper.api.client.RealmAwareZkClient; +import org.apache.helix.zookeeper.datamodel.serializer.ChainedPathZkSerializer; +import org.apache.helix.zookeeper.datamodel.serializer.ZNRecordSerializer; +import org.apache.helix.zookeeper.zkclient.serialize.PathBasedZkSerializer; +import org.apache.pinot.common.utils.helix.CompactZNRecordSerializer; +import org.apache.pinot.spi.env.PinotConfiguration; +import org.apache.pinot.spi.utils.CommonConstants.Broker; +import org.apache.pinot.spi.utils.CommonConstants.Helix; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + + +/// A spectator [ZKHelixManager] that reads table ideal states and external views with +/// [CompactZNRecordSerializer]. +/// +/// Helix creates the ZooKeeper client of the manager inside [#connect], with a default +/// [ChainedPathZkSerializer] that sends every path to [ZNRecordSerializer]. Right after that, this class replaces the +/// serializer of the client with one that sends `/<cluster>/IDEALSTATES/...` and `/<cluster>/EXTERNALVIEW/...` to +/// [CompactZNRecordSerializer]. The broker resource keeps [ZNRecordSerializer]: it is small and its readers are +/// outside of routing. The client keeps the serializer across session expiry, because Helix creates a new client only +/// in [#connect], and this class installs the serializer again after each [#connect]. +/// +/// The serializer applies to everything that reads through this manager: its data accessor, its `HelixAdmin`, its +/// property store (no path match) and the Helix callback handlers. All of them only read the records under those +/// paths. Writes still go through [ZNRecordSerializer]. +/// +/// The map fields of a compact record are immutable, and the `IdealState` and `ExternalView` constructors share the +/// inner instance-state maps. So every table ideal state or external view read through this manager, by any +/// reader, has immutable instance-state maps. A read-modify-write of a table ideal state through this manager (for +/// example `HelixHelper.updateIdealState`) fails with [UnsupportedOperationException]. No broker code does that; the +/// participant manager of the broker keeps its own client and the default serializer. +/// +/// Thread-safety: same as [ZKHelixManager]. The serializer is installed from the thread that calls [#connect], +/// before the caller registers any listener. +public class CompactAssignmentZkHelixManager extends ZKHelixManager { + private static final Logger LOGGER = LoggerFactory.getLogger(CompactAssignmentZkHelixManager.class); + + public CompactAssignmentZkHelixManager(String clusterName, String instanceName, InstanceType instanceType, + String zkAddress) { + super(clusterName, instanceName, instanceType, zkAddress); + } + + /// Creates the spectator Helix manager for the broker. Returns a [CompactAssignmentZkHelixManager] when + /// [Broker#CONFIG_OF_ROUTING_COMPACT_ASSIGNMENT_READER_ENABLED] is on, otherwise the default Helix manager. + public static HelixManager createSpectatorHelixManager(PinotConfiguration brokerConf, String clusterName, + String instanceId, String zkServers) { + if (brokerConf.getProperty(Broker.CONFIG_OF_ROUTING_COMPACT_ASSIGNMENT_READER_ENABLED, + Broker.DEFAULT_ROUTING_COMPACT_ASSIGNMENT_READER_ENABLED)) { + LOGGER.info("Using the compact ideal state and external view reader for cluster: {}", clusterName); + return new CompactAssignmentZkHelixManager(clusterName, instanceId, InstanceType.SPECTATOR, zkServers); + } + return HelixManagerFactory.getZKHelixManager(clusterName, instanceId, InstanceType.SPECTATOR, zkServers); + } + + @Override + public void connect() + throws Exception { + super.connect(); + installCompactSerializer(); + } + + /// Replaces the serializer of the ZooKeeper client created by [#connect]. + @VisibleForTesting + void installCompactSerializer() { Review Comment: A few opportunities to reduce the API surface: neither `installCompactSerializer()` nor `createPathBasedSerializer()` has test callers, so both can be private without `@VisibleForTesting`. `getInstalledZkSerializer()` duplicates `getZkClient().getZkSerializer()`. The constructor could also hardcode `InstanceType.SPECTATOR`, matching its contract and every caller, instead of accepting an arbitrary instance type. ########## 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: Consider adding one successful-fallback case. The current malformed-input and bad-trailer cases are rejected by both parsers, so they do not demonstrate recovery through the default deserializer. A gzip record containing `"rawPayload":[1,2,3]` is one candidate: Jackson's byte[] binding accepts integer arrays, while the compact parser's `getBinaryValue()` on START_ARRAY rejects it. Verify the direct compact parse fails but public `deserialize()` returns the same valid record as the default serializer. ########## 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: The assertions verify copy contents and ordering, not linear complexity. Please rename this to something like `testTreeMapCopyCompatibility`, or add an observable check if complexity is intended to be the tested guarantee. ########## pinot-broker/src/test/java/org/apache/pinot/broker/broker/helix/CompactAssignmentZkHelixManagerTest.java: ########## @@ -0,0 +1,233 @@ +/** + * 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.broker.helix; + +import java.lang.reflect.Field; +import java.lang.reflect.Modifier; +import java.util.Map; +import java.util.TreeMap; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; +import org.apache.helix.AccessOption; +import org.apache.helix.BaseDataAccessor; +import org.apache.helix.HelixAdmin; +import org.apache.helix.HelixManager; +import org.apache.helix.InstanceType; +import org.apache.helix.PropertyKey; +import org.apache.helix.manager.zk.ZKHelixAdmin; +import org.apache.helix.manager.zk.ZKHelixManager; +import org.apache.helix.model.ExternalView; +import org.apache.helix.model.IdealState; +import org.apache.helix.zookeeper.api.client.RealmAwareZkClient; +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.helix.zookeeper.zkclient.ZkConnection; +import org.apache.helix.zookeeper.zkclient.serialize.PathBasedZkSerializer; +import org.apache.pinot.common.utils.ZkStarter; +import org.apache.pinot.common.utils.helix.ImmutableSortedArrayMap; +import org.apache.pinot.spi.env.PinotConfiguration; +import org.apache.pinot.spi.utils.CommonConstants.Broker; +import org.apache.pinot.spi.utils.CommonConstants.Helix; +import org.apache.pinot.util.TestUtils; +import org.apache.zookeeper.Watcher; +import org.apache.zookeeper.ZooKeeper; +import org.testng.annotations.AfterClass; +import org.testng.annotations.BeforeClass; +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.assertSame; +import static org.testng.Assert.assertTrue; + + +/// Checks that [CompactAssignmentZkHelixManager] installs +/// [org.apache.pinot.common.utils.helix.CompactZNRecordSerializer] on the client that Helix creates in `connect()`, for +/// table ideal states and external views only. +/// +/// This test fails when a Helix upgrade changes how [ZKHelixManager] creates its ZooKeeper client or its serializer, +/// for example by renaming the protected `_zkclient` field (compile error) or by creating the client outside of +/// `connect()` (the records read through the manager would no longer be compact). +public class CompactAssignmentZkHelixManagerTest { + private static final String CLUSTER_NAME = "CompactAssignmentZkHelixManagerTest"; + private static final String TABLE_NAME = "myTable_OFFLINE"; + + private ZkStarter.ZookeeperInstance _zookeeperInstance; + private ZkClient _zkClient; + + @BeforeClass + public void setUp() { + _zookeeperInstance = ZkStarter.startLocalZkServer(); + String zkUrl = _zookeeperInstance.getZkUrl(); + HelixAdmin helixAdmin = new ZKHelixAdmin.Builder().setZkAddress(zkUrl).build(); + try { + helixAdmin.addCluster(CLUSTER_NAME); + } finally { + helixAdmin.close(); + } + _zkClient = new ZkClient.Builder().setZkServer(zkUrl).setZkSerializer(new ZNRecordSerializer()).build(); + PropertyKey.Builder keyBuilder = new PropertyKey.Builder(CLUSTER_NAME); + _zkClient.createPersistent(keyBuilder.idealStates(TABLE_NAME).getPath(), assignment(TABLE_NAME)); + _zkClient.createPersistent(keyBuilder.externalView(TABLE_NAME).getPath(), assignment(TABLE_NAME)); + _zkClient.createPersistent(keyBuilder.idealStates(Helix.BROKER_RESOURCE_INSTANCE).getPath(), + assignment(Helix.BROKER_RESOURCE_INSTANCE)); + _zkClient.createPersistent(keyBuilder.externalView(Helix.BROKER_RESOURCE_INSTANCE).getPath(), + assignment(Helix.BROKER_RESOURCE_INSTANCE)); + } + + @AfterClass + public void tearDown() { + if (_zkClient != null) { + _zkClient.close(); + } + ZkStarter.stopLocalZkServer(_zookeeperInstance); + } + + private static ZNRecord assignment(String id) { + ZNRecord record = new ZNRecord(id); + // Compressed, because the compact serializer only parses gzip-compressed znodes itself + record.setBooleanField(ZNRecord.ENABLE_COMPRESSION_BOOLEAN_FIELD, true); + record.setSimpleField("REBALANCE_MODE", "CUSTOMIZED"); + record.setMapField("segment_1", Map.of("Server_b_8098", "ONLINE", "Server_a_8098", "ONLINE")); + record.setMapField("segment_0", Map.of("Server_b_8098", "ONLINE", "Server_a_8098", "ONLINE")); + return record; + } + + @Test + public void testZkClientFieldIsProtected() + throws NoSuchFieldException { + Field field = ZKHelixManager.class.getDeclaredField("_zkclient"); + assertTrue(Modifier.isProtected(field.getModifiers())); + assertEquals(field.getType(), RealmAwareZkClient.class); Review Comment: This reflection test duplicates compile-time protection from direct `_zkclient` access in the subclass. Its exact declaring-class, visibility, and type assertions can also reject compatible Helix refactors, such as making the field public or moving it to a superclass. I'd remove this test and retain the behavioral compact-read, reconnect, and session-expiry tests. -- 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]
