Repository: apache/pinot (2025-07-03T10:20:19Z)

Issue:
Throttle Task Submission in MSE if heap usage is at critical level



Discussion:

Patch:
diff --git a/pinot-core/src/main/java/org/apache/pinot/core/accounting/PerQueryCPUMemAccountantFactory.java b/pinot-core/src/main/java/org/apache/pinot/core/accounting/PerQueryCPUMemAccountantFactory.java
index aa387807d5bc..5f243e03bffa 100644
--- a/pinot-core/src/main/java/org/apache/pinot/core/accounting/PerQueryCPUMemAccountantFactory.java
+++ b/pinot-core/src/main/java/org/apache/pinot/core/accounting/PerQueryCPUMemAccountantFactory.java
@@ -249,6 +249,11 @@ public void sampleUsageMSE() {
       }
     }
 
+    @Override
+    public boolean throttleQuerySubmission() {
+      return getWatcherTask().getHeapUsageBytes() > getWatcherTask().getQueryMonitorConfig().getAlarmingLevel();
+    }
+
     @Override
     public boolean isAnchorThreadInterrupted() {
       ThreadExecutionContext context = _threadLocalEntry.get().getCurrentThreadTaskStatus();
@@ -632,6 +637,10 @@ public QueryMonitorConfig getQueryMonitorConfig() {
         return _queryMonitorConfig.get();
       }
 
+      public long getHeapUsageBytes() {
+        return _usedBytes;
+      }
+
       @Override
       public synchronized void onChange(Set<String> changedConfigs, Map<String, String> clusterConfigs) {
         // Filter configs that have CommonConstants.PREFIX_SCHEDULER_PREFIX
diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/scheduler/BinaryWorkloadScheduler.java b/pinot-core/src/main/java/org/apache/pinot/core/query/scheduler/BinaryWorkloadScheduler.java
index d5c971472014..b4bedb3ff64b 100644
--- a/pinot-core/src/main/java/org/apache/pinot/core/query/scheduler/BinaryWorkloadScheduler.java
+++ b/pinot-core/src/main/java/org/apache/pinot/core/query/scheduler/BinaryWorkloadScheduler.java
@@ -35,6 +35,7 @@
 import org.apache.pinot.core.query.request.ServerQueryRequest;
 import org.apache.pinot.core.query.scheduler.resources.BinaryWorkloadResourceManager;
 import org.apache.pinot.core.query.scheduler.resources.QueryExecutorService;
+import org.apache.pinot.spi.accounting.ThreadResourceUsageAccountant;
 import org.apache.pinot.spi.env.PinotConfiguration;
 import org.apache.pinot.spi.query.QueryThreadContext;
 import org.slf4j.Logger;
@@ -73,8 +74,9 @@ public class BinaryWorkloadScheduler extends QueryScheduler {
   Thread _scheduler;
 
   public BinaryWorkloadScheduler(PinotConfiguration config, QueryExecutor queryExecutor, ServerMetrics metrics,
-      LongAccumulator latestQueryTime) {
-    super(config, queryExecutor, new BinaryWorkloadResourceManager(config), metrics, latestQueryTime);
+      LongAccumulator latestQueryTime, ThreadResourceUsageAccountant resourceUsageAccountant) {
+    super(config, queryExecutor, new BinaryWorkloadResourceManager(config, resourceUsageAccountant), metrics,
+        latestQueryTime);
 
     _secondaryQueryQ = new SecondaryWorkloadQueue(config, _resourceManager);
     _numSecondaryRunners = config.getProperty(MAX_SECONDARY_QUERIES, DEFAULT_MAX_SECONDARY_QUERIES);
diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/scheduler/QuerySchedulerFactory.java b/pinot-core/src/main/java/org/apache/pinot/core/query/scheduler/QuerySchedulerFactory.java
index 69dd8bed3650..e72255987320 100644
--- a/pinot-core/src/main/java/org/apache/pinot/core/query/scheduler/QuerySchedulerFactory.java
+++ b/pinot-core/src/main/java/org/apache/pinot/core/query/scheduler/QuerySchedulerFactory.java
@@ -27,6 +27,7 @@
 import org.apache.pinot.core.query.scheduler.fcfs.BoundedFCFSScheduler;
 import org.apache.pinot.core.query.scheduler.fcfs.FCFSQueryScheduler;
 import org.apache.pinot.core.query.scheduler.tokenbucket.TokenPriorityScheduler;
+import org.apache.pinot.spi.accounting.ThreadResourceUsageAccountant;
 import org.apache.pinot.spi.env.PinotConfiguration;
 import org.apache.pinot.spi.plugin.PluginManager;
 import org.slf4j.Logger;
@@ -58,7 +59,8 @@ private QuerySchedulerFactory() {
    * @return returns an instance of query scheduler
    */
   public static QueryScheduler create(PinotConfiguration schedulerConfig, QueryExecutor queryExecutor,
-      ServerMetrics serverMetrics, LongAccumulator latestQueryTime) {
+      ServerMetrics serverMetrics, LongAccumulator latestQueryTime,
+      ThreadResourceUsageAccountant resourceUsageAccountant) {
     Preconditions.checkNotNull(schedulerConfig);
     Preconditions.checkNotNull(queryExecutor);
 
@@ -66,16 +68,20 @@ public static QueryScheduler create(PinotConfiguration schedulerConfig, QueryExe
     QueryScheduler scheduler;
     switch (schedulerName.toLowerCase()) {
       case FCFS_ALGORITHM:
-        scheduler = new FCFSQueryScheduler(schedulerConfig, queryExecutor, serverMetrics, latestQueryTime);
+        scheduler = new FCFSQueryScheduler(schedulerConfig, queryExecutor, serverMetrics, latestQueryTime,
+            resourceUsageAccountant);
         break;
       case TOKEN_BUCKET_ALGORITHM:
-        scheduler = TokenPriorityScheduler.create(schedulerConfig, queryExecutor, serverMetrics, latestQueryTime);
+        scheduler = TokenPriorityScheduler.create(schedulerConfig, queryExecutor, serverMetrics, latestQueryTime,
+            resourceUsageAccountant);
         break;
       case BOUNDED_FCFS_ALGORITHM:
-        scheduler = BoundedFCFSScheduler.create(schedulerConfig, queryExecutor, serverMetrics, latestQueryTime);
+        scheduler = BoundedFCFSScheduler.create(schedulerConfig, queryExecutor, serverMetrics, latestQueryTime,
+            resourceUsageAccountant);
         break;
       case BINARY_WORKLOAD_ALGORITHM:
-        scheduler = new BinaryWorkloadScheduler(schedulerConfig, queryExecutor, serverMetrics, latestQueryTime);
+        scheduler = new BinaryWorkloadScheduler(schedulerConfig, queryExecutor, serverMetrics, latestQueryTime,
+            resourceUsageAccountant);
         break;
       default:
         scheduler =
@@ -93,7 +99,8 @@ public static QueryScheduler create(PinotConfiguration schedulerConfig, QueryExe
     // Failure on bad configuration will cause outage vs an inferior algorithm that
     // will provide degraded service
     LOGGER.warn("Scheduler {} not found. Using default FCFS query scheduler", schedulerName);
-    return new FCFSQueryScheduler(schedulerConfig, queryExecutor, serverMetrics, latestQueryTime);
+    return new FCFSQueryScheduler(schedulerConfig, queryExecutor, serverMetrics, latestQueryTime,
+        resourceUsageAccountant);
   }
 
   @Nullable
diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/scheduler/fcfs/BoundedFCFSScheduler.java b/pinot-core/src/main/java/org/apache/pinot/core/query/scheduler/fcfs/BoundedFCFSScheduler.java
index 60495ad2f903..30f273e892b0 100644
--- a/pinot-core/src/main/java/org/apache/pinot/core/query/scheduler/fcfs/BoundedFCFSScheduler.java
+++ b/pinot-core/src/main/java/org/apache/pinot/core/query/scheduler/fcfs/BoundedFCFSScheduler.java
@@ -29,6 +29,7 @@
 import org.apache.pinot.core.query.scheduler.TableBasedGroupMapper;
 import org.apache.pinot.core.query.scheduler.resources.PolicyBasedResourceManager;
 import org.apache.pinot.core.query.scheduler.resources.ResourceManager;
+import org.apache.pinot.spi.accounting.ThreadResourceUsageAccountant;
 import org.apache.pinot.spi.env.PinotConfiguration;
 
 
@@ -39,8 +40,9 @@
  */
 public class BoundedFCFSScheduler extends PriorityScheduler {
   public static BoundedFCFSScheduler create(PinotConfiguration config, QueryExecutor queryExecutor,
-      ServerMetrics serverMetrics, LongAccumulator latestQueryTime) {
-    final ResourceManager rm = new PolicyBasedResourceManager(config);
+      ServerMetrics serverMetrics, LongAccumulator latestQueryTime,
+      ThreadResourceUsageAccountant resourceUsageAccountant) {
+    final ResourceManager rm = new PolicyBasedResourceManager(config, resourceUsageAccountant);
     final SchedulerGroupFactory groupFactory = new SchedulerGroupFactory() {
       @Override
       public SchedulerGroup create(PinotConfiguration config, String groupName) {
diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/scheduler/fcfs/FCFSQueryScheduler.java b/pinot-core/src/main/java/org/apache/pinot/core/query/scheduler/fcfs/FCFSQueryScheduler.java
index d61972d347aa..035b795f7376 100644
--- a/pinot-core/src/main/java/org/apache/pinot/core/query/scheduler/fcfs/FCFSQueryScheduler.java
+++ b/pinot-core/src/main/java/org/apache/pinot/core/query/scheduler/fcfs/FCFSQueryScheduler.java
@@ -29,6 +29,7 @@
 import org.apache.pinot.core.query.scheduler.QueryScheduler;
 import org.apache.pinot.core.query.scheduler.resources.QueryExecutorService;
 import org.apache.pinot.core.query.scheduler.resources.UnboundedResourceManager;
+import org.apache.pinot.spi.accounting.ThreadResourceUsageAccountant;
 import org.apache.pinot.spi.env.PinotConfiguration;
 import org.apache.pinot.spi.query.QueryThreadContext;
 
@@ -41,8 +42,9 @@
 public class FCFSQueryScheduler extends QueryScheduler {
 
   public FCFSQueryScheduler(PinotConfiguration config, QueryExecutor queryExecutor, ServerMetrics serverMetrics,
-      LongAccumulator latestQueryTime) {
-    super(config, queryExecutor, new UnboundedResourceManager(config), serverMetrics, latestQueryTime);
+      LongAccumulator latestQueryTime, ThreadResourceUsageAccountant resourceUsageAccountant) {
+    super(config, queryExecutor, new UnboundedResourceManager(config, resourceUsageAccountant), serverMetrics,
+        latestQueryTime);
   }
 
   @Override
diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/scheduler/resources/BinaryWorkloadResourceManager.java b/pinot-core/src/main/java/org/apache/pinot/core/query/scheduler/resources/BinaryWorkloadResourceManager.java
index fb7ffd93e3f9..feab2930bdba 100644
--- a/pinot-core/src/main/java/org/apache/pinot/core/query/scheduler/resources/BinaryWorkloadResourceManager.java
+++ b/pinot-core/src/main/java/org/apache/pinot/core/query/scheduler/resources/BinaryWorkloadResourceManager.java
@@ -22,6 +22,7 @@
 import org.apache.pinot.common.utils.config.QueryOptionsUtils;
 import org.apache.pinot.core.query.request.ServerQueryRequest;
 import org.apache.pinot.core.query.scheduler.SchedulerGroupAccountant;
+import org.apache.pinot.spi.accounting.ThreadResourceUsageAccountant;
 import org.apache.pinot.spi.env.PinotConfiguration;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -34,8 +35,9 @@ public class BinaryWorkloadResourceManager extends ResourceManager {
   private static final Logger LOGGER = LoggerFactory.getLogger(BinaryWorkloadResourceManager.class);
   private final ResourceLimitPolicy _secondaryWorkloadPolicy;
 
-  public BinaryWorkloadResourceManager(PinotConfiguration config) {
-    super(config);
+  public BinaryWorkloadResourceManager(PinotConfiguration config,
+      ThreadResourceUsageAccountant resourceUsageAccountant) {
+    super(config, resourceUsageAccountant);
     _secondaryWorkloadPolicy = new ResourceLimitPolicy(config, _numQueryWorkerThreads);
   }
 
diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/scheduler/resources/PolicyBasedResourceManager.java b/pinot-core/src/main/java/org/apache/pinot/core/query/scheduler/resources/PolicyBasedResourceManager.java
index dd2c942c565d..1fa16063469f 100644
--- a/pinot-core/src/main/java/org/apache/pinot/core/query/scheduler/resources/PolicyBasedResourceManager.java
+++ b/pinot-core/src/main/java/org/apache/pinot/core/query/scheduler/resources/PolicyBasedResourceManager.java
@@ -21,6 +21,7 @@
 import com.google.common.base.Preconditions;
 import org.apache.pinot.core.query.request.ServerQueryRequest;
 import org.apache.pinot.core.query.scheduler.SchedulerGroupAccountant;
+import org.apache.pinot.spi.accounting.ThreadResourceUsageAccountant;
 import org.apache.pinot.spi.env.PinotConfiguration;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -34,8 +35,8 @@ public class PolicyBasedResourceManager extends ResourceManager {
 
   private final ResourceLimitPolicy _resourcePolicy;
 
-  public PolicyBasedResourceManager(PinotConfiguration config) {
-    super(config);
+  public PolicyBasedResourceManager(PinotConfiguration config, ThreadResourceUsageAccountant resourceUsageAccountant) {
+    super(config, resourceUsageAccountant);
     _resourcePolicy = new ResourceLimitPolicy(config, _numQueryWorkerThreads);
   }
 
diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/scheduler/resources/ResourceManager.java b/pinot-core/src/main/java/org/apache/pinot/core/query/scheduler/resources/ResourceManager.java
index 9414dc11edde..a705b735459c 100644
--- a/pinot-core/src/main/java/org/apache/pinot/core/query/scheduler/resources/ResourceManager.java
+++ b/pinot-core/src/main/java/org/apache/pinot/core/query/scheduler/resources/ResourceManager.java
@@ -27,7 +27,9 @@
 import org.apache.pinot.core.query.request.ServerQueryRequest;
 import org.apache.pinot.core.query.scheduler.SchedulerGroupAccountant;
 import org.apache.pinot.core.util.trace.TracedThreadFactory;
+import org.apache.pinot.spi.accounting.ThreadResourceUsageAccountant;
 import org.apache.pinot.spi.env.PinotConfiguration;
+import org.apache.pinot.spi.executor.ThrottleOnCriticalHeapUsageExecutor;
 import org.apache.pinot.spi.utils.CommonConstants;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -76,7 +78,7 @@ public abstract class ResourceManager {
   /**
    * @param config configuration for initializing resource manager
    */
-  public ResourceManager(PinotConfiguration config) {
+  public ResourceManager(PinotConfiguration config, ThreadResourceUsageAccountant resourceUsageAccountant) {
     _numQueryRunnerThreads = config.getProperty(QUERY_RUNNER_CONFIG_KEY, DEFAULT_QUERY_RUNNER_THREADS);
     _numQueryWorkerThreads = config.getProperty(QUERY_WORKER_CONFIG_KEY, DEFAULT_QUERY_WORKER_THREADS);
 
@@ -85,14 +87,23 @@ public ResourceManager(PinotConfiguration config) {
     // pqr -> pinot query runner (to give short names)
     ThreadFactory queryRunnerFactory = new TracedThreadFactory(QUERY_RUNNER_THREAD_PRIORITY, false,
         CommonConstants.ExecutorService.PINOT_QUERY_RUNNER_NAME_FORMAT);
-    _queryRunners =
-        MoreExecutors.listeningDecorator(Executors.newFixedThreadPool(_numQueryRunnerThreads, queryRunnerFactory));
+
+    ExecutorService runnerService = Executors.newFixedThreadPool(_numQueryRunnerThreads, queryRunnerFactory);
+    if (config.getProperty(CommonConstants.Server.CONFIG_OF_ENABLE_QUERY_SCHEDULER_THROTTLING_ON_HEAP_USAGE,
+        CommonConstants.Server.DEFAULT_ENABLE_QUERY_SCHEDULER_THROTTLING_ON_HEAP_USAGE)) {
+      runnerService = new ThrottleOnCriticalHeapUsageExecutor(runnerService, resourceUsageAccountant);
+    }
+    _queryRunners = MoreExecutors.listeningDecorator(runnerService);
 
     // pqw -> pinot query workers
     ThreadFactory queryWorkersFactory = new TracedThreadFactory(Thread.NORM_PRIORITY, false,
         CommonConstants.ExecutorService.PINOT_QUERY_WORKER_NAME_FORMAT);
-    _queryWorkers =
-        MoreExecutors.listeningDecorator(Executors.newFixedThreadPool(_numQueryWorkerThreads, queryWorkersFactory));
+    ExecutorService workerService = Executors.newFixedThreadPool(_numQueryWorkerThreads, queryWorkersFactory);
+    if (config.getProperty(CommonConstants.Server.CONFIG_OF_ENABLE_QUERY_SCHEDULER_THROTTLING_ON_HEAP_USAGE,
+        CommonConstants.Server.DEFAULT_ENABLE_QUERY_SCHEDULER_THROTTLING_ON_HEAP_USAGE)) {
+      workerService = new ThrottleOnCriticalHeapUsageExecutor(workerService, resourceUsageAccountant);
+    }
+    _queryWorkers = MoreExecutors.listeningDecorator(workerService);
   }
 
   public void stop() {
diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/scheduler/resources/UnboundedResourceManager.java b/pinot-core/src/main/java/org/apache/pinot/core/query/scheduler/resources/UnboundedResourceManager.java
index 075e35dd021e..ff0c2eb2835c 100644
--- a/pinot-core/src/main/java/org/apache/pinot/core/query/scheduler/resources/UnboundedResourceManager.java
+++ b/pinot-core/src/main/java/org/apache/pinot/core/query/scheduler/resources/UnboundedResourceManager.java
@@ -20,6 +20,7 @@
 
 import org.apache.pinot.core.query.request.ServerQueryRequest;
 import org.apache.pinot.core.query.scheduler.SchedulerGroupAccountant;
+import org.apache.pinot.spi.accounting.ThreadResourceUsageAccountant;
 import org.apache.pinot.spi.env.PinotConfiguration;
 
 
@@ -30,8 +31,8 @@
  */
 public class UnboundedResourceManager extends ResourceManager {
 
-  public UnboundedResourceManager(PinotConfiguration config) {
-    super(config);
+  public UnboundedResourceManager(PinotConfiguration config, ThreadResourceUsageAccountant resourceUsageAccountant) {
+    super(config, resourceUsageAccountant);
   }
 
   @Override
diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/scheduler/tokenbucket/TokenPriorityScheduler.java b/pinot-core/src/main/java/org/apache/pinot/core/query/scheduler/tokenbucket/TokenPriorityScheduler.java
index cac4c2131317..25f02c6de89a 100644
--- a/pinot-core/src/main/java/org/apache/pinot/core/query/scheduler/tokenbucket/TokenPriorityScheduler.java
+++ b/pinot-core/src/main/java/org/apache/pinot/core/query/scheduler/tokenbucket/TokenPriorityScheduler.java
@@ -28,6 +28,7 @@
 import org.apache.pinot.core.query.scheduler.TableBasedGroupMapper;
 import org.apache.pinot.core.query.scheduler.resources.PolicyBasedResourceManager;
 import org.apache.pinot.core.query.scheduler.resources.ResourceManager;
+import org.apache.pinot.spi.accounting.ThreadResourceUsageAccountant;
 import org.apache.pinot.spi.env.PinotConfiguration;
 
 
@@ -42,8 +43,8 @@ public class TokenPriorityScheduler extends PriorityScheduler {
   private static final int DEFAULT_TOKEN_LIFETIME_MS = 100;
 
   public static TokenPriorityScheduler create(PinotConfiguration config, QueryExecutor queryExecutor,
-      ServerMetrics metrics, LongAccumulator latestQueryTime) {
-    final ResourceManager rm = new PolicyBasedResourceManager(config);
+      ServerMetrics metrics, LongAccumulator latestQueryTime, ThreadResourceUsageAccountant resourceUsageAccountant) {
+    final ResourceManager rm = new PolicyBasedResourceManager(config, resourceUsageAccountant);
     final SchedulerGroupFactory groupFactory = new SchedulerGroupFactory() {
       @Override
       public SchedulerGroup create(PinotConfiguration config, String groupName) {
diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/QueryRunner.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/QueryRunner.java
index d98f919f29eb..054a80ef14d1 100644
--- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/QueryRunner.java
+++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/QueryRunner.java
@@ -75,10 +75,12 @@
 import org.apache.pinot.query.runtime.timeseries.TimeSeriesExecutionContext;
 import org.apache.pinot.query.runtime.timeseries.serde.TimeSeriesBlockSerde;
 import org.apache.pinot.spi.accounting.ThreadExecutionContext;
+import org.apache.pinot.spi.accounting.ThreadResourceUsageAccountant;
 import org.apache.pinot.spi.env.PinotConfiguration;
 import org.apache.pinot.spi.executor.ExecutorServiceUtils;
 import org.apache.pinot.spi.executor.HardLimitExecutor;
 import org.apache.pinot.spi.executor.MetricsExecutor;
+import org.apache.pinot.spi.executor.ThrottleOnCriticalHeapUsageExecutor;
 import org.apache.pinot.spi.query.QueryThreadContext;
 import org.apache.pinot.spi.utils.CommonConstants;
 import org.apache.pinot.spi.utils.CommonConstants.Broker.Request.QueryOptionKey;
@@ -144,7 +146,7 @@ public class QueryRunner {
    * <p>Should be called only once and before calling any other method.
    */
   public void init(PinotConfiguration serverConf, InstanceDataManager instanceDataManager,
-      @Nullable TlsConfig tlsConfig, BooleanSupplier sendStats) {
+      @Nullable TlsConfig tlsConfig, BooleanSupplier sendStats, ThreadResourceUsageAccountant resourceUsageAccountant) {
     String hostname = serverConf.getProperty(CommonConstants.MultiStageQueryRunner.KEY_OF_QUERY_RUNNER_HOSTNAME);
     if (hostname.startsWith(CommonConstants.Helix.PREFIX_OF_SERVER_INSTANCE)) {
       hostname = hostname.substring(CommonConstants.Helix.SERVER_INSTANCE_PREFIX_LENGTH);
@@ -208,6 +210,12 @@ public void init(PinotConfiguration serverConf, InstanceDataManager instanceData
       _executorService = new HardLimitExecutor(hardLimit, _executorService);
     }
 
+    if (serverConf.getProperty(Server.CONFIG_OF_ENABLE_QUERY_SCHEDULER_THROTTLING_ON_HEAP_USAGE,
+        Server.DEFAULT_ENABLE_QUERY_SCHEDULER_THROTTLING_ON_HEAP_USAGE)) {
+      LOGGER.info("Enable OOM Throttling on critical heap usage for multi-stage executor");
+      _executorService = new ThrottleOnCriticalHeapUsageExecutor(_executorService, resourceUsageAccountant);
+    }
+
     _opChainScheduler = new OpChainSchedulerService(_executorService, serverConf);
     _mailboxService = new MailboxService(hostname, port, serverConf, tlsConfig);
     try {
diff --git a/pinot-server/src/main/java/org/apache/pinot/server/starter/ServerInstance.java b/pinot-server/src/main/java/org/apache/pinot/server/starter/ServerInstance.java
index 2ba34b61bdc0..516fb720fe1d 100644
--- a/pinot-server/src/main/java/org/apache/pinot/server/starter/ServerInstance.java
+++ b/pinot-server/src/main/java/org/apache/pinot/server/starter/ServerInstance.java
@@ -51,6 +51,7 @@
 import org.apache.pinot.server.conf.ServerConf;
 import org.apache.pinot.server.starter.helix.SendStatsPredicate;
 import org.apache.pinot.server.worker.WorkerQueryServer;
+import org.apache.pinot.spi.accounting.ThreadResourceUsageAccountant;
 import org.apache.pinot.spi.env.PinotConfiguration;
 import org.apache.pinot.spi.metrics.PinotMetricUtils;
 import org.apache.pinot.spi.metrics.PinotMetricsRegistry;
@@ -85,7 +86,8 @@ public class ServerInstance {
   private boolean _queryServerStarted = false;
 
   public ServerInstance(ServerConf serverConf, HelixManager helixManager, AccessControlFactory accessControlFactory,
-      @Nullable SegmentOperationsThrottler segmentOperationsThrottler, SendStatsPredicate sendStatsPredicate)
+      @Nullable SegmentOperationsThrottler segmentOperationsThrottler, SendStatsPredicate sendStatsPredicate,
+      ThreadResourceUsageAccountant resourceUsageAccountant)
       throws Exception {
     LOGGER.info("Initializing server instance");
     _helixManager = helixManager;
@@ -122,7 +124,8 @@ public ServerInstance(ServerConf serverConf, HelixManager helixManager, AccessCo
     LOGGER.info("Initializing query scheduler");
     _latestQueryTime = new LongAccumulator(Long::max, 0);
     _queryScheduler =
-        QuerySchedulerFactory.create(serverConf.getSchedulerConfig(), _queryExecutor, _serverMetrics, _latestQueryTime);
+        QuerySchedulerFactory.create(serverConf.getSchedulerConfig(), _queryExecutor, _serverMetrics, _latestQueryTime,
+            resourceUsageAccountant);
 
     TlsConfig tlsConfig =
         TlsUtils.extractTlsConfig(serverConf.getPinotConfig(), CommonConstants.Server.SERVER_TLS_PREFIX);
@@ -135,7 +138,7 @@ public ServerInstance(ServerConf serverConf, HelixManager helixManager, AccessCo
     if (serverConf.isMultiStageServerEnabled()) {
       LOGGER.info("Initializing Multi-stage query engine");
       _workerQueryServer = new WorkerQueryServer(serverConf.getPinotConfig(), _instanceDataManager,
-          serverConf.isMultiStageEngineTlsEnabled() ? tlsConfig : null, sendStatsPredicate);
+          serverConf.isMultiStageEngineTlsEnabled() ? tlsConfig : null, sendStatsPredicate, resourceUsageAccountant);
     } else {
       _workerQueryServer = null;
     }
diff --git a/pinot-server/src/main/java/org/apache/pinot/server/starter/helix/BaseServerStarter.java b/pinot-server/src/main/java/org/apache/pinot/server/starter/helix/BaseServerStarter.java
index fa8b8fa6bc93..044d33bf7cb9 100644
--- a/pinot-server/src/main/java/org/apache/pinot/server/starter/helix/BaseServerStarter.java
+++ b/pinot-server/src/main/java/org/apache/pinot/server/starter/helix/BaseServerStarter.java
@@ -663,14 +663,6 @@ public void start()
               segmentDownloadThrottler, segmentMultiColTextIndexPreprocessThrottler);
     }
 
-    SendStatsPredicate sendStatsPredicate = SendStatsPredicate.create(_serverConf, _helixManager);
-    ServerConf serverConf = new ServerConf(_serverConf);
-    _serverInstance = new ServerInstance(serverConf, _helixManager, _accessControlFactory, _segmentOperationsThrottler,
-        sendStatsPredicate);
-    ServerMetrics serverMetrics = _serverInstance.getServerMetrics();
-
-    InstanceDataManager instanceDataManager = _serverInstance.getInstanceDataManager();
-    instanceDataManager.setSupplierOfIsServerReadyToServeQueries(() -> _isServerReadyToServeQueries);
     // initialize the thread accountant for query killing
     Tracing.ThreadAccountantOps.initializeThreadAccountant(
         _serverConf.subset(CommonConstants.PINOT_QUERY_SCHEDULER_PREFIX), _instanceId,
@@ -680,6 +672,15 @@ public void start()
           Tracing.getThreadAccountant().getClusterConfigChangeListener());
     }
 
+    SendStatsPredicate sendStatsPredicate = SendStatsPredicate.create(_serverConf, _helixManager);
+    ServerConf serverConf = new ServerConf(_serverConf);
+    _serverInstance = new ServerInstance(serverConf, _helixManager, _accessControlFactory, _segmentOperationsThrottler,
+        sendStatsPredicate, Tracing.getThreadAccountant());
+    ServerMetrics serverMetrics = _serverInstance.getServerMetrics();
+
+    InstanceDataManager instanceDataManager = _serverInstance.getInstanceDataManager();
+    instanceDataManager.setSupplierOfIsServerReadyToServeQueries(() -> _isServerReadyToServeQueries);
+
     initSegmentFetcher(_serverConf);
     StateModelFactory<?> stateModelFactory =
         new SegmentOnlineOfflineStateModelFactory(_instanceId, instanceDataManager);
diff --git a/pinot-server/src/main/java/org/apache/pinot/server/worker/WorkerQueryServer.java b/pinot-server/src/main/java/org/apache/pinot/server/worker/WorkerQueryServer.java
index f8372d9960c3..9528ec79af0c 100644
--- a/pinot-server/src/main/java/org/apache/pinot/server/worker/WorkerQueryServer.java
+++ b/pinot-server/src/main/java/org/apache/pinot/server/worker/WorkerQueryServer.java
@@ -24,6 +24,7 @@
 import org.apache.pinot.query.runtime.QueryRunner;
 import org.apache.pinot.query.service.server.QueryServer;
 import org.apache.pinot.server.starter.helix.SendStatsPredicate;
+import org.apache.pinot.spi.accounting.ThreadResourceUsageAccountant;
 import org.apache.pinot.spi.env.PinotConfiguration;
 import org.apache.pinot.spi.utils.CommonConstants;
 import org.apache.pinot.spi.utils.NetUtils;
@@ -34,13 +35,14 @@ public class WorkerQueryServer {
   private final QueryServer _queryWorkerService;
 
   public WorkerQueryServer(PinotConfiguration serverConf, InstanceDataManager instanceDataManager,
-      @Nullable TlsConfig tlsConfig, SendStatsPredicate sendStats) {
+      @Nullable TlsConfig tlsConfig, SendStatsPredicate sendStats,
+      ThreadResourceUsageAccountant resourceUsageAccountant) {
     serverConf = toWorkerQueryConfig(serverConf);
     String instanceId = serverConf.getProperty(CommonConstants.Server.CONFIG_OF_INSTANCE_ID);
     _queryServicePort = serverConf.getProperty(CommonConstants.MultiStageQueryRunner.KEY_OF_QUERY_SERVER_PORT,
         CommonConstants.MultiStageQueryRunner.DEFAULT_QUERY_SERVER_PORT);
     QueryRunner queryRunner = new QueryRunner();
-    queryRunner.init(serverConf, instanceDataManager, tlsConfig, sendStats::isSendStats);
+    queryRunner.init(serverConf, instanceDataManager, tlsConfig, sendStats::isSendStats, resourceUsageAccountant);
     _queryWorkerService = new QueryServer(instanceId, _queryServicePort, queryRunner, tlsConfig, serverConf);
   }
 
diff --git a/pinot-spi/src/main/java/org/apache/pinot/spi/accounting/ThreadResourceUsageAccountant.java b/pinot-spi/src/main/java/org/apache/pinot/spi/accounting/ThreadResourceUsageAccountant.java
index b3a270aea1ea..1d6d590e0fa3 100644
--- a/pinot-spi/src/main/java/org/apache/pinot/spi/accounting/ThreadResourceUsageAccountant.java
+++ b/pinot-spi/src/main/java/org/apache/pinot/spi/accounting/ThreadResourceUsageAccountant.java
@@ -84,6 +84,10 @@ void setupWorker(int taskId, ThreadExecutionContext.TaskType taskType,
    */
   void sampleUsageMSE();
 
+  default boolean throttleQuerySubmission() {
+    return false;
+  }
+
   /**
    * special interface to aggregate usage to the stats store only once, it is used for response
    * ser/de threads where the thread execution context cannot be setup before hands as
diff --git a/pinot-spi/src/main/java/org/apache/pinot/spi/executor/ThrottleOnCriticalHeapUsageExecutor.java b/pinot-spi/src/main/java/org/apache/pinot/spi/executor/ThrottleOnCriticalHeapUsageExecutor.java
new file mode 100644
index 000000000000..091fb5e3fe0d
--- /dev/null
+++ b/pinot-spi/src/main/java/org/apache/pinot/spi/executor/ThrottleOnCriticalHeapUsageExecutor.java
@@ -0,0 +1,62 @@
+/**
+ * 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.executor;
+
+import java.util.concurrent.Callable;
+import java.util.concurrent.ExecutorService;
+import org.apache.pinot.spi.accounting.ThreadResourceUsageAccountant;
+
+
+/**
+ * An Executor that throttles task submission when the heap usage is critical.
+ * Heap Usage level is obtained from {@link ThreadResourceUsageAccountant#throttleQuerySubmission()}.
+ */
+public class ThrottleOnCriticalHeapUsageExecutor extends DecoratorExecutorService {
+  ThreadResourceUsageAccountant _threadResourceUsageAccountant;
+
+  public ThrottleOnCriticalHeapUsageExecutor(ExecutorService executorService,
+      ThreadResourceUsageAccountant threadResourceUsageAccountant) {
+    super(executorService);
+    _threadResourceUsageAccountant = threadResourceUsageAccountant;
+  }
+
+  protected void checkTaskAllowed() {
+    if (_threadResourceUsageAccountant.throttleQuerySubmission()) {
+      throw new IllegalStateException("Tasks throttled due to high heap usage.");
+    }
+  }
+
+  @Override
+  protected <T> Callable<T> decorate(Callable<T> task) {
+    checkTaskAllowed();
+    return () -> {
+      checkTaskAllowed();
+      return task.call();
+    };
+  }
+
+  @Override
+  protected Runnable decorate(Runnable task) {
+    checkTaskAllowed();
+    return () -> {
+      checkTaskAllowed();
+      task.run();
+    };
+  }
+}
diff --git a/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java b/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java
index 3a2b703edc1c..b8a4d16152e3 100644
--- a/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java
+++ b/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java
@@ -1099,6 +1099,11 @@ public static class Server {
     @Deprecated
     public static final String DEFAULT_QUERY_EXECUTOR_OPCHAIN_EXECUTOR = DEFAULT_MULTISTAGE_EXECUTOR_TYPE;
 
+    // Enable SSE & MSE task throttling on critical heap usage.
+    public static final String CONFIG_OF_ENABLE_QUERY_SCHEDULER_THROTTLING_ON_HEAP_USAGE =
+        QUERY_EXECUTOR_CONFIG_PREFIX + ".enableThrottlingOnHeapUsage";
+    public static final boolean DEFAULT_ENABLE_QUERY_SCHEDULER_THROTTLING_ON_HEAP_USAGE = false;
+
     /**
      * The ExecutorServiceProvider to be used for timeseries threads.
      *


Test patch:
diff --git a/pinot-core/src/test/java/org/apache/pinot/core/accounting/ResourceManagerAccountingTest.java b/pinot-core/src/test/java/org/apache/pinot/core/accounting/ResourceManagerAccountingTest.java
index 7ffa5f342d13..c2e2f0cef3ce 100644
--- a/pinot-core/src/test/java/org/apache/pinot/core/accounting/ResourceManagerAccountingTest.java
+++ b/pinot-core/src/test/java/org/apache/pinot/core/accounting/ResourceManagerAccountingTest.java
@@ -550,7 +550,7 @@ public void testThreadMemory()
   private ResourceManager getResourceManager(int runners, int workers, final int softLimit, final int hardLimit,
       Map<String, Object> map) {
 
-    return new ResourceManager(getConfig(runners, workers, map)) {
+    return new ResourceManager(getConfig(runners, workers, map), new Tracing.DefaultThreadResourceUsageAccountant()) {
 
       @Override
       public QueryExecutorService getExecutorService(ServerQueryRequest query, SchedulerGroupAccountant accountant) {
diff --git a/pinot-core/src/test/java/org/apache/pinot/core/query/scheduler/MultiLevelPriorityQueueTest.java b/pinot-core/src/test/java/org/apache/pinot/core/query/scheduler/MultiLevelPriorityQueueTest.java
index 2006aa7992f5..84ab91ed92d9 100644
--- a/pinot-core/src/test/java/org/apache/pinot/core/query/scheduler/MultiLevelPriorityQueueTest.java
+++ b/pinot-core/src/test/java/org/apache/pinot/core/query/scheduler/MultiLevelPriorityQueueTest.java
@@ -33,6 +33,7 @@
 import org.apache.pinot.core.query.scheduler.resources.UnboundedResourceManager;
 import org.apache.pinot.spi.env.PinotConfiguration;
 import org.apache.pinot.spi.metrics.PinotMetricUtils;
+import org.apache.pinot.spi.trace.Tracing;
 import org.testng.annotations.BeforeMethod;
 import org.testng.annotations.Test;
 
@@ -81,7 +82,8 @@ public void testPutOutOfCapacity()
 
     PinotConfiguration configuration = new PinotConfiguration(properties);
 
-    ResourceManager rm = new UnboundedResourceManager(configuration);
+    ResourceManager rm =
+        new UnboundedResourceManager(configuration, new Tracing.DefaultThreadResourceUsageAccountant());
     MultiLevelPriorityQueue queue = createQueue(configuration, rm);
     queue.put(createQueryRequest(GROUP_ONE, METRICS));
     GROUP_FACTORY._groupMap.get(GROUP_ONE).addReservedThreads(rm.getTableThreadsHardLimit());
@@ -129,7 +131,8 @@ public void testTakeWithLimits()
 
     PinotConfiguration configuration = new PinotConfiguration(properties);
 
-    PolicyBasedResourceManager rm = new PolicyBasedResourceManager(configuration);
+    PolicyBasedResourceManager rm =
+        new PolicyBasedResourceManager(configuration, new Tracing.DefaultThreadResourceUsageAccountant());
     MultiLevelPriorityQueue queue = createQueue(configuration, rm);
 
     queue.put(createQueryRequest(GROUP_ONE, METRICS));
@@ -214,7 +217,7 @@ public void testNoPendingAfterTrim()
 
   private MultiLevelPriorityQueue createQueue() {
     PinotConfiguration conf = new PinotConfiguration();
-    return createQueue(conf, new UnboundedResourceManager(conf));
+    return createQueue(conf, new UnboundedResourceManager(conf, new Tracing.DefaultThreadResourceUsageAccountant()));
   }
 
   private MultiLevelPriorityQueue createQueue(PinotConfiguration config, ResourceManager rm) {
diff --git a/pinot-core/src/test/java/org/apache/pinot/core/query/scheduler/PrioritySchedulerTest.java b/pinot-core/src/test/java/org/apache/pinot/core/query/scheduler/PrioritySchedulerTest.java
index 0e7a4d674766..94d91dd80e6c 100644
--- a/pinot-core/src/test/java/org/apache/pinot/core/query/scheduler/PrioritySchedulerTest.java
+++ b/pinot-core/src/test/java/org/apache/pinot/core/query/scheduler/PrioritySchedulerTest.java
@@ -52,6 +52,7 @@
 import org.apache.pinot.spi.env.PinotConfiguration;
 import org.apache.pinot.spi.exception.QueryErrorCode;
 import org.apache.pinot.spi.metrics.PinotMetricUtils;
+import org.apache.pinot.spi.trace.Tracing;
 import org.testng.annotations.AfterMethod;
 import org.testng.annotations.Test;
 
@@ -241,7 +242,7 @@ public TestPriorityScheduler(PinotConfiguration config, ResourceManager resource
     }
 
     public static TestPriorityScheduler create(PinotConfiguration config) {
-      ResourceManager rm = new PolicyBasedResourceManager(config);
+      ResourceManager rm = new PolicyBasedResourceManager(config, new Tracing.DefaultThreadResourceUsageAccountant());
       QueryExecutor qe = new TestQueryExecutor();
       _groupFactory = new TestSchedulerGroupFactory();
       MultiLevelPriorityQueue queue =
diff --git a/pinot-core/src/test/java/org/apache/pinot/core/query/scheduler/QuerySchedulerFactoryTest.java b/pinot-core/src/test/java/org/apache/pinot/core/query/scheduler/QuerySchedulerFactoryTest.java
index 9c4c59ea1efc..b28bf4ac8ec0 100644
--- a/pinot-core/src/test/java/org/apache/pinot/core/query/scheduler/QuerySchedulerFactoryTest.java
+++ b/pinot-core/src/test/java/org/apache/pinot/core/query/scheduler/QuerySchedulerFactoryTest.java
@@ -27,7 +27,9 @@
 import org.apache.pinot.core.query.scheduler.fcfs.FCFSQueryScheduler;
 import org.apache.pinot.core.query.scheduler.resources.UnboundedResourceManager;
 import org.apache.pinot.core.query.scheduler.tokenbucket.TokenPriorityScheduler;
+import org.apache.pinot.spi.accounting.ThreadResourceUsageAccountant;
 import org.apache.pinot.spi.env.PinotConfiguration;
+import org.apache.pinot.spi.trace.Tracing;
 import org.testng.annotations.Test;
 
 import static org.mockito.Mockito.mock;
@@ -43,25 +45,27 @@ public void testQuerySchedulerFactory() {
     LongAccumulator latestQueryTime = mock(LongAccumulator.class);
 
     PinotConfiguration config = new PinotConfiguration();
+    ThreadResourceUsageAccountant accountant = new Tracing.DefaultThreadResourceUsageAccountant();
     config.setProperty(QuerySchedulerFactory.ALGORITHM_NAME_CONFIG_KEY, QuerySchedulerFactory.FCFS_ALGORITHM);
-    QueryScheduler queryScheduler = QuerySchedulerFactory.create(config, queryExecutor, serverMetrics, latestQueryTime);
+    QueryScheduler queryScheduler =
+        QuerySchedulerFactory.create(config, queryExecutor, serverMetrics, latestQueryTime, accountant);
     assertTrue(queryScheduler instanceof FCFSQueryScheduler);
 
     config.setProperty(QuerySchedulerFactory.ALGORITHM_NAME_CONFIG_KEY, QuerySchedulerFactory.TOKEN_BUCKET_ALGORITHM);
-    queryScheduler = QuerySchedulerFactory.create(config, queryExecutor, serverMetrics, latestQueryTime);
+    queryScheduler = QuerySchedulerFactory.create(config, queryExecutor, serverMetrics, latestQueryTime, accountant);
     assertTrue(queryScheduler instanceof TokenPriorityScheduler);
 
     config.setProperty(QuerySchedulerFactory.ALGORITHM_NAME_CONFIG_KEY, QuerySchedulerFactory.BOUNDED_FCFS_ALGORITHM);
-    queryScheduler = QuerySchedulerFactory.create(config, queryExecutor, serverMetrics, latestQueryTime);
+    queryScheduler = QuerySchedulerFactory.create(config, queryExecutor, serverMetrics, latestQueryTime, accountant);
     assertTrue(queryScheduler instanceof BoundedFCFSScheduler);
 
     config.setProperty(QuerySchedulerFactory.ALGORITHM_NAME_CONFIG_KEY,
         QuerySchedulerFactory.BINARY_WORKLOAD_ALGORITHM);
-    queryScheduler = QuerySchedulerFactory.create(config, queryExecutor, serverMetrics, latestQueryTime);
+    queryScheduler = QuerySchedulerFactory.create(config, queryExecutor, serverMetrics, latestQueryTime, accountant);
     assertTrue(queryScheduler instanceof BinaryWorkloadScheduler);
 
     config.setProperty(QuerySchedulerFactory.ALGORITHM_NAME_CONFIG_KEY, TestQueryScheduler.class.getName());
-    queryScheduler = QuerySchedulerFactory.create(config, queryExecutor, serverMetrics, latestQueryTime);
+    queryScheduler = QuerySchedulerFactory.create(config, queryExecutor, serverMetrics, latestQueryTime, accountant);
     assertTrue(queryScheduler instanceof TestQueryScheduler);
   }
 
@@ -69,7 +73,9 @@ public static final class TestQueryScheduler extends QueryScheduler {
 
     public TestQueryScheduler(PinotConfiguration config, QueryExecutor queryExecutor, ServerMetrics serverMetrics,
         LongAccumulator latestQueryTime) {
-      super(config, queryExecutor, new UnboundedResourceManager(config), serverMetrics, latestQueryTime);
+      super(config, queryExecutor,
+          new UnboundedResourceManager(config, new Tracing.DefaultThreadResourceUsageAccountant()), serverMetrics,
+          latestQueryTime);
     }
 
     @Override
diff --git a/pinot-core/src/test/java/org/apache/pinot/core/query/scheduler/resources/ResourceManagerTest.java b/pinot-core/src/test/java/org/apache/pinot/core/query/scheduler/resources/ResourceManagerTest.java
index 6e6b01cf431f..23f7b9f7b1f3 100644
--- a/pinot-core/src/test/java/org/apache/pinot/core/query/scheduler/resources/ResourceManagerTest.java
+++ b/pinot-core/src/test/java/org/apache/pinot/core/query/scheduler/resources/ResourceManagerTest.java
@@ -23,6 +23,7 @@
 import org.apache.pinot.core.query.request.ServerQueryRequest;
 import org.apache.pinot.core.query.scheduler.SchedulerGroupAccountant;
 import org.apache.pinot.spi.env.PinotConfiguration;
+import org.apache.pinot.spi.trace.Tracing;
 import org.testng.annotations.Test;
 
 import static org.mockito.Mockito.mock;
@@ -48,7 +49,7 @@ public void testCanSchedule()
 
   private ResourceManager getResourceManager(int runners, int workers, final int softLimit, final int hardLimit) {
 
-    return new ResourceManager(getConfig(runners, workers)) {
+    return new ResourceManager(getConfig(runners, workers), new Tracing.DefaultThreadResourceUsageAccountant()) {
 
       @Override
       public QueryExecutorService getExecutorService(ServerQueryRequest query, SchedulerGroupAccountant accountant) {
diff --git a/pinot-core/src/test/java/org/apache/pinot/core/query/scheduler/resources/UnboundedResourceManagerTest.java b/pinot-core/src/test/java/org/apache/pinot/core/query/scheduler/resources/UnboundedResourceManagerTest.java
index 7031a8ea2277..0baeec19e524 100644
--- a/pinot-core/src/test/java/org/apache/pinot/core/query/scheduler/resources/UnboundedResourceManagerTest.java
+++ b/pinot-core/src/test/java/org/apache/pinot/core/query/scheduler/resources/UnboundedResourceManagerTest.java
@@ -22,6 +22,7 @@
 import java.util.Map;
 import org.apache.pinot.core.query.scheduler.SchedulerGroupAccountant;
 import org.apache.pinot.spi.env.PinotConfiguration;
+import org.apache.pinot.spi.trace.Tracing;
 import org.testng.annotations.Test;
 
 import static org.mockito.Mockito.mock;
@@ -35,7 +36,8 @@ public class UnboundedResourceManagerTest {
 
   @Test
   public void testDefault() {
-    UnboundedResourceManager rm = new UnboundedResourceManager(new PinotConfiguration());
+    UnboundedResourceManager rm =
+        new UnboundedResourceManager(new PinotConfiguration(), new Tracing.DefaultThreadResourceUsageAccountant());
     assertTrue(rm.getNumQueryRunnerThreads() > 1);
     assertTrue(rm.getNumQueryWorkerThreads() >= 1);
     assertEquals(rm.getTableThreadsHardLimit(), rm.getNumQueryRunnerThreads() + rm.getNumQueryWorkerThreads());
@@ -51,7 +53,8 @@ public void testWithConfig() {
     properties.put(ResourceManager.QUERY_RUNNER_CONFIG_KEY, runners);
     properties.put(ResourceManager.QUERY_WORKER_CONFIG_KEY, workers);
 
-    UnboundedResourceManager rm = new UnboundedResourceManager(new PinotConfiguration(properties));
+    UnboundedResourceManager rm = new UnboundedResourceManager(new PinotConfiguration(properties),
+        new Tracing.DefaultThreadResourceUsageAccountant());
     assertEquals(rm.getNumQueryWorkerThreads(), workers);
     assertEquals(rm.getNumQueryRunnerThreads(), runners);
     assertEquals(rm.getTableThreadsHardLimit(), runners + workers);
diff --git a/pinot-query-runtime/src/test/java/org/apache/pinot/query/QueryServerEnclosure.java b/pinot-query-runtime/src/test/java/org/apache/pinot/query/QueryServerEnclosure.java
index 43e99f1c0cbc..5df708551bec 100644
--- a/pinot-query-runtime/src/test/java/org/apache/pinot/query/QueryServerEnclosure.java
+++ b/pinot-query-runtime/src/test/java/org/apache/pinot/query/QueryServerEnclosure.java
@@ -30,6 +30,7 @@
 import org.apache.pinot.spi.accounting.ThreadExecutionContext;
 import org.apache.pinot.spi.env.PinotConfiguration;
 import org.apache.pinot.spi.query.QueryThreadContext;
+import org.apache.pinot.spi.trace.Tracing;
 import org.apache.pinot.spi.utils.CommonConstants;
 
 
@@ -61,7 +62,8 @@ public QueryServerEnclosure(MockInstanceDataManagerFactory factory, Map<String,
     runnerConfig.put(CommonConstants.MultiStageQueryRunner.KEY_OF_QUERY_RUNNER_PORT, _queryRunnerPort);
     InstanceDataManager instanceDataManager = factory.buildInstanceDataManager();
     _queryRunner = new QueryRunner();
-    _queryRunner.init(new PinotConfiguration(runnerConfig), instanceDataManager, null, () -> true);
+    _queryRunner.init(new PinotConfiguration(runnerConfig), instanceDataManager, null, () -> true,
+        new Tracing.DefaultThreadResourceUsageAccountant());
   }
 
   public int getPort() {
diff --git a/pinot-spi/src/test/java/org/apache/pinot/spi/executor/ThrottleOnCriticalHeapUsageExecutorTest.java b/pinot-spi/src/test/java/org/apache/pinot/spi/executor/ThrottleOnCriticalHeapUsageExecutorTest.java
new file mode 100644
index 000000000000..82bc2e808df1
--- /dev/null
+++ b/pinot-spi/src/test/java/org/apache/pinot/spi/executor/ThrottleOnCriticalHeapUsageExecutorTest.java
@@ -0,0 +1,150 @@
+/**
+ * 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.executor;
+
+import java.util.Collection;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.BrokenBarrierException;
+import java.util.concurrent.CyclicBarrier;
+import java.util.concurrent.Executors;
+import java.util.concurrent.atomic.AtomicLong;
+import javax.annotation.Nullable;
+import org.apache.pinot.spi.accounting.QueryResourceTracker;
+import org.apache.pinot.spi.accounting.ThreadExecutionContext;
+import org.apache.pinot.spi.accounting.ThreadResourceTracker;
+import org.apache.pinot.spi.accounting.ThreadResourceUsageAccountant;
+import org.apache.pinot.spi.accounting.ThreadResourceUsageProvider;
+import org.testng.annotations.Test;
+
+import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.fail;
+
+
+public class ThrottleOnCriticalHeapUsageExecutorTest {
+  @Test
+  void testThrottle() throws Exception {
+    ThreadResourceUsageAccountant accountant = new ThreadResourceUsageAccountant() {
+      final AtomicLong _numCalls = new AtomicLong(0);
+      @Override
+      public void clear() {
+      }
+
+      @Override
+      public boolean isAnchorThreadInterrupted() {
+        return false;
+      }
+
+      @Override
+      public void createExecutionContext(String queryId, int taskId, ThreadExecutionContext.TaskType taskType,
+          @Nullable ThreadExecutionContext parentContext) {
+      }
+
+      @Override
+      public void setupRunner(String queryId, int taskId, ThreadExecutionContext.TaskType taskType) {
+      }
+
+      @Override
+      public void setupWorker(int taskId, ThreadExecutionContext.TaskType taskType,
+          @Nullable ThreadExecutionContext parentContext) {
+      }
+
+      @Nullable
+      @Override
+      public ThreadExecutionContext getThreadExecutionContext() {
+        return null;
+      }
+
+      @Override
+      public void setThreadResourceUsageProvider(ThreadResourceUsageProvider threadResourceUsageProvider) {
+      }
+
+      @Override
+      public void sampleUsage() {
+      }
+
+      @Override
+      public void sampleUsageMSE() {
+      }
+
+      @Override
+      public boolean throttleQuerySubmission() {
+        return _numCalls.getAndIncrement() > 1;
+      }
+
+      @Override
+      public void updateQueryUsageConcurrently(String queryId, long cpuTimeNs, long allocatedBytes) {
+      }
+
+      @Override
+      public void updateQueryUsageConcurrently(String queryId) {
+      }
+
+      @Override
+      public void startWatcherTask() {
+      }
+
+      @Override
+      public Exception getErrorStatus() {
+        return null;
+      }
+
+      @Override
+      public Collection<? extends ThreadResourceTracker> getThreadResources() {
+        return List.of();
+      }
+
+      @Override
+      public Map<String, ? extends QueryResourceTracker> getQueryResources() {
+        return Map.of();
+      }
+    };
+
+    ThrottleOnCriticalHeapUsageExecutor executor = new ThrottleOnCriticalHeapUsageExecutor(
+        Executors.newCachedThreadPool(), accountant);
+
+    CyclicBarrier barrier = new CyclicBarrier(2);
+
+    try {
+      executor.execute(() -> {
+        try {
+          barrier.await();
+          Thread.sleep(Long.MAX_VALUE);
+        } catch (InterruptedException | BrokenBarrierException e) {
+          // do nothing
+        }
+      });
+
+      barrier.await();
+      try {
+        executor.execute(() -> {
+          // do nothing
+        });
+        fail("Should not allow more than 1 task");
+      } catch (Exception e) {
+        // as expected
+        assertEquals(e.getMessage(), "Tasks throttled due to high heap usage.");
+      }
+    } catch (BrokenBarrierException | InterruptedException e) {
+      throw new RuntimeException(e);
+    } finally {
+      executor.shutdownNow();
+    }
+  }
+}