Fix OOM Cancellation in MSE

[vrajat]

QueryServer in MSE submits work asynchronously.

OOM protection framework assumes that there is a runner thread which is waiting for all the worker threads. The framework cancels the runner thread if required. Worker threads check regularly if the anchor thread is canceled and abort if it is.

Since QueryServer submits asynchronously, the above assumptions break down. It is unclear if cancellation by OOM protection works for MSE. If it does, it is by luck.

[gortiz]

Probably instead of registering the query/thread at QueryServer.submitWorker, we should do that at OpChainSchedulerService.register. But that method is called more than once on the same server for the same query and even worker. In fact, it is called once for every tuple of (query, worker, pipeline breaker or main stage). Would that break another invariant?

[vrajat]

Is it OK if that thread is interrupted ?

[gortiz]

This diagram is a bird's-eye view of the thread model on a server when handling MSE requests.

The thread that

sequenceDiagram

  box Grpc thread

  participant GrpcServer

  end

  box Submission Thread
  participant QueryServer
  end

  box ExecutionThread 1
  participant QueryRunner
  participant PipelineBreakerExecutor
  end

  box ExecutionThread 2
  participant OpChainSchedulerService
  end 

  participant MultiStageOperator

  activate GrpcServer
  GrpcServer ->> QueryServer: submit
  activate QueryServer
  QueryServer ->> QueryRunner: processQueryAsync
  activate QueryRunner
  QueryRunner -->> QueryServer: return not waiting 
  QueryServer -->> GrpcServer: return 
  deactivate QueryServer
  deactivate GrpcServer
  QueryRunner ->> QueryRunner: processQueryBlocking
  QueryRunner ->> PipelineBreakerExecutor: executePipelineBreakers
  PipelineBreakerExecutor ->> MultiStageOperator: nextBlock
  Note right of MultiStageOperator: Starts new threads,<br/> but wait for them
  MultiStageOperator -->> PipelineBreakerExecutor: return
  PipelineBreakerExecutor -->> QueryRunner: return
  
  QueryRunner ->> OpChainSchedulerService: register
  activate OpChainSchedulerService
OpChainSchedulerService -->> QueryRunner: return not waiting 
  deactivate QueryRunner
  
  OpChainSchedulerService ->> MultiStageOperator: nextBlock
  Note right of MultiStageOperator: Starts new threads,<br/> but wait for them
  MultiStageOperator -->> OpChainSchedulerService: return

  deactivate OpChainSchedulerService

Loading

MultiStageOperators periodically call sampleAndCheckInterruption, which could interrupt the thread and modify the operator to return an error block.

[vrajat]

TIL, in SSE there are two actions taken to interrupt. The code in watcher task when it decides to cancel a query is:

          if (shouldKill) {
            maxUsageTuple._exceptionAtomicReference
                .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());
          }

_exceptionAtomicReference is a field in the anchor thread's thread local.

Then as InterruptedException propogates, exceptionAtomicReference != null is the check for termination vs cancel. Example code from InstanceResponseOperator.java

  protected BaseResultsBlock getCombinedResults() {
    try {
      prefetchAll();
      return _combineOperator.nextBlock();
    } catch (EarlyTerminationException e) {
      Exception killedErrorMsg = Tracing.getThreadAccountant().getErrorStatus();
      QueryErrorMessage errMsg = QueryErrorMessage.safeMsg(QueryErrorCode.QUERY_CANCELLATION,
          "Cancelled while combining results" + (killedErrorMsg == null ? StringUtils.EMPTY : " " + killedErrorMsg));
      return new ExceptionResultsBlock(errMsg);
    } finally {
      releaseAll();
    }
  }

IMO, the right way seems to be to propogate an extra piece of info on the reason for termination/cancel. Maybe we should not try to shoe-horn the SSE model onto MSE.