gortiz commented on code in PR #20104:
URL: https://github.com/apache/pinot/pull/20104#discussion_r4426887038
##########
pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/MultiStageBrokerRequestHandler.java:
##########
@@ -660,6 +660,12 @@ protected void
prepareCompiledQueryForPlanning(QueryEnvironment.CompiledQuery co
RequestContext requestContext, HttpHeaders httpHeaders) {
}
+ /// Extension hook for preparing a stage right before the broker sends it to
its workers, to run the query or to
+ /// explain it on the servers. Subclasses can add entries to the custom
properties of the stage
+ /// ([DispatchablePlanFragment#getCustomProperties()]), which its workers
get in their stage metadata.
+ protected void prepareStageForDispatch(DispatchablePlanFragment stage,
Map<String, String> queryOptions) {
Review Comment:
Two things I would add to the Javadoc. First, the stages the hook sees are
not the same in both paths. When the query runs, it sees every stage, including
stage 0 (the broker reduce). When we explain asking servers, it only sees leaf
stages that have segments, because `AskingServerStageExplainer` returns early
for fragments without segments. Second, the custom properties map also holds
`tableName` and `timeBoundaryInfo.*`, which the servers read for leaf-stage
routing. I think it is worth saying that implementations add keys and must not
replace the existing ones.
##########
pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/MultiStageBrokerRequestHandlerTest.java:
##########
@@ -153,6 +157,66 @@ protected void onQueryCompletion(RequestContext
requestContext, BrokerResponse b
"onQueryCompletion hook must be called with the BrokerResponse from
handleRequest for MSE");
}
+ /// The broker calls
[MultiStageBrokerRequestHandler#prepareStageForDispatch] for each stage before
it sends the stage
+ /// to its workers, both to run the query and to explain it on the servers.
+ @Test
+ public void testPrepareStageForDispatchHook()
+ throws Exception {
+ int port1 = NetUtils.findOpenPort();
+ int port2 = NetUtils.findOpenPort();
+ MockRoutingManagerFactory factory = new MockRoutingManagerFactory(port1,
port2);
+ TABLE_SCHEMAS.forEach((tableName, schema) -> factory.registerTable(schema,
tableName));
+ SERVER1_SEGMENTS.forEach(
+ (tableName, segments) -> segments.forEach(segment ->
factory.registerSegment(port1, tableName, segment)));
+ SERVER2_SEGMENTS.forEach(
+ (tableName, segments) -> segments.forEach(segment ->
factory.registerSegment(port2, tableName, segment)));
+ RoutingManager routingManager = factory.buildRoutingManager(null);
+ WorkerManager workerManager =
+ new WorkerManager("Broker_localhost", "localhost",
NetUtils.findOpenPort(), routingManager);
+ PinotConfiguration config = new PinotConfiguration();
+ config.setProperty(MultiStageQueryRunner.KEY_OF_QUERY_RUNNER_HOSTNAME,
"localhost");
+ config.setProperty(MultiStageQueryRunner.KEY_OF_QUERY_RUNNER_PORT,
Integer.toString(NetUtils.findOpenPort()));
+ BrokerQueryEventListenerFactory.init(config);
+ QueryQuotaManager queryQuotaManager = mock(QueryQuotaManager.class);
+ when(queryQuotaManager.acquire(anyString())).thenReturn(true);
+ when(queryQuotaManager.acquireDatabase(anyString())).thenReturn(true);
+ when(queryQuotaManager.acquireApplication(anyString())).thenReturn(true);
+ AtomicReference<DispatchableSubPlan> throttledPlan = new
AtomicReference<>();
+ try (MockedConstruction<QueryDispatcher> dispatchers =
mockConstruction(QueryDispatcher.class)) {
+ MultiStageBrokerRequestHandler handler = new
MultiStageBrokerRequestHandler(config, "testBrokerId",
Review Comment:
The handler is never shut down. The constructor creates the
`multi-stage-query-compile-executor` pool and a `MailboxService`, so this test
leaks threads. Can we wrap the body in a try/finally that calls
`handler.shutDown()`? The dispatcher is mocked, so it is safe.
##########
pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/MultiStageBrokerRequestHandlerTest.java:
##########
@@ -153,6 +157,66 @@ protected void onQueryCompletion(RequestContext
requestContext, BrokerResponse b
"onQueryCompletion hook must be called with the BrokerResponse from
handleRequest for MSE");
}
+ /// The broker calls
[MultiStageBrokerRequestHandler#prepareStageForDispatch] for each stage before
it sends the stage
+ /// to its workers, both to run the query and to explain it on the servers.
+ @Test
+ public void testPrepareStageForDispatchHook()
+ throws Exception {
+ int port1 = NetUtils.findOpenPort();
+ int port2 = NetUtils.findOpenPort();
+ MockRoutingManagerFactory factory = new MockRoutingManagerFactory(port1,
port2);
+ TABLE_SCHEMAS.forEach((tableName, schema) -> factory.registerTable(schema,
tableName));
+ SERVER1_SEGMENTS.forEach(
+ (tableName, segments) -> segments.forEach(segment ->
factory.registerSegment(port1, tableName, segment)));
+ SERVER2_SEGMENTS.forEach(
+ (tableName, segments) -> segments.forEach(segment ->
factory.registerSegment(port2, tableName, segment)));
+ RoutingManager routingManager = factory.buildRoutingManager(null);
+ WorkerManager workerManager =
+ new WorkerManager("Broker_localhost", "localhost",
NetUtils.findOpenPort(), routingManager);
+ PinotConfiguration config = new PinotConfiguration();
+ config.setProperty(MultiStageQueryRunner.KEY_OF_QUERY_RUNNER_HOSTNAME,
"localhost");
+ config.setProperty(MultiStageQueryRunner.KEY_OF_QUERY_RUNNER_PORT,
Integer.toString(NetUtils.findOpenPort()));
+ BrokerQueryEventListenerFactory.init(config);
+ QueryQuotaManager queryQuotaManager = mock(QueryQuotaManager.class);
+ when(queryQuotaManager.acquire(anyString())).thenReturn(true);
+ when(queryQuotaManager.acquireDatabase(anyString())).thenReturn(true);
+ when(queryQuotaManager.acquireApplication(anyString())).thenReturn(true);
+ AtomicReference<DispatchableSubPlan> throttledPlan = new
AtomicReference<>();
+ try (MockedConstruction<QueryDispatcher> dispatchers =
mockConstruction(QueryDispatcher.class)) {
+ MultiStageBrokerRequestHandler handler = new
MultiStageBrokerRequestHandler(config, "testBrokerId",
+ new BrokerRequestIdGenerator(), routingManager, new
AllowAllAccessControlFactory(), queryQuotaManager,
+ factory.buildTableCache(), mock(MultiStageQueryThrottler.class),
mock(FailureDetector.class),
+ ThreadAccountantUtils.getNoOpAccountant(), null, workerManager,
workerManager) {
+ @Override
+ protected void prepareStageForDispatch(DispatchablePlanFragment stage,
Map<String, String> queryOptions) {
+ stage.getCustomProperties().put("testKey", "stage" +
stage.getPlanFragment().getFragmentId());
+ }
+
+ @Override
+ protected int getEstimatedNumQueryThreads(DispatchableSubPlan
dispatchableSubPlan) {
+ throttledPlan.set(dispatchableSubPlan);
+ return super.getEstimatedNumQueryThreads(dispatchableSubPlan);
+ }
+ };
+ String query = "SELECT col1, COUNT(*) FROM a GROUP BY col1";
+
+ // The mocked throttler stops the query before it runs
+ handler.handleRequest(query);
Review Comment:
This works because the mocked throttler returns false from `tryAcquire`. Can
we assert that the response has `EXECUTION_TIMEOUT`? If the setup ever fails
before throttling, the test fails today with an NPE on `throttledPlan.get()`,
which is confusing.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
