Fix TestEvictableCache.testLoadFailure flakiness

core/trino-server-rpm/pom.xml

@@ -55,6 +55,12 @@
             <scope>test</scope>
         </dependency>
 
+        <dependency>
+            <groupId>io.trino</groupId>
+            <artifactId>trino-testing-services</artifactId>
+            <scope>test</scope>
+        </dependency>
+
         <dependency>
             <groupId>org.assertj</groupId>
             <artifactId>assertj-core</artifactId>

lib/trino-cache/src/test/java/io/trino/cache/TestEmptyCache.java

@@ -16,24 +16,27 @@
 import com.google.common.cache.Cache;
 import com.google.common.cache.CacheLoader;
 import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.Timeout;
 
 import java.util.ArrayList;
 import java.util.List;
-import java.util.concurrent.CyclicBarrier;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.Exchanger;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Future;
-import java.util.concurrent.atomic.AtomicBoolean;
 
+import static io.trino.testing.assertions.Assert.assertEventually;
 import static java.util.concurrent.Executors.newFixedThreadPool;
 import static java.util.concurrent.TimeUnit.SECONDS;
 import static org.assertj.core.api.Assertions.assertThat;
 
 public class TestEmptyCache
 {
-    private static final int TEST_TIMEOUT_MILLIS = 10_000;
+    private static final int TEST_TIMEOUT_SECONDS = 10;
 
     @Test
+    @Timeout(TEST_TIMEOUT_SECONDS)
     public void testLoadFailure()
             throws Exception
     {
@@ -46,17 +49,25 @@ public void testLoadFailure()
 
         ExecutorService executor = newFixedThreadPool(2);
         try {
-            AtomicBoolean first = new AtomicBoolean(true);
-            CyclicBarrier barrier = new CyclicBarrier(2);
+            Exchanger<Thread> exchanger = new Exchanger<>();
+            CountDownLatch secondUnblocked = new CountDownLatch(1);
 
             List<Future<String>> futures = new ArrayList<>();
             for (int i = 0; i < 2; i++) {
+                boolean first = i == 0;
                 futures.add(executor.submit(() -> {
-                    barrier.await(10, SECONDS);
+                    if (!first) {
+                        // Wait for the first one to start the call
+                        exchanger.exchange(Thread.currentThread(), 10, SECONDS);
+                        // Prove that we are back in RUNNABLE state.
+                        secondUnblocked.countDown();
+                    }
                     return cache.get(key, () -> {
-                        if (first.compareAndSet(true, false)) {
-                            // first
-                            Thread.sleep(1); // increase chances that second thread calls cache.get before we return
+                        if (first) {
+                            Thread secondThread = exchanger.exchange(null, 10, SECONDS);
+                            assertThat(secondUnblocked.await(10, SECONDS)).isTrue();
+                            // Wait for the second one to hang inside the cache.get call.
+                            assertEventually(() -> assertThat(secondThread.getState()).isNotEqualTo(Thread.State.RUNNABLE));
                             throw new RuntimeException("first attempt is poised to fail");
                         }
                         return "success";

lib/trino-cache/src/test/java/io/trino/cache/TestEvictableCache.java

@@ -31,11 +31,11 @@
 import java.util.concurrent.Callable;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.CyclicBarrier;
+import java.util.concurrent.Exchanger;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Future;
 import java.util.concurrent.TimeUnit;
-import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicLong;
 import java.util.stream.IntStream;
@@ -44,6 +44,7 @@
 import static com.google.common.collect.ImmutableList.toImmutableList;
 import static com.google.common.collect.ImmutableSet.toImmutableSet;
 import static io.trino.cache.CacheStatsAssertions.assertCacheStats;
+import static io.trino.testing.assertions.Assert.assertEventually;
 import static java.lang.Math.toIntExact;
 import static java.lang.String.format;
 import static java.util.concurrent.Executors.newFixedThreadPool;
@@ -300,17 +301,25 @@ public void testLoadFailure()
 
         ExecutorService executor = newFixedThreadPool(2);
         try {
-            AtomicBoolean first = new AtomicBoolean(true);
-            CyclicBarrier barrier = new CyclicBarrier(2);
+            Exchanger<Thread> exchanger = new Exchanger<>();
+            CountDownLatch secondUnblocked = new CountDownLatch(1);
 
             List<Future<String>> futures = new ArrayList<>();
             for (int i = 0; i < 2; i++) {
+                boolean first = i == 0;
                 futures.add(executor.submit(() -> {
-                    barrier.await(10, SECONDS);
+                    if (!first) {
+                        // Wait for the first one to start the call
+                        exchanger.exchange(Thread.currentThread(), 10, SECONDS);
+                        // Prove that we are back in RUNNABLE state.
+                        secondUnblocked.countDown();
+                    }
                     return cache.get(key, () -> {
-                        if (first.compareAndSet(true, false)) {
-                            // first
-                            Thread.sleep(1); // increase chances that second thread calls cache.get before we return
+                        if (first) {
+                            Thread secondThread = exchanger.exchange(null, 10, SECONDS);
+                            assertThat(secondUnblocked.await(10, SECONDS)).isTrue();
+                            // Wait for the second one to hang inside the cache.get call.
+                            assertEventually(() -> assertThat(secondThread.getState()).isNotEqualTo(Thread.State.RUNNABLE));
                             throw new RuntimeException("first attempt is poised to fail");
                         }
                         return "success";

plugin/trino-google-sheets/pom.xml

@@ -191,6 +191,12 @@
             <scope>test</scope>
         </dependency>
 
+        <dependency>
+            <groupId>io.trino</groupId>
+            <artifactId>trino-testing-services</artifactId>
+            <scope>test</scope>
+        </dependency>
+
         <dependency>
             <groupId>io.trino</groupId>
             <artifactId>trino-tpch</artifactId>