private long getMemoryAlreadyBeingRevoked(Collection<SqlTask> sqlTasks, MemoryPool memoryPool) { return sqlTasks.stream() .filter(task -> task.getTaskStatus().getState() == TaskState.RUNNING) .filter(task -> task.getQueryContext().getMemoryPool() == memoryPool) .mapToLong(task -> task.getQueryContext().accept(new TraversingQueryContextVisitor<Void, Long>() { @Override public Long visitOperatorContext(OperatorContext operatorContext, Void context) { if (operatorContext.isMemoryRevokingRequested()) { return operatorContext.getReservedRevocableBytes(); } return 0L; } @Override public Long mergeResults(List<Long> childrenResults) { return childrenResults.stream() .mapToLong(i -> i).sum(); } }, null)) .sum(); }
sqlTasks.stream() .filter(task -> task.getTaskStatus().getState() == TaskState.RUNNING) .filter(task -> task.getQueryContext().getMemoryPool() == memoryPool) .sorted(ORDER_BY_CREATE_TIME) .forEach(task -> task.getQueryContext().accept(new VoidTraversingQueryContextVisitor<AtomicLong>()
@BeforeMethod public void setUp() { executor = newCachedThreadPool(daemonThreadsNamed("test-executor-%s")); scheduledExecutor = newScheduledThreadPool(2, daemonThreadsNamed("test-scheduledExecutor-%s")); TaskContext taskContext = TestingTaskContext.builder(executor, scheduledExecutor, TEST_SESSION) .setQueryMaxMemory(DataSize.valueOf("100MB")) .setMemoryPoolSize(DataSize.valueOf("10B")) .setQueryId(QUERY_ID) .build(); memoryPool = taskContext.getQueryContext().getMemoryPool(); driverContext = taskContext .addPipelineContext(0, true, true, false) .addDriverContext(); }