Add a timeout for the backend stats collection query.

Default the stats collection query to use EXECUTE IMMEDIATE by default, with the option to use explicit PREPARE if desired

Author: willmostly

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,6 +13,7 @@
  */
 package io.trino.gateway.ha.config;
 
+import io.airlift.units.Duration;
 import io.trino.gateway.ha.clustermonitor.ActiveClusterMonitor;
 
 public class MonitorConfiguration
@@ -21,6 +22,10 @@ public class MonitorConfiguration
 
     private int retries;
 
+    private Duration queryTimeout = Duration.valueOf("10s");
+
+    private boolean explicitPrepare;
+
     public MonitorConfiguration() {}
 
     public int getTaskDelaySeconds()
@@ -42,4 +47,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;
@@ -58,7 +59,9 @@ void testHttpMonitor()
     @Test
     void testJdbcMonitor()
     {
-        testClusterStatsMonitor(backendStateConfiguration -> new ClusterStatsJdbcMonitor(backendStateConfiguration, new MonitorConfiguration()));
+        MonitorConfiguration monitorConfigurationWithTimeout = new MonitorConfiguration();
+        monitorConfigurationWithTimeout.setQueryTimeout(Duration.valueOf("30s"));
+        testClusterStatsMonitor(backendStateConfiguration -> new ClusterStatsJdbcMonitor(backendStateConfiguration, monitorConfigurationWithTimeout));
     }
 
     @Test