Emit QueryCountMetrics for SQL Queries

extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/http/DartSqlResource.java

@@ -33,6 +33,7 @@
 import org.apache.druid.msq.dart.controller.sql.DartSqlClients;
 import org.apache.druid.msq.dart.controller.sql.DartSqlEngine;
 import org.apache.druid.query.DefaultQueryConfig;
+import org.apache.druid.server.BaseQueryCountResource;
 import org.apache.druid.server.DruidNode;
 import org.apache.druid.server.ResponseContextConfig;
 import org.apache.druid.server.initialization.ServerConfig;
@@ -98,7 +99,8 @@ public DartSqlResource(
       final ServerConfig serverConfig,
       final ResponseContextConfig responseContextConfig,
       @Self final DruidNode selfNode,
-      @Dart final DefaultQueryConfig dartQueryConfig
+      @Dart final DefaultQueryConfig dartQueryConfig,
+      final BaseQueryCountResource baseQueryResource
   )
   {
     super(
@@ -108,7 +110,8 @@ public DartSqlResource(
         sqlLifecycleManager,
         serverConfig,
         responseContextConfig,
-        selfNode
+        selfNode,
+        baseQueryResource
     );
     this.controllerRegistry = controllerRegistry;
     this.sqlLifecycleManager = sqlLifecycleManager;

extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/dart/controller/http/DartSqlResourceTest.java

@@ -51,6 +51,7 @@
 import org.apache.druid.query.DefaultQueryConfig;
 import org.apache.druid.query.QueryContext;
 import org.apache.druid.query.QueryContexts;
+import org.apache.druid.server.BaseQueryCountResource;
 import org.apache.druid.server.DruidNode;
 import org.apache.druid.server.QueryStackTests;
 import org.apache.druid.server.ResponseContextConfig;
@@ -239,7 +240,8 @@ public void register(ControllerHolder holder)
         new ServerConfig() /* currently only used for error transform strategy */,
         ResponseContextConfig.newConfig(false),
         SELF_NODE,
-        new DefaultQueryConfig(ImmutableMap.of("foo", "bar"))
+        new DefaultQueryConfig(ImmutableMap.of("foo", "bar")),
+        new BaseQueryCountResource()
     );
 
     // Setup mocks

quidem-ut/src/main/java/org/apache/druid/quidem/ExposedAsBrokerQueryComponentSupplierWrapper.java

@@ -74,6 +74,7 @@
 import org.apache.druid.query.RetryQueryRunnerConfig;
 import org.apache.druid.rpc.guice.ServiceClientModule;
 import org.apache.druid.segment.writeout.SegmentWriteOutMediumModule;
+import org.apache.druid.server.BaseQueryCountResource;
 import org.apache.druid.server.BrokerQueryResource;
 import org.apache.druid.server.ClientInfoResource;
 import org.apache.druid.server.DruidNode;
@@ -242,7 +243,7 @@ static List<? extends Module> brokerModules()
           binder.bind(BrokerQueryResource.class).in(LazySingleton.class);
           Jerseys.addResource(binder, BrokerQueryResource.class);
           binder.bind(SubqueryGuardrailHelper.class).toProvider(SubqueryGuardrailHelperProvider.class);
-          binder.bind(QueryCountStatsProvider.class).to(BrokerQueryResource.class).in(LazySingleton.class);
+          binder.bind(QueryCountStatsProvider.class).to(BaseQueryCountResource.class).in(LazySingleton.class);
           binder.bind(SubqueryCountStatsProvider.class).toInstance(new SubqueryCountStatsProvider());
           Jerseys.addResource(binder, BrokerResource.class);
           Jerseys.addResource(binder, ClientInfoResource.class);

server/src/main/java/org/apache/druid/server/BaseQueryCountResource.java

@@ -0,0 +1,82 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you 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 org.apache.druid.server;
+
+import org.apache.druid.guice.LazySingleton;
+import org.apache.druid.server.metrics.QueryCountStatsProvider;
+
+import java.util.concurrent.atomic.AtomicLong;
+
+@LazySingleton
+public class BaseQueryCountResource implements QueryCountStatsProvider, QueryMetricCounter
+{
+  private final AtomicLong successfulQueryCount = new AtomicLong();
+  private final AtomicLong failedQueryCount = new AtomicLong();
+  private final AtomicLong interruptedQueryCount = new AtomicLong();
+  private final AtomicLong timedOutQueryCount = new AtomicLong();
+
+  @Override
+  public long getSuccessfulQueryCount()
+  {
+    return successfulQueryCount.get();
+  }
+
+  @Override
+  public long getFailedQueryCount()
+  {
+    return failedQueryCount.get();
+  }
+
+  @Override
+  public long getInterruptedQueryCount()
+  {
+    return interruptedQueryCount.get();
+  }
+
+  @Override
+  public long getTimedOutQueryCount()
+  {
+    return timedOutQueryCount.get();
+  }
+
+  @Override
+  public void incrementSuccess()
+  {
+    successfulQueryCount.incrementAndGet();
+  }
+
+  @Override
+  public void incrementFailed()
+  {
+    failedQueryCount.incrementAndGet();
+  }
+
+  @Override
+  public void incrementInterrupted()
+  {
+    interruptedQueryCount.incrementAndGet();
+  }
+
+  @Override
+  public void incrementTimedOut()
+  {
+    timedOutQueryCount.incrementAndGet();
+  }
+}

server/src/main/java/org/apache/druid/server/BrokerQueryResource.java

@@ -63,7 +63,8 @@ public BrokerQueryResource(
       AuthorizerMapper authorizerMapper,
       ResponseContextConfig responseContextConfig,
       @Self DruidNode selfNode,
-      TimelineServerView brokerServerView
+      TimelineServerView brokerServerView,
+      BaseQueryCountResource counter
   )
   {
     super(
@@ -74,7 +75,8 @@ public BrokerQueryResource(
         authConfig,
         authorizerMapper,
         responseContextConfig,
-        selfNode
+        selfNode,
+        counter
     );
     this.brokerServerView = brokerServerView;
   }

server/src/main/java/org/apache/druid/server/QueryMetricCounter.java

@@ -0,0 +1,31 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you 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 org.apache.druid.server;
+
+public interface QueryMetricCounter
+{
+  void incrementSuccess();
+
+  void incrementFailed();
+
+  void incrementInterrupted();
+
+  void incrementTimedOut();
+}

server/src/main/java/org/apache/druid/server/QueryResource.java

@@ -48,7 +48,6 @@
 import org.apache.druid.query.QueryToolChest;
 import org.apache.druid.query.context.ResponseContext;
 import org.apache.druid.query.context.ResponseContext.Keys;
-import org.apache.druid.server.metrics.QueryCountStatsProvider;
 import org.apache.druid.server.security.Access;
 import org.apache.druid.server.security.AuthConfig;
 import org.apache.druid.server.security.AuthorizationUtils;
@@ -76,11 +75,10 @@
 import java.io.OutputStream;
 import java.util.Set;
 import java.util.TreeSet;
-import java.util.concurrent.atomic.AtomicLong;
 
 @LazySingleton
 @Path("/druid/v2/")
-public class QueryResource implements QueryCountStatsProvider
+public class QueryResource
 {
   protected static final EmittingLogger log = new EmittingLogger(QueryResource.class);
   public static final EmittingLogger NO_STACK_LOGGER = log.noStackTrace();
@@ -109,11 +107,7 @@ public class QueryResource implements QueryCountStatsProvider
   private final ResponseContextConfig responseContextConfig;
   private final DruidNode selfNode;
 
-  private final AtomicLong successfulQueryCount = new AtomicLong();
-  private final AtomicLong failedQueryCount = new AtomicLong();
-  private final AtomicLong interruptedQueryCount = new AtomicLong();
-  private final AtomicLong timedOutQueryCount = new AtomicLong();
-  private final QueryResourceQueryMetricCounter counter = new QueryResourceQueryMetricCounter();
+  final BaseQueryCountResource counter;
 
   @Inject
   public QueryResource(
@@ -124,7 +118,8 @@ public QueryResource(
       AuthConfig authConfig,
       AuthorizerMapper authorizerMapper,
       ResponseContextConfig responseContextConfig,
-      @Self DruidNode selfNode
+      @Self DruidNode selfNode,
+      final BaseQueryCountResource counter
   )
   {
     this.queryLifecycleFactory = queryLifecycleFactory;
@@ -136,6 +131,7 @@ public QueryResource(
     this.authorizerMapper = authorizerMapper;
     this.responseContextConfig = responseContextConfig;
     this.selfNode = selfNode;
+    this.counter = counter;
   }
 
   @DELETE
@@ -258,17 +254,6 @@ public Response doPost(
     }
   }
 
-  public interface QueryMetricCounter
-  {
-    void incrementSuccess();
-
-    void incrementFailed();
-
-    void incrementInterrupted();
-
-    void incrementTimedOut();
-  }
-
   private Query<?> readQuery(
       final HttpServletRequest req,
       final InputStream in,
@@ -416,30 +401,6 @@ Response buildNonOkResponse(int status, Exception e) throws JsonProcessingExcept
     }
   }
 
-  @Override
-  public long getSuccessfulQueryCount()
-  {
-    return successfulQueryCount.get();
-  }
-
-  @Override
-  public long getFailedQueryCount()
-  {
-    return failedQueryCount.get();
-  }
-
-  @Override
-  public long getInterruptedQueryCount()
-  {
-    return interruptedQueryCount.get();
-  }
-
-  @Override
-  public long getTimedOutQueryCount()
-  {
-    return timedOutQueryCount.get();
-  }
-
   @VisibleForTesting
   public static void transferEntityTag(ResponseContext context, Response.ResponseBuilder builder)
   {
@@ -449,33 +410,6 @@ public static void transferEntityTag(ResponseContext context, Response.ResponseB
     }
   }
 
-  private class QueryResourceQueryMetricCounter implements QueryMetricCounter
-  {
-    @Override
-    public void incrementSuccess()
-    {
-      successfulQueryCount.incrementAndGet();
-    }
-
-    @Override
-    public void incrementFailed()
-    {
-      failedQueryCount.incrementAndGet();
-    }
-
-    @Override
-    public void incrementInterrupted()
-    {
-      interruptedQueryCount.incrementAndGet();
-    }
-
-    @Override
-    public void incrementTimedOut()
-    {
-      timedOutQueryCount.incrementAndGet();
-    }
-  }
-
   private class QueryResourceQueryResultPusher extends QueryResultPusher
   {
     private final HttpServletRequest req;

server/src/main/java/org/apache/druid/server/QueryResultPusher.java

@@ -63,7 +63,7 @@ public abstract class QueryResultPusher
   private final ObjectMapper jsonMapper;
   private final ResponseContextConfig responseContextConfig;
   private final DruidNode selfNode;
-  private final QueryResource.QueryMetricCounter counter;
+  private final QueryMetricCounter counter;
   private final MediaType contentType;
   private final Map<String, String> extraHeaders;
   private final HttpFields trailerFields;
@@ -77,7 +77,7 @@ public QueryResultPusher(
       ObjectMapper jsonMapper,
       ResponseContextConfig responseContextConfig,
       DruidNode selfNode,
-      QueryResource.QueryMetricCounter counter,
+      QueryMetricCounter counter,
       String queryId,
       MediaType contentType,
       Map<String, String> extraHeaders