Linearizability checker memory reduction (#40149)

The cache used in linearizability checker now uses approximately 6x less
memory by changing the cache from a set of (bits, state) tuples into a
map from bits -> { state }.

Each combination of states is kept once only, building on the
assumption that the number of state permutations is small compared to
the number of bits permutations. For those histories that are difficult
to check we will have many bits combinations that use the same state
permutations.

We end up now using approximately 15 bytes per entry compared to 101
bytes before, ie. a 6x improvement, allowing us to linearizability check
significantly longer histories.

Re-enabled linearizability checker in CoordinatorTests, hoping above
ensures we no longer run out of memory.

Resolves #39437

server/src/test/java/org/elasticsearch/cluster/coordination/CoordinatorTests.java

@@ -1474,8 +1474,7 @@ public class CoordinatorTests extends ESTestCase {
                 leader.improveConfiguration(lastAcceptedState), sameInstance(lastAcceptedState));
 
             logger.info("checking linearizability of history with size {}: {}", history.size(), history);
-            // See https://github.com/elastic/elasticsearch/issues/39437
+            assertTrue("history not linearizable: " + history, linearizabilityChecker.isLinearizable(spec, history, i -> null));
-            //assertTrue("history not linearizable: " + history, linearizabilityChecker.isLinearizable(spec, history, i -> null));
             logger.info("linearizability check completed");
         }
 

server/src/test/java/org/elasticsearch/cluster/coordination/LinearizabilityChecker.java

@@ -18,6 +18,7 @@
  */
 package org.elasticsearch.cluster.coordination;
 
+import com.carrotsearch.hppc.LongObjectHashMap;
 import org.apache.lucene.util.FixedBitSet;
 import org.elasticsearch.common.collect.Tuple;
 
@@ -217,7 +218,7 @@ public class LinearizabilityChecker {
         Object state = spec.initialState(); // the current state of the datatype
         final FixedBitSet linearized = new FixedBitSet(history.size() / 2); // the linearized prefix of the history
 
-        final Set<Tuple<Object, FixedBitSet>> cache = new HashSet<>(); // cache of explored <state, linearized prefix> pairs
+        final Cache cache = new Cache();
         final Deque<Tuple<Entry, Object>> calls = new LinkedList<>(); // path we're currently exploring
 
         final Entry headEntry = createLinkedEntries(history);
@@ -231,7 +232,7 @@ public class LinearizabilityChecker {
                     // check if we have already explored this linearization
                     final FixedBitSet updatedLinearized = linearized.clone();
                     updatedLinearized.set(entry.id);
-                    shouldExploreNextState = cache.add(new Tuple<>(maybeNextState.get(), updatedLinearized));
+                    shouldExploreNextState = cache.add(maybeNextState.get(), updatedLinearized);
                 }
                 if (shouldExploreNextState) {
                     calls.push(new Tuple<>(entry, state));
@@ -373,4 +374,79 @@ public class LinearizabilityChecker {
         }
     }
 
+
+    /**
+     * A cache optimized for small bit-counts (less than 64) and small number of unique permutations of state objects.
+     *
+     * Each combination of states is kept once only, building on the
+     * assumption that the number of permutations is small compared to the
+     * number of bits permutations. For those histories that are difficult to check
+     * we will have many bits combinations that use the same state permutations.
+     *
+     * The smallMap optimization allows us to avoid object overheads for bit-sets up to 64 bit large.
+     *
+     * Comparing set of (bits, state) to smallMap:
+     * (bits, state) : 24 (tuple) + 24 (FixedBitSet) + 24 (bits) + 5 (hash buckets) + 24 (hashmap node).
+     * smallMap bits to {state} : 10 (bits) + 5 (hash buckets) + avg-size of unique permutations.
+     *
+     * The avg-size of the unique permutations part is very small compared to the
+     * sometimes large number of bits combinations (which are the cases where
+     * we run into trouble).
+     *
+     * set of (bits, state) totals 101 bytes compared to smallMap bits to { state }
+     * which totals 15 bytes, ie. a 6x improvement in memory usage.
+     */
+    private static class Cache {
+        private final Map<Object, Set<FixedBitSet>> largeMap = new HashMap<>();
+        private final LongObjectHashMap<Set<Object>> smallMap = new LongObjectHashMap<>();
+        private final Map<Object, Object> internalizeStateMap = new HashMap<>();
+        private final Map<Set<Object>, Set<Object>> statePermutations = new HashMap<>();
+
+        /**
+         * Add state, bits combination
+         * @return true if added, false if already registered.
+         */
+        public boolean add(Object state, FixedBitSet bitSet) {
+            return addInternal(internalizeStateMap.computeIfAbsent(state, k -> state), bitSet);
+        }
+
+        private boolean addInternal(Object state, FixedBitSet bitSet) {
+            long[] bits = bitSet.getBits();
+            if (bits.length == 1)
+                return addSmall(state, bits[0]);
+            else
+                return addLarge(state, bitSet);
+        }
+
+        private boolean addSmall(Object state, long bits) {
+            int index = smallMap.indexOf(bits);
+            Set<Object> states;
+            if (index < 0) {
+                states = Collections.singleton(state);
+            } else {
+                Set<Object> oldStates = smallMap.indexGet(index);
+                if (oldStates.contains(state))
+                    return false;
+                states = new HashSet<>(oldStates.size() + 1);
+                states.addAll(oldStates);
+                states.add(state);
+            }
+
+            // Get a unique set object per state permutation. We assume that the number of permutations of states are small.
+            // We thus avoid the overhead of the set data structure.
+            states = statePermutations.computeIfAbsent(states, k -> k);
+
+            if (index < 0) {
+                smallMap.indexInsert(index, bits, states);
+            } else {
+                smallMap.indexReplace(index, states);
+            }
+
+            return true;
+        }
+
+        private boolean addLarge(Object state, FixedBitSet bitSet) {
+            return largeMap.computeIfAbsent(state, k -> new HashSet<>()).add(bitSet);
+        }
+    }
 }