From f096ddea648a7c6c71d3c500665591bb82129c44 Mon Sep 17 00:00:00 2001 From: Jesse Tuglu Date: Sat, 2 Nov 2024 19:02:10 -0700 Subject: [PATCH 1/2] Create base counter query count class sharable by query resource types This also fixes a bug in the QueryResultPusher code regarding handling of exceptions. --- .../dart/controller/http/DartSqlResource.java | 7 +- .../controller/http/DartSqlResourceTest.java | 4 +- ...AsBrokerQueryComponentSupplierWrapper.java | 3 +- .../druid/server/BaseQueryCountResource.java | 63 +++++++++++++++ .../druid/server/BrokerQueryResource.java | 7 +- .../druid/server/QueryMetricCounter.java | 12 +++ .../apache/druid/server/QueryResource.java | 76 ++----------------- .../druid/server/QueryResultPusher.java | 6 +- .../druid/server/QueryResourceTest.java | 40 ++++++---- .../druid/server/QueryResultPusherTest.java | 3 +- .../java/org/apache/druid/cli/CliBroker.java | 3 +- .../org/apache/druid/cli/CliHistorical.java | 3 +- .../druid/guice/QueryablePeonModule.java | 3 +- .../apache/druid/sql/http/SqlResource.java | 38 ++-------- .../druid/sql/http/SqlResourceTest.java | 12 ++- 15 files changed, 148 insertions(+), 132 deletions(-) create mode 100644 server/src/main/java/org/apache/druid/server/BaseQueryCountResource.java create mode 100644 server/src/main/java/org/apache/druid/server/QueryMetricCounter.java diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/http/DartSqlResource.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/http/DartSqlResource.java index 65d770a29c55..5c89f7eb7267 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/http/DartSqlResource.java +++ b/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; diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/dart/controller/http/DartSqlResourceTest.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/dart/controller/http/DartSqlResourceTest.java index 51e17235203d..f9b1cdad5e29 100644 --- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/dart/controller/http/DartSqlResourceTest.java +++ b/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 diff --git a/quidem-ut/src/main/java/org/apache/druid/quidem/ExposedAsBrokerQueryComponentSupplierWrapper.java b/quidem-ut/src/main/java/org/apache/druid/quidem/ExposedAsBrokerQueryComponentSupplierWrapper.java index d1fa6a349ba4..65bc1dfc90fa 100644 --- a/quidem-ut/src/main/java/org/apache/druid/quidem/ExposedAsBrokerQueryComponentSupplierWrapper.java +++ b/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 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); diff --git a/server/src/main/java/org/apache/druid/server/BaseQueryCountResource.java b/server/src/main/java/org/apache/druid/server/BaseQueryCountResource.java new file mode 100644 index 000000000000..bde2be41c7fb --- /dev/null +++ b/server/src/main/java/org/apache/druid/server/BaseQueryCountResource.java @@ -0,0 +1,63 @@ +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(); + } +} diff --git a/server/src/main/java/org/apache/druid/server/BrokerQueryResource.java b/server/src/main/java/org/apache/druid/server/BrokerQueryResource.java index b4ebc8c6b576..0250acb765e6 100644 --- a/server/src/main/java/org/apache/druid/server/BrokerQueryResource.java +++ b/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,9 @@ public BrokerQueryResource( authConfig, authorizerMapper, responseContextConfig, - selfNode + selfNode, + counter + ); this.brokerServerView = brokerServerView; } diff --git a/server/src/main/java/org/apache/druid/server/QueryMetricCounter.java b/server/src/main/java/org/apache/druid/server/QueryMetricCounter.java new file mode 100644 index 000000000000..84d2881f790a --- /dev/null +++ b/server/src/main/java/org/apache/druid/server/QueryMetricCounter.java @@ -0,0 +1,12 @@ +package org.apache.druid.server; + +public interface QueryMetricCounter +{ + void incrementSuccess(); + + void incrementFailed(); + + void incrementInterrupted(); + + void incrementTimedOut(); +} diff --git a/server/src/main/java/org/apache/druid/server/QueryResource.java b/server/src/main/java/org/apache/druid/server/QueryResource.java index 06104000b1ca..db0f1e056889 100644 --- a/server/src/main/java/org/apache/druid/server/QueryResource.java +++ b/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; diff --git a/server/src/main/java/org/apache/druid/server/QueryResultPusher.java b/server/src/main/java/org/apache/druid/server/QueryResultPusher.java index 710c8ccc9199..bce9ed565e52 100644 --- a/server/src/main/java/org/apache/druid/server/QueryResultPusher.java +++ b/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 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 extraHeaders @@ -233,7 +233,6 @@ private Response handleDruidException(ResultsWriter resultsWriter, DruidExceptio { if (resultsWriter != null) { resultsWriter.recordFailure(e); - counter.incrementFailed(); if (accumulator != null && accumulator.isInitialized()) { // We already started sending a response when we got the error message. In this case we just give up @@ -243,6 +242,7 @@ private Response handleDruidException(ResultsWriter resultsWriter, DruidExceptio // the future. trailerFields.put(QueryResource.ERROR_MESSAGE_TRAILER_HEADER, e.getMessage()); trailerFields.put(QueryResource.RESPONSE_COMPLETE_TRAILER_HEADER, "false"); + counter.incrementFailed(); return null; } } diff --git a/server/src/test/java/org/apache/druid/server/QueryResourceTest.java b/server/src/test/java/org/apache/druid/server/QueryResourceTest.java index bf2c1933d082..7ac613c2bba8 100644 --- a/server/src/test/java/org/apache/druid/server/QueryResourceTest.java +++ b/server/src/test/java/org/apache/druid/server/QueryResourceTest.java @@ -253,7 +253,8 @@ private QueryResource createQueryResource(ResponseContextConfig responseContextC new AuthConfig(), null, responseContextConfig, - DRUID_NODE + DRUID_NODE, + new BaseQueryCountResource() ); } @@ -288,7 +289,8 @@ public void testGoodQueryWithQueryConfigOverrideDefault() throws IOException new AuthConfig(), null, ResponseContextConfig.newConfig(true), - DRUID_NODE + DRUID_NODE, + new BaseQueryCountResource() ); expectPermissiveHappyPathAuth(); @@ -362,7 +364,8 @@ public QueryRunner getQueryRunnerForSegments( new AuthConfig(), null, ResponseContextConfig.newConfig(true), - DRUID_NODE + DRUID_NODE, + new BaseQueryCountResource() ); expectPermissiveHappyPathAuth(); @@ -450,7 +453,8 @@ public QueryRunner getQueryRunnerForSegments( new AuthConfig(), null, ResponseContextConfig.newConfig(true), - DRUID_NODE + DRUID_NODE, + new BaseQueryCountResource() ); expectPermissiveHappyPathAuth(); @@ -492,7 +496,8 @@ public void testSuccessResponseWithTrailerHeader() throws IOException new AuthConfig(), null, ResponseContextConfig.newConfig(true), - DRUID_NODE + DRUID_NODE, + new BaseQueryCountResource() ); expectPermissiveHappyPathAuth(); @@ -574,7 +579,8 @@ public void emitLogsAndMetrics(@Nullable Throwable e, @Nullable String remoteAdd new AuthConfig(), null, ResponseContextConfig.newConfig(true), - DRUID_NODE + DRUID_NODE, + new BaseQueryCountResource() ); expectPermissiveHappyPathAuth(); @@ -616,7 +622,8 @@ public void testGoodQueryWithQueryConfigDoesNotOverrideQueryContext() throws IOE new AuthConfig(), null, ResponseContextConfig.newConfig(true), - DRUID_NODE + DRUID_NODE, + new BaseQueryCountResource() ); expectPermissiveHappyPathAuth(); @@ -659,7 +666,7 @@ public void testTruncatedResponseContextShouldFail() throws IOException SIMPLE_TIMESERIES_QUERY.getBytes(StandardCharsets.UTF_8), queryResource ); - Assert.assertEquals(1, queryResource.getInterruptedQueryCount()); + Assert.assertEquals(1, queryResource.counter.getInterruptedQueryCount()); Assert.assertEquals(HttpStatus.SC_INTERNAL_SERVER_ERROR, response.getStatus()); final String expectedException = new QueryInterruptedException( new TruncatedResponseContextException("Serialized response context exceeds the max size[0]"), @@ -853,7 +860,8 @@ public Access authorize(AuthenticationResult authenticationResult, Resource reso new AuthConfig(), authMapper, ResponseContextConfig.newConfig(true), - DRUID_NODE + DRUID_NODE, + new BaseQueryCountResource() ); @@ -928,7 +936,8 @@ public QueryRunner getQueryRunnerForSegments(Query query, Iterable QueryRunner getQueryRunnerForSegments(Query query, Iterable QueryRunner getQueryRunnerForSegments(Query query, Iterable extraHeaders = new HashMap(); diff --git a/services/src/main/java/org/apache/druid/cli/CliBroker.java b/services/src/main/java/org/apache/druid/cli/CliBroker.java index 94160f557791..655c59a261f5 100644 --- a/services/src/main/java/org/apache/druid/cli/CliBroker.java +++ b/services/src/main/java/org/apache/druid/cli/CliBroker.java @@ -56,6 +56,7 @@ import org.apache.druid.query.QuerySegmentWalker; import org.apache.druid.query.RetryQueryRunnerConfig; import org.apache.druid.query.lookup.LookupModule; +import org.apache.druid.server.BaseQueryCountResource; import org.apache.druid.server.BrokerQueryResource; import org.apache.druid.server.ClientInfoResource; import org.apache.druid.server.ClientQuerySegmentWalker; @@ -154,7 +155,7 @@ protected List getModules() 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); diff --git a/services/src/main/java/org/apache/druid/cli/CliHistorical.java b/services/src/main/java/org/apache/druid/cli/CliHistorical.java index ea8bbd994348..0bf1d5b52b57 100644 --- a/services/src/main/java/org/apache/druid/cli/CliHistorical.java +++ b/services/src/main/java/org/apache/druid/cli/CliHistorical.java @@ -47,6 +47,7 @@ import org.apache.druid.java.util.common.logger.Logger; import org.apache.druid.query.QuerySegmentWalker; import org.apache.druid.query.lookup.LookupModule; +import org.apache.druid.server.BaseQueryCountResource; import org.apache.druid.server.QueryResource; import org.apache.druid.server.ResponseContextConfig; import org.apache.druid.server.SegmentManager; @@ -119,7 +120,7 @@ protected List getModules() binder.bind(ServerTypeConfig.class).toInstance(new ServerTypeConfig(ServerType.HISTORICAL)); binder.bind(JettyServerInitializer.class).to(QueryJettyServerInitializer.class).in(LazySingleton.class); - binder.bind(QueryCountStatsProvider.class).to(QueryResource.class); + binder.bind(QueryCountStatsProvider.class).to(BaseQueryCountResource.class); Jerseys.addResource(binder, QueryResource.class); Jerseys.addResource(binder, SegmentListerResource.class); Jerseys.addResource(binder, HistoricalResource.class); diff --git a/services/src/main/java/org/apache/druid/guice/QueryablePeonModule.java b/services/src/main/java/org/apache/druid/guice/QueryablePeonModule.java index 0d9c316b3ecf..64c1ea39bf4a 100644 --- a/services/src/main/java/org/apache/druid/guice/QueryablePeonModule.java +++ b/services/src/main/java/org/apache/druid/guice/QueryablePeonModule.java @@ -21,6 +21,7 @@ import com.google.inject.Binder; import org.apache.druid.initialization.DruidModule; +import org.apache.druid.server.BaseQueryCountResource; import org.apache.druid.server.QueryResource; import org.apache.druid.server.metrics.QueryCountStatsProvider; @@ -29,7 +30,7 @@ public class QueryablePeonModule implements DruidModule @Override public void configure(Binder binder) { - binder.bind(QueryCountStatsProvider.class).to(QueryResource.class); + binder.bind(QueryCountStatsProvider.class).to(BaseQueryCountResource.class); Jerseys.addResource(binder, QueryResource.class); LifecycleModule.register(binder, QueryResource.class); } diff --git a/sql/src/main/java/org/apache/druid/sql/http/SqlResource.java b/sql/src/main/java/org/apache/druid/sql/http/SqlResource.java index d957e7155b5e..5c716b413753 100644 --- a/sql/src/main/java/org/apache/druid/sql/http/SqlResource.java +++ b/sql/src/main/java/org/apache/druid/sql/http/SqlResource.java @@ -27,8 +27,8 @@ import org.apache.druid.guice.annotations.Self; import org.apache.druid.java.util.common.StringUtils; import org.apache.druid.java.util.common.logger.Logger; +import org.apache.druid.server.BaseQueryCountResource; import org.apache.druid.server.DruidNode; -import org.apache.druid.server.QueryResource; import org.apache.druid.server.QueryResponse; import org.apache.druid.server.QueryResultPusher; import org.apache.druid.server.ResponseContextConfig; @@ -71,7 +71,6 @@ public class SqlResource public static final String SQL_HEADER_RESPONSE_HEADER = "X-Druid-SQL-Header-Included"; public static final String SQL_HEADER_VALUE = "yes"; private static final Logger log = new Logger(SqlResource.class); - public static final SqlResourceQueryMetricCounter QUERY_METRIC_COUNTER = new SqlResourceQueryMetricCounter(); private final ObjectMapper jsonMapper; private final AuthorizerMapper authorizerMapper; @@ -81,6 +80,8 @@ public class SqlResource private final ResponseContextConfig responseContextConfig; private final DruidNode selfNode; + final BaseQueryCountResource counter; + @Inject protected SqlResource( final ObjectMapper jsonMapper, @@ -89,7 +90,8 @@ protected SqlResource( final SqlLifecycleManager sqlLifecycleManager, final ServerConfig serverConfig, ResponseContextConfig responseContextConfig, - @Self DruidNode selfNode + @Self DruidNode selfNode, + final BaseQueryCountResource counter ) { this.jsonMapper = Preconditions.checkNotNull(jsonMapper, "jsonMapper"); @@ -99,6 +101,7 @@ protected SqlResource( this.serverConfig = Preconditions.checkNotNull(serverConfig, "serverConfig"); this.responseContextConfig = responseContextConfig; this.selfNode = selfNode; + this.counter = counter; } @POST @@ -152,33 +155,6 @@ public Response cancelQuery( } } - /** - * The SqlResource only generates metrics and doesn't keep track of aggregate counts of successful/failed/interrupted - * queries, so this implementation is effectively just a noop. - */ - private static class SqlResourceQueryMetricCounter implements QueryResource.QueryMetricCounter - { - @Override - public void incrementSuccess() - { - } - - @Override - public void incrementFailed() - { - } - - @Override - public void incrementInterrupted() - { - } - - @Override - public void incrementTimedOut() - { - } - } - private SqlResourceQueryResultPusher makePusher(HttpServletRequest req, HttpStatement stmt, SqlQuery sqlQuery) { final String sqlQueryId = stmt.sqlQueryId(); @@ -211,7 +187,7 @@ public SqlResourceQueryResultPusher( jsonMapper, responseContextConfig, selfNode, - SqlResource.QUERY_METRIC_COUNTER, + SqlResource.this.counter, sqlQueryId, MediaType.APPLICATION_JSON_TYPE, headers diff --git a/sql/src/test/java/org/apache/druid/sql/http/SqlResourceTest.java b/sql/src/test/java/org/apache/druid/sql/http/SqlResourceTest.java index b906fe731781..ae03ed37d755 100644 --- a/sql/src/test/java/org/apache/druid/sql/http/SqlResourceTest.java +++ b/sql/src/test/java/org/apache/druid/sql/http/SqlResourceTest.java @@ -64,6 +64,7 @@ import org.apache.druid.query.ResourceLimitExceededException; import org.apache.druid.query.context.ResponseContext; import org.apache.druid.query.groupby.GroupByQueryConfig; +import org.apache.druid.server.BaseQueryCountResource; import org.apache.druid.server.DruidNode; import org.apache.druid.server.QueryResource; import org.apache.druid.server.QueryResponse; @@ -182,6 +183,7 @@ public class SqlResourceTest extends CalciteTestBase private NativeSqlEngine engine; private SqlStatementFactory sqlStatementFactory; private StubServiceEmitter stubServiceEmitter; + private BaseQueryCountResource baseQueryResource; private CountDownLatch lifecycleAddLatch; private final SettableSupplier> validateAndAuthorizeLatchSupplier = new SettableSupplier<>(); @@ -248,6 +250,7 @@ public void setUp() throws Exception req = request(); testRequestLogger = new TestRequestLogger(); + baseQueryResource = new BaseQueryCountResource(); final PlannerFactory plannerFactory = new PlannerFactory( rootSchema, @@ -330,7 +333,8 @@ public PreparedStatement preparedStatement(SqlQueryPlus sqlRequest) lifecycleManager, new ServerConfig(), TEST_RESPONSE_CONTEXT_CONFIG, - DUMMY_DRUID_NODE + DUMMY_DRUID_NODE, + baseQueryResource ); } @@ -1547,7 +1551,8 @@ public ErrorResponseTransformStrategy getErrorResponseTransformStrategy() } }, TEST_RESPONSE_CONTEXT_CONFIG, - DUMMY_DRUID_NODE + DUMMY_DRUID_NODE, + baseQueryResource ); String errorMessage = "This will be supported in Druid 9999"; @@ -1701,7 +1706,9 @@ public void testTooManyRequestsAfterTotalLaning() throws Exception } } Assert.assertEquals(2, success); + Assert.assertEquals(2, resource.counter.getSuccessfulQueryCount()); Assert.assertEquals(1, limited); + Assert.assertEquals(1, resource.counter.getFailedQueryCount()); Assert.assertEquals(3, testRequestLogger.getSqlQueryLogs().size()); Assert.assertTrue(lifecycleManager.getAll(sqlQueryId).isEmpty()); } @@ -1737,6 +1744,7 @@ public void testQueryTimeoutException() throws Exception "" ); Assert.assertTrue(lifecycleManager.getAll(sqlQueryId).isEmpty()); + Assert.assertEquals(1, resource.counter.getTimedOutQueryCount()); } @Test From 03d0fd323361c57d59b0bdb4745592b4cad5fd7a Mon Sep 17 00:00:00 2001 From: Jesse Tuglu Date: Tue, 5 Nov 2024 17:36:59 -0800 Subject: [PATCH 2/2] add licenses --- .../druid/server/BaseQueryCountResource.java | 19 +++++++++++++++++++ .../druid/server/BrokerQueryResource.java | 1 - .../druid/server/QueryMetricCounter.java | 19 +++++++++++++++++++ .../druid/server/QueryResultPusher.java | 2 +- 4 files changed, 39 insertions(+), 2 deletions(-) diff --git a/server/src/main/java/org/apache/druid/server/BaseQueryCountResource.java b/server/src/main/java/org/apache/druid/server/BaseQueryCountResource.java index bde2be41c7fb..6228318e70c9 100644 --- a/server/src/main/java/org/apache/druid/server/BaseQueryCountResource.java +++ b/server/src/main/java/org/apache/druid/server/BaseQueryCountResource.java @@ -1,3 +1,22 @@ +/* + * 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; diff --git a/server/src/main/java/org/apache/druid/server/BrokerQueryResource.java b/server/src/main/java/org/apache/druid/server/BrokerQueryResource.java index 0250acb765e6..9ffadb91b1b0 100644 --- a/server/src/main/java/org/apache/druid/server/BrokerQueryResource.java +++ b/server/src/main/java/org/apache/druid/server/BrokerQueryResource.java @@ -77,7 +77,6 @@ public BrokerQueryResource( responseContextConfig, selfNode, counter - ); this.brokerServerView = brokerServerView; } diff --git a/server/src/main/java/org/apache/druid/server/QueryMetricCounter.java b/server/src/main/java/org/apache/druid/server/QueryMetricCounter.java index 84d2881f790a..c62fe18bf6c1 100644 --- a/server/src/main/java/org/apache/druid/server/QueryMetricCounter.java +++ b/server/src/main/java/org/apache/druid/server/QueryMetricCounter.java @@ -1,3 +1,22 @@ +/* + * 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 diff --git a/server/src/main/java/org/apache/druid/server/QueryResultPusher.java b/server/src/main/java/org/apache/druid/server/QueryResultPusher.java index bce9ed565e52..ebe5f819a7a9 100644 --- a/server/src/main/java/org/apache/druid/server/QueryResultPusher.java +++ b/server/src/main/java/org/apache/druid/server/QueryResultPusher.java @@ -233,6 +233,7 @@ private Response handleDruidException(ResultsWriter resultsWriter, DruidExceptio { if (resultsWriter != null) { resultsWriter.recordFailure(e); + counter.incrementFailed(); if (accumulator != null && accumulator.isInitialized()) { // We already started sending a response when we got the error message. In this case we just give up @@ -242,7 +243,6 @@ private Response handleDruidException(ResultsWriter resultsWriter, DruidExceptio // the future. trailerFields.put(QueryResource.ERROR_MESSAGE_TRAILER_HEADER, e.getMessage()); trailerFields.put(QueryResource.RESPONSE_COMPLETE_TRAILER_HEADER, "false"); - counter.incrementFailed(); return null; } }