You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@pinot.apache.org by go...@apache.org on 2025/08/06 11:02:23 UTC
(pinot) branch master updated: Add tests for OOM Termination of MSE queries. (#16514)
This is an automated email from the ASF dual-hosted git repository.
gortiz pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/pinot.git
The following commit(s) were added to refs/heads/master by this push:
new e50947a53f9 Add tests for OOM Termination of MSE queries. (#16514)
e50947a53f9 is described below
commit e50947a53f9eec910d3f7a8aa80b082c5f6367fd
Author: Rajat Venkatesh <16...@users.noreply.github.com>
AuthorDate: Wed Aug 6 16:32:18 2025 +0530
Add tests for OOM Termination of MSE queries. (#16514)
---
...flineClusterMemBasedServerQueryKillingTest.java | 138 +++++++++++++++++----
1 file changed, 113 insertions(+), 25 deletions(-)
diff --git a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/OfflineClusterMemBasedServerQueryKillingTest.java b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/OfflineClusterMemBasedServerQueryKillingTest.java
index b39cb237504..efa34e69eb5 100644
--- a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/OfflineClusterMemBasedServerQueryKillingTest.java
+++ b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/OfflineClusterMemBasedServerQueryKillingTest.java
@@ -37,7 +37,9 @@ import org.apache.commons.lang3.StringUtils;
import org.apache.log4j.Level;
import org.apache.log4j.LogManager;
import org.apache.pinot.core.accounting.PerQueryCPUMemAccountantFactory;
+import org.apache.pinot.spi.accounting.QueryResourceTracker;
import org.apache.pinot.spi.accounting.ThreadResourceUsageProvider;
+import org.apache.pinot.spi.config.instance.InstanceType;
import org.apache.pinot.spi.config.table.TableConfig;
import org.apache.pinot.spi.config.table.TableType;
import org.apache.pinot.spi.data.FieldSpec;
@@ -48,6 +50,7 @@ import org.apache.pinot.spi.trace.Tracing;
import org.apache.pinot.spi.utils.CommonConstants;
import org.apache.pinot.spi.utils.builder.TableConfigBuilder;
import org.apache.pinot.util.TestUtils;
+import org.testng.annotations.AfterMethod;
import org.testng.annotations.BeforeClass;
import org.testng.annotations.Test;
@@ -100,6 +103,57 @@ public class OfflineClusterMemBasedServerQueryKillingTest extends BaseClusterInt
return NUM_SERVERS;
}
+ /**
+ * Keeps track of metadata of queries that have been terminated due to OOM.
+ * This is used to verify that the query was killed and the method used to kill the query.
+ */
+ public static class TestAccountantFactory extends PerQueryCPUMemAccountantFactory {
+ @Override
+ public PerQueryCPUMemResourceUsageAccountant init(PinotConfiguration config, String instanceId,
+ InstanceType instanceType) {
+ return new TestAccountant(config, instanceId, instanceType);
+ }
+
+ public static class TestAccountant extends PerQueryCPUMemAccountantFactory.PerQueryCPUMemResourceUsageAccountant {
+ private QueryResourceTracker _queryResourceTracker = null;
+ private long _totalHeapMemoryUsage = 0L;
+ private boolean _hasCallback = false;
+
+ public TestAccountant(PinotConfiguration config, String instanceId, InstanceType instanceType) {
+ super(config, instanceId, instanceType);
+ }
+
+ @Override
+ public void logTerminatedQuery(QueryResourceTracker queryResourceTracker, long totalHeapMemoryUsage,
+ boolean hasCallback) {
+ super.logTerminatedQuery(queryResourceTracker, totalHeapMemoryUsage, hasCallback);
+ _queryResourceTracker = queryResourceTracker;
+ _totalHeapMemoryUsage = totalHeapMemoryUsage;
+ _hasCallback = hasCallback;
+ }
+
+ public void reset() {
+ _queryResourceTracker = null;
+ _totalHeapMemoryUsage = 0L;
+ _hasCallback = false;
+ }
+
+ public QueryResourceTracker getQueryResourceTracker() {
+ return _queryResourceTracker;
+ }
+
+ public long getTotalHeapMemoryUsage() {
+ return _totalHeapMemoryUsage;
+ }
+
+ public boolean hasCallback() {
+ return _hasCallback;
+ }
+ }
+ }
+
+ private TestAccountantFactory.TestAccountant _testAccountant = null;
+
@BeforeClass
public void setUp()
throws Exception {
@@ -114,9 +168,9 @@ public class OfflineClusterMemBasedServerQueryKillingTest extends BaseClusterInt
startZk();
startController();
startServers();
- while (!Tracing.isAccountantRegistered()) {
- Thread.sleep(100L);
- }
+ TestUtils.waitForCondition(aVoid -> Tracing.isAccountantRegistered(), 100L, 60_000L,
+ "Waiting for accountant to be registered");
+ _testAccountant = (TestAccountantFactory.TestAccountant) Tracing.getThreadAccountant();
startBrokers();
@@ -158,6 +212,13 @@ public class OfflineClusterMemBasedServerQueryKillingTest extends BaseClusterInt
startServers(getNumServers());
}
+ @AfterMethod
+ public void resetTestAccountant() {
+ if (_testAccountant != null) {
+ _testAccountant.reset();
+ }
+ }
+
protected void overrideServerConf(PinotConfiguration serverConf) {
serverConf.setProperty(CommonConstants.PINOT_QUERY_SCHEDULER_PREFIX + "."
+ CommonConstants.Accounting.CONFIG_OF_ALARMING_LEVEL_HEAP_USAGE_RATIO, 0.0f);
@@ -165,7 +226,7 @@ public class OfflineClusterMemBasedServerQueryKillingTest extends BaseClusterInt
+ CommonConstants.Accounting.CONFIG_OF_CRITICAL_LEVEL_HEAP_USAGE_RATIO, 0.15f);
serverConf.setProperty(
CommonConstants.PINOT_QUERY_SCHEDULER_PREFIX + "." + CommonConstants.Accounting.CONFIG_OF_FACTORY_NAME,
- "org.apache.pinot.core.accounting.PerQueryCPUMemAccountantFactory");
+ TestAccountantFactory.class.getName());
serverConf.setProperty(CommonConstants.PINOT_QUERY_SCHEDULER_PREFIX + "."
+ CommonConstants.Accounting.CONFIG_OF_ENABLE_THREAD_MEMORY_SAMPLING, true);
serverConf.setProperty(
@@ -185,7 +246,7 @@ public class OfflineClusterMemBasedServerQueryKillingTest extends BaseClusterInt
+ CommonConstants.Accounting.CONFIG_OF_PANIC_LEVEL_HEAP_USAGE_RATIO, 1.1f);
brokerConf.setProperty(CommonConstants.PINOT_QUERY_SCHEDULER_PREFIX + "."
+ CommonConstants.Accounting.CONFIG_OF_FACTORY_NAME,
- "org.apache.pinot.core.accounting.PerQueryCPUMemAccountantFactory");
+ TestAccountantFactory.class.getName());
brokerConf.setProperty(CommonConstants.PINOT_QUERY_SCHEDULER_PREFIX + "."
+ CommonConstants.Accounting.CONFIG_OF_ENABLE_THREAD_MEMORY_SAMPLING, true);
brokerConf.setProperty(CommonConstants.PINOT_QUERY_SCHEDULER_PREFIX + "."
@@ -214,22 +275,30 @@ public class OfflineClusterMemBasedServerQueryKillingTest extends BaseClusterInt
.build();
}
- @Test(dataProvider = "useBothQueryEngines")
- public void testDigestOOM(boolean useMultiStageQueryEngine)
+ @Test
+ public void testDigestOOM()
throws Exception {
- setUseMultiStageQueryEngine(useMultiStageQueryEngine);
- notSupportedInV2();
JsonNode queryResponse = postQuery(OOM_QUERY);
String exceptionsNode = queryResponse.get("exceptions").toString();
assertTrue(exceptionsNode.contains("\"errorCode\":" + QueryErrorCode.QUERY_CANCELLATION.getId()), exceptionsNode);
assertTrue(exceptionsNode.contains("got killed because"), exceptionsNode);
}
- @Test(dataProvider = "useBothQueryEngines")
- public void testMemoryAllocationStats(boolean useMultiStageQueryEngine)
+ @Test
+ public void testDigestOOMMSE()
+ throws Exception {
+ setUseMultiStageQueryEngine(true);
+ JsonNode queryResponse = postQuery(OOM_QUERY);
+ String exceptionsNode = queryResponse.get("exceptions").toString();
+ assertTrue(exceptionsNode.contains("\"errorCode\":" + QueryErrorCode.INTERNAL.getId()), exceptionsNode);
+ assertTrue(exceptionsNode.contains("Received 1 error from stage 1"), exceptionsNode);
+ assertTrue(_testAccountant.hasCallback());
+ assertEquals(queryResponse.get("requestId").asText(), _testAccountant.getQueryResourceTracker().getQueryId());
+ }
+
+ @Test
+ public void testMemoryAllocationStats()
throws Exception {
- setUseMultiStageQueryEngine(useMultiStageQueryEngine);
- notSupportedInV2();
JsonNode queryResponse = postQuery(COUNT_STAR_QUERY);
long offlineThreadMemAllocatedBytes = queryResponse.get("offlineThreadMemAllocatedBytes").asLong();
long offlineResponseSerMemAllocatedBytes = queryResponse.get("offlineResponseSerMemAllocatedBytes").asLong();
@@ -240,11 +309,9 @@ public class OfflineClusterMemBasedServerQueryKillingTest extends BaseClusterInt
assertEquals(offlineThreadMemAllocatedBytes + offlineResponseSerMemAllocatedBytes, offlineTotalMemAllocatedBytes);
}
- @Test(dataProvider = "useBothQueryEngines")
- public void testSelectionOnlyOOM(boolean useMultiStageQueryEngine)
+ @Test
+ public void testSelectionOnlyOOM()
throws Exception {
- setUseMultiStageQueryEngine(useMultiStageQueryEngine);
- notSupportedInV2();
JsonNode queryResponse = postQuery(OOM_QUERY_SELECTION_ONLY);
String exceptionsNode = queryResponse.get("exceptions").toString();
@@ -252,21 +319,42 @@ public class OfflineClusterMemBasedServerQueryKillingTest extends BaseClusterInt
assertTrue(exceptionsNode.contains("got killed because"), exceptionsNode);
}
- @Test(dataProvider = "useBothQueryEngines")
- public void testDigestOOM2(boolean useMultiStageQueryEngine)
+ @Test
+ public void testSelectionOnlyOOMMSE()
+ throws Exception {
+ setUseMultiStageQueryEngine(true);
+ JsonNode queryResponse = postQuery(OOM_QUERY_SELECTION_ONLY);
+
+ String exceptionsNode = queryResponse.get("exceptions").toString();
+ assertTrue(exceptionsNode.contains("\"errorCode\":" + QueryErrorCode.INTERNAL.getId()), exceptionsNode);
+ assertTrue(exceptionsNode.contains("Received 1 error from stage 1"), exceptionsNode);
+ assertTrue(_testAccountant.hasCallback());
+ assertEquals(queryResponse.get("requestId").asText(), _testAccountant.getQueryResourceTracker().getQueryId());
+ }
+
+ @Test
+ public void testDigestOOM2()
throws Exception {
- setUseMultiStageQueryEngine(useMultiStageQueryEngine);
- notSupportedInV2();
JsonNode queryResponse = postQuery(OOM_QUERY_2);
String exceptionsNode = queryResponse.get("exceptions").toString();
assertTrue(exceptionsNode.contains("got killed because"), exceptionsNode);
}
- @Test(dataProvider = "useBothQueryEngines")
- public void testDigestOOMMultipleQueries(boolean useMultiStageQueryEngine)
+ @Test
+ public void testDigestOOM2MSE()
+ throws Exception {
+ setUseMultiStageQueryEngine(true);
+ JsonNode queryResponse = postQuery(OOM_QUERY_2);
+ String exceptionsNode = queryResponse.get("exceptions").toString();
+ assertTrue(exceptionsNode.contains("\"errorCode\":" + QueryErrorCode.INTERNAL.getId()), exceptionsNode);
+ assertTrue(exceptionsNode.contains("Received 1 error from stage 1"), exceptionsNode);
+ assertTrue(_testAccountant.hasCallback());
+ assertEquals(queryResponse.get("requestId").asText(), _testAccountant.getQueryResourceTracker().getQueryId());
+ }
+
+ @Test
+ public void testDigestOOMMultipleQueries()
throws Exception {
- setUseMultiStageQueryEngine(useMultiStageQueryEngine);
- notSupportedInV2();
AtomicReference<JsonNode> queryResponse1 = new AtomicReference<>();
AtomicReference<JsonNode> queryResponse2 = new AtomicReference<>();
AtomicReference<JsonNode> queryResponse3 = new AtomicReference<>();
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org