add protection against flushing during a test to avoid flakiness for PartitionRowAccountingTest

(cherry picked from commit 451246e17aac5b0c00918666a27c7b8e4d992601)

test/unit/org/apache/cassandra/db/partitions/PartitionRowAccountingTest.java

@@ -20,6 +20,7 @@ package org.apache.cassandra.db.partitions;
 
 import java.lang.reflect.Field;
 import java.util.concurrent.atomic.AtomicLong;
+import java.util.function.Supplier;
 
 import org.junit.BeforeClass;
 import org.junit.Test;
@@ -38,7 +39,6 @@ import org.apache.cassandra.db.Memtable;
 import org.apache.cassandra.utils.btree.BTree;
 
 import static org.assertj.core.api.Assertions.assertThat;
-import static org.junit.Assert.assertSame;
 
 /**
  * Row-level analogue of {@link SetCellAccountingTest}. Where that test grows/resets a single {@code set<text>}
@@ -81,6 +81,70 @@ public class PartitionRowAccountingTest extends CQLTester
                 c.forceBlockingFlush();
     }
 
+    private static final int MAX_ROUND_ATTEMPTS = 10;
+
+    /**
+     * Runs one round of a test -- the writes plus the ownership measurement taken at their end -- and returns that
+     * measurement, guaranteeing it was read from the very memtables the round started writing into.
+     * <p>
+     * Ownership is only comparable while everything the round wrote is still held by the memtable it is measured from,
+     * but a flush can roll a memtable at any moment (memtable pressure, the commitlog, another test's flush). So
+     * instead of failing when that happens we discard the round and repeat it: every attempt starts from a freshly
+     * flushed baseline, and replaying the identical writes (identical timestamps) rebuilds the identical state in the
+     * new memtables. The measurement is taken inside {@code round}, before the memtables are re-checked, so a roll
+     * between the writes and the measurement is caught too.
+     */
+    private <T> T onStableMemtables(Supplier<T> round, ColumnFamilyStore... cfss)
+    {
+        for (int attempt = 1; attempt <= MAX_ROUND_ATTEMPTS; attempt++)
+        {
+            flushAllToBaseline();
+            Memtable[] started = currentMemtables(cfss);
+
+            T measurement = round.get();
+
+            if (stillCurrent(started, cfss))
+                return measurement;
+
+            logger.info("attempt " + attempt + '/' + MAX_ROUND_ATTEMPTS +
+                        " discarded: a memtable rolled (flushed) mid-round, repeating the round");
+        }
+        throw new AssertionError("a memtable rolled (flushed) mid-round in each of the " + MAX_ROUND_ATTEMPTS +
+                                 " attempts -- ownership measurement never valid");
+    }
+
+    private static Memtable[] currentMemtables(ColumnFamilyStore... cfss)
+    {
+        Memtable[] memtables = new Memtable[cfss.length];
+        for (int i = 0; i < cfss.length; i++)
+            memtables[i] = cfss[i].getTracker().getView().getCurrentMemtable();
+        return memtables;
+    }
+
+    private static boolean stillCurrent(Memtable[] started, ColumnFamilyStore... cfss)
+    {
+        for (int i = 0; i < cfss.length; i++)
+            if (started[i] != cfss[i].getTracker().getView().getCurrentMemtable())
+                return false;
+        return true;
+    }
+
+    /**
+     * The on/off-heap ownership of the retain- and removeShadowed-path memtables, captured in one measurement.
+     */
+    private static class Owns
+    {
+        final long onHeapRetain, offHeapRetain, onHeapShadow, offHeapShadow;
+
+        Owns(ColumnFamilyStore cfsRetain, ColumnFamilyStore cfsShadow)
+        {
+            onHeapRetain = ownsOnHeapNow(cfsRetain);
+            offHeapRetain = ownsOffHeapNow(cfsRetain);
+            onHeapShadow = ownsOnHeapNow(cfsShadow);
+            offHeapShadow = ownsOffHeapNow(cfsShadow);
+        }
+    }
+
     private ColumnFamilyStore createTestTable()
     {
         createTable("CREATE TABLE %s (" +
@@ -175,10 +239,7 @@ public class PartitionRowAccountingTest extends CQLTester
         final long insertTs = 1000L;
         final long deleteTs = 2000L; // newer than insertTs, so the tombstone shadows the cells
 
-        flushAllToBaseline();
+        Owns owns = onStableMemtables(() -> {
-        Memtable mtRetain = cfsRetain.getTracker().getView().getCurrentMemtable();
-        Memtable mtShadow = cfsShadow.getTracker().getView().getCurrentMemtable();
-
             for (int ck = 0; ck < rows; ck++)
             {
                 // retain path: write the row, then delete it -> the tombstone shadows existing (owned) cells
@@ -189,16 +250,12 @@ public class PartitionRowAccountingTest extends CQLTester
                 QueryProcessor.executeInternal(rowDelete(tblShadow, ck, deleteTs));
                 QueryProcessor.executeInternal(wideInsert(tblShadow, ck, insertTs));
             }
+            return new Owns(cfsRetain, cfsShadow);
+        }, cfsRetain, cfsShadow);
 
-        assertSame("retain memtable rolled (flushed) mid-test -- on-heap ownership measurement invalid",
+        long onHeapRetain = owns.onHeapRetain, onHeapShadow = owns.onHeapShadow;
-                   mtRetain, cfsRetain.getTracker().getView().getCurrentMemtable());
+        logger.info("retain-path onHeap=" + onHeapRetain + " offHeap=" + owns.offHeapRetain +
-        assertSame("removeShadowed memtable rolled (flushed) mid-test -- on-heap ownership measurement invalid",
+                    "; removeShadowed-path onHeap=" + onHeapShadow + " offHeap=" + owns.offHeapShadow);
-                   mtShadow, cfsShadow.getTracker().getView().getCurrentMemtable());
-
-        long onHeapRetain = ownsOnHeapNow(cfsRetain), onHeapShadow = ownsOnHeapNow(cfsShadow);
-        long offHeapRetain = ownsOffHeapNow(cfsRetain), offHeapShadow = ownsOffHeapNow(cfsShadow);
-        logger.info("retain-path onHeap=" + onHeapRetain + " offHeap=" + offHeapRetain +
-                    "; removeShadowed-path onHeap=" + onHeapShadow + " offHeap=" + offHeapShadow);
 
         assertOwnsNonNegative(cfsRetain, "retain path");
         assertOwnsNonNegative(cfsShadow, "removeShadowed path");
@@ -233,14 +290,11 @@ public class PartitionRowAccountingTest extends CQLTester
         final long insertTs = 1000L;
         final long deleteTs = 2000L; // newer than insertTs, so the tombstone shadows the collection cells
 
-        flushAllToBaseline();
+        Owns owns = onStableMemtables(() -> {
-        Memtable mtRetain = cfsRetain.getTracker().getView().getCurrentMemtable();
-        Memtable mtShadow = cfsShadow.getTracker().getView().getCurrentMemtable();
-
             for (int ck = 0; ck < rows; ck++)
             {
-            // retain path: write the row (large collection), then delete it -> the tombstone shadows existing (owned)
+                // retain path: write the row (large collection), then delete it -> the tombstone shadows existing
-            // cells and the collection's multi-node internal tree
+                // (owned) cells and the collection's multi-node internal tree
                 QueryProcessor.executeInternal(setInsert(tblRetain, ck, insertTs));
                 QueryProcessor.executeInternal(rowDelete(tblRetain, ck, deleteTs));
 
@@ -248,16 +302,12 @@ public class PartitionRowAccountingTest extends CQLTester
                 QueryProcessor.executeInternal(rowDelete(tblShadow, ck, deleteTs));
                 QueryProcessor.executeInternal(setInsert(tblShadow, ck, insertTs));
             }
+            return new Owns(cfsRetain, cfsShadow);
+        }, cfsRetain, cfsShadow);
 
-        assertSame("retain memtable rolled (flushed) mid-test -- on-heap ownership measurement invalid",
+        long onHeapRetain = owns.onHeapRetain, onHeapShadow = owns.onHeapShadow;
-                   mtRetain, cfsRetain.getTracker().getView().getCurrentMemtable());
+        logger.info("collection retain-path onHeap=" + onHeapRetain + " offHeap=" + owns.offHeapRetain +
-        assertSame("removeShadowed memtable rolled (flushed) mid-test -- on-heap ownership measurement invalid",
+                    "; removeShadowed-path onHeap=" + onHeapShadow + " offHeap=" + owns.offHeapShadow);
-                   mtShadow, cfsShadow.getTracker().getView().getCurrentMemtable());
-
-        long onHeapRetain = ownsOnHeapNow(cfsRetain), onHeapShadow = ownsOnHeapNow(cfsShadow);
-        long offHeapRetain = ownsOffHeapNow(cfsRetain), offHeapShadow = ownsOffHeapNow(cfsShadow);
-        logger.info("collection retain-path onHeap=" + onHeapRetain + " offHeap=" + offHeapRetain +
-                    "; removeShadowed-path onHeap=" + onHeapShadow + " offHeap=" + offHeapShadow);
 
         assertOwnsNonNegative(cfsRetain, "retain path");
         assertOwnsNonNegative(cfsShadow, "removeShadowed path");
@@ -295,10 +345,7 @@ public class PartitionRowAccountingTest extends CQLTester
         final long collectionDeleteTs = 1000L;
         final long rowDeleteTs = 2000L; // newer, so the row tombstone supersedes the collection tombstone
 
-        flushAllToBaseline();
+        Owns owns = onStableMemtables(() -> {
-        Memtable mtRetain = cfsRetain.getTracker().getView().getCurrentMemtable();
-        Memtable mtShadow = cfsShadow.getTracker().getView().getCurrentMemtable();
-
             for (int ck = 0; ck < rows; ck++)
             {
                 // retain path: write a collection tombstone (an owned complex deletion over an empty cell tree),
@@ -310,13 +357,10 @@ public class PartitionRowAccountingTest extends CQLTester
                 QueryProcessor.executeInternal(rowDelete(tblShadow, ck, rowDeleteTs));
                 QueryProcessor.executeInternal(collectionDelete(tblShadow, ck, collectionDeleteTs));
             }
+            return new Owns(cfsRetain, cfsShadow);
+        }, cfsRetain, cfsShadow);
 
-        assertSame("retain memtable rolled (flushed) mid-test -- on-heap ownership measurement invalid",
+        long onHeapRetain = owns.onHeapRetain, onHeapShadow = owns.onHeapShadow;
-                   mtRetain, cfsRetain.getTracker().getView().getCurrentMemtable());
-        assertSame("removeShadowed memtable rolled (flushed) mid-test -- on-heap ownership measurement invalid",
-                   mtShadow, cfsShadow.getTracker().getView().getCurrentMemtable());
-
-        long onHeapRetain = ownsOnHeapNow(cfsRetain), onHeapShadow = ownsOnHeapNow(cfsShadow);
         logger.info("collection-tombstone retain-path onHeap=" + onHeapRetain +
                     "; removeShadowed-path onHeap=" + onHeapShadow);