From 8de33336623fbdda97d91a81f584d0c6a7c83ba4 Mon Sep 17 00:00:00 2001
From: Piotr Findeisen <piotr.findeisen@gmail.com>
Date: Tue, 21 Nov 2023 10:30:00 +0100
Subject: [PATCH] Fix TestShardOrganizer.testShardOrganizerInProgress flakiness

There was a race, the test job was sleeping for only 10ms.
However, since GC pause is possible, sleeping "long enough" should be
measured in seconds. Use a latch for synchronization instead of fixed
sleep time.
---
 .../TestShardOrganizationManager.java         | 17 +++++++++++-
 .../organization/TestShardOrganizer.java      | 26 +++++++------------
 2 files changed, 25 insertions(+), 18 deletions(-)

diff --git a/plugin/trino-raptor-legacy/src/test/java/io/trino/plugin/raptor/legacy/storage/organization/TestShardOrganizationManager.java b/plugin/trino-raptor-legacy/src/test/java/io/trino/plugin/raptor/legacy/storage/organization/TestShardOrganizationManager.java
index 9115d1b3b222..7ee6ae0061a6 100644
--- a/plugin/trino-raptor-legacy/src/test/java/io/trino/plugin/raptor/legacy/storage/organization/TestShardOrganizationManager.java
+++ b/plugin/trino-raptor-legacy/src/test/java/io/trino/plugin/raptor/legacy/storage/organization/TestShardOrganizationManager.java
@@ -35,13 +35,13 @@
 import java.util.UUID;
 
 import static com.google.common.collect.Iterables.getOnlyElement;
+import static com.google.common.util.concurrent.Uninterruptibles.sleepUninterruptibly;
 import static io.airlift.units.Duration.nanosSince;
 import static io.trino.plugin.raptor.legacy.DatabaseTesting.createTestingJdbi;
 import static io.trino.plugin.raptor.legacy.metadata.SchemaDaoUtil.createTablesWithRetry;
 import static io.trino.plugin.raptor.legacy.metadata.TestDatabaseShardManager.createShardManager;
 import static io.trino.plugin.raptor.legacy.storage.organization.ShardOrganizationManager.createOrganizationSets;
 import static io.trino.plugin.raptor.legacy.storage.organization.TestCompactionSetCreator.extractIndexes;
-import static io.trino.plugin.raptor.legacy.storage.organization.TestShardOrganizer.createShardOrganizer;
 import static io.trino.spi.type.BigintType.BIGINT;
 import static io.trino.spi.type.DateType.DATE;
 import static io.trino.spi.type.TimestampType.TIMESTAMP_MILLIS;
@@ -208,4 +208,19 @@ private ShardOrganizationManager createShardOrganizationManager(long intervalMil
                 new Duration(intervalMillis, MILLISECONDS),
                 new Duration(5, MINUTES));
     }
+
+    private static class MockJobFactory
+            implements JobFactory
+    {
+        @Override
+        public Runnable create(OrganizationSet organizationSet)
+        {
+            return () -> sleepUninterruptibly(10, MILLISECONDS);
+        }
+    }
+
+    private static ShardOrganizer createShardOrganizer()
+    {
+        return new ShardOrganizer(new MockJobFactory(), 1);
+    }
 }
diff --git a/plugin/trino-raptor-legacy/src/test/java/io/trino/plugin/raptor/legacy/storage/organization/TestShardOrganizer.java b/plugin/trino-raptor-legacy/src/test/java/io/trino/plugin/raptor/legacy/storage/organization/TestShardOrganizer.java
index 9ba90eace90f..f72f4f464a55 100644
--- a/plugin/trino-raptor-legacy/src/test/java/io/trino/plugin/raptor/legacy/storage/organization/TestShardOrganizer.java
+++ b/plugin/trino-raptor-legacy/src/test/java/io/trino/plugin/raptor/legacy/storage/organization/TestShardOrganizer.java
@@ -20,10 +20,13 @@
 import java.util.OptionalInt;
 import java.util.Set;
 import java.util.UUID;
+import java.util.concurrent.CountDownLatch;
 
+import static com.google.common.base.Preconditions.checkState;
 import static com.google.common.collect.Iterables.getOnlyElement;
-import static com.google.common.util.concurrent.Uninterruptibles.sleepUninterruptibly;
+import static com.google.common.util.concurrent.Uninterruptibles.awaitUninterruptibly;
 import static java.util.concurrent.TimeUnit.MILLISECONDS;
+import static java.util.concurrent.TimeUnit.SECONDS;
 import static org.assertj.core.api.Assertions.assertThat;
 
 public class TestShardOrganizer
@@ -33,7 +36,10 @@ public class TestShardOrganizer
     public void testShardOrganizerInProgress()
             throws Exception
     {
-        ShardOrganizer organizer = createShardOrganizer();
+        CountDownLatch canComplete = new CountDownLatch(1);
+        ShardOrganizer organizer = new ShardOrganizer(
+                organizationSet -> () -> checkState(awaitUninterruptibly(canComplete, 10, SECONDS)),
+                1);
 
         Set<UUID> shards = ImmutableSet.of(UUID.randomUUID());
         OrganizationSet organizationSet = new OrganizationSet(1L, shards, OptionalInt.empty());
@@ -43,6 +49,7 @@ public void testShardOrganizerInProgress()
         assertThat(organizer.inProgress(getOnlyElement(shards))).isTrue();
         assertThat(organizer.getShardsInProgress()).isEqualTo(1);
 
+        canComplete.countDown();
         while (organizer.inProgress(getOnlyElement(shards))) {
             MILLISECONDS.sleep(10);
         }
@@ -50,19 +57,4 @@ public void testShardOrganizerInProgress()
         assertThat(organizer.getShardsInProgress()).isEqualTo(0);
         organizer.shutdown();
     }
-
-    private static class MockJobFactory
-            implements JobFactory
-    {
-        @Override
-        public Runnable create(OrganizationSet organizationSet)
-        {
-            return () -> sleepUninterruptibly(10, MILLISECONDS);
-        }
-    }
-
-    static ShardOrganizer createShardOrganizer()
-    {
-        return new ShardOrganizer(new MockJobFactory(), 1);
-    }
 }