From 0a82c25739a512c71d445069941427901f6b75b7 Mon Sep 17 00:00:00 2001
From: Lari Hotari <lari@hotari.net>
Date: Sat, 15 May 2021 15:58:21 +0300
Subject: [PATCH] [Tests] Fix flaky test GracefulExecutorServicesShutdownTest

- fix race condition in test by adding a CountDownLatch
  to verify that execution has entered the awaitTermination
  method before the future is cancelled
---
 .../GracefulExecutorServicesShutdownTest.java      | 14 +++++++++-----
 1 file changed, 9 insertions(+), 5 deletions(-)

diff --git a/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/GracefulExecutorServicesShutdownTest.java b/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/GracefulExecutorServicesShutdownTest.java
index 726f37d12cee4..a78460871045d 100644
--- a/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/GracefulExecutorServicesShutdownTest.java
+++ b/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/GracefulExecutorServicesShutdownTest.java
@@ -29,6 +29,7 @@
 import static org.testng.Assert.assertTrue;
 import java.time.Duration;
 import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.TimeUnit;
@@ -120,25 +121,28 @@ public void shouldWaitForExecutorToTerminate() throws ExecutionException, Interr
 
 
     @Test
-    public void shouldTerminateWhenFutureIsCancelled() throws InterruptedException {
+    public void shouldTerminateWhenFutureIsCancelled() throws InterruptedException, ExecutionException {
         // given
         GracefulExecutorServicesShutdown shutdown = GracefulExecutorServicesShutdown.initiate();
         shutdown.timeout(Duration.ofMillis(15000));
         ExecutorService executorService = mock(ExecutorService.class);
         when(executorService.isShutdown()).thenReturn(true);
         AtomicBoolean terminated = new AtomicBoolean();
-        AtomicBoolean awaitTerminationInterrupted = new AtomicBoolean();
+        CompletableFuture<Boolean> awaitTerminationInterrupted = new CompletableFuture<>();
         when(executorService.isTerminated()).thenAnswer(invocation -> terminated.get());
+        CountDownLatch awaitingTerminationEntered = new CountDownLatch(1);
         when(executorService.awaitTermination(anyLong(), any())).thenAnswer(invocation  -> {
             long timeout = invocation.getArgument(0);
             TimeUnit unit = invocation.getArgument(1);
+            awaitingTerminationEntered.countDown();
             try {
                 Thread.sleep(unit.toMillis(timeout));
             } catch (InterruptedException e) {
-                awaitTerminationInterrupted.set(true);
+                awaitTerminationInterrupted.complete(true);
                 Thread.currentThread().interrupt();
                 throw e;
             }
+            awaitTerminationInterrupted.complete(false);
             throw new IllegalStateException("Thread.sleep should have been interrupted");
         });
         when(executorService.shutdownNow()).thenAnswer(invocation -> {
@@ -149,11 +153,11 @@ public void shouldTerminateWhenFutureIsCancelled() throws InterruptedException {
         // when
         shutdown.shutdown(executorService);
         CompletableFuture<Void> future = shutdown.handle();
+        awaitingTerminationEntered.await();
         future.cancel(false);
 
         // then
-        Awaitility.await().untilAsserted(() -> assertTrue(awaitTerminationInterrupted.get(),
-                "awaitTermination should have been interrupted"));
+        assertTrue(awaitTerminationInterrupted.get(), "awaitTermination should have been interrupted");
         verify(executorService, times(1)).awaitTermination(anyLong(), any());
         verify(executorService, times(1)).shutdownNow();
     }