Fix flakiness of TopPartitionsTest

patch by Dmitry Konstantinov; reviewed by Stefan Miklosovic, Berenguer Blasi for CASSANDRA-19991

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

@@ -42,7 +42,7 @@ public abstract class Sampler<T>
     MonotonicClock clock = MonotonicClock.approxTime;
 
     @VisibleForTesting
-    static final ThreadPoolExecutor samplerExecutor = new JMXEnabledThreadPoolExecutor(1, 1,
+    public static final ThreadPoolExecutor samplerExecutor = new JMXEnabledThreadPoolExecutor(1, 1,
             TimeUnit.SECONDS,
             new ArrayBlockingQueue<Runnable>(1000),
             new NamedThreadFactory("Sampler"),

test/unit/org/apache/cassandra/Util.java

@@ -86,6 +86,7 @@ import org.apache.cassandra.utils.CounterId;
 import org.apache.cassandra.utils.FBUtilities;
 import org.apache.cassandra.utils.FilterFactory;
 import org.awaitility.Awaitility;
+import org.hamcrest.Matcher;
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
@@ -626,12 +627,17 @@ public class Util
     }
 
     public static <T> void spinAssertEquals(String message, T expected, Supplier<? extends T> actualSupplier, long timeout, TimeUnit timeUnit)
+    {
+        spinAssert(message, equalTo(expected), actualSupplier, timeout, timeUnit);
+    }
+
+    public static <T> void spinAssert(String message, Matcher<T> matcher, Supplier<? extends T> actualSupplier, long timeout, TimeUnit timeUnit)
     {
         Awaitility.await()
                   .pollInterval(Duration.ofMillis(100))
                   .pollDelay(0, TimeUnit.MILLISECONDS)
                   .atMost(timeout, timeUnit)
-                  .untilAsserted(() -> assertThat(message, actualSupplier.get(), equalTo(expected)));
+                  .untilAsserted(() -> assertThat(message, actualSupplier.get(), matcher));
     }
 
     public static void joinThread(Thread thread) throws InterruptedException

test/unit/org/apache/cassandra/tools/TopPartitionsTest.java

@@ -33,13 +33,16 @@ import org.junit.BeforeClass;
 import org.junit.Test;
 
 import org.apache.cassandra.SchemaLoader;
+import org.apache.cassandra.Util;
 import org.apache.cassandra.db.ColumnFamilyStore;
 import org.apache.cassandra.db.SystemKeyspace;
 import org.apache.cassandra.exceptions.ConfigurationException;
+import org.apache.cassandra.metrics.Sampler;
 import org.apache.cassandra.service.StorageService;
 
 import static java.lang.String.format;
 import static org.apache.cassandra.cql3.QueryProcessor.executeInternal;
+import static org.hamcrest.Matchers.greaterThan;
 import static org.junit.Assert.assertEquals;
 
 public class TopPartitionsTest
@@ -76,10 +79,35 @@ public class TopPartitionsTest
     @Test
     public void testServiceTopPartitionsSingleTable() throws Exception
     {
-        ColumnFamilyStore.getIfExists("system", "local").beginLocalSampling("READS", 5, 240000);
+        ColumnFamilyStore columnFamilyStore = ColumnFamilyStore.getIfExists("system", "local");
+        String samplerName = "READS";
+        long executedBefore = Sampler.samplerExecutor.getCompletedTaskCount();
+        columnFamilyStore.beginLocalSampling(samplerName, 5, 240_000);
+
         String req = "SELECT * FROM system.%s WHERE key='%s'";
         executeInternal(format(req, SystemKeyspace.LOCAL, SystemKeyspace.LOCAL));
-        List<CompositeData> result = ColumnFamilyStore.getIfExists("system", "local").finishLocalSampling("READS", 5);
+        ensureThatSamplerExecutorProcessedAllSamples(executedBefore);
+
+        List<CompositeData> result = columnFamilyStore.finishLocalSampling(samplerName, 5);
         assertEquals("If this failed you probably have to raise the beginLocalSampling duration", 1, result.size());
     }
+
+    private static void ensureThatSamplerExecutorProcessedAllSamples(long executedBefore)
+    {
+        Util.spinAssertEquals("samplerExecutor should not have pending tasks",
+                              0,
+                              () -> Sampler.samplerExecutor.getQueue().size(),
+                              60,
+                              TimeUnit.SECONDS);
+        Util.spinAssertEquals("samplerExecutor should not have active tasks",
+                              0,
+                              Sampler.samplerExecutor::getActiveCount,
+                              60,
+                              TimeUnit.SECONDS);
+        Util.spinAssert("samplerExecutor has not completed any new tasks after beginLocalSampling",
+                        greaterThan(executedBefore),
+                        Sampler.samplerExecutor::getCompletedTaskCount,
+                        60,
+                        TimeUnit.SECONDS);
+    }
 }