CNDB-12899: Add a cached version of the estimatedPartitionCount metric (#1559)

### What is the issue
[CNDB-12899](riptano/cndb#12899)
`CompactionRealm.estimatedPartitionCount()` is very expensive

### What does this PR fix and why was it fixed
Adds a cached version of the metric and removes the memtable partitions
from the calculation to make it more precise for the compaction use
case.

Also makes sure that the `estimatedPartitionCount` metric is not
recalculated if the table's data view (i.e. sstable and memtable set)
has not changed.

---------

Co-authored-by: Szymon Miężał <[email protected]>

Author: blambov

src/java/org/apache/cassandra/db/compaction/CompactionRealm.java

@@ -31,7 +31,6 @@
 import org.apache.cassandra.db.Directories;
 import org.apache.cassandra.db.DiskBoundaries;
 import org.apache.cassandra.db.compaction.unified.Environment;
-import org.apache.cassandra.db.compaction.unified.RealEnvironment;
 import org.apache.cassandra.db.lifecycle.LifecycleTransaction;
 import org.apache.cassandra.db.lifecycle.SSTableSet;
 import org.apache.cassandra.db.memtable.Memtable;
@@ -126,12 +125,12 @@ default IPartitioner getPartitioner()
      * Return the estimated partition count, used when the number of partitions in an sstable is not sufficient to give
      * a sensible range estimation.
      */
-    default long estimatedPartitionCount()
+    default long estimatedPartitionCountInSSTables()
     {
         final long INITIAL_ESTIMATED_PARTITION_COUNT = 1 << 16; // If we don't yet have a count, use a sensible default.
         if (metrics() == null)
             return INITIAL_ESTIMATED_PARTITION_COUNT;
-        final Long estimation = metrics().estimatedPartitionCount.getValue();
+        final Long estimation = metrics().estimatedPartitionCountInSSTablesCached.getValue();
         if (estimation == null || estimation == 0)
             return INITIAL_ESTIMATED_PARTITION_COUNT;
         return estimation;

src/java/org/apache/cassandra/db/compaction/DelegatingShardManager.java

@@ -60,7 +60,7 @@ public double shardSetCoverage()
     @Override
     public double minimumPerPartitionSpan()
     {
-        return localSpaceCoverage() / Math.max(1, realm.estimatedPartitionCount());
+        return localSpaceCoverage() / Math.max(1, realm.estimatedPartitionCountInSSTables());
     }
 
     @Override

src/java/org/apache/cassandra/db/compaction/ShardManagerNoDisks.java

@@ -87,7 +87,7 @@ public double shardSetCoverage()
 
     public double minimumPerPartitionSpan()
     {
-        return localSpaceCoverage() / Math.max(1, localRanges.getRealm().estimatedPartitionCount());
+        return localSpaceCoverage() / Math.max(1, localRanges.getRealm().estimatedPartitionCountInSSTables());
     }
 
     @Override

src/java/org/apache/cassandra/db/compaction/ShardManagerReplicaAware.java

@@ -20,21 +20,15 @@
 
 import java.util.ArrayList;
 import java.util.Arrays;
-import java.util.Set;
 import java.util.concurrent.ConcurrentHashMap;
 
-import javax.annotation.Nullable;
-
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
-import org.apache.cassandra.db.PartitionPosition;
 import org.apache.cassandra.dht.IPartitioner;
 import org.apache.cassandra.dht.Range;
 import org.apache.cassandra.dht.Token;
 import org.apache.cassandra.dht.tokenallocator.IsolatedTokenAllocator;
-import org.apache.cassandra.io.sstable.format.SSTableReader;
-import org.apache.cassandra.io.sstable.format.SSTableWriter;
 import org.apache.cassandra.locator.AbstractReplicationStrategy;
 import org.apache.cassandra.locator.TokenMetadata;
 
@@ -94,7 +88,7 @@ public double shardSetCoverage()
     @Override
     public double minimumPerPartitionSpan()
     {
-        return localSpaceCoverage() / Math.max(1, realm.estimatedPartitionCount());
+        return localSpaceCoverage() / Math.max(1, realm.estimatedPartitionCountInSSTables());
     }
 
     @Override

src/java/org/apache/cassandra/io/sstable/format/SSTableReader.java

@@ -332,7 +332,7 @@ public static long getApproximateKeyCount(Iterable<? extends SSTableReader> ssta
         long count = -1;
 
         if (Iterables.isEmpty(sstables))
-            return count;
+            return 0;
 
         boolean failed = false;
         ICardinality cardinality = null;

src/java/org/apache/cassandra/metrics/TableMetrics.java

@@ -17,6 +17,7 @@
  */
 package org.apache.cassandra.metrics;
 
+import java.lang.ref.WeakReference;
 import java.nio.BufferUnderflowException;
 import java.nio.ByteBuffer;
 import java.util.ArrayList;
@@ -30,6 +31,8 @@
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentMap;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicReference;
+import java.util.function.LongSupplier;
 import java.util.function.Predicate;
 import java.util.stream.Stream;
 
@@ -44,6 +47,7 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import com.codahale.metrics.CachedGauge;
 import com.codahale.metrics.Counter;
 import com.codahale.metrics.Gauge;
 import com.codahale.metrics.Histogram;
@@ -74,6 +78,7 @@
 import org.apache.cassandra.utils.Hex;
 import org.apache.cassandra.utils.MovingAverage;
 import org.apache.cassandra.utils.Pair;
+import org.apache.cassandra.utils.concurrent.Refs;
 
 import static org.apache.cassandra.io.sstable.format.SSTableReader.selectOnlyBigTableReaders;
 import static org.apache.cassandra.metrics.CassandraMetricsRegistry.Metrics;
@@ -174,6 +179,12 @@ public String asCQLString()
     public final Gauge<long[]> estimatedPartitionSizeHistogram;
     /** Approximate number of keys in table. */
     public final Gauge<Long> estimatedPartitionCount;
+    /** This function is used to calculate estimated partition count in sstables and store the calculated value for the
+     *  current set of sstables. */
+    public final LongSupplier estimatedPartitionCountInSSTables;
+    /** A cached version of the estimated partition count in sstables, used by compaction. This value will be more
+     *  precise when the table has a small number of partitions that keep getting written to. */
+    public final Gauge<Long> estimatedPartitionCountInSSTablesCached;
     /** Histogram of estimated number of columns. */
     public final Gauge<long[]> estimatedColumnCountHistogram;
 
@@ -605,20 +616,53 @@ public Long getValue()
         estimatedPartitionSizeHistogram = createTableGauge("EstimatedPartitionSizeHistogram", "EstimatedRowSizeHistogram",
                                                            () -> combineHistograms(cfs.getSSTables(SSTableSet.CANONICAL),
                                                                                    SSTableReader::getEstimatedPartitionSize), null);
-        
+
+        estimatedPartitionCountInSSTables = new LongSupplier()
+        {
+            // Since the sstables only change when the tracker view changes, we can cache the value.
+            AtomicReference<Pair<WeakReference<View>, Long>> collected = new AtomicReference<>(Pair.create(new WeakReference<>(null), 0L));
+
+            public long getAsLong()
+            {
+                final View currentView = cfs.getTracker().getView();
+                final Pair<WeakReference<View>, Long> currentCollected = collected.get();
+                if (currentView != currentCollected.left.get())
+                {
+                    Refs<SSTableReader> refs = Refs.tryRef(currentView.select(SSTableSet.CANONICAL));
+                    if (refs != null)
+                    {
+                        try (refs)
+                        {
+                            long collectedValue = SSTableReader.getApproximateKeyCount(refs);
+                            final Pair<WeakReference<View>, Long> newCollected = Pair.create(new WeakReference<>(currentView), collectedValue);
+                            collected.compareAndSet(currentCollected, newCollected); // okay if failed, a different thread did it
+                            return collectedValue;
+                        }
+                    }
+                    // If we can't reference, simply return the previous collected value; it can only result in a delay
+                    // in reporting the correct key count.
+                }
+                return currentCollected.right;
+            }
+        };
         estimatedPartitionCount = createTableGauge("EstimatedPartitionCount", "EstimatedRowCount", new Gauge<Long>()
         {
             public Long getValue()
             {
-                long memtablePartitions = 0;
+                long estimatedPartitions = estimatedPartitionCountInSSTables.getAsLong();
                 for (Memtable memtable : cfs.getTracker().getView().getAllMemtables())
-                   memtablePartitions += memtable.partitionCount();
-                try(ColumnFamilyStore.RefViewFragment refViewFragment = cfs.selectAndReference(View.selectFunction(SSTableSet.CANONICAL)))
-                {
-                    return SSTableReader.getApproximateKeyCount(refViewFragment.sstables) + memtablePartitions;
-                }
+                    estimatedPartitions += memtable.partitionCount();
+                return estimatedPartitions;
             }
         }, null);
+        estimatedPartitionCountInSSTablesCached = new CachedGauge<Long>(1, TimeUnit.SECONDS)
+        {
+            public Long loadValue()
+            {
+                return estimatedPartitionCountInSSTables.getAsLong();
+            }
+        };
+
         estimatedColumnCountHistogram = createTableGauge("EstimatedColumnCountHistogram", "EstimatedColumnCountHistogram",
                                                          () -> combineHistograms(cfs.getSSTables(SSTableSet.CANONICAL), 
                                                                                  SSTableReader::getEstimatedCellPerPartitionCount), null);

test/unit/org/apache/cassandra/db/compaction/DelegatingShardManagerTest.java

@@ -42,7 +42,7 @@ public void testWrappingShardManagerNoDisks()
     {
         CompactionRealm realm = Mockito.mock(CompactionRealm.class);
         when(realm.getPartitioner()).thenReturn(partitioner);
-        when(realm.estimatedPartitionCount()).thenReturn(1L << 16);
+        when(realm.estimatedPartitionCountInSSTables()).thenReturn(1L << 16);
         SortedLocalRanges localRanges = SortedLocalRanges.forTestingFull(realm);
         ShardManager delegate = new ShardManagerNoDisks(localRanges);
 

test/unit/org/apache/cassandra/db/compaction/ShardManagerReplicaAwareTest.java

@@ -49,7 +49,7 @@ public class ShardManagerReplicaAwareTest
     public void testRangeEndsForShardCountEqualtToNumTokensPlusOne() throws UnknownHostException
     {
         var mockCompationRealm = mock(CompactionRealm.class);
-        when(mockCompationRealm.estimatedPartitionCount()).thenReturn(1L<<16);
+        when(mockCompationRealm.estimatedPartitionCountInSSTables()).thenReturn(1L<<16);
 
         for (int numTokens = 1; numTokens < 32; numTokens++)
         {
@@ -75,7 +75,7 @@ public void testRangeEndsForShardCountEqualtToNumTokensPlusOne() throws UnknownH
     public void testRangeEndsAreFromTokenListAndContainLowerRangeEnds() throws UnknownHostException
     {
         var mockCompationRealm = mock(CompactionRealm.class);
-        when(mockCompationRealm.estimatedPartitionCount()).thenReturn(1L<<16);
+        when(mockCompationRealm.estimatedPartitionCountInSSTables()).thenReturn(1L<<16);
 
         for (int nodeCount = 1; nodeCount <= 6; nodeCount++)
         {

test/unit/org/apache/cassandra/db/compaction/ShardManagerTest.java

@@ -27,7 +27,6 @@
 import java.util.stream.Collectors;
 
 import org.junit.Assert;
-import org.junit.Assume;
 import org.junit.Before;
 import org.junit.Test;
 
@@ -69,7 +68,7 @@ public void setUp()
         localRanges = Mockito.mock(SortedLocalRanges.class, Mockito.withSettings().defaultAnswer(Mockito.CALLS_REAL_METHODS));
         Mockito.when(localRanges.getRanges()).thenAnswer(invocation -> weightedRanges);
         Mockito.when(localRanges.getRealm()).thenReturn(realm);
-        Mockito.when(realm.estimatedPartitionCount()).thenReturn(10000L);
+        Mockito.when(realm.estimatedPartitionCountInSSTables()).thenReturn(10000L);
     }
 
     @Test

test/unit/org/apache/cassandra/metrics/TableMetricsTest.java

@@ -19,6 +19,7 @@
 package org.apache.cassandra.metrics;
 
 import java.io.IOException;
+import java.util.concurrent.TimeUnit;
 import java.util.function.Supplier;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
@@ -37,6 +38,7 @@
 import org.apache.cassandra.db.Keyspace;
 import org.apache.cassandra.exceptions.ConfigurationException;
 import org.apache.cassandra.service.EmbeddedCassandraService;
+import org.awaitility.Awaitility;
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertTrue;
@@ -363,6 +365,55 @@ public void testViewMetricsCleanupOnDrop()
         assertEquals(metrics.get().collect(Collectors.joining(",")), 0, metrics.get().count());
     }
 
+    @Test
+    public void testEstimatedPartitionCount() throws InterruptedException
+    {
+        ColumnFamilyStore cfs = recreateTable();
+        assertEquals(0L, cfs.metric.estimatedPartitionCount.getValue().longValue());
+        assertEquals(0L, cfs.metric.estimatedPartitionCountInSSTablesCached.getValue().longValue());
+        long startTime = System.currentTimeMillis();
+
+        int partitionCount = 10;
+        int numRows = 100;
+
+        for (int i = 0; i < numRows; i++)
+            session.execute(String.format("INSERT INTO %s.%s (id, val1, val2) VALUES (%d, '%s', '%s')", KEYSPACE, TABLE, i % partitionCount, "val" + i, "val" + i));
+
+        assertEquals(partitionCount, cfs.metric.estimatedPartitionCount.getValue().longValue());
+        cfs.forceBlockingFlush(ColumnFamilyStore.FlushReason.UNIT_TESTS);
+        assertEquals(partitionCount, cfs.metric.estimatedPartitionCount.getValue().longValue());
+
+        long estimatedPartitionCountInSSTables = cfs.metric.estimatedPartitionCountInSSTablesCached.getValue().longValue();
+        long elapsedTime = System.currentTimeMillis() - startTime;
+        // the caching time is one second; avoid flakiness by only checking if a long time has not passed
+        // (Because we take the time after calling the method, elapsedTime < 1000 should also be stable, but let's also
+        // accommodate the possibility that the cache uses a different timer with different tick times.)
+        if (elapsedTime < 980)
+            assertEquals(0, estimatedPartitionCountInSSTables);
+
+        for (int i = 0; i < numRows; i++)
+            session.execute(String.format("INSERT INTO %s.%s (id, val1, val2) VALUES (%d, '%s', '%s')", KEYSPACE, TABLE, i % partitionCount, "val" + i, "val" + i));
+
+        estimatedPartitionCountInSSTables = cfs.metric.estimatedPartitionCountInSSTablesCached.getValue().longValue();
+        elapsedTime = System.currentTimeMillis() - startTime;
+        if (elapsedTime < 980)
+            assertEquals(0, estimatedPartitionCountInSSTables);
+        else if (elapsedTime >= 1020)
+            assertEquals(partitionCount, estimatedPartitionCountInSSTables);
+
+        // The answer below is incorrect but what the metric currently returns.
+        assertEquals(partitionCount * 2, cfs.metric.estimatedPartitionCount.getValue().longValue());
+        cfs.forceBlockingFlush(ColumnFamilyStore.FlushReason.UNIT_TESTS);
+        // Recalculation for the new sstable set will correct it.
+        assertEquals(partitionCount, cfs.metric.estimatedPartitionCount.getValue().longValue());
+
+        // The cached estimatedPartitionCountInSSTables lags one second, check that.
+        // Assert that the metric will return a correct value after at least a second passes
+        Awaitility.await()
+                  .atMost(2, TimeUnit.SECONDS)
+                  .untilAsserted(() -> assertEquals(partitionCount, (long) cfs.metric.estimatedPartitionCountInSSTablesCached.getValue()));
+    }
+
 
     @AfterClass
     public static void tearDown()