yashmayya commented on code in PR #20104:
URL: https://github.com/apache/pinot/pull/20104#discussion_r4427095141
##########
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:
Done in 997ad70. The Javadoc now says which stages each path gets, and that
subclasses add entries 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:
Done in 997ad70: the test now calls `handler.shutDown()` in a `finally`.
##########
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:
Done in 997ad70: the test now checks for `EXECUTION_TIMEOUT` before it reads
`throttledPlan`.
--
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]
