You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@pinot.apache.org by "vrajat (via GitHub)" <gi...@apache.org> on 2025/06/18 09:56:19 UTC
[PR] ThreadResourceUsageAccountant uses QueryCancelManager to cancel queries. [pinot]
vrajat opened a new pull request, #16142:
URL: https://github.com/apache/pinot/pull/16142
Instructions:
1. The PR has to be tagged with at least one of the following labels (*):
1. `feature`
2. `bugfix`
3. `performance`
4. `ui`
5. `backward-incompat`
6. `release-notes` (**)
2. Remove these instructions before publishing the PR.
(*) Other labels to consider:
- `testing`
- `dependencies`
- `docker`
- `kubernetes`
- `observability`
- `security`
- `code-style`
- `extension-point`
- `refactor`
- `cleanup`
(**) Use `release-notes` label for scenarios like:
- New configuration options
- Deprecation of configurations
- Signature changes to public methods/interfaces
- New plugins added or old plugins removed
--
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: commits-unsubscribe@pinot.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org
Re: [PR] Use QueryCancelCallback to cancel queries from ThreadResourceUsageAccountant [pinot]
Posted by "Copilot (via GitHub)" <gi...@apache.org>.
Copilot commented on code in PR #16142:
URL: https://github.com/apache/pinot/pull/16142#discussion_r2197179730
##########
pinot-core/src/main/java/org/apache/pinot/core/accounting/PerQueryCPUMemAccountantFactory.java:
##########
@@ -814,14 +837,21 @@ void killAllQueries() {
if (config.isOomKillQueryEnabled()) {
int killedCount = 0;
+ Set<String> terminatedQueryIds = new HashSet<>();
for (Map.Entry<Thread, CPUMemThreadLevelAccountingObjects.ThreadEntry> entry : _threadEntriesMap.entrySet()) {
CPUMemThreadLevelAccountingObjects.ThreadEntry threadEntry = entry.getValue();
CPUMemThreadLevelAccountingObjects.TaskEntry taskEntry = threadEntry.getCurrentThreadTaskStatus();
- if (taskEntry != null && taskEntry.isAnchorThread()) {
- threadEntry._errorStatus
- .set(new RuntimeException(String.format("Query killed due to %s out of memory!", _instanceType)));
- taskEntry.getAnchorThread().interrupt();
- killedCount += 1;
+ if (taskEntry != null && !terminatedQueryIds.contains(taskEntry.getQueryId())) {
+ if (_queryCancelCallbacks.containsKey(taskEntry.getQueryId())) {
Review Comment:
When cancelling via callback, `killedCount` is not incremented, so the query-killed metric may undercount. Consider incrementing `killedCount` in the callback branch as well.
--
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: commits-unsubscribe@pinot.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org
Re: [PR] Use QueryCancelCallback to cancel queries from ThreadResourceUsageAccountant [pinot]
Posted by "gortiz (via GitHub)" <gi...@apache.org>.
gortiz commented on PR #16142:
URL: https://github.com/apache/pinot/pull/16142#issuecomment-3061494402
There is a conflict that has to be resolved, mark it as ready for review and add that nullable notation. Apart from these minor tasks, I think it is ready to merge
--
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: commits-unsubscribe@pinot.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org
Re: [PR] Use QueryCancelCallback to cancel queries from ThreadResourceUsageAccountant [pinot]
Posted by "gortiz (via GitHub)" <gi...@apache.org>.
gortiz commented on code in PR #16142:
URL: https://github.com/apache/pinot/pull/16142#discussion_r2200204526
##########
pinot-core/src/main/java/org/apache/pinot/core/accounting/PerQueryCPUMemAccountantFactory.java:
##########
@@ -254,6 +268,15 @@ public boolean throttleQuerySubmission() {
return getWatcherTask().getHeapUsageBytes() > getWatcherTask().getQueryMonitorConfig().getAlarmingLevel();
}
+ @Override
+ public void registerQueryCancelCallback(String queryId, MseCancelCallback callback) {
+ _queryCancelCallbacks.put(queryId, callback);
+ }
+
+ public MseCancelCallback getQueryCancelCallback(String queryId) {
Review Comment:
This should be marked as nullable
--
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: commits-unsubscribe@pinot.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org
Re: [PR] Use QueryCancelCallback to cancel queries from ThreadResourceUsageAccountant [pinot]
Posted by "gortiz (via GitHub)" <gi...@apache.org>.
gortiz commented on code in PR #16142:
URL: https://github.com/apache/pinot/pull/16142#discussion_r2197712214
##########
pinot-core/src/main/java/org/apache/pinot/core/accounting/PerQueryCPUMemAccountantFactory.java:
##########
@@ -106,6 +107,8 @@ public static class PerQueryCPUMemResourceUsageAccountant implements ThreadResou
protected final HashMap<String, Long> _finishedTaskCPUStatsAggregator = new HashMap<>();
protected final HashMap<String, Long> _finishedTaskMemStatsAggregator = new HashMap<>();
+ protected final ConcurrentHashMap<String, QueryCancelCallback> _queryCancelCallbacks = new ConcurrentHashMap<>();
Review Comment:
Are we 100% we are going to clean up this even in case of error? Otherwise it would be better to use a Guava Cache, which will be cleaned after a while by other threads when new values are added
##########
pinot-spi/src/main/java/org/apache/pinot/spi/accounting/QueryCancelCallback.java:
##########
@@ -0,0 +1,26 @@
+/**
+ * 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.pinot.spi.accounting;
+
+public interface QueryCancelCallback {
+ /**
+ * Cancels the query with the given queryId.
+ */
+ void cancelQuery();
Review Comment:
I like the idea of using an interface here, but I think we need to specify whether this method is blocking (and therefore once the call finishes the query is guarantee to be killed) or could be async
##########
pinot-core/src/main/java/org/apache/pinot/core/accounting/PerQueryCPUMemAccountantFactory.java:
##########
@@ -814,14 +837,21 @@ void killAllQueries() {
if (config.isOomKillQueryEnabled()) {
int killedCount = 0;
+ Set<String> terminatedQueryIds = new HashSet<>();
for (Map.Entry<Thread, CPUMemThreadLevelAccountingObjects.ThreadEntry> entry : _threadEntriesMap.entrySet()) {
CPUMemThreadLevelAccountingObjects.ThreadEntry threadEntry = entry.getValue();
CPUMemThreadLevelAccountingObjects.TaskEntry taskEntry = threadEntry.getCurrentThreadTaskStatus();
- if (taskEntry != null && taskEntry.isAnchorThread()) {
- threadEntry._errorStatus
- .set(new RuntimeException(String.format("Query killed due to %s out of memory!", _instanceType)));
- taskEntry.getAnchorThread().interrupt();
- killedCount += 1;
+ if (taskEntry != null && !terminatedQueryIds.contains(taskEntry.getQueryId())) {
+ if (_queryCancelCallbacks.containsKey(taskEntry.getQueryId())) {
+ _queryCancelCallbacks.get(taskEntry.getQueryId()).cancelQuery();
+ terminatedQueryIds.add(taskEntry.getQueryId());
+ } else if (taskEntry.isAnchorThread()) {
Review Comment:
Do we still need this older option?
##########
pinot-core/src/main/java/org/apache/pinot/core/accounting/PerQueryCPUMemAccountantFactory.java:
##########
@@ -879,7 +909,11 @@ private void killMostExpensiveQuery() {
.set(new RuntimeException(String.format(
" Query %s got killed because using %d bytes of memory on %s: %s, exceeding the quota",
maxUsageTuple._queryId, maxUsageTuple.getAllocatedBytes(), _instanceType, _instanceId)));
- interruptRunnerThread(maxUsageTuple.getAnchorThread());
+ if (_queryCancelCallbacks.containsKey(maxUsageTuple.getQueryId())) {
+ _queryCancelCallbacks.get(maxUsageTuple.getQueryId()).cancelQuery();
Review Comment:
Same race condition here
##########
pinot-core/src/main/java/org/apache/pinot/core/accounting/PerQueryCPUMemAccountantFactory.java:
##########
@@ -905,7 +939,11 @@ private void killCPUTimeExceedQueries() {
String.format("Query %s got killed on %s: %s because using %d "
+ "CPU time exceeding limit of %d ns CPU time", value._queryId, _instanceType, _instanceId,
value.getCpuTimeNs(), config.getCpuTimeBasedKillingThresholdNS())));
- interruptRunnerThread(value.getAnchorThread());
+ if (_queryCancelCallbacks.containsKey(value.getQueryId())) {
+ _queryCancelCallbacks.get(value.getQueryId()).cancelQuery();
Review Comment:
And here
##########
pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/QueryRunner.java:
##########
@@ -267,6 +269,11 @@ public void shutDown() {
/// If any error happened during the asynchronous execution, an error block will be sent to all receiver mailboxes.
public CompletableFuture<Void> processQuery(WorkerMetadata workerMetadata, StagePlan stagePlan,
Map<String, String> requestMetadata, @Nullable ThreadExecutionContext parentContext) {
+ String requestIdStr = Long.toString(QueryThreadContext.getRequestId());
+ long requestId = QueryThreadContext.getRequestId();
+
+ _resourceUsageAccountant.registerQueryCancelCallback(requestIdStr, () -> _opChainScheduler.cancel(requestId));
Review Comment:
Not super important, but we need to allocate a new instance per query here.
Instead we could change the signature of the QueryCancelCallback to receive the request id as argument and then we could create a single instance callback and reuse it.
##########
pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/queries/QueryAccountantCancelTest.java:
##########
@@ -0,0 +1,117 @@
+/**
+ * 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.pinot.query.runtime.queries;
+
+import java.util.HashMap;
+import java.util.Map;
+import org.apache.pinot.common.response.broker.ResultTable;
+import org.apache.pinot.core.accounting.PerQueryCPUMemAccountantFactory;
+import org.apache.pinot.query.QueryEnvironmentTestBase;
+import org.apache.pinot.query.QueryServerEnclosure;
+import org.apache.pinot.query.mailbox.MailboxService;
+import org.apache.pinot.query.routing.QueryServerInstance;
+import org.apache.pinot.query.testutils.MockInstanceDataManagerFactory;
+import org.apache.pinot.query.testutils.QueryTestUtils;
+import org.apache.pinot.spi.accounting.ThreadResourceUsageProvider;
+import org.apache.pinot.spi.config.instance.InstanceType;
+import org.apache.pinot.spi.env.PinotConfiguration;
+import org.apache.pinot.spi.trace.Tracing;
+import org.apache.pinot.spi.utils.CommonConstants;
+import org.mockito.MockedStatic;
+import org.mockito.Mockito;
+import org.testng.Assert;
+import org.testng.annotations.AfterClass;
+import org.testng.annotations.BeforeClass;
+import org.testng.annotations.Test;
+
+import static org.testng.Assert.assertNotNull;
+
+
+public class QueryAccountantCancelTest extends QueryRunnerTestBase {
+ private PerQueryCPUMemAccountantFactory.PerQueryCPUMemResourceUsageAccountant _accountant;
+
+ @BeforeClass
+ public void setUp()
+ throws Exception {
+ MockInstanceDataManagerFactory factory1 = new MockInstanceDataManagerFactory("server1");
+ factory1.registerTable(QueryRunnerTest.SCHEMA_BUILDER.setSchemaName("a").build(), "a_REALTIME");
+ factory1.addSegment("a_REALTIME", QueryRunnerTest.buildRows("a_REALTIME"));
+
+ MockInstanceDataManagerFactory factory2 = new MockInstanceDataManagerFactory("server2");
+ factory2.registerTable(QueryRunnerTest.SCHEMA_BUILDER.setSchemaName("b").build(), "b_REALTIME");
+ factory2.addSegment("b_REALTIME", QueryRunnerTest.buildRows("b_REALTIME"));
+
+ _reducerHostname = "localhost";
+ _reducerPort = QueryTestUtils.getAvailablePort();
+ Map<String, Object> reducerConfig = new HashMap<>();
+ reducerConfig.put(CommonConstants.MultiStageQueryRunner.KEY_OF_QUERY_RUNNER_HOSTNAME, _reducerHostname);
+ reducerConfig.put(CommonConstants.MultiStageQueryRunner.KEY_OF_QUERY_RUNNER_PORT, _reducerPort);
+ _mailboxService = new MailboxService(_reducerHostname, _reducerPort, new PinotConfiguration(reducerConfig));
+ _mailboxService.start();
+
+ HashMap<String, Object> configs = getAccountingConfig();
+
+ ThreadResourceUsageProvider.setThreadMemoryMeasurementEnabled(true);
+ _accountant =
+ new PerQueryCPUMemAccountantFactory.PerQueryCPUMemResourceUsageAccountant(new PinotConfiguration(configs),
+ "testWithPerQueryAccountantFactory", InstanceType.SERVER);
+
+ QueryServerEnclosure server1 = new QueryServerEnclosure(factory1, Map.of(), _accountant);
+ server1.start();
+ // Start server1 to ensure the next server will have a different port.
+ QueryServerEnclosure server2 = new QueryServerEnclosure(factory2, Map.of(), _accountant);
+ server2.start();
+ // this doesn't test the QueryServer functionality so the server port can be the same as the mailbox port.
+ // this is only use for test identifier purpose.
+ int port1 = server1.getPort();
+ int port2 = server2.getPort();
+ _servers.put(new QueryServerInstance("Server_localhost_" + port1, "localhost", port1, port1), server1);
+ _servers.put(new QueryServerInstance("Server_localhost_" + port2, "localhost", port2, port2), server2);
+
+ _queryEnvironment = QueryEnvironmentTestBase.getQueryEnvironment(_reducerPort, server1.getPort(), server2.getPort(),
+ factory1.getRegisteredSchemaMap(), factory1.buildTableSegmentNameMap(), factory2.buildTableSegmentNameMap(),
+ null);
+ }
+
+ @AfterClass
+ public void tearDown() {
+ for (QueryServerEnclosure server : _servers.values()) {
+ server.shutDown();
+ }
+ _mailboxService.shutdown();
+ }
+
+ @Test
+ void testWithPerQueryAccountantFactory() {
+ try (MockedStatic<Tracing> tracing = Mockito.mockStatic(Tracing.class, Mockito.CALLS_REAL_METHODS)) {
+ tracing.when(Tracing::getThreadAccountant).thenReturn(_accountant);
+
+ ResultTable resultTable = queryRunner("SELECT * FROM a LIMIT 2", false).getResultTable();
+ Assert.assertEquals(resultTable.getRows().size(), 2);
+ // TODO: How to programmatically get the request id ?
+ assertNotNull(_accountant.getQueryCancelCallback("0"));
Review Comment:
We are using corretlations id here, right? in that case you can set the CID by using query options. See https://docs.pinot.apache.org/users/user-guide-query/query-correlation-id
##########
pinot-core/src/main/java/org/apache/pinot/core/accounting/PerQueryCPUMemAccountantFactory.java:
##########
@@ -491,6 +504,16 @@ public Map<String, AggregatedStats> aggregate(boolean isTriggered) {
public void postAggregation(Map<String, AggregatedStats> aggregatedUsagePerActiveQuery) {
}
+ public boolean cancelQuery(String queryId) {
+ if (_queryCancelCallbacks.containsKey(queryId)) {
+ QueryCancelCallback callback = _queryCancelCallbacks.get(queryId);
+ callback.cancelQuery();
+ return true;
Review Comment:
Could we remove the entry here? or we expect to try to cancel them again in future?
##########
pinot-core/src/main/java/org/apache/pinot/core/accounting/PerQueryCPUMemAccountantFactory.java:
##########
@@ -814,14 +837,21 @@ void killAllQueries() {
if (config.isOomKillQueryEnabled()) {
int killedCount = 0;
+ Set<String> terminatedQueryIds = new HashSet<>();
for (Map.Entry<Thread, CPUMemThreadLevelAccountingObjects.ThreadEntry> entry : _threadEntriesMap.entrySet()) {
CPUMemThreadLevelAccountingObjects.ThreadEntry threadEntry = entry.getValue();
CPUMemThreadLevelAccountingObjects.TaskEntry taskEntry = threadEntry.getCurrentThreadTaskStatus();
- if (taskEntry != null && taskEntry.isAnchorThread()) {
- threadEntry._errorStatus
- .set(new RuntimeException(String.format("Query killed due to %s out of memory!", _instanceType)));
- taskEntry.getAnchorThread().interrupt();
- killedCount += 1;
+ if (taskEntry != null && !terminatedQueryIds.contains(taskEntry.getQueryId())) {
+ if (_queryCancelCallbacks.containsKey(taskEntry.getQueryId())) {
+ _queryCancelCallbacks.get(taskEntry.getQueryId()).cancelQuery();
Review Comment:
I think there is a race condition here. It seems strange copilot doesn't detect it.
The entry could be removed after this if but before we can get the value.
The solution is simple, you can always execute `get` and only call cancel if the returned value is not null.
--
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: commits-unsubscribe@pinot.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org
Re: [PR] Use QueryCancelCallback to cancel queries from ThreadResourceUsageAccountant [pinot]
Posted by "vrajat (via GitHub)" <gi...@apache.org>.
vrajat commented on code in PR #16142:
URL: https://github.com/apache/pinot/pull/16142#discussion_r2200698230
##########
pinot-core/src/main/java/org/apache/pinot/core/accounting/PerQueryCPUMemAccountantFactory.java:
##########
@@ -254,6 +268,15 @@ public boolean throttleQuerySubmission() {
return getWatcherTask().getHeapUsageBytes() > getWatcherTask().getQueryMonitorConfig().getAlarmingLevel();
}
+ @Override
+ public void registerQueryCancelCallback(String queryId, MseCancelCallback callback) {
+ _queryCancelCallbacks.put(queryId, callback);
+ }
+
+ public MseCancelCallback getQueryCancelCallback(String queryId) {
Review Comment:
Done
--
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: commits-unsubscribe@pinot.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org
Re: [PR] Use QueryCancelCallback to cancel queries from ThreadResourceUsageAccountant [pinot]
Posted by "vrajat (via GitHub)" <gi...@apache.org>.
vrajat commented on PR #16142:
URL: https://github.com/apache/pinot/pull/16142#issuecomment-3062009914
> There is a conflict that has to be resolved, mark it as ready for review and add that nullable notation. Apart from these minor tasks, I think it is ready to merge
Will fix. I'll get https://github.com/apache/pinot/pull/16299 merged first and then fix conflicts and get the branch ready for merge.
--
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: commits-unsubscribe@pinot.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org
Re: [PR] Use QueryCancelCallback to cancel queries from ThreadResourceUsageAccountant [pinot]
Posted by "gortiz (via GitHub)" <gi...@apache.org>.
gortiz merged PR #16142:
URL: https://github.com/apache/pinot/pull/16142
--
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: commits-unsubscribe@pinot.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org
Re: [PR] Use QueryCancelCallback to cancel queries from ThreadResourceUsageAccountant [pinot]
Posted by "vrajat (via GitHub)" <gi...@apache.org>.
vrajat commented on code in PR #16142:
URL: https://github.com/apache/pinot/pull/16142#discussion_r2199603364
##########
pinot-core/src/main/java/org/apache/pinot/core/accounting/PerQueryCPUMemAccountantFactory.java:
##########
@@ -814,14 +837,21 @@ void killAllQueries() {
if (config.isOomKillQueryEnabled()) {
int killedCount = 0;
+ Set<String> terminatedQueryIds = new HashSet<>();
for (Map.Entry<Thread, CPUMemThreadLevelAccountingObjects.ThreadEntry> entry : _threadEntriesMap.entrySet()) {
CPUMemThreadLevelAccountingObjects.ThreadEntry threadEntry = entry.getValue();
CPUMemThreadLevelAccountingObjects.TaskEntry taskEntry = threadEntry.getCurrentThreadTaskStatus();
- if (taskEntry != null && taskEntry.isAnchorThread()) {
- threadEntry._errorStatus
- .set(new RuntimeException(String.format("Query killed due to %s out of memory!", _instanceType)));
- taskEntry.getAnchorThread().interrupt();
- killedCount += 1;
+ if (taskEntry != null && !terminatedQueryIds.contains(taskEntry.getQueryId())) {
+ if (_queryCancelCallbacks.containsKey(taskEntry.getQueryId())) {
+ _queryCancelCallbacks.get(taskEntry.getQueryId()).cancelQuery();
Review Comment:
This part of the code - `WatcherTask` is single-threaded.
* There are multiple writers to the hash table. Although every writer will create a new entry.
* Get and delete will be by only the `WatcherTask` thread.
So each key is write once. Read/Delete once.
##########
pinot-core/src/main/java/org/apache/pinot/core/accounting/PerQueryCPUMemAccountantFactory.java:
##########
@@ -106,6 +107,8 @@ public static class PerQueryCPUMemResourceUsageAccountant implements ThreadResou
protected final HashMap<String, Long> _finishedTaskCPUStatsAggregator = new HashMap<>();
protected final HashMap<String, Long> _finishedTaskMemStatsAggregator = new HashMap<>();
+ protected final ConcurrentHashMap<String, QueryCancelCallback> _queryCancelCallbacks = new ConcurrentHashMap<>();
Review Comment:
I'll check. At the minimum, I'll use a data structure that is bounded in size.
##########
pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/queries/QueryAccountantCancelTest.java:
##########
@@ -0,0 +1,117 @@
+/**
+ * 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.pinot.query.runtime.queries;
+
+import java.util.HashMap;
+import java.util.Map;
+import org.apache.pinot.common.response.broker.ResultTable;
+import org.apache.pinot.core.accounting.PerQueryCPUMemAccountantFactory;
+import org.apache.pinot.query.QueryEnvironmentTestBase;
+import org.apache.pinot.query.QueryServerEnclosure;
+import org.apache.pinot.query.mailbox.MailboxService;
+import org.apache.pinot.query.routing.QueryServerInstance;
+import org.apache.pinot.query.testutils.MockInstanceDataManagerFactory;
+import org.apache.pinot.query.testutils.QueryTestUtils;
+import org.apache.pinot.spi.accounting.ThreadResourceUsageProvider;
+import org.apache.pinot.spi.config.instance.InstanceType;
+import org.apache.pinot.spi.env.PinotConfiguration;
+import org.apache.pinot.spi.trace.Tracing;
+import org.apache.pinot.spi.utils.CommonConstants;
+import org.mockito.MockedStatic;
+import org.mockito.Mockito;
+import org.testng.Assert;
+import org.testng.annotations.AfterClass;
+import org.testng.annotations.BeforeClass;
+import org.testng.annotations.Test;
+
+import static org.testng.Assert.assertNotNull;
+
+
+public class QueryAccountantCancelTest extends QueryRunnerTestBase {
+ private PerQueryCPUMemAccountantFactory.PerQueryCPUMemResourceUsageAccountant _accountant;
+
+ @BeforeClass
+ public void setUp()
+ throws Exception {
+ MockInstanceDataManagerFactory factory1 = new MockInstanceDataManagerFactory("server1");
+ factory1.registerTable(QueryRunnerTest.SCHEMA_BUILDER.setSchemaName("a").build(), "a_REALTIME");
+ factory1.addSegment("a_REALTIME", QueryRunnerTest.buildRows("a_REALTIME"));
+
+ MockInstanceDataManagerFactory factory2 = new MockInstanceDataManagerFactory("server2");
+ factory2.registerTable(QueryRunnerTest.SCHEMA_BUILDER.setSchemaName("b").build(), "b_REALTIME");
+ factory2.addSegment("b_REALTIME", QueryRunnerTest.buildRows("b_REALTIME"));
+
+ _reducerHostname = "localhost";
+ _reducerPort = QueryTestUtils.getAvailablePort();
+ Map<String, Object> reducerConfig = new HashMap<>();
+ reducerConfig.put(CommonConstants.MultiStageQueryRunner.KEY_OF_QUERY_RUNNER_HOSTNAME, _reducerHostname);
+ reducerConfig.put(CommonConstants.MultiStageQueryRunner.KEY_OF_QUERY_RUNNER_PORT, _reducerPort);
+ _mailboxService = new MailboxService(_reducerHostname, _reducerPort, new PinotConfiguration(reducerConfig));
+ _mailboxService.start();
+
+ HashMap<String, Object> configs = getAccountingConfig();
+
+ ThreadResourceUsageProvider.setThreadMemoryMeasurementEnabled(true);
+ _accountant =
+ new PerQueryCPUMemAccountantFactory.PerQueryCPUMemResourceUsageAccountant(new PinotConfiguration(configs),
+ "testWithPerQueryAccountantFactory", InstanceType.SERVER);
+
+ QueryServerEnclosure server1 = new QueryServerEnclosure(factory1, Map.of(), _accountant);
+ server1.start();
+ // Start server1 to ensure the next server will have a different port.
+ QueryServerEnclosure server2 = new QueryServerEnclosure(factory2, Map.of(), _accountant);
+ server2.start();
+ // this doesn't test the QueryServer functionality so the server port can be the same as the mailbox port.
+ // this is only use for test identifier purpose.
+ int port1 = server1.getPort();
+ int port2 = server2.getPort();
+ _servers.put(new QueryServerInstance("Server_localhost_" + port1, "localhost", port1, port1), server1);
+ _servers.put(new QueryServerInstance("Server_localhost_" + port2, "localhost", port2, port2), server2);
+
+ _queryEnvironment = QueryEnvironmentTestBase.getQueryEnvironment(_reducerPort, server1.getPort(), server2.getPort(),
+ factory1.getRegisteredSchemaMap(), factory1.buildTableSegmentNameMap(), factory2.buildTableSegmentNameMap(),
+ null);
+ }
+
+ @AfterClass
+ public void tearDown() {
+ for (QueryServerEnclosure server : _servers.values()) {
+ server.shutDown();
+ }
+ _mailboxService.shutdown();
+ }
+
+ @Test
+ void testWithPerQueryAccountantFactory() {
+ try (MockedStatic<Tracing> tracing = Mockito.mockStatic(Tracing.class, Mockito.CALLS_REAL_METHODS)) {
+ tracing.when(Tracing::getThreadAccountant).thenReturn(_accountant);
+
+ ResultTable resultTable = queryRunner("SELECT * FROM a LIMIT 2", false).getResultTable();
+ Assert.assertEquals(resultTable.getRows().size(), 2);
+ // TODO: How to programmatically get the request id ?
+ assertNotNull(_accountant.getQueryCancelCallback("0"));
Review Comment:
Not yet. The feature works with correlation id. However the test dont work with correlation id - sigh
##########
pinot-spi/src/main/java/org/apache/pinot/spi/accounting/QueryCancelCallback.java:
##########
@@ -0,0 +1,26 @@
+/**
+ * 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.pinot.spi.accounting;
+
+public interface QueryCancelCallback {
+ /**
+ * Cancels the query with the given queryId.
+ */
+ void cancelQuery();
Review Comment:
It will be async. I'll add javadocs.
##########
pinot-core/src/main/java/org/apache/pinot/core/accounting/PerQueryCPUMemAccountantFactory.java:
##########
@@ -491,6 +504,16 @@ public Map<String, AggregatedStats> aggregate(boolean isTriggered) {
public void postAggregation(Map<String, AggregatedStats> aggregatedUsagePerActiveQuery) {
}
+ public boolean cancelQuery(String queryId) {
+ if (_queryCancelCallbacks.containsKey(queryId)) {
+ QueryCancelCallback callback = _queryCancelCallbacks.get(queryId);
+ callback.cancelQuery();
+ return true;
Review Comment:
I can remove the entry here. However it still has to be removed in the `reapFinishedTasks` for cases where cancel is not called.
##########
pinot-core/src/main/java/org/apache/pinot/core/accounting/PerQueryCPUMemAccountantFactory.java:
##########
@@ -814,14 +837,21 @@ void killAllQueries() {
if (config.isOomKillQueryEnabled()) {
int killedCount = 0;
+ Set<String> terminatedQueryIds = new HashSet<>();
for (Map.Entry<Thread, CPUMemThreadLevelAccountingObjects.ThreadEntry> entry : _threadEntriesMap.entrySet()) {
CPUMemThreadLevelAccountingObjects.ThreadEntry threadEntry = entry.getValue();
CPUMemThreadLevelAccountingObjects.TaskEntry taskEntry = threadEntry.getCurrentThreadTaskStatus();
- if (taskEntry != null && taskEntry.isAnchorThread()) {
- threadEntry._errorStatus
- .set(new RuntimeException(String.format("Query killed due to %s out of memory!", _instanceType)));
- taskEntry.getAnchorThread().interrupt();
- killedCount += 1;
+ if (taskEntry != null && !terminatedQueryIds.contains(taskEntry.getQueryId())) {
+ if (_queryCancelCallbacks.containsKey(taskEntry.getQueryId())) {
+ _queryCancelCallbacks.get(taskEntry.getQueryId()).cancelQuery();
+ terminatedQueryIds.add(taskEntry.getQueryId());
+ } else if (taskEntry.isAnchorThread()) {
Review Comment:
I'll keep it for now.
* I removed it and reverted the change because the diff got too big.
* Previously I had added checks to make sure anchor thread/parent context exists. Fixed NPE etc.
I'll undo these changes esp. for MSE in a separate thread with unit tests. It will be safer that way.
--
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: commits-unsubscribe@pinot.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org
Re: [PR] ThreadResourceUsageAccountant uses QueryCancelManager to cancel queries. [pinot]
Posted by "codecov-commenter (via GitHub)" <gi...@apache.org>.
codecov-commenter commented on PR #16142:
URL: https://github.com/apache/pinot/pull/16142#issuecomment-3052962673
## [Codecov](https://app.codecov.io/gh/apache/pinot/pull/16142?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) Report
Attention: Patch coverage is `41.93548%` with `18 lines` in your changes missing coverage. Please review.
> Project coverage is 56.33%. Comparing base [(`1a476de`)](https://app.codecov.io/gh/apache/pinot/commit/1a476dea72ebf1c15e14ded28ae60c385f71b2c3?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) to head [(`4aff5c4`)](https://app.codecov.io/gh/apache/pinot/commit/4aff5c477127a97ea37a3952b5e8a2f7846674ea?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache).
> Report is 425 commits behind head on master.
| [Files with missing lines](https://app.codecov.io/gh/apache/pinot/pull/16142?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | Patch % | Lines |
|---|---|---|
| [...re/accounting/PerQueryCPUMemAccountantFactory.java](https://app.codecov.io/gh/apache/pinot/pull/16142?src=pr&el=tree&filepath=pinot-core%2Fsrc%2Fmain%2Fjava%2Forg%2Fapache%2Fpinot%2Fcore%2Faccounting%2FPerQueryCPUMemAccountantFactory.java&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache#diff-cGlub3QtY29yZS9zcmMvbWFpbi9qYXZhL29yZy9hcGFjaGUvcGlub3QvY29yZS9hY2NvdW50aW5nL1BlclF1ZXJ5Q1BVTWVtQWNjb3VudGFudEZhY3RvcnkuamF2YQ==) | 14.28% | [17 Missing and 1 partial :warning: ](https://app.codecov.io/gh/apache/pinot/pull/16142?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) |
> :exclamation: There is a different number of reports uploaded between BASE (1a476de) and HEAD (4aff5c4). Click for more details.
>
> <details><summary>HEAD has 68 uploads less than BASE</summary>
>
>| Flag | BASE (1a476de) | HEAD (4aff5c4) |
>|------|------|------|
>|java-21|10|1|
>|unittests|6|1|
>|unittests1|3|1|
>|skip-bytebuffers-true|4|0|
>|temurin|15|1|
>|unittests2|3|0|
>|skip-bytebuffers-false|8|0|
>|java-11|5|0|
>|integration|9|0|
>|integration2|3|0|
>|integration1|3|0|
>|custom-integration1|3|0|
></details>
<details><summary>Additional details and impacted files</summary>
```diff
@@ Coverage Diff @@
## master #16142 +/- ##
============================================
- Coverage 62.90% 56.33% -6.57%
+ Complexity 1386 699 -687
============================================
Files 2867 2313 -554
Lines 163354 127469 -35885
Branches 24952 20317 -4635
============================================
- Hits 102755 71810 -30945
+ Misses 52847 49739 -3108
+ Partials 7752 5920 -1832
```
| [Flag](https://app.codecov.io/gh/apache/pinot/pull/16142/flags?src=pr&el=flags&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | Coverage Δ | |
|---|---|---|
| [custom-integration1](https://app.codecov.io/gh/apache/pinot/pull/16142/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | `?` | |
| [integration](https://app.codecov.io/gh/apache/pinot/pull/16142/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | `?` | |
| [integration1](https://app.codecov.io/gh/apache/pinot/pull/16142/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | `?` | |
| [integration2](https://app.codecov.io/gh/apache/pinot/pull/16142/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | `?` | |
| [java-11](https://app.codecov.io/gh/apache/pinot/pull/16142/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | `?` | |
| [java-21](https://app.codecov.io/gh/apache/pinot/pull/16142/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | `56.33% <41.93%> (-6.49%)` | :arrow_down: |
| [skip-bytebuffers-false](https://app.codecov.io/gh/apache/pinot/pull/16142/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | `?` | |
| [skip-bytebuffers-true](https://app.codecov.io/gh/apache/pinot/pull/16142/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | `?` | |
| [temurin](https://app.codecov.io/gh/apache/pinot/pull/16142/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | `56.33% <41.93%> (-6.57%)` | :arrow_down: |
| [unittests](https://app.codecov.io/gh/apache/pinot/pull/16142/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | `56.33% <41.93%> (-6.57%)` | :arrow_down: |
| [unittests1](https://app.codecov.io/gh/apache/pinot/pull/16142/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | `56.33% <41.93%> (+0.51%)` | :arrow_up: |
| [unittests2](https://app.codecov.io/gh/apache/pinot/pull/16142/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | `?` | |
Flags with carried forward coverage won't be shown. [Click here](https://docs.codecov.io/docs/carryforward-flags?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache#carryforward-flags-in-the-pull-request-comment) to find out more.
</details>
[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/apache/pinot/pull/16142?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache).
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache).
<details><summary> :rocket: New features to boost your workflow: </summary>
- :package: [JS Bundle Analysis](https://docs.codecov.com/docs/javascript-bundle-analysis): Save yourself from yourself by tracking and limiting bundle sizes in JS merges.
</details>
--
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: commits-unsubscribe@pinot.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org
Re: [PR] Use QueryCancelCallback to cancel queries from ThreadResourceUsageAccountant [pinot]
Posted by "gortiz (via GitHub)" <gi...@apache.org>.
gortiz commented on code in PR #16142:
URL: https://github.com/apache/pinot/pull/16142#discussion_r2200891308
##########
pinot-core/src/main/java/org/apache/pinot/core/accounting/PerQueryCPUMemAccountantFactory.java:
##########
@@ -106,6 +107,8 @@ public static class PerQueryCPUMemResourceUsageAccountant implements ThreadResou
protected final HashMap<String, Long> _finishedTaskCPUStatsAggregator = new HashMap<>();
protected final HashMap<String, Long> _finishedTaskMemStatsAggregator = new HashMap<>();
+ protected final ConcurrentHashMap<String, QueryCancelCallback> _queryCancelCallbacks = new ConcurrentHashMap<>();
Review Comment:
Well, the good thing is that each entry is very cheap. I had issues in the past with the map we keep to return stats because I was keeping a reference to the mse operators, which can be huge. Yours literally only references the OpChainScheduler and the query id.
--
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: commits-unsubscribe@pinot.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org