Fix QueueResizableOpenSearchThreadPoolExecutorTests

There was a race condition in testResizeQueueDown() where depending on random
parameters we could submit up to 1002 tasks into an executor with a queue size
of 900. That introduced a race condition where if the tasks didn't execute fast
enough then a rejected execution exception could happen and fail the test.

And finally I refactored the tests to reduce duplication of code and ensure the
executor gets shutdown properly even in case of a test failure.

Signed-off-by: Andrew Ross <andrross@amazon.com>

Author: andrross

server/src/test/java/org/opensearch/common/util/concurrent/QueueResizableOpenSearchThreadPoolExecutorTests.java

@@ -10,7 +10,10 @@
 
 import org.opensearch.common.settings.Settings;
 import org.opensearch.test.OpenSearchTestCase;
+import org.opensearch.threadpool.ThreadPool;
+import org.junit.After;
 
+import java.util.Objects;
 import java.util.concurrent.TimeUnit;
 import java.util.function.Function;
 
@@ -23,14 +26,19 @@
  * based on the time taken for each event.
  */
 public class QueueResizableOpenSearchThreadPoolExecutorTests extends OpenSearchTestCase {
-    public void testResizeQueueSameSize() throws Exception {
-        ThreadContext context = new ThreadContext(Settings.EMPTY);
-        ResizableBlockingQueue<Runnable> queue = new ResizableBlockingQueue<>(ConcurrentCollections.<Runnable>newBlockingQueue(), 2000);
+    private QueueResizableOpenSearchThreadPoolExecutor executor;
+    private ResizableBlockingQueue<Runnable> queue;
+    private int measureWindow;
 
+    private void createExecutor(int queueSize) {
         int threads = randomIntBetween(1, 10);
-        int measureWindow = randomIntBetween(100, 200);
-        logger.info("--> auto-queue with a measurement window of {} tasks", measureWindow);
-        QueueResizableOpenSearchThreadPoolExecutor executor = new QueueResizableOpenSearchThreadPoolExecutor(
+        measureWindow = randomIntBetween(100, 200);
+        ThreadContext context = new ThreadContext(Settings.EMPTY);
+        this.queue = new ResizableBlockingQueue<>(
+            ConcurrentCollections.newBlockingQueue(),
+            Objects.requireNonNull(queueSize, "All tests must set a queue size")
+        );
+        this.executor = new QueueResizableOpenSearchThreadPoolExecutor(
             "test-threadpool",
             threads,
             threads,
@@ -44,152 +52,67 @@ public void testResizeQueueSameSize() throws Exception {
         );
         executor.prestartAllCoreThreads();
         logger.info("--> executor: {}", executor);
+    }
+
+    @After
+    public void stopExecutor() {
+        ThreadPool.terminate(executor, 10, TimeUnit.SECONDS);
+    }
+
+    public void testResizeQueueSameSize() throws Exception {
+        createExecutor(2000);
 
         // Execute a task multiple times that takes 1ms
         assertThat(executor.resize(1000), equalTo(1000));
         executeTask(executor, (measureWindow * 5) + 2);
 
-        assertBusy(() -> { assertThat(queue.capacity(), lessThanOrEqualTo(1000)); });
-        executor.shutdown();
-        executor.awaitTermination(10, TimeUnit.SECONDS);
+        assertBusy(() -> assertThat(queue.capacity(), lessThanOrEqualTo(1000)));
     }
 
     public void testResizeQueueUp() throws Exception {
-        ThreadContext context = new ThreadContext(Settings.EMPTY);
-        ResizableBlockingQueue<Runnable> queue = new ResizableBlockingQueue<>(ConcurrentCollections.<Runnable>newBlockingQueue(), 2000);
-
-        int threads = randomIntBetween(1, 10);
-        int measureWindow = randomIntBetween(100, 200);
-        logger.info("--> auto-queue with a measurement window of {} tasks", measureWindow);
-        QueueResizableOpenSearchThreadPoolExecutor executor = new QueueResizableOpenSearchThreadPoolExecutor(
-            "test-threadpool",
-            threads,
-            threads,
-            1000,
-            TimeUnit.MILLISECONDS,
-            queue,
-            fastWrapper(),
-            OpenSearchExecutors.daemonThreadFactory("queuetest"),
-            new OpenSearchAbortPolicy(),
-            context
-        );
-        executor.prestartAllCoreThreads();
-        logger.info("--> executor: {}", executor);
-
+        createExecutor(2000);
         // Execute a task multiple times that takes 1ms
         assertThat(executor.resize(3000), equalTo(3000));
         executeTask(executor, (measureWindow * 5) + 2);
 
-        assertBusy(() -> { assertThat(queue.capacity(), greaterThanOrEqualTo(2000)); });
-        executor.shutdown();
-        executor.awaitTermination(10, TimeUnit.SECONDS);
+        assertBusy(() -> assertThat(queue.capacity(), greaterThanOrEqualTo(2000)));
     }
 
     public void testResizeQueueDown() throws Exception {
-        ThreadContext context = new ThreadContext(Settings.EMPTY);
-        ResizableBlockingQueue<Runnable> queue = new ResizableBlockingQueue<>(ConcurrentCollections.<Runnable>newBlockingQueue(), 2000);
-
-        int threads = randomIntBetween(1, 10);
-        int measureWindow = randomIntBetween(100, 200);
-        logger.info("--> auto-queue with a measurement window of {} tasks", measureWindow);
-        QueueResizableOpenSearchThreadPoolExecutor executor = new QueueResizableOpenSearchThreadPoolExecutor(
-            "test-threadpool",
-            threads,
-            threads,
-            1000,
-            TimeUnit.MILLISECONDS,
-            queue,
-            fastWrapper(),
-            OpenSearchExecutors.daemonThreadFactory("queuetest"),
-            new OpenSearchAbortPolicy(),
-            context
-        );
-        executor.prestartAllCoreThreads();
-        logger.info("--> executor: {}", executor);
-
+        createExecutor(2000);
         // Execute a task multiple times that takes 1ms
-        assertThat(executor.resize(900), equalTo(900));
+        assertThat(executor.resize(1500), equalTo(1500));
         executeTask(executor, (measureWindow * 5) + 2);
 
-        assertBusy(() -> { assertThat(queue.capacity(), lessThanOrEqualTo(900)); });
-        executor.shutdown();
-        executor.awaitTermination(10, TimeUnit.SECONDS);
+        assertBusy(() -> assertThat(queue.capacity(), lessThanOrEqualTo(1500)));
     }
 
     public void testExecutionEWMACalculation() throws Exception {
-        ThreadContext context = new ThreadContext(Settings.EMPTY);
-        ResizableBlockingQueue<Runnable> queue = new ResizableBlockingQueue<>(ConcurrentCollections.<Runnable>newBlockingQueue(), 100);
-
-        QueueResizableOpenSearchThreadPoolExecutor executor = new QueueResizableOpenSearchThreadPoolExecutor(
-            "test-threadpool",
-            1,
-            1,
-            1000,
-            TimeUnit.MILLISECONDS,
-            queue,
-            fastWrapper(),
-            OpenSearchExecutors.daemonThreadFactory("queuetest"),
-            new OpenSearchAbortPolicy(),
-            context
-        );
-        executor.prestartAllCoreThreads();
-        logger.info("--> executor: {}", executor);
-
+        createExecutor(100);
         assertThat((long) executor.getTaskExecutionEWMA(), equalTo(0L));
         executeTask(executor, 1);
-        assertBusy(() -> { assertThat((long) executor.getTaskExecutionEWMA(), equalTo(30L)); });
+        assertBusy(() -> assertThat((long) executor.getTaskExecutionEWMA(), equalTo(30L)));
         executeTask(executor, 1);
-        assertBusy(() -> { assertThat((long) executor.getTaskExecutionEWMA(), equalTo(51L)); });
+        assertBusy(() -> assertThat((long) executor.getTaskExecutionEWMA(), equalTo(51L)));
         executeTask(executor, 1);
-        assertBusy(() -> { assertThat((long) executor.getTaskExecutionEWMA(), equalTo(65L)); });
+        assertBusy(() -> assertThat((long) executor.getTaskExecutionEWMA(), equalTo(65L)));
         executeTask(executor, 1);
-        assertBusy(() -> { assertThat((long) executor.getTaskExecutionEWMA(), equalTo(75L)); });
+        assertBusy(() -> assertThat((long) executor.getTaskExecutionEWMA(), equalTo(75L)));
         executeTask(executor, 1);
-        assertBusy(() -> { assertThat((long) executor.getTaskExecutionEWMA(), equalTo(83L)); });
-
-        executor.shutdown();
-        executor.awaitTermination(10, TimeUnit.SECONDS);
+        assertBusy(() -> assertThat((long) executor.getTaskExecutionEWMA(), equalTo(83L)));
     }
 
     /** Use a runnable wrapper that simulates a task with unknown failures. */
-    public void testExceptionThrowingTask() throws Exception {
-        ThreadContext context = new ThreadContext(Settings.EMPTY);
-        ResizableBlockingQueue<Runnable> queue = new ResizableBlockingQueue<>(ConcurrentCollections.<Runnable>newBlockingQueue(), 100);
-
-        QueueResizableOpenSearchThreadPoolExecutor executor = new QueueResizableOpenSearchThreadPoolExecutor(
-            "test-threadpool",
-            1,
-            1,
-            1000,
-            TimeUnit.MILLISECONDS,
-            queue,
-            exceptionalWrapper(),
-            OpenSearchExecutors.daemonThreadFactory("queuetest"),
-            new OpenSearchAbortPolicy(),
-            context
-        );
-        executor.prestartAllCoreThreads();
-        logger.info("--> executor: {}", executor);
-
+    public void testExceptionThrowingTask() {
+        createExecutor(100);
         assertThat((long) executor.getTaskExecutionEWMA(), equalTo(0L));
         executeTask(executor, 1);
-        executor.shutdown();
-        executor.awaitTermination(10, TimeUnit.SECONDS);
     }
 
     private Function<Runnable, WrappedRunnable> fastWrapper() {
         return (runnable) -> new SettableTimedRunnable(TimeUnit.NANOSECONDS.toNanos(100), false);
     }
 
-    /**
-     * The returned function outputs a WrappedRunnabled that simulates the case
-     * where {@link TimedRunnable#getTotalExecutionNanos()} returns -1 because
-     * the job failed or was rejected before it finished.
-     */
-    private Function<Runnable, WrappedRunnable> exceptionalWrapper() {
-        return (runnable) -> new SettableTimedRunnable(TimeUnit.NANOSECONDS.toNanos(-1), true);
-    }
-
     /** Execute a blank task {@code times} times for the executor */
     private void executeTask(QueueResizableOpenSearchThreadPoolExecutor executor, int times) {
         logger.info("--> executing a task [{}] times", times);
@@ -198,7 +121,7 @@ private void executeTask(QueueResizableOpenSearchThreadPoolExecutor executor, in
         }
     }
 
-    public class SettableTimedRunnable extends TimedRunnable {
+    private static class SettableTimedRunnable extends TimedRunnable {
         private final long timeTaken;
         private final boolean testFailedOrRejected;