Add airlift stats

gateway-ha/pom.xml

@@ -121,6 +121,11 @@
             <artifactId>log-manager</artifactId>
         </dependency>
 
+        <dependency>
+            <groupId>io.airlift</groupId>
+            <artifactId>stats</artifactId>
+        </dependency>
+
         <dependency>
             <groupId>io.dropwizard</groupId>
             <artifactId>dropwizard-assets</artifactId>
@@ -343,6 +348,11 @@
             <artifactId>slf4j-api</artifactId>
         </dependency>
 
+        <dependency>
+            <groupId>org.weakref</groupId>
+            <artifactId>jmxutils</artifactId>
+        </dependency>
+
         <dependency>
             <groupId>io.dropwizard</groupId>
             <artifactId>dropwizard-views-freemarker</artifactId>
@@ -529,6 +539,8 @@
                 <configuration>
                     <!-- necessary for integration tests, also see /.mvn/jvm.config -->
                     <argLine>--add-opens=java.base/java.lang=ALL-UNNAMED --add-opens=java.base/java.net=ALL-UNNAMED</argLine>
+                    <!-- This is needed because we call main() in test setups. Otherwise Guice singleton won't work. -->
+                    <reuseForks>false</reuseForks>
                 </configuration>
             </plugin>
 

gateway-ha/src/main/java/io/trino/gateway/ha/handler/ProxyHandlerStats.java

@@ -0,0 +1,63 @@
+/*
+ * 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.handler;
+
+import com.codahale.metrics.Meter;
+import io.airlift.stats.CounterStat;
+import io.dropwizard.core.setup.Environment;
+import org.weakref.jmx.MBeanExporter;
+import org.weakref.jmx.Managed;
+import org.weakref.jmx.Nested;
+
+import java.lang.management.ManagementFactory;
+
+import static java.util.Objects.requireNonNull;
+
+public final class ProxyHandlerStats
+{
+    // Airlift
+    private final CounterStat requestCount = new CounterStat();
+
+    // Dropwizard
+    private final Meter requestMeter;
+
+    private ProxyHandlerStats(Environment environment, String metricsName)
+    {
+        this.requestMeter = requireNonNull(environment, "environment is null")
+                .metrics()
+                .meter(requireNonNull(metricsName, "metricsName is null"));
+    }
+
+    public void recordRequest()
+    {
+        requestCount.update(1);
+        requestMeter.mark();
+    }
+
+    // Replace this with Guice bind after migrated to Airlift
+    public static ProxyHandlerStats create(Environment environment, String metricsName)
+    {
+        ProxyHandlerStats proxyHandlerStats = new ProxyHandlerStats(environment, metricsName);
+        MBeanExporter exporter = new MBeanExporter(ManagementFactory.getPlatformMBeanServer());
+        exporter.exportWithGeneratedName(proxyHandlerStats);
+        return proxyHandlerStats;
+    }
+
+    @Managed
+    @Nested
+    public CounterStat getRequestCount()
+    {
+        return requestCount;
+    }
+}

gateway-ha/src/main/java/io/trino/gateway/ha/handler/QueryIdCachingProxyHandler.java

@@ -13,7 +13,6 @@
  */
 package io.trino.gateway.ha.handler;
 
-import com.codahale.metrics.Meter;
 import com.fasterxml.jackson.databind.ObjectMapper;
 import com.google.common.base.Splitter;
 import com.google.common.io.CharStreams;
@@ -82,7 +81,7 @@ public class QueryIdCachingProxyHandler
     private final RoutingGroupSelector routingGroupSelector;
     private final QueryHistoryManager queryHistoryManager;
 
-    private final Meter requestMeter;
+    private final ProxyHandlerStats proxyHandlerStats;
     private final List<String> extraWhitelistPaths;
     private final String applicationEndpoint;
 
@@ -91,9 +90,10 @@ public QueryIdCachingProxyHandler(
             RoutingManager routingManager,
             RoutingGroupSelector routingGroupSelector,
             int serverApplicationPort,
-            Meter requestMeter, List<String> extraWhitelistPaths)
+            ProxyHandlerStats proxyHandlerStats,
+            List<String> extraWhitelistPaths)
     {
-        this.requestMeter = requestMeter;
+        this.proxyHandlerStats = proxyHandlerStats;
         this.routingManager = routingManager;
         this.routingGroupSelector = routingGroupSelector;
         this.queryHistoryManager = queryHistoryManager;
@@ -238,7 +238,7 @@ public void preConnectionHook(HttpServletRequest request, Request proxyRequest)
     {
         if (request.getMethod().equals(HttpMethod.POST)
                 && request.getRequestURI().startsWith(V1_STATEMENT_PATH)) {
-            requestMeter.mark();
+            proxyHandlerStats.recordRequest();
             try {
                 String requestBody = CharStreams.toString(request.getReader());
                 log.info(

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

@@ -13,7 +13,6 @@
  */
 package io.trino.gateway.ha.module;
 
-import com.codahale.metrics.Meter;
 import com.google.inject.Provides;
 import com.google.inject.Singleton;
 import io.dropwizard.auth.AuthFilter;
@@ -28,6 +27,7 @@
 import io.trino.gateway.ha.config.RequestRouterConfiguration;
 import io.trino.gateway.ha.config.RoutingRulesConfiguration;
 import io.trino.gateway.ha.config.UserConfiguration;
+import io.trino.gateway.ha.handler.ProxyHandlerStats;
 import io.trino.gateway.ha.handler.QueryIdCachingProxyHandler;
 import io.trino.gateway.ha.router.BackendStateManager;
 import io.trino.gateway.ha.router.QueryHistoryManager;
@@ -133,10 +133,9 @@ private ChainedAuthFilter getAuthenticationFilters(AuthenticationConfiguration c
     protected ProxyHandler getProxyHandler(QueryHistoryManager queryHistoryManager,
                                          RoutingManager routingManager)
     {
-        Meter requestMeter =
-                getEnvironment()
-                        .metrics()
-                        .meter(getConfiguration().getRequestRouter().getName() + ".requests");
+        ProxyHandlerStats proxyHandlerStats = ProxyHandlerStats.create(
+                getEnvironment(),
+                getConfiguration().getRequestRouter().getName() + ".requests");
 
         // By default, use routing group header to route
         RoutingGroupSelector routingGroupSelector = RoutingGroupSelector.byRoutingGroupHeader();
@@ -152,7 +151,7 @@ protected ProxyHandler getProxyHandler(QueryHistoryManager queryHistoryManager,
                 routingManager,
                 routingGroupSelector,
                 getApplicationPort(),
-                requestMeter,
+                proxyHandlerStats,
                 extraWhitelistPaths);
     }