Fix flaky testFailStuckSplitTasks unit test

core/trino-main/src/test/java/io/trino/execution/TestSqlTaskManager.java

@@ -49,11 +49,15 @@
 import org.testng.annotations.AfterClass;
 import org.testng.annotations.Test;
 
+import javax.annotation.concurrent.GuardedBy;
+
 import java.net.URI;
 import java.util.List;
 import java.util.Optional;
 import java.util.OptionalInt;
+import java.util.concurrent.ExecutionException;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
 import java.util.function.Predicate;
 
 import static com.google.common.util.concurrent.Futures.immediateVoidFuture;
@@ -249,6 +253,7 @@ public void testRemoveOldTasks()
 
     @Test
     public void testFailStuckSplitTasks()
+            throws InterruptedException, ExecutionException, TimeoutException
     {
         TestingTicker ticker = new TestingTicker();
 
@@ -265,19 +270,32 @@ public void testFailStuckSplitTasks()
         taskExecutor.enqueueSplits(taskHandle, false, ImmutableList.of(mockSplitRunner));
         taskExecutor.start();
 
+        // wait for the task executor to start processing the split
+        mockSplitRunner.waitForStart();
+
         TaskManagerConfig taskManagerConfig = new TaskManagerConfig()
                 .setInterruptStuckSplitTasksEnabled(true)
                 .setInterruptStuckSplitTasksDetectionInterval(new Duration(10, SECONDS))
                 .setInterruptStuckSplitTasksWarningThreshold(new Duration(10, SECONDS))
                 .setInterruptStuckSplitTasksTimeout(new Duration(10, SECONDS));
 
         try (SqlTaskManager sqlTaskManager = createSqlTaskManager(taskManagerConfig, new NodeMemoryConfig(), taskExecutor, stackTraceElements -> true)) {
+            sqlTaskManager.addStateChangeListener(TASK_ID, (state) -> {
+                if (state.isDone()) {
+                    taskExecutor.removeTask(taskHandle);
+                }
+            });
+
             ticker.increment(30, SECONDS);
             sqlTaskManager.failStuckSplitTasks();
 
+            mockSplitRunner.waitForFinish();
             assertEquals(sqlTaskManager.getAllTaskInfo().size(), 1);
             assertEquals(sqlTaskManager.getAllTaskInfo().get(0).getTaskStatus().getState(), TaskState.FAILED);
         }
+        finally {
+            taskExecutor.stop();
+        }
     }
 
     @Test
@@ -444,17 +462,45 @@ public URI createMemoryInfoLocation(InternalNode node)
     private static class MockSplitRunner
             implements SplitRunner
     {
-        private SettableFuture<Boolean> interrupted = SettableFuture.create();
+        private final SettableFuture<Void> startedFuture = SettableFuture.create();
+        private final SettableFuture<Void> finishedFuture = SettableFuture.create();
+
+        @GuardedBy("this")
+        private Thread runnerThread;
+        @GuardedBy("this")
+        private boolean closed;
+
+        public void waitForStart()
+                throws ExecutionException, InterruptedException, TimeoutException
+        {
+            startedFuture.get(10, SECONDS);
+        }
+
+        public void waitForFinish()
+                throws ExecutionException, InterruptedException, TimeoutException
+        {
+            finishedFuture.get(10, SECONDS);
+        }
 
         @Override
-        public boolean isFinished()
+        public synchronized boolean isFinished()
         {
-            return interrupted.isDone();
+            return closed;
         }
 
         @Override
         public ListenableFuture<Void> processFor(Duration duration)
         {
+            startedFuture.set(null);
+            synchronized (this) {
+                runnerThread = Thread.currentThread();
+
+                if (closed) {
+                    finishedFuture.set(null);
+                    return immediateVoidFuture();
+                }
+            }
+
             while (true) {
                 try {
                     Thread.sleep(100000);
@@ -463,19 +509,29 @@ public ListenableFuture<Void> processFor(Duration duration)
                     break;
                 }
             }
-            interrupted.set(true);
+
+            synchronized (this) {
+                closed = true;
+            }
+            finishedFuture.set(null);
+
             return immediateVoidFuture();
         }
 
         @Override
         public String getInfo()
         {
-            return "";
+            return "MockSplitRunner";
         }
 
         @Override
-        public void close()
+        public synchronized void close()
         {
+            closed = true;
+
+            if (runnerThread != null) {
+                runnerThread.interrupt();
+            }
         }
     }
 }