Fix testDirectTrinoClientLongQuery

Record heartbeats more frequently to prevent query timeouts during large
result sets or when the query state changes more often than the heartbeat
interval.

Author: piotrrzysko

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

@@ -110,6 +110,8 @@ public DispatchQuery execute(SessionContext sessionContext, @Language("SQL") Str
                             !(dispatchQuery.getState() == FINISHING && dispatchQuery.getFullQueryInfo().getStages().isEmpty());
                     state = queryManager.getQueryState(queryId)) {
                 for (Slice serializedPage = exchangeClient.pollPage(); serializedPage != null; serializedPage = exchangeClient.pollPage()) {
+                    // record heartbeat for each page to avoid query timeout during large result sets
+                    dispatchQuery.recordHeartbeat();
                     Page page = pageDeserializer.deserialize(serializedPage);
                     queryResultsListener.consumeOutputPage(page);
                 }
@@ -148,9 +150,7 @@ private void getQueryFutureWithHeartbeats(ListenableFuture<Object> anyCompleteFu
                 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();
+                // ignore, heartbeat will be recorded in finally block
             }
             catch (InterruptedException e) {
                 Thread.currentThread().interrupt();
@@ -159,6 +159,11 @@ private void getQueryFutureWithHeartbeats(ListenableFuture<Object> anyCompleteFu
             catch (ExecutionException e) {
                 throw new TrinoException(GENERIC_INTERNAL_ERROR, "Error processing query", e.getCause());
             }
+            finally {
+                // 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();
+            }
         }
     }