Add a timeout for the backend stats collection query

docs/installation.md

@@ -365,17 +365,26 @@ method of backend clusters. Configure a username and password by adding
 `backendState` to your configuration. The username and password must be valid 
 across all backends.
 
+Trino Gateway uses `explicitPrepare=false` by default. This property was introduced
+in Trino 431, and uses a single query for prepared statements, instead of a 
+`PREPARE/EXECUTE` pair. If you are using the JDBC health check option with older 
+versions of Trino, set
+```yaml
+monitorConfiguration:
+   explicitPrepare: false
+```
+
 ```yaml
 backendState:
   username: "user"
   password: "password"
 ```
 
-The request timeout can be set through
+The query timeout can be set through
 
 ```yaml
 monitor:
-  requestTimeoutSeconds: 10
+    queryTimeout: 10
 ```
 
 Other timeout parameters are not applicable to the JDBC connection.

gateway-ha/src/main/java/io/trino/gateway/ha/clustermonitor/ClusterStatsJdbcMonitor.java

@@ -15,6 +15,7 @@
 
 import com.google.common.util.concurrent.SimpleTimeLimiter;
 import io.airlift.log.Logger;
+import io.airlift.units.Duration;
 import io.trino.gateway.ha.config.BackendStateConfiguration;
 import io.trino.gateway.ha.config.MonitorConfiguration;
 import io.trino.gateway.ha.config.ProxyBackendConfiguration;
@@ -29,15 +30,17 @@
 import java.util.Map;
 import java.util.Properties;
 import java.util.concurrent.Executors;
-import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
 
+import static java.util.concurrent.TimeUnit.SECONDS;
+
 public class ClusterStatsJdbcMonitor
         implements ClusterStatsMonitor
 {
     private static final Logger log = Logger.get(ClusterStatsJdbcMonitor.class);
 
     private final Properties properties; // TODO Avoid using a mutable field
+    private final Duration queryTimeout;
 
     private static final String STATE_QUERY = "SELECT state, COUNT(*) as count "
             + "FROM runtime.queries "
@@ -50,6 +53,12 @@ public ClusterStatsJdbcMonitor(BackendStateConfiguration backendStateConfigurati
         properties.setProperty("user", backendStateConfiguration.getUsername());
         properties.setProperty("password", backendStateConfiguration.getPassword());
         properties.setProperty("SSL", String.valueOf(backendStateConfiguration.getSsl()));
+        // explicitPrepare is a valid property for Trino versions >= 431. To avoid compatibility
+        // issues with versions < 431, this property is left unset when explicitPrepare=true, which is the default
+        if (!monitorConfiguration.isExplicitPrepare()) {
+            properties.setProperty("explicitPrepare", "false");
+        }
+        queryTimeout = monitorConfiguration.getQueryTimeout();
         log.info("state check configured");
     }
 
@@ -75,8 +84,9 @@ public ClusterStats monitor(ProxyBackendConfiguration backend)
 
         try (Connection conn = DriverManager.getConnection(jdbcUrl, properties);
                 PreparedStatement statement = SimpleTimeLimiter.create(Executors.newSingleThreadExecutor()).callWithTimeout(
-                        () -> conn.prepareStatement(STATE_QUERY), 10, TimeUnit.SECONDS)) {
+                        () -> conn.prepareStatement(STATE_QUERY), 10, SECONDS)) {
             statement.setString(1, (String) properties.get("user"));
+            statement.setQueryTimeout((int) queryTimeout.roundTo(SECONDS));
             Map<String, Integer> partialState = new HashMap<>();
             ResultSet rs = statement.executeQuery();
             while (rs.next()) {

gateway-ha/src/main/java/io/trino/gateway/ha/config/MonitorConfiguration.java

@@ -13,14 +13,21 @@
  */
 package io.trino.gateway.ha.config;
 
+import io.airlift.units.Duration;
 import io.trino.gateway.ha.clustermonitor.ActiveClusterMonitor;
 
+import static java.util.concurrent.TimeUnit.SECONDS;
+
 public class MonitorConfiguration
 {
     private int taskDelaySeconds = ActiveClusterMonitor.MONITOR_TASK_DELAY_SECONDS;
 
     private int retries;
 
+    private Duration queryTimeout = new Duration(10, SECONDS);
+
+    private boolean explicitPrepare;
+
     public MonitorConfiguration() {}
 
     public int getTaskDelaySeconds()
@@ -42,4 +49,24 @@ public void setRetries(int retries)
     {
         this.retries = retries;
     }
+
+    public Duration getQueryTimeout()
+    {
+        return queryTimeout;
+    }
+
+    public void setQueryTimeout(Duration queryTimeout)
+    {
+        this.queryTimeout = queryTimeout;
+    }
+
+    public boolean isExplicitPrepare()
+    {
+        return explicitPrepare;
+    }
+
+    public void setExplicitPrepare(boolean explicitPrepare)
+    {
+        this.explicitPrepare = explicitPrepare;
+    }
 }

gateway-ha/src/test/java/io/trino/gateway/ha/clustermonitor/TestClusterStatsMonitor.java

@@ -15,6 +15,7 @@
 
 import io.airlift.http.client.HttpClientConfig;
 import io.airlift.http.client.jetty.JettyHttpClient;
+import io.airlift.units.Duration;
 import io.trino.gateway.ha.config.BackendStateConfiguration;
 import io.trino.gateway.ha.config.MonitorConfiguration;
 import io.trino.gateway.ha.config.ProxyBackendConfiguration;
@@ -26,6 +27,7 @@
 
 import java.util.function.Function;
 
+import static java.util.concurrent.TimeUnit.SECONDS;
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.junit.jupiter.api.TestInstance.Lifecycle.PER_CLASS;
 import static org.testcontainers.utility.MountableFile.forClasspathResource;
@@ -58,7 +60,9 @@ void testHttpMonitor()
     @Test
     void testJdbcMonitor()
     {
-        testClusterStatsMonitor(backendStateConfiguration -> new ClusterStatsJdbcMonitor(backendStateConfiguration, new MonitorConfiguration()));
+        MonitorConfiguration monitorConfigurationWithTimeout = new MonitorConfiguration();
+        monitorConfigurationWithTimeout.setQueryTimeout(new Duration(30, SECONDS));
+        testClusterStatsMonitor(backendStateConfiguration -> new ClusterStatsJdbcMonitor(backendStateConfiguration, monitorConfigurationWithTimeout));
     }
 
     @Test