From 8eca03ddc08416a36deb07ed09f71afdf62a8d4b Mon Sep 17 00:00:00 2001
From: Piotr Findeisen <piotr.findeisen@gmail.com>
Date: Wed, 6 Apr 2022 11:07:56 +0200
Subject: [PATCH] Pick failing query faster in testOutOfMemoryKiller

Within the test, one query should get killed, but the other may not be
killed and run to completion. If the first submitted query happens to
run successfully, the test takes longer to complete.

Inspect queries in order of their completion, not submission.
---
 .../io/trino/memory/TestMemoryManager.java     | 18 +++++++++++-------
 1 file changed, 11 insertions(+), 7 deletions(-)

diff --git a/testing/trino-tests/src/test/java/io/trino/memory/TestMemoryManager.java b/testing/trino-tests/src/test/java/io/trino/memory/TestMemoryManager.java
index 28994e6d7609..f3cd57a459d2 100644
--- a/testing/trino-tests/src/test/java/io/trino/memory/TestMemoryManager.java
+++ b/testing/trino-tests/src/test/java/io/trino/memory/TestMemoryManager.java
@@ -21,6 +21,7 @@
 import io.trino.server.testing.TestingTrinoServer;
 import io.trino.spi.QueryId;
 import io.trino.testing.DistributedQueryRunner;
+import io.trino.testing.MaterializedResult;
 import io.trino.testing.QueryRunner;
 import io.trino.tests.tpch.TpchQueryRunnerBuilder;
 import org.intellij.lang.annotations.Language;
@@ -32,7 +33,9 @@
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
+import java.util.concurrent.CompletionService;
 import java.util.concurrent.ExecutionException;
+import java.util.concurrent.ExecutorCompletionService;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Future;
 
@@ -129,16 +132,17 @@ public void testOutOfMemoryKiller()
                 assertTrue(memoryPool.tryReserve(fakeQueryId, "test", memoryPool.getMaxBytes()));
             }
 
-            List<Future<?>> queryFutures = new ArrayList<>();
-            for (int i = 0; i < 2; i++) {
-                queryFutures.add(executor.submit(() -> queryRunner.execute("" +
+            int queries = 2;
+            CompletionService<MaterializedResult> completionService = new ExecutorCompletionService<>(executor);
+            for (int i = 0; i < queries; i++) {
+                completionService.submit(() -> queryRunner.execute("" +
                         "SELECT COUNT(*), clerk " +
                         "FROM (SELECT clerk FROM orders UNION ALL SELECT dummy FROM blackhole.default.take_30s)" +
-                        "GROUP BY clerk")));
+                        "GROUP BY clerk"));
             }
 
             // Wait for queries to start
-            assertEventually(() -> assertThat(queryRunner.getCoordinator().getQueryManager().getQueries()).hasSize(1 + queryFutures.size()));
+            assertEventually(() -> assertThat(queryRunner.getCoordinator().getQueryManager().getQueries()).hasSize(1 + queries));
 
             // Wait for one of the queries to die
             waitForQueryToBeKilled(queryRunner);
@@ -152,8 +156,8 @@ public void testOutOfMemoryKiller()
             }
 
             assertThatThrownBy(() -> {
-                for (Future<?> query : queryFutures) {
-                    query.get();
+                for (int i = 0; i < queries; i++) {
+                    completionService.take().get();
                 }
             })
                     .isInstanceOf(ExecutionException.class)