The PR add the following improvements to OOM protection for Multi-Stage Engine.

Add a config parameter to turn off sampling for MSE

pinot.query.scheduler.accounting.enable.thread.sampling.mse.debug
Since the feature is new, a debug config has been added to turn off sampling for MSE when there are mixed workloads.

As part of this change, a new fn Tracing::sampleMSE has been added to not add an if condition in the sampling path. The sample function is in the critical path wrt performance and SSE will not be affected by the extra check.

Assign query id to workers

Worker threads are not assigned query ids. Refer:

Tracing.getThreadAccountant().createExecutionContext(null, taskId, taskType, threadExecutionContext);

Therefore, the resource used by these threads was not attributed to a query. Now, if there is a parentContext available, the query id is set. During aggregation, memory used by workers is also added to a query's aggregate statistics.

  String queryId = null;
      if (threadExecutionContext != null) {
        queryId = threadExecutionContext.getQueryId();
      } else {
        LOGGER.warn("Request ID not available. ParentContext not set for query worker thread.");
      }
      Tracing.getThreadAccountant()
          .createExecutionContext(queryId, taskId, taskType, threadExecutionContext);
    }

Instrument threads for Leaf Stage Operator.

Threads that executed Leaf Stage operators in MSE were not tracked. Now runners and workers are instrumented correctly.

Fix annotations

Nullable/NotNull annotations have been added. The most important is CPUMemThreadLevelAccountingObjects.::setThreadTaskStatus . Tracing::setupWorker can pass null for queryId.

tags: bugfix