8296405: java/util/concurrent/forkjoin/AsyncShutdownNow.java is too slow

Reviewed-by: jpai

test/jdk/java/util/concurrent/forkjoin/AsyncShutdownNow.java

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2020, 2021, Oracle and/or its affiliates. All rights reserved.
+ * Copyright (c) 2020, 2022, Oracle and/or its affiliates. All rights reserved.
  * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
  *
  * This code is free software; you can redistribute it and/or modify it
@@ -28,8 +28,7 @@
  *          invokeAll, and invokeAny
  */
 
-// TODO: this test is far too slow
-
+import java.util.Arrays;
 import java.util.List;
 import java.util.concurrent.Callable;
 import java.util.concurrent.CancellationException;
@@ -38,12 +37,9 @@ import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import java.util.concurrent.ForkJoinPool;
 import java.util.concurrent.Future;
-import java.util.concurrent.RejectedExecutionException;
-import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.TimeUnit;
+import static java.lang.Thread.State.*;
 
-import org.testng.annotations.AfterClass;
-import org.testng.annotations.BeforeClass;
 import org.testng.annotations.DataProvider;
 import org.testng.annotations.Test;
 import static org.testng.Assert.*;
@@ -56,29 +52,6 @@ public class AsyncShutdownNow {
         return null;
     };
 
-    private ScheduledExecutorService scheduledExecutor;
-
-    @BeforeClass
-    public void setup() {
-        scheduledExecutor = Executors.newScheduledThreadPool(1);
-    }
-
-    @AfterClass
-    public void teardown() {
-        scheduledExecutor.shutdown();
-    }
-
-    /**
-     * Schedule the given executor service to be shutdown abruptly after the given
-     * delay, in seconds.
-     */
-    private void scheduleShutdownNow(ExecutorService executor, int delayInSeconds) {
-        scheduledExecutor.schedule(() -> {
-            executor.shutdownNow();
-            return null;
-        }, delayInSeconds, TimeUnit.SECONDS);
-    }
-
     /**
      * The executors to test.
      */
@@ -96,18 +69,17 @@ public class AsyncShutdownNow {
     @Test(dataProvider = "executors")
     public void testFutureGet(ExecutorService executor) throws Exception {
         System.out.format("testFutureGet: %s%n", executor);
-        scheduleShutdownNow(executor, 5);
-        try {
-            // submit long running task, the task should be cancelled
+        try (executor) {
             Future<?> future = executor.submit(SLEEP_FOR_A_DAY);
+
+            // shutdownNow when main thread waits in ForkJoinTask.get
+            onWait("java.util.concurrent.ForkJoinTask.get", executor::shutdownNow);
             try {
                 future.get();
-                assertTrue(false);
-            } catch (ExecutionException | RejectedExecutionException e) {
+                fail();
+            } catch (ExecutionException | CancellationException e) {
                 // expected
             }
-        } finally {
-            executor.shutdown();
         }
     }
 
@@ -117,18 +89,17 @@ public class AsyncShutdownNow {
     @Test(dataProvider = "executors")
     public void testTimedFutureGet(ExecutorService executor) throws Exception {
         System.out.format("testTimedFutureGet: %s%n", executor);
-        scheduleShutdownNow(executor, 5);
-        try {
-            // submit long running task, the task should be cancelled
+        try (executor) {
             Future<?> future = executor.submit(SLEEP_FOR_A_DAY);
+
+            // shutdownNow when main thread waits in ForkJoinTask.get
+            onWait("java.util.concurrent.ForkJoinTask.get", executor::shutdownNow);
             try {
                 future.get(1, TimeUnit.HOURS);
-                assertTrue(false);
-            } catch (ExecutionException | RejectedExecutionException e) {
+                fail();
+            } catch (ExecutionException | CancellationException e) {
                 // expected
             }
-        } finally {
-            executor.shutdown();
         }
     }
 
@@ -138,41 +109,79 @@ public class AsyncShutdownNow {
     @Test(dataProvider = "executors")
     public void testInvokeAll(ExecutorService executor) throws Exception {
         System.out.format("testInvokeAll: %s%n", executor);
-        scheduleShutdownNow(executor, 5);
-        try {
-            // execute long running tasks
+        try (executor) {
+            // shutdownNow when main thread waits in ForkJoinTask.quietlyJoin
+            onWait("java.util.concurrent.ForkJoinTask.quietlyJoin", executor::shutdownNow);
             List<Future<Void>> futures = executor.invokeAll(List.of(SLEEP_FOR_A_DAY, SLEEP_FOR_A_DAY));
             for (Future<Void> f : futures) {
                 assertTrue(f.isDone());
                 try {
                     Object result = f.get();
-                    assertTrue(false);
+                    fail();
                 } catch (ExecutionException | CancellationException e) {
                     // expected
                 }
             }
-        } finally {
-            executor.shutdown();
         }
     }
 
     /**
      * Test shutdownNow with thread blocked in invokeAny.
      */
-    @Test(dataProvider = "executors")
+    @Test(dataProvider = "executors", enabled = false)
     public void testInvokeAny(ExecutorService executor) throws Exception {
         System.out.format("testInvokeAny: %s%n", executor);
-        scheduleShutdownNow(executor, 5);
-        try {
+        try (executor) {
+            // shutdownNow when main thread waits in ForkJoinTask.get
+            onWait("java.util.concurrent.ForkJoinTask.get", executor::shutdownNow);
             try {
-                // execute long running tasks
                 executor.invokeAny(List.of(SLEEP_FOR_A_DAY, SLEEP_FOR_A_DAY));
-                assertTrue(false);
-            } catch (ExecutionException | RejectedExecutionException e) {
+                fail();
+            } catch (ExecutionException e) {
                 // expected
             }
-        } finally {
-            executor.shutdown();
         }
     }
+
+    /**
+     * Runs the given action when the current thread is sampled as waiting (timed or
+     * untimed) at the given location. The location takes the form "{@code c.m}" where
+     * {@code c} is the fully qualified class name and {@code m} is the method name.
+     */
+    private void onWait(String location, Runnable action) {
+        int index = location.lastIndexOf('.');
+        String className = location.substring(0, index);
+        String methodName = location.substring(index + 1);
+        Thread target = Thread.currentThread();
+        var thread = new Thread(() -> {
+            try {
+                boolean found = false;
+                while (!found) {
+                    Thread.State state = target.getState();
+                    assertTrue(state != TERMINATED);
+                    if ((state == WAITING || state == TIMED_WAITING)
+                            && contains(target.getStackTrace(), className, methodName)) {
+                        found = true;
+                    } else {
+                        Thread.sleep(20);
+                    }
+                }
+                action.run();
+            } catch (Exception e) {
+                e.printStackTrace();
+            }
+        });
+        thread.setDaemon(true);
+        thread.start();
+    }
+
+    /**
+     * Returns true if the given stack trace contains an element for the given class
+     * and method name.
+     */
+    private boolean contains(StackTraceElement[] stack, String className, String methodName) {
+        return Arrays.stream(stack)
+                .anyMatch(e -> className.equals(e.getClassName())
+                        && methodName.equals(e.getMethodName()));
+    }
 }