Handle long running queries in DirectTrinoClient

core/trino-main/src/main/java/io/trino/client/direct/DirectTrinoClient.java

@@ -21,6 +21,7 @@
 import io.trino.dispatcher.DispatchQuery;
 import io.trino.exchange.DirectExchangeInput;
 import io.trino.execution.QueryManager;
+import io.trino.execution.QueryManagerConfig;
 import io.trino.execution.QueryState;
 import io.trino.execution.buffer.PageDeserializer;
 import io.trino.memory.context.SimpleLocalMemoryContext;
@@ -40,6 +41,8 @@
 import java.util.List;
 import java.util.Optional;
 import java.util.concurrent.ExecutionException;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
 
 import static io.airlift.concurrent.MoreFutures.whenAnyComplete;
 import static io.trino.SystemSessionProperties.getRetryPolicy;
@@ -56,13 +59,20 @@ public class DirectTrinoClient
     private final QueryManager queryManager;
     private final DirectExchangeClientSupplier directExchangeClientSupplier;
     private final BlockEncodingSerde blockEncodingSerde;
+    private final long heartBeatIntervalMillis;
 
-    public DirectTrinoClient(DispatchManager dispatchManager, QueryManager queryManager, DirectExchangeClientSupplier directExchangeClientSupplier, BlockEncodingSerde blockEncodingSerde)
+    public DirectTrinoClient(
+            DispatchManager dispatchManager,
+            QueryManager queryManager,
+            QueryManagerConfig queryManagerConfig,
+            DirectExchangeClientSupplier directExchangeClientSupplier,
+            BlockEncodingSerde blockEncodingSerde)
     {
         this.dispatchManager = requireNonNull(dispatchManager, "dispatchManager is null");
         this.queryManager = requireNonNull(queryManager, "queryManager is null");
         this.directExchangeClientSupplier = requireNonNull(directExchangeClientSupplier, "directExchangeClientSupplier is null");
         this.blockEncodingSerde = requireNonNull(blockEncodingSerde, "blockEncodingSerde is null");
+        this.heartBeatIntervalMillis = requireNonNull(queryManagerConfig, "queryManagerConfig is null").getClientTimeout().toMillis() / 2;
     }
 
     public DispatchQuery execute(SessionContext sessionContext, @Language("SQL") String sql, QueryResultsListener queryResultsListener)
@@ -103,7 +113,27 @@ public DispatchQuery execute(SessionContext sessionContext, @Language("SQL") Str
                     Page page = pageDeserializer.deserialize(serializedPage);
                     queryResultsListener.consumeOutputPage(page);
                 }
-                getQueryFuture(whenAnyComplete(ImmutableList.of(queryManager.getStateChange(queryId, state), exchangeClient.isBlocked())));
+
+                ListenableFuture<Object> anyCompleteFuture = whenAnyComplete(ImmutableList.of(
+                        queryManager.getStateChange(queryId, state),
+                        exchangeClient.isBlocked()));
+                while (!anyCompleteFuture.isDone()) {
+                    try {
+                        anyCompleteFuture.get(heartBeatIntervalMillis, TimeUnit.MILLISECONDS);
+                    }
+                    catch (TimeoutException e) {
+                        // continue waiting until the query state changes or the exchange client is blocked.
+                        // we need to periodically record the heartbeat to prevent the query from being canceled
+                        dispatchQuery.recordHeartbeat();
+                    }
+                    catch (InterruptedException e) {
+                        Thread.currentThread().interrupt();
+                        throw new TrinoException(GENERIC_INTERNAL_ERROR, "Thread interrupted", e);
+                    }
+                    catch (ExecutionException e) {
+                        throw new TrinoException(GENERIC_INTERNAL_ERROR, "Error processing query", e.getCause());
+                    }
+                }
             }
         }
 

core/trino-main/src/main/java/io/trino/server/testing/TestingTrinoServer.java

@@ -758,6 +758,19 @@ public Builder addProperty(String name, String value)
             return this;
         }
 
+        public Builder overrideProperties(Map<String, String> properties)
+        {
+            ImmutableMap.Builder<String, String> builder = ImmutableMap.builder();
+            this.properties.forEach((k, v) -> {
+                if (!properties.containsKey(k)) {
+                    builder.put(k, v);
+                }
+            });
+            builder.putAll(properties);
+            this.properties = builder.buildOrThrow();
+            return this;
+        }
+
         public Builder setProperties(Map<String, String> properties)
         {
             this.properties = ImmutableMap.copyOf(requireNonNull(properties, "properties is null"));

core/trino-main/src/main/java/io/trino/testing/StandaloneQueryRunner.java

@@ -22,6 +22,7 @@
 import io.trino.Session;
 import io.trino.cost.StatsCalculator;
 import io.trino.execution.FailureInjector.InjectedFailureType;
+import io.trino.execution.QueryManagerConfig;
 import io.trino.execution.warnings.WarningCollector;
 import io.trino.metadata.FunctionBundle;
 import io.trino.metadata.MetadataUtil;
@@ -91,6 +92,7 @@ public StandaloneQueryRunner(Session defaultSession, Consumer<TestingTrinoServer
         this.trinoClient = new TestingDirectTrinoClient(
                 server.getDispatchManager(),
                 server.getQueryManager(),
+                server.getInstance(Key.get(QueryManagerConfig.class)),
                 server.getInstance(Key.get(DirectExchangeClientSupplier.class)),
                 server.getInstance(Key.get(BlockEncodingSerde.class)));
     }

core/trino-main/src/main/java/io/trino/testing/TestingDirectTrinoClient.java

@@ -20,6 +20,7 @@
 import io.trino.dispatcher.DispatchQuery;
 import io.trino.execution.QueryInfo;
 import io.trino.execution.QueryManager;
+import io.trino.execution.QueryManagerConfig;
 import io.trino.operator.DirectExchangeClientSupplier;
 import io.trino.server.ResultQueryInfo;
 import io.trino.server.SessionContext;
@@ -49,9 +50,9 @@ public class TestingDirectTrinoClient
 {
     private final DirectTrinoClient directTrinoClient;
 
-    public TestingDirectTrinoClient(DispatchManager dispatchManager, QueryManager queryManager, DirectExchangeClientSupplier directExchangeClientSupplier, BlockEncodingSerde blockEncodingSerde)
+    public TestingDirectTrinoClient(DispatchManager dispatchManager, QueryManager queryManager, QueryManagerConfig queryManagerConfig, DirectExchangeClientSupplier directExchangeClientSupplier, BlockEncodingSerde blockEncodingSerde)
     {
-        directTrinoClient = new DirectTrinoClient(dispatchManager, queryManager, directExchangeClientSupplier, blockEncodingSerde);
+        directTrinoClient = new DirectTrinoClient(dispatchManager, queryManager, queryManagerConfig, directExchangeClientSupplier, blockEncodingSerde);
     }
 
     public Result execute(Session session, @Language("SQL") String sql)

testing/trino-tests/src/test/java/io/trino/client/direct/TestDirectTrinoClient.java

@@ -0,0 +1,64 @@
+/*
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package io.trino.client.direct;
+
+import com.google.common.collect.ImmutableMap;
+import io.trino.plugin.blackhole.BlackHolePlugin;
+import io.trino.testing.QueryRunner;
+import io.trino.testing.StandaloneQueryRunner;
+import org.junit.jupiter.api.BeforeAll;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.TestInstance;
+import org.junit.jupiter.api.Timeout;
+import org.junit.jupiter.api.parallel.Execution;
+
+import java.util.concurrent.TimeUnit;
+
+import static io.trino.SessionTestUtils.TEST_SESSION;
+import static org.junit.jupiter.api.TestInstance.Lifecycle.PER_CLASS;
+import static org.junit.jupiter.api.parallel.ExecutionMode.CONCURRENT;
+
+@TestInstance(PER_CLASS)
+@Execution(CONCURRENT)
+public class TestDirectTrinoClient
+{
+    private QueryRunner queryRunner;
+
+    @BeforeAll
+    public void setup()
+            throws Exception
+    {
+        queryRunner = new StandaloneQueryRunner(
+                TEST_SESSION,
+                builder -> builder.overrideProperties(ImmutableMap.of(
+                        "query.client.timeout", "1s")));
+        queryRunner.installPlugin(new BlackHolePlugin());
+        queryRunner.createCatalog("blackhole", "blackhole");
+        queryRunner.execute("CREATE SCHEMA blackhole.test_schema");
+        queryRunner.execute("CREATE TABLE blackhole.test_schema.slow_test_table (col1 VARCHAR, col2 INTEGER)" +
+                "WITH (" +
+                "   split_count = 1, " +
+                "   pages_per_split = 1, " +
+                "   rows_per_page = 1, " +
+                "   page_processing_delay = '3s'" +
+                ")");
+    }
+
+    @Test
+    @Timeout(value = 20, unit = TimeUnit.SECONDS)
+    public void testDirectTrinoClientLongQuery()
+    {
+        queryRunner.execute(TEST_SESSION, "SELECT * FROM blackhole.test_schema.slow_test_table");
+    }
+}