HCD-75: Fix and improve EstimatedRowCount metric (#1604)

This patch introduces two changes:
- it adds a reading group to guard against sweeping the memtable which
the metric
is going to potentially iterate through (preventing crashes).
- changes the metric calculation by using an estimate (used already by
SAI query planner) instead of iterating
through the whole memtable (which is quite a heavy operation).

src/java/org/apache/cassandra/db/memtable/TrieMemtable.java

@@ -300,22 +300,6 @@ public int getShardCount()
         return shards.length;
     }
 
-    public long rowCount(final ColumnFilter columnFilter, final DataRange dataRange)
-    {
-        int total = 0;
-        for (MemtableUnfilteredPartitionIterator iter = makePartitionIterator(columnFilter, dataRange); iter.hasNext(); )
-        {
-            for (UnfilteredRowIterator it = iter.next(); it.hasNext(); )
-            {
-                Unfiltered uRow = it.next();
-                if (uRow.isRow())
-                    total++;
-            }
-        }
-
-        return total;
-    }
-
     @Override
     public long getEstimatedAverageRowSize()
     {

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

@@ -78,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.OpOrder;
 import org.apache.cassandra.utils.concurrent.Refs;
 
 import static org.apache.cassandra.io.sstable.format.SSTableReader.selectOnlyBigTableReaders;
@@ -667,28 +668,38 @@ public Long loadValue()
                                                          () -> combineHistograms(cfs.getSSTables(SSTableSet.CANONICAL), 
                                                                                  SSTableReader::getEstimatedCellPerPartitionCount), null);
 
-        estimatedRowCount = createTableGauge("EstimatedRowCount", "EstimatedRowCount", new Gauge<Long>()
+        estimatedRowCount = createTableGauge("EstimatedRowCount", "EstimatedRowCount", new CachedGauge<>(1, TimeUnit.SECONDS)
         {
-            public Long getValue()
+            public Long loadValue()
             {
                 long memtableRows = 0;
-                for (Memtable memtable : cfs.getTracker().getView().getAllMemtables())
+                OpOrder.Group readGroup = null;
+                try
                 {
-                    if (memtable instanceof TrieMemtable)
+                    for (Memtable memtable : cfs.getTracker().getView().getAllMemtables())
                     {
-                        ColumnFilter.Builder builder = ColumnFilter.allRegularColumnsBuilder(cfs.metadata(), true);
-                        memtableRows += ((TrieMemtable) memtable).rowCount(builder.build(), DataRange.allData(cfs.getPartitioner()));
+                        if (readGroup == null)
+                        {
+                            readGroup = memtable.readOrdering().start();
+                        }
+                        memtableRows += Memtable.estimateRowCount(memtable);
                     }
                 }
+                finally
+                {
+                    if (readGroup != null)
+                        readGroup.close();
+                }
+
+                long sstableRows = 0;
                 try(ColumnFamilyStore.RefViewFragment refViewFragment = cfs.selectAndReference(View.selectFunction(SSTableSet.CANONICAL)))
                 {
-                    long total = 0;
                     for (SSTableReader reader: refViewFragment.sstables)
                     {
-                        total += reader.getTotalRows();
+                        sstableRows += reader.getTotalRows();
                     }
-                    return total + memtableRows;
                 }
+                return sstableRows + memtableRows;
             }
         }, null);
 

test/unit/org/apache/cassandra/db/memtable/MemtableSizeTestBase.java

@@ -31,15 +31,10 @@
 import org.apache.cassandra.config.DatabaseDescriptor;
 import org.apache.cassandra.cql3.CQLTester;
 import org.apache.cassandra.db.ColumnFamilyStore;
-import org.apache.cassandra.db.DataRange;
 import org.apache.cassandra.db.Keyspace;
-import org.apache.cassandra.db.filter.ColumnFilter;
 import org.apache.cassandra.utils.FBUtilities;
-import org.apache.cassandra.utils.memory.SlabAllocator;
 import org.github.jamm.MemoryMeter;
 
-import static org.assertj.core.api.Assertions.assertThat;
-
 // Note: This test can be run in idea with the allocation type configured in the test yaml and memtable using the
 // value memtableClass is initialized with.
 public abstract class MemtableSizeTestBase extends CQLTester
@@ -213,27 +208,6 @@ public void testSize() throws Throwable
         }
     }
 
-    @Test
-    public void testRowCountInTrieMemtable() throws Throwable
-    {
-        buildAndFillTable("TrieMemtable");
-
-        String writeStatement = "INSERT INTO " + table + "(userid,picid,commentid)VALUES(?,?,?)";
-
-        Memtable memtable = cfs.getTracker().getView().getCurrentMemtable();
-        System.out.println("Writing " + partitions + " partitions of " + rowsPerPartition + " rows");
-        for (long i = 0; i < partitions; ++i)
-        {
-            for (long j = 0; j < rowsPerPartition; ++j)
-                execute(writeStatement, i, j, i + j);
-        }
-
-        assertThat(memtable).isExactlyInstanceOf(TrieMemtable.class);
-        ColumnFilter.Builder builder = ColumnFilter.allRegularColumnsBuilder(cfs.metadata(), true);
-        long rowCount = ((TrieMemtable)cfs.getTracker().getView().getCurrentMemtable()).rowCount(builder.build(), DataRange.allData(cfs.getPartitioner()));
-        Assert.assertEquals(rowCount, partitions*rowsPerPartition);
-    }
-
     @Test
     public void testRowSize() throws Throwable
     {
@@ -253,4 +227,4 @@ public void testRowSize() throws Throwable
         double expectedRowSize = (double) memtable.getLiveDataSize() / (partitions * rowsPerPartition);
         Assert.assertEquals(expectedRowSize, rowSize, (partitions * rowsPerPartition) * 0.05);  // 5% accuracy
     }
-}
+}