Examples with TaskContext com.facebook.presto.operator.TaskContext used on opensource projects

Example 21 with TaskContext

use of com.facebook.presto.operator.TaskContext in project presto by prestodb.

the class MemoryRevokingScheduler method revokeQueryMemory.

private void revokeQueryMemory(QueryContext queryContext, long maxTotalMemory) {
    QueryId queryId = queryContext.getQueryId();
    MemoryPool memoryPool = queryContext.getMemoryPool();
    // get a fresh value for queryTotalMemory in case it's changed (e.g. by a previous revocation request)
    long queryTotalMemory = getTotalQueryMemoryReservation(queryId, memoryPool);
    // order tasks by decreasing revocableMemory so that we don't spill more tasks than needed
    SortedMap<Long, TaskContext> queryTaskContextsMap = new TreeMap<>(Comparator.reverseOrder());
    queryContext.getAllTaskContexts().forEach(taskContext -> queryTaskContextsMap.put(taskContext.getTaskMemoryContext().getRevocableMemory(), taskContext));
    AtomicLong remainingBytesToRevoke = new AtomicLong(queryTotalMemory - maxTotalMemory);
    Collection<TaskContext> queryTaskContexts = queryTaskContextsMap.values();
    remainingBytesToRevoke.addAndGet(-MemoryRevokingSchedulerUtils.getMemoryAlreadyBeingRevoked(queryTaskContexts, remainingBytesToRevoke.get()));
    for (TaskContext taskContext : queryTaskContexts) {
        if (remainingBytesToRevoke.get() <= 0) {
            break;
        }
        taskContext.accept(new VoidTraversingQueryContextVisitor<AtomicLong>() {

            @Override
            public Void visitOperatorContext(OperatorContext operatorContext, AtomicLong remainingBytesToRevoke) {
                if (remainingBytesToRevoke.get() > 0) {
                    long revokedBytes = operatorContext.requestMemoryRevoking();
                    if (revokedBytes > 0) {
                        remainingBytesToRevoke.addAndGet(-revokedBytes);
                        log.debug("taskId=%s: requested revoking %s; remaining %s", taskContext.getTaskId(), revokedBytes, remainingBytesToRevoke);
                    }
                }
                return null;
            }
        }, remainingBytesToRevoke);
    }
}

Example 22 with TaskContext

use of com.facebook.presto.operator.TaskContext in project presto by prestodb.

the class MemoryRevokingSchedulerUtils method getMemoryAlreadyBeingRevoked.

public static long getMemoryAlreadyBeingRevoked(Collection<TaskContext> taskContexts, long targetRevokingLimit) {
    TraversingQueryContextVisitor<Void, Long> visitor = 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();
        }
    };
    long currentRevoking = 0;
    for (TaskContext taskContext : taskContexts) {
        currentRevoking += taskContext.accept(visitor, null);
        if (currentRevoking > targetRevokingLimit) {
            // Return early, target value exceeded and revoking will not occur
            return currentRevoking;
        }
    }
    return currentRevoking;
}

Example 23 with TaskContext

use of com.facebook.presto.operator.TaskContext in project presto by prestodb.

the class TestMemoryRevokingScheduler method createContexts.

private OperatorContext createContexts(SqlTask sqlTask) {
    TaskContext taskContext = getOrCreateTaskContext(sqlTask);
    PipelineContext pipelineContext = taskContext.addPipelineContext(0, false, false, false);
    DriverContext driverContext = pipelineContext.addDriverContext();
    return driverContext.addOperatorContext(1, new PlanNodeId("na"), "na");
}

Example 24 with TaskContext

use of com.facebook.presto.operator.TaskContext in project presto by prestodb.

the class TestMemoryRevokingScheduler method createTestingOperatorContexts.

private TestOperatorContext createTestingOperatorContexts(SqlTask sqlTask, String operatorName) {
    // update task to update underlying taskHolderReference with taskExecution + create a new taskContext
    sqlTask.updateTask(TEST_SESSION, Optional.of(PLAN_FRAGMENT), ImmutableList.of(new TaskSource(TABLE_SCAN_NODE_ID, ImmutableSet.of(SPLIT), false)), createInitialEmptyOutputBuffers(PARTITIONED).withBuffer(OUT, 0).withNoMoreBufferIds(), Optional.of(new TableWriteInfo(Optional.empty(), Optional.empty(), Optional.empty())));
    // use implicitly created task context from updateTask. It should be the only task in this QueryContext's tasks
    TaskContext taskContext = sqlTask.getQueryContext().getTaskContextByTaskId(sqlTask.getTaskId());
    PipelineContext pipelineContext = taskContext.addPipelineContext(0, false, false, false);
    DriverContext driverContext = pipelineContext.addDriverContext();
    TestOperatorContext testOperatorContext = new TestOperatorContext(1, new PlanNodeId("na"), "na", driverContext, singleThreadedExecutor, driverContext.getDriverMemoryContext().newMemoryTrackingContext(), operatorName);
    driverContext.addOperatorContext(testOperatorContext);
    return testOperatorContext;
}

Example 25 with TaskContext

use of com.facebook.presto.operator.TaskContext in project presto by prestodb.

the class LocalQueryRunner method executeInternal.

private MaterializedResultWithPlan executeInternal(Session session, @Language("SQL") String sql, WarningCollector warningCollector) {
    lock.readLock().lock();
    try (Closer closer = Closer.create()) {
        AtomicReference<MaterializedResult.Builder> builder = new AtomicReference<>();
        PageConsumerOutputFactory outputFactory = new PageConsumerOutputFactory(types -> {
            builder.compareAndSet(null, MaterializedResult.resultBuilder(session, types));
            return builder.get()::page;
        });
        Plan plan = createPlan(session, sql, warningCollector);
        TaskContext taskContext = TestingTaskContext.builder(notificationExecutor, yieldExecutor, session).setMaxSpillSize(nodeSpillConfig.getMaxSpillPerNode()).setQueryMaxSpillSize(nodeSpillConfig.getQueryMaxSpillPerNode()).setQueryMaxTotalMemory(getQueryMaxTotalMemoryPerNode(session)).setTaskPlan(plan.getRoot()).build();
        taskContext.getQueryContext().setVerboseExceededMemoryLimitErrorsEnabled(isVerboseExceededMemoryLimitErrorsEnabled(session));
        taskContext.getQueryContext().setHeapDumpOnExceededMemoryLimitEnabled(isHeapDumpOnExceededMemoryLimitEnabled(session));
        String heapDumpFilePath = Paths.get(getHeapDumpFileDirectory(session), format("%s_%s.hprof", session.getQueryId().getId(), taskContext.getTaskId().getStageExecutionId().getStageId().getId())).toString();
        taskContext.getQueryContext().setHeapDumpFilePath(heapDumpFilePath);
        List<Driver> drivers = createDrivers(session, plan, outputFactory, taskContext);
        drivers.forEach(closer::register);
        boolean done = false;
        while (!done) {
            boolean processed = false;
            for (Driver driver : drivers) {
                if (alwaysRevokeMemory) {
                    driver.getDriverContext().getOperatorContexts().stream().filter(operatorContext -> operatorContext.getOperatorStats().getRevocableMemoryReservation().getValue() > 0).forEach(OperatorContext::requestMemoryRevoking);
                }
                if (!driver.isFinished()) {
                    driver.process();
                    processed = true;
                }
            }
            done = !processed;
        }
        verify(builder.get() != null, "Output operator was not created");
        return new MaterializedResultWithPlan(builder.get().build(), plan);
    } catch (IOException e) {
        throw new UncheckedIOException(e);
    } finally {
        lock.readLock().unlock();
    }
}