From 36e994a20f680baf52060be28c9dad78f33e1dd1 Mon Sep 17 00:00:00 2001
From: Stuart Douglas <stuart.w.douglas@gmail.com>
Date: Thu, 4 Feb 2021 12:16:02 +1100
Subject: [PATCH] Fix flakey test

Just because the executor is shutdown it does not mean that
Thread.isAlive() will return false.
---
 .../runtime/ConnectionPoolsClosedTest.java    | 31 +++++++++++++------
 1 file changed, 21 insertions(+), 10 deletions(-)

diff --git a/extensions/reactive-datasource/deployment/src/test/java/io/quarkus/reactive/datasource/runtime/ConnectionPoolsClosedTest.java b/extensions/reactive-datasource/deployment/src/test/java/io/quarkus/reactive/datasource/runtime/ConnectionPoolsClosedTest.java
index 7dc8bec3e8777..83c9c67ad3386 100644
--- a/extensions/reactive-datasource/deployment/src/test/java/io/quarkus/reactive/datasource/runtime/ConnectionPoolsClosedTest.java
+++ b/extensions/reactive-datasource/deployment/src/test/java/io/quarkus/reactive/datasource/runtime/ConnectionPoolsClosedTest.java
@@ -1,10 +1,12 @@
 package io.quarkus.reactive.datasource.runtime;
 
+import java.util.HashSet;
+import java.util.Set;
+import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import java.util.concurrent.Future;
-import java.util.concurrent.atomic.AtomicReference;
 
 import org.junit.jupiter.api.Test;
 import org.wildfly.common.Assert;
@@ -48,13 +50,14 @@ public void connectionsFromOtherThreadsGetClosed() throws ExecutionException, In
     public void connectionsThatAreNotTheLastGetClosedSuccessfully()
             throws ExecutionException, InterruptedException {
         ExecutorService e = Executors.newFixedThreadPool(3);
+        Set<Thread> threads = new HashSet<>();
         try {
             TestableThreadLocalPool globalPool = new TestableThreadLocalPool();
             Assert.assertTrue(globalPool.trackedSize() == 0);
-            final TestPoolInterface p1 = grabPoolFromLongLivingThread(globalPool, e);
+            final TestPoolInterface p1 = grabPoolFromLongLivingThread(globalPool, e, threads);
             Assert.assertFalse(p1.isClosed());
             Assert.assertTrue(globalPool.trackedSize() == 1);
-            final TestPoolInterface p2 = grabPoolFromLongLivingThread(globalPool, e);
+            final TestPoolInterface p2 = grabPoolFromLongLivingThread(globalPool, e, threads);
             Assert.assertFalse(p1.isClosed());
             Assert.assertFalse(p2.isClosed());
             Assert.assertFalse(p1.isClosed());
@@ -64,6 +67,9 @@ public void connectionsThatAreNotTheLastGetClosedSuccessfully()
             while (!e.isTerminated()) {
                 Thread.sleep(1);
             }
+            for (Thread thread : threads) {
+                waitForThreadStop(thread);
+            }
             final TestPoolInterface p3 = grabPoolFromOtherThread(globalPool);
             Assert.assertTrue(p1.isClosed());
             Assert.assertTrue(p2.isClosed());
@@ -107,30 +113,35 @@ public void plainClose() {
     private TestPoolInterface grabPoolFromOtherThread(TestableThreadLocalPool globalPool)
             throws ExecutionException, InterruptedException {
         final ExecutorService e = Executors.newSingleThreadExecutor();
-        AtomicReference<Thread> thread = new AtomicReference<>();
+        CompletableFuture<Thread> thread = new CompletableFuture<>();
         final Future<TestPoolInterface> creation = e.submit(() -> {
-            thread.set(Thread.currentThread());
+            thread.complete(Thread.currentThread());
             return globalPool.pool();
         });
         final TestPoolInterface poolInstance = creation.get();
         e.shutdown();
-        while (thread.get().isAlive()) {
+        waitForThreadStop(thread.get());
+        return poolInstance;
+    }
+
+    private void waitForThreadStop(Thread thread) throws InterruptedException {
+        while (thread.isAlive()) {
             //even after the pool is shutdown the thread may not actually report itself
             //as being dead yet, which can lead to race conditions
             //so we wait just to be sure
             Thread.sleep(1);
         }
-        return poolInstance;
     }
 
     private TestPoolInterface grabPoolFromLongLivingThread(TestableThreadLocalPool globalPool,
-            ExecutorService e)
+            ExecutorService e, Set<Thread> allThreads)
             throws InterruptedException, ExecutionException {
-        AtomicReference<Thread> thread = new AtomicReference<>();
+        CompletableFuture<Thread> thread = new CompletableFuture<>();
         final Future<TestPoolInterface> creation = e.submit(() -> {
-            thread.set(Thread.currentThread());
+            thread.complete(Thread.currentThread());
             return globalPool.pool();
         });
+        allThreads.add(thread.get());
         return creation.get();
     }