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
trinodb · Jan 2, 2025 · b005207 · b005207
1 parent 5e8d877
commit b005207
Show file tree

Hide file tree

Showing 3 changed files with 36 additions and 0 deletions.
diff --git a/gateway-ha/src/main/java/io/trino/gateway/ha/clustermonitor/ClusterStatsJdbcMonitor.java b/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;
@@ -38,6 +39,7 @@ public class ClusterStatsJdbcMonitor
     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 +52,10 @@ public ClusterStatsJdbcMonitor(BackendStateConfiguration backendStateConfigurati
         properties.setProperty("user", backendStateConfiguration.getUsername());
         properties.setProperty("password", backendStateConfiguration.getPassword());
         properties.setProperty("SSL", String.valueOf(backendStateConfiguration.getSsl()));
+        if (!monitorConfiguration.isExplicitPrepare()) { //do not set property if true (default) to avoid issues with older Trinos
+            properties.setProperty("explicitPrepare", "false");
+        }
+        queryTimeout = monitorConfiguration.getQueryTimeout();
         log.info("state check configured");
     }
 
@@ -77,6 +83,7 @@ public ClusterStats monitor(ProxyBackendConfiguration backend)
                 PreparedStatement statement = SimpleTimeLimiter.create(Executors.newSingleThreadExecutor()).callWithTimeout(
                         () -> conn.prepareStatement(STATE_QUERY), 10, TimeUnit.SECONDS)) {
             statement.setString(1, (String) properties.get("user"));
+            statement.setQueryTimeout((int)queryTimeout.toMillis()/1000);
             Map<String, Integer> partialState = new HashMap<>();
             ResultSet rs = statement.executeQuery();
             while (rs.next()) {

diff --git a/gateway-ha/src/main/java/io/trino/gateway/ha/config/MonitorConfiguration.java b/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;
+    }
 }
diff --git a/gateway-ha/src/test/java/io/trino/gateway/ha/clustermonitor/TestClusterStatsMonitor.java b/gateway-ha/src/test/java/io/trino/gateway/ha/clustermonitor/TestClusterStatsMonitor.java
@@ -13,8 +13,10 @@
  */
 package io.trino.gateway.ha.clustermonitor;
 
+import com.google.common.collect.ImmutableMap;
 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,6 +60,8 @@ void testHttpMonitor()
     @Test
     void testJdbcMonitor()
     {
+        MonitorConfiguration monitorConfigurationWithTimeout = new MonitorConfiguration();
+        monitorConfigurationWithTimeout.setQueryTimeout(Duration.valueOf("30s"));
         testClusterStatsMonitor(backendStateConfiguration -> new ClusterStatsJdbcMonitor(backendStateConfiguration, new MonitorConfiguration()));
     }