Add a monitor for the OpenMetrics endpoint

docs/installation.md

@@ -387,6 +387,38 @@ monitor:
 
 All timeout parameters are optional.
 
+#### METRICS
+
+This pulls statistics from Trino's [OpenMetrics](https://openmetrics.io/) endpoint.
+It retrieves the number of running and queued queries for use with
+the `QueryCountBasedRouter` (either `METRICS` or `JDBC` must be enabled if
+`QueryCountBasedRouter` is used).
+
+This monitor allows customizing health definitions by comparing metrics to fixed
+values. This is configured through two maps: `metricMinimumValues` and 
+`metricMaximumValues`. The keys of these maps are the metric names, and the values
+are the minimum or maximum values (inclusive) that are considered healthy. By default,
+the only metric populated is:
+
+```yaml
+monitorConfiguration:
+    metricMinimumValues:
+        trino_metadata_name_DiscoveryNodeManager_ActiveNodeCount: 1
+```
+
+This requires the cluster to have at least one active worker node in order to be considered
+healthy. The map is overwritten if configured explicitly. For example, to increase the minimum
+worker count to 10 and disqualify clusters that have been experiencing frequent major Garbage
+Collections, set
+
+```yaml
+monitorConfiguration:
+    metricMinimumValues:
+        trino_metadata_name_DiscoveryNodeManager_ActiveNodeCount: 10
+    metricMaximumValues:
+        io_airlift_stats_name_GcMonitor_MajorGc_FiveMinutes_count: 2
+```
+
 #### JDBC
 
 This uses a JDBC connection to query `system.runtime` tables for cluster 

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

@@ -27,9 +27,7 @@
 import static io.airlift.http.client.JsonResponseHandler.createJsonResponseHandler;
 import static io.airlift.http.client.Request.Builder.prepareGet;
 import static io.airlift.json.JsonCodec.jsonCodec;
-import static java.net.HttpURLConnection.HTTP_BAD_GATEWAY;
-import static java.net.HttpURLConnection.HTTP_GATEWAY_TIMEOUT;
-import static java.net.HttpURLConnection.HTTP_UNAVAILABLE;
+import static io.trino.gateway.ha.clustermonitor.MonitorUtils.shouldRetry;
 import static java.util.Objects.requireNonNull;
 
 public class ClusterStatsInfoApiMonitor
@@ -88,16 +86,4 @@ private TrinoStatus checkStatus(String baseUrl, int retriesRemaining)
         }
         return TrinoStatus.UNHEALTHY;
     }
-
-    public static boolean shouldRetry(int statusCode)
-    {
-        switch (statusCode) {
-            case HTTP_BAD_GATEWAY:
-            case HTTP_UNAVAILABLE:
-            case HTTP_GATEWAY_TIMEOUT:
-                return true;
-            default:
-                return false;
-        }
-    }
 }

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

@@ -0,0 +1,200 @@
+/*
+ * 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.gateway.ha.clustermonitor;
+
+import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.ImmutableSet;
+import io.airlift.http.client.HttpClient;
+import io.airlift.http.client.HttpUriBuilder;
+import io.airlift.http.client.Request;
+import io.airlift.http.client.Response;
+import io.airlift.http.client.ResponseHandler;
+import io.airlift.http.client.UnexpectedResponseException;
+import io.airlift.log.Logger;
+import io.trino.gateway.ha.config.BackendStateConfiguration;
+import io.trino.gateway.ha.config.MonitorConfiguration;
+import io.trino.gateway.ha.config.ProxyBackendConfiguration;
+import io.trino.gateway.ha.security.util.BasicCredentials;
+
+import java.io.IOException;
+import java.net.URI;
+import java.util.Arrays;
+import java.util.Map;
+import java.util.Set;
+
+import static com.google.common.base.Strings.isNullOrEmpty;
+import static com.google.common.collect.ImmutableMap.toImmutableMap;
+import static io.airlift.http.client.HttpUriBuilder.uriBuilderFrom;
+import static io.airlift.http.client.Request.Builder.prepareGet;
+import static io.airlift.http.client.ResponseHandlerUtils.propagate;
+import static io.trino.gateway.ha.clustermonitor.MonitorUtils.shouldRetry;
+import static java.lang.String.format;
+import static java.nio.charset.StandardCharsets.UTF_8;
+import static java.util.Objects.requireNonNull;
+
+public class ClusterStatsMetricsMonitor
+        implements ClusterStatsMonitor
+{
+    public static final String RUNNING_QUERIES_METRIC = "trino_execution_name_QueryManager_RunningQueries";
+    public static final String QUEUED_QUERIES_METRIC = "trino_execution_name_QueryManager_QueuedQueries";
+    private static final Logger log = Logger.get(ClusterStatsMetricsMonitor.class);
+
+    private final HttpClient httpClient;
+    private final int retries;
+    private final MetricsResponseHandler metricsResponseHandler;
+    private final Header identityHeader;
+    private final String metricsEndpoint;
+    private final ImmutableSet<String> metricNames;
+    private final Map<String, Float> metricMinimumValues;
+    private final Map<String, Float> metricMaximumValues;
+
+    public ClusterStatsMetricsMonitor(HttpClient httpClient, BackendStateConfiguration backendStateConfiguration, MonitorConfiguration monitorConfiguration)
+    {
+        this.httpClient = requireNonNull(httpClient, "httpClient is null");
+        retries = monitorConfiguration.getRetries();
+        if (!isNullOrEmpty(backendStateConfiguration.getPassword())) {
+            identityHeader = new Header("Authorization",
+                    new BasicCredentials(backendStateConfiguration.getUsername(), backendStateConfiguration.getPassword()).getBasicAuthHeader());
+        }
+        else {
+            identityHeader = new Header("X-Trino-User", backendStateConfiguration.getUsername());
+        }
+        metricsEndpoint = monitorConfiguration.getMetricsEndpoint();
+        metricMinimumValues = ImmutableMap.copyOf(monitorConfiguration.getMetricMinimumValues());
+        metricMaximumValues = ImmutableMap.copyOf(monitorConfiguration.getMetricMaximumValues());
+        metricNames = ImmutableSet.<String>builder()
+                .add(RUNNING_QUERIES_METRIC, QUEUED_QUERIES_METRIC)
+                .addAll(metricMinimumValues.keySet())
+                .addAll(metricMaximumValues.keySet())
+                .build();
+        metricsResponseHandler = new MetricsResponseHandler(metricNames);
+    }
+
+    private static ClusterStats getUnhealthyStats(ProxyBackendConfiguration backend)
+    {
+        return ClusterStats.builder(backend.getName())
+                .trinoStatus(TrinoStatus.UNHEALTHY)
+                .proxyTo(backend.getProxyTo())
+                .externalUrl(backend.getExternalUrl())
+                .routingGroup(backend.getRoutingGroup())
+                .build();
+    }
+
+    @Override
+    public ClusterStats monitor(ProxyBackendConfiguration backend)
+    {
+        Map<String, String> metrics = getMetrics(backend.getProxyTo(), retries);
+        if (metrics.isEmpty()) {
+            log.error("No metrics available for %s!", backend.getName());
+            return getUnhealthyStats(backend);
+        }
+
+        for (Map.Entry<String, Float> entry : metricMinimumValues.entrySet()) {
+            if (!metrics.containsKey(entry.getKey())
+                    || Float.parseFloat(metrics.get(entry.getKey())) < entry.getValue()) {
+                log.warn("Health metric value below min for cluster %s: %s=%s", backend.getName(), entry.getKey(), metrics.get(entry.getKey()));
+                return getUnhealthyStats(backend);
+            }
+        }
+
+        for (Map.Entry<String, Float> entry : metricMaximumValues.entrySet()) {
+            if (!metrics.containsKey(entry.getKey())
+                    || Float.parseFloat(metrics.get(entry.getKey())) > entry.getValue()) {
+                log.warn("Health metric value over max for cluster %s: %s=%s", backend.getName(), entry.getKey(), metrics.get(entry.getKey()));
+                return getUnhealthyStats(backend);
+            }
+        }
+        return ClusterStats.builder(backend.getName())
+                .trinoStatus(TrinoStatus.HEALTHY)
+                .runningQueryCount((int) Float.parseFloat(metrics.get(RUNNING_QUERIES_METRIC)))
+                .queuedQueryCount((int) Float.parseFloat(metrics.get(QUEUED_QUERIES_METRIC)))
+                .proxyTo(backend.getProxyTo())
+                .externalUrl(backend.getExternalUrl())
+                .routingGroup(backend.getRoutingGroup())
+                .build();
+    }
+
+    private Map<String, String> getMetrics(String baseUrl, int retriesRemaining)
+    {
+        HttpUriBuilder uri = uriBuilderFrom(URI.create(baseUrl)).appendPath(metricsEndpoint);
+        for (String metric : metricNames) {
+            uri.addParameter("name[]", metric);
+        }
+
+        Request request = prepareGet()
+                .setUri(uri.build())
+                .addHeader(identityHeader.name, identityHeader.value)
+                .addHeader("Content-Type", "application/openmetrics-text; version=1.0.0; charset=utf-8")
+                .build();
+        try {
+            return httpClient.execute(request, metricsResponseHandler);
+        }
+        catch (UnexpectedResponseException e) {
+            if (shouldRetry(e.getStatusCode())) {
+                if (retriesRemaining > 0) {
+                    log.warn("Retrying health check on error: %s, ", e.toString());
+                    return getMetrics(baseUrl, retriesRemaining - 1);
+                }
+                log.error("Encountered error %s, no retries remaining", e.toString());
+            }
+            log.error(e, "Health check failed with non-retryable response.\n%s", e.toString());
+        }
+        catch (Exception e) {
+            log.error(e, "Exception checking %s for health", request.getUri());
+        }
+        return ImmutableMap.of();
+    }
+
+    private static class MetricsResponseHandler
+            implements ResponseHandler<Map<String, String>, RuntimeException>
+    {
+        private final ImmutableSet<String> requiredKeys;
+
+        public MetricsResponseHandler(Set<String> requiredKeys)
+        {
+            this.requiredKeys = ImmutableSet.copyOf(requiredKeys);
+        }
+
+        @Override
+        public Map<String, String> handleException(Request request, Exception exception)
+                throws RuntimeException
+        {
+            throw propagate(request, exception);
+        }
+
+        @Override
+        public Map<String, String> handle(Request request, Response response)
+                throws RuntimeException
+        {
+            try {
+                String responseBody = new String(response.getInputStream().readAllBytes(), UTF_8);
+                Map<String, String> metrics = Arrays.stream(responseBody.split("\n"))
+                        .filter(line -> !line.startsWith("#"))
+                        .collect(toImmutableMap(s -> s.split(" ")[0], s -> s.split(" ")[1]));
+                if (!metrics.keySet().containsAll(requiredKeys)) {
+                    throw new UnexpectedResponseException(
+                            format("Request is missing required keys: \n%s\nin response: '%s'", String.join("\n", requiredKeys), responseBody),
+                            request,
+                            response);
+                }
+                return metrics;
+            }
+            catch (IOException e) {
+                throw new UnexpectedResponseException(request, response);
+            }
+        }
+    }
+
+    private record Header(String name, String value) {}
+}

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

@@ -0,0 +1,31 @@
+/*
+ * 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.gateway.ha.clustermonitor;
+
+import static java.net.HttpURLConnection.HTTP_BAD_GATEWAY;
+import static java.net.HttpURLConnection.HTTP_GATEWAY_TIMEOUT;
+import static java.net.HttpURLConnection.HTTP_UNAVAILABLE;
+
+public final class MonitorUtils
+{
+    private MonitorUtils() {}
+
+    public static boolean shouldRetry(int statusCode)
+    {
+        return switch (statusCode) {
+            case HTTP_BAD_GATEWAY, HTTP_UNAVAILABLE, HTTP_GATEWAY_TIMEOUT -> true;
+            default -> false;
+        };
+    }
+}

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

@@ -19,5 +19,6 @@ public enum ClusterStatsMonitorType
     INFO_API,
     UI_API,
     JDBC,
-    JMX
+    JMX,
+    METRICS
 }

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

@@ -13,9 +13,12 @@
  */
 package io.trino.gateway.ha.config;
 
+import com.google.common.collect.ImmutableMap;
 import io.airlift.units.Duration;
 import io.trino.gateway.ha.clustermonitor.ActiveClusterMonitor;
 
+import java.util.Map;
+
 import static java.util.concurrent.TimeUnit.SECONDS;
 
 public class MonitorConfiguration
@@ -28,6 +31,13 @@ public class MonitorConfiguration
 
     private boolean explicitPrepare;
 
+    private String metricsEndpoint = "/metrics";
+
+    // Require 1 node for health by default. This configuration only applies to the ClusterStatsMetricsMonitor
+    private Map<String, Float> metricMinimumValues = ImmutableMap.of("trino_metadata_name_DiscoveryNodeManager_ActiveNodeCount", 1f);
+
+    private Map<String, Float> metricMaximumValues = ImmutableMap.of();
+
     public MonitorConfiguration() {}
 
     public int getTaskDelaySeconds()
@@ -69,4 +79,34 @@ public void setExplicitPrepare(boolean explicitPrepare)
     {
         this.explicitPrepare = explicitPrepare;
     }
+
+    public String getMetricsEndpoint()
+    {
+        return metricsEndpoint;
+    }
+
+    public void setMetricsEndpoint(String metricsEndpoint)
+    {
+        this.metricsEndpoint = metricsEndpoint;
+    }
+
+    public Map<String, Float> getMetricMinimumValues()
+    {
+        return metricMinimumValues;
+    }
+
+    public void setMetricMinimumValues(Map<String, Float> metricMinimumValues)
+    {
+        this.metricMinimumValues = ImmutableMap.copyOf(metricMinimumValues);
+    }
+
+    public Map<String, Float> getMetricMaximumValues()
+    {
+        return ImmutableMap.copyOf(metricMaximumValues);
+    }
+
+    public void setMetricMaximumValues(Map<String, Float> metricMaximumValues)
+    {
+        this.metricMaximumValues = metricMaximumValues;
+    }
 }

gateway-ha/src/main/java/io/trino/gateway/ha/module/HaGatewayProviderModule.java

@@ -23,6 +23,7 @@
 import io.trino.gateway.ha.clustermonitor.ClusterStatsInfoApiMonitor;
 import io.trino.gateway.ha.clustermonitor.ClusterStatsJdbcMonitor;
 import io.trino.gateway.ha.clustermonitor.ClusterStatsJmxMonitor;
+import io.trino.gateway.ha.clustermonitor.ClusterStatsMetricsMonitor;
 import io.trino.gateway.ha.clustermonitor.ClusterStatsMonitor;
 import io.trino.gateway.ha.clustermonitor.ClusterStatsObserver;
 import io.trino.gateway.ha.clustermonitor.ForMonitor;
@@ -230,6 +231,7 @@ public ClusterStatsMonitor getClusterStatsMonitor(@ForMonitor HttpClient httpCli
             case UI_API -> new ClusterStatsHttpMonitor(configuration.getBackendState());
             case JDBC -> new ClusterStatsJdbcMonitor(configuration.getBackendState(), configuration.getMonitor());
             case JMX -> new ClusterStatsJmxMonitor(httpClient, configuration.getBackendState());
+            case METRICS -> new ClusterStatsMetricsMonitor(httpClient, configuration.getBackendState(), configuration.getMonitor());
             case NOOP -> new NoopClusterStatsMonitor();
         };
     }