gortiz commented on code in PR #19786: URL: https://github.com/apache/pinot/pull/19786#discussion_r4280901302
########## 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: Agreed, the production caller already filters plain input, so the direct method is gzip-only now with an assert. The direct-parser tests compress their input, and the public `deserialize()` tests still cover plain 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]. +/// +/// 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: Good catch. `copyOf()` is package-private now (only tests in this package use it), and the class javadoc says values must be immutable because the hash code is cached. ########## 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: You are right, the javadoc promised more than the code does. It now says range methods return copies without view bounds, and that some inherited view mutators return normally when they have nothing to change. ########## 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: Done. Both methods are private without `@VisibleForTesting`, `getInstalledZkSerializer()` is gone, and the constructor always creates a spectator. ########## 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: Removed the reflection test. The compact-read, reconnect and session-expiry tests remain. ########## 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: Removed. Occupancy is just `_maps[slot] == null`, so a zero hash needs no special handling. ########## 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: Removed `put()`, the view `add()` overrides and `EntrySet.equals`, together with the cached `EntrySet.hashCode()`, so the entry set uses the `AbstractSet` versions. -- 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]
