Add OOM Protection Support for Multi-Stage Queries (apache#13598)

track cpu and memory usage in multi-stage queries if query resource usage tracking is enabled

Author: vrajat

pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/MultiStageBrokerRequestHandler.java

@@ -64,10 +64,12 @@
 import org.apache.pinot.query.runtime.MultiStageStatsTreeBuilder;
 import org.apache.pinot.query.runtime.plan.MultiStageQueryStats;
 import org.apache.pinot.query.service.dispatch.QueryDispatcher;
+import org.apache.pinot.spi.accounting.ThreadExecutionContext;
 import org.apache.pinot.spi.auth.TableAuthorizationResult;
 import org.apache.pinot.spi.env.PinotConfiguration;
 import org.apache.pinot.spi.exception.DatabaseConflictException;
 import org.apache.pinot.spi.trace.RequestContext;
+import org.apache.pinot.spi.trace.Tracing;
 import org.apache.pinot.spi.utils.CommonConstants;
 import org.apache.pinot.spi.utils.builder.TableNameBuilder;
 import org.apache.pinot.sql.parsers.SqlNodeAndOptions;
@@ -210,6 +212,8 @@ protected BrokerResponse handleRequest(long requestId, String query, @Nullable S
       return new BrokerResponseNative(QueryException.getException(QueryException.QUOTA_EXCEEDED_ERROR, errorMessage));
     }
 
+    Tracing.ThreadAccountantOps.setupRunner(String.valueOf(requestId), ThreadExecutionContext.TaskType.MSE);
+
     long executionStartTimeNs = System.nanoTime();
     QueryDispatcher.QueryResult queryResults;
     try {
@@ -228,6 +232,8 @@ protected BrokerResponse handleRequest(long requestId, String query, @Nullable S
       requestContext.setErrorCode(QueryException.QUERY_EXECUTION_ERROR_CODE);
       return new BrokerResponseNative(
           QueryException.getException(QueryException.QUERY_EXECUTION_ERROR, consolidatedMessage));
+    } finally {
+      Tracing.getThreadAccountant().clear();
     }
     long executionEndTimeNs = System.nanoTime();
     updatePhaseTimingForTables(tableNames, BrokerQueryPhase.QUERY_EXECUTION, executionEndTimeNs - executionStartTimeNs);

pinot-core/src/main/java/org/apache/pinot/core/accounting/CPUMemThreadLevelAccountingObjects.java

@@ -103,8 +103,15 @@ public int getTaskId() {
       return taskEntry == null ? -1 : taskEntry.getTaskId();
     }
 
-    public void setThreadTaskStatus(@Nonnull String queryId, int taskId, @Nonnull Thread anchorThread) {
-      _currentThreadTaskStatus.set(new TaskEntry(queryId, taskId, anchorThread));
+    @Override
+    public ThreadExecutionContext.TaskType getTaskType() {
+      TaskEntry taskEntry = _currentThreadTaskStatus.get();
+      return taskEntry == null ? ThreadExecutionContext.TaskType.UNKNOWN : taskEntry.getTaskType();
+    }
+
+    public void setThreadTaskStatus(@Nonnull String queryId, int taskId, ThreadExecutionContext.TaskType taskType,
+        @Nonnull Thread anchorThread) {
+      _currentThreadTaskStatus.set(new TaskEntry(queryId, taskId, taskType, anchorThread));
     }
   }
 
@@ -117,15 +124,17 @@ public static class TaskEntry implements ThreadExecutionContext {
     private final String _queryId;
     private final int _taskId;
     private final Thread _anchorThread;
+    private final TaskType _taskType;
 
     public boolean isAnchorThread() {
       return _taskId == CommonConstants.Accounting.ANCHOR_TASK_ID;
     }
 
-    public TaskEntry(String queryId, int taskId, Thread anchorThread) {
+    public TaskEntry(String queryId, int taskId, TaskType taskType, Thread anchorThread) {
       _queryId = queryId;
       _taskId = taskId;
       _anchorThread = anchorThread;
+      _taskType = taskType;
     }
 
     public String getQueryId() {
@@ -140,6 +149,11 @@ public Thread getAnchorThread() {
       return _anchorThread;
     }
 
+    @Override
+    public TaskType getTaskType() {
+      return _taskType;
+    }
+
     @Override
     public String toString() {
       return "TaskEntry{" + "_queryId='" + _queryId + '\'' + ", _taskId=" + _taskId + ", _rootThread=" + _anchorThread

pinot-core/src/main/java/org/apache/pinot/core/accounting/PerQueryCPUMemAccountantFactory.java

@@ -293,10 +293,10 @@ public void createExecutionContextInner(@Nullable String queryId, int taskId, @N
         // is anchor thread
         assert queryId != null;
         _threadLocalEntry.get().setThreadTaskStatus(queryId, CommonConstants.Accounting.ANCHOR_TASK_ID,
-            Thread.currentThread());
+            ThreadExecutionContext.TaskType.UNKNOWN, Thread.currentThread());
       } else {
         // not anchor thread
-        _threadLocalEntry.get().setThreadTaskStatus(parentContext.getQueryId(), taskId,
+        _threadLocalEntry.get().setThreadTaskStatus(parentContext.getQueryId(), taskId, parentContext.getTaskType(),
             parentContext.getAnchorThread());
       }
     }

pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/QueryRunner.java

@@ -50,6 +50,7 @@
 import org.apache.pinot.query.runtime.plan.pipeline.PipelineBreakerExecutor;
 import org.apache.pinot.query.runtime.plan.pipeline.PipelineBreakerResult;
 import org.apache.pinot.query.runtime.plan.server.ServerPlanRequestUtils;
+import org.apache.pinot.spi.accounting.ThreadExecutionContext;
 import org.apache.pinot.spi.env.PinotConfiguration;
 import org.apache.pinot.spi.utils.CommonConstants;
 import org.apache.pinot.spi.utils.CommonConstants.Broker.Request.QueryOptionKey;
@@ -152,7 +153,8 @@ public void shutDown() {
    * <p>This execution entry point should be asynchronously called by the request handler and caller should not wait
    * for results/exceptions.</p>
    */
-  public void processQuery(WorkerMetadata workerMetadata, StagePlan stagePlan, Map<String, String> requestMetadata) {
+  public void processQuery(WorkerMetadata workerMetadata, StagePlan stagePlan, Map<String, String> requestMetadata,
+      @Nullable ThreadExecutionContext parentContext) {
     long requestId = Long.parseLong(requestMetadata.get(CommonConstants.Query.Request.MetadataKeys.REQUEST_ID));
     long timeoutMs = Long.parseLong(requestMetadata.get(CommonConstants.Broker.Request.QueryOptionKey.TIMEOUT_MS));
     long deadlineMs = System.currentTimeMillis() + timeoutMs;
@@ -163,7 +165,7 @@ public void processQuery(WorkerMetadata workerMetadata, StagePlan stagePlan, Map
     // run pre-stage execution for all pipeline breakers
     PipelineBreakerResult pipelineBreakerResult =
         PipelineBreakerExecutor.executePipelineBreakers(_opChainScheduler, _mailboxService, workerMetadata, stagePlan,
-            opChainMetadata, requestId, deadlineMs);
+            opChainMetadata, requestId, deadlineMs, parentContext);
 
     // Send error block to all the receivers if pipeline breaker fails
     if (pipelineBreakerResult != null && pipelineBreakerResult.getErrorBlock() != null) {
@@ -196,7 +198,7 @@ public void processQuery(WorkerMetadata workerMetadata, StagePlan stagePlan, Map
     // run OpChain
     OpChainExecutionContext executionContext =
         new OpChainExecutionContext(_mailboxService, requestId, deadlineMs, opChainMetadata, stageMetadata,
-            workerMetadata, pipelineBreakerResult);
+            workerMetadata, pipelineBreakerResult, parentContext);
     OpChain opChain;
     if (workerMetadata.isLeafStageWorker()) {
       opChain = ServerPlanRequestUtils.compileLeafStage(executionContext, stagePlan, _helixManager, _serverMetrics,

pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/executor/OpChainSchedulerService.java

@@ -27,6 +27,9 @@
 import org.apache.pinot.query.runtime.blocks.TransferableBlock;
 import org.apache.pinot.query.runtime.operator.OpChain;
 import org.apache.pinot.query.runtime.operator.OpChainId;
+import org.apache.pinot.spi.accounting.ThreadExecutionContext;
+import org.apache.pinot.spi.accounting.ThreadResourceUsageProvider;
+import org.apache.pinot.spi.trace.Tracing;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -50,6 +53,10 @@ public void runJob() {
         TransferableBlock returnedErrorBlock = null;
         Throwable thrown = null;
         try {
+          ThreadResourceUsageProvider threadResourceUsageProvider = new ThreadResourceUsageProvider();
+          Tracing.ThreadAccountantOps.setupWorker(operatorChain.getId().getStageId(),
+              ThreadExecutionContext.TaskType.MSE, threadResourceUsageProvider,
+              operatorChain.getParentContext());
           LOGGER.trace("({}): Executing", operatorChain);
           TransferableBlock result = operatorChain.getRoot().nextBlock();
           while (!result.isEndOfStreamBlock()) {
@@ -76,6 +83,7 @@ public void runJob() {
           } else if (isFinished) {
             operatorChain.close();
           }
+          Tracing.ThreadAccountantOps.clear();
         }
       }
     });

pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/AggregateOperator.java

@@ -173,6 +173,7 @@ private TransferableBlock consumeGroupBy() {
     TransferableBlock block = _input.nextBlock();
     while (block.isDataBlock()) {
       _groupByExecutor.processBlock(block);
+      sampleAndCheckInterruption();
       block = _input.nextBlock();
     }
     return block;
@@ -187,6 +188,7 @@ private TransferableBlock consumeAggregation() {
     TransferableBlock block = _input.nextBlock();
     while (block.isDataBlock()) {
       _aggregationExecutor.processBlock(block);
+      sampleAndCheckInterruption();
       block = _input.nextBlock();
     }
     return block;

pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/HashJoinOperator.java

@@ -253,6 +253,7 @@ private void buildBroadcastHashTable()
         hashCollection.add(row);
       }
       _currentRowsInHashTable += container.size();
+      sampleAndCheckInterruption();
       rightBlock = _rightInput.nextBlock();
     }
     if (rightBlock.isErrorBlock()) {
@@ -297,6 +298,7 @@ private TransferableBlock buildJoinedDataBlock() throws ProcessingException {
       }
       assert leftBlock.isDataBlock();
       List<Object[]> rows = buildJoinedRows(leftBlock);
+      sampleAndCheckInterruption();
       if (!rows.isEmpty()) {
         return new TransferableBlock(rows, _resultSchema, DataBlock.Type.ROW);
       }

pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/MailboxReceiveOperator.java

@@ -60,6 +60,8 @@ protected TransferableBlock getNextBlock() {
     }
     if (block.isSuccessfulEndOfStreamBlock()) {
       updateEosBlock(block, _statMap);
+    } else if (block.isDataBlock()) {
+      sampleAndCheckInterruption();
     }
     return block;
   }

pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/MailboxSendOperator.java

@@ -143,6 +143,7 @@ protected TransferableBlock getNextBlock() {
           earlyTerminate();
         }
       }
+      sampleAndCheckInterruption();
       return block;
     } catch (QueryCancelledException e) {
       LOGGER.debug("Query was cancelled! for opChain: {}", _context.getId());

pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/MultiStageOperator.java

@@ -66,6 +66,15 @@ public MultiStageOperator(OpChainExecutionContext context) {
 
   public abstract void registerExecution(long time, int numRows);
 
+  // Samples resource usage of the operator. The operator should call this function for every block of data or
+  // assuming the block holds 10000 rows or more.
+  protected void sampleAndCheckInterruption() {
+    Tracing.ThreadAccountantOps.sample();
+    if (Tracing.ThreadAccountantOps.isInterrupted()) {
+      earlyTerminate();
+    }
+  }
+
   /**
    * Returns the next block from the operator. It should return non-empty data blocks followed by an end-of-stream (EOS)
    * block when all the data is processed, or an error block if an error occurred. After it returns EOS or error block,