private TaskInfo createTask(SqlTaskManager sqlTaskManager, TaskId taskId, OutputBuffers outputBuffers) { sqlTaskManager.getQueryContext(taskId.getQueryId()) .addTaskContext(new TaskStateMachine(taskId, directExecutor()), testSessionBuilder().build(), false, false, OptionalInt.empty()); return sqlTaskManager.updateTask(TEST_SESSION, taskId, Optional.of(PLAN_FRAGMENT), ImmutableList.of(), outputBuffers, OptionalInt.empty()); }
@Override public synchronized void updateMemoryPoolAssignments(MemoryPoolAssignmentsRequest assignments) { if (coordinatorId != null && coordinatorId.equals(assignments.getCoordinatorId()) && assignments.getVersion() <= currentMemoryPoolAssignmentVersion) { return; } currentMemoryPoolAssignmentVersion = assignments.getVersion(); if (coordinatorId != null && !coordinatorId.equals(assignments.getCoordinatorId())) { log.warn("Switching coordinator affinity from " + coordinatorId + " to " + assignments.getCoordinatorId()); } coordinatorId = assignments.getCoordinatorId(); for (MemoryPoolAssignment assignment : assignments.getAssignments()) { if (assignment.getPoolId().equals(GENERAL_POOL)) { queryContexts.getUnchecked(assignment.getQueryId()).setMemoryPool(localMemoryManager.getGeneralPool()); } else if (assignment.getPoolId().equals(RESERVED_POOL)) { MemoryPool reservedPool = localMemoryManager.getReservedPool() .orElseThrow(() -> new IllegalArgumentException(format("Cannot move %s to the reserved pool as the reserved pool is not enabled", assignment.getQueryId()))); queryContexts.getUnchecked(assignment.getQueryId()).setMemoryPool(reservedPool); } else { new IllegalArgumentException(format("Cannot move %s to %s as the target memory pool id is invalid", assignment.getQueryId(), assignment.getPoolId())); } } }
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>()
MemoryPool systemMemoryPool = new MemoryPool(new MemoryPoolId("testSystem"), new DataSize(1, GIGABYTE)); TaskContext taskContext = new QueryContext(new QueryId("test"), new DataSize(256, MEGABYTE), memoryPool, systemMemoryPool, executor) .addTaskContext(new TaskStateMachine(new TaskId("query", "stage", "task"), executor), session, new DataSize(1, MEGABYTE),
@Override public QueryContext load(QueryId key) throws Exception { return new QueryContext(key, maxQueryMemoryPerNode, localMemoryManager.getPool(LocalMemoryManager.GENERAL_POOL), localMemoryManager.getPool(LocalMemoryManager.SYSTEM_POOL), taskNotificationExecutor); } });
@Override public TaskInfo updateTask(Session session, TaskId taskId, Optional<PlanFragment> fragment, List<TaskSource> sources, OutputBuffers outputBuffers, OptionalInt totalPartitions) { requireNonNull(session, "session is null"); requireNonNull(taskId, "taskId is null"); requireNonNull(fragment, "fragment is null"); requireNonNull(sources, "sources is null"); requireNonNull(outputBuffers, "outputBuffers is null"); if (resourceOvercommit(session)) { // TODO: This should have been done when the QueryContext was created. However, the session isn't available at that point. queryContexts.getUnchecked(taskId.getQueryId()).setResourceOvercommit(); } SqlTask sqlTask = tasks.getUnchecked(taskId); sqlTask.recordHeartbeat(); return sqlTask.updateTask(session, fragment, sources, outputBuffers, totalPartitions); }
@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(); }
public MockRemoteTask(TaskId taskId, PlanFragment fragment, String nodeId, int partition, Executor executor, Multimap<PlanNodeId, Split> initialSplits, PartitionedSplitCountTracker partitionedSplitCountTracker) { this.taskStateMachine = new TaskStateMachine(requireNonNull(taskId, "taskId is null"), requireNonNull(executor, "executor is null")); MemoryPool memoryPool = new MemoryPool(new MemoryPoolId("test"), new DataSize(1, GIGABYTE)); MemoryPool memorySystemPool = new MemoryPool(new MemoryPoolId("testSystem"), new DataSize(1, GIGABYTE)); this.taskContext = new QueryContext(taskId.getQueryId(), new DataSize(1, MEGABYTE), memoryPool, memorySystemPool, executor).addTaskContext(taskStateMachine, TEST_SESSION, new DataSize(1, MEGABYTE), true, true); this.location = URI.create("fake://task/" + taskId); this.sharedBuffer = new SharedBuffer(taskId, TASK_INSTANCE_ID, executor, requireNonNull(new DataSize(1, BYTE), "maxBufferSize is null")); this.fragment = requireNonNull(fragment, "fragment is null"); this.nodeId = requireNonNull(nodeId, "nodeId is null"); checkArgument(partition >= 0, "partition is negative"); this.partition = partition; splits.putAll(initialSplits); this.partitionedSplitCountTracker = requireNonNull(partitionedSplitCountTracker, "partitionedSplitCountTracker is null"); partitionedSplitCountTracker.setPartitionedSplitCount(getPartitionedSplitCount()); }
public static TaskContext createTaskContext(Executor executor, Session session, DataSize maxMemory, DataSize preallocated) { MemoryPool memoryPool = new MemoryPool(new MemoryPoolId("test"), new DataSize(1, GIGABYTE)); MemoryPool systemMemoryPool = new MemoryPool(new MemoryPoolId("testSystem"), new DataSize(1, GIGABYTE)); QueryContext queryContext = new QueryContext(new QueryId("test_query"), maxMemory, memoryPool, systemMemoryPool, executor); return createTaskContext(queryContext, executor, session, preallocated); }
@Override public TaskInfo updateTask(Session session, TaskId taskId, Optional<PlanFragment> fragment, List<TaskSource> sources, OutputBuffers outputBuffers) { requireNonNull(session, "session is null"); requireNonNull(taskId, "taskId is null"); requireNonNull(fragment, "fragment is null"); requireNonNull(sources, "sources is null"); requireNonNull(outputBuffers, "outputBuffers is null"); if (resourceOvercommit(session)) { // TODO: This should have been done when the QueryContext was created. However, the session isn't available at that point. queryContexts.getUnchecked(taskId.getQueryId()).setResourceOvercommit(); } SqlTask sqlTask = tasks.getUnchecked(taskId); sqlTask.recordHeartbeat(); return sqlTask.updateTask(session, fragment, sources, outputBuffers); }
public SqlTaskExecution create(Session session, QueryContext queryContext, TaskStateMachine taskStateMachine, OutputBuffer outputBuffer, PlanFragment fragment, List<TaskSource> sources, OptionalInt totalPartitions) TaskContext taskContext = queryContext.addTaskContext( taskStateMachine, session,
public SqlTask createInitialTask() { TaskId taskId = new TaskId("query", "stage", "task" + nextTaskId.incrementAndGet()); URI location = URI.create("fake://task/" + taskId); return new SqlTask( taskId, location, new QueryContext(new QueryId("query"), new DataSize(1, MEGABYTE), new MemoryPool(new MemoryPoolId("test"), new DataSize(1, GIGABYTE)), new MemoryPool(new MemoryPoolId("testSystem"), new DataSize(1, GIGABYTE)), taskNotificationExecutor), sqlTaskExecutionFactory, taskNotificationExecutor, Functions.<SqlTask>identity(), new DataSize(32, MEGABYTE)); } }
@Test public void testCountAlreadyRevokedMemoryWithinAPool() throws Exception { // Given SqlTask sqlTask1 = newSqlTask(); MemoryPool anotherMemoryPool = new MemoryPool(new MemoryPoolId("test"), new DataSize(10, BYTE)); sqlTask1.getQueryContext().setMemoryPool(anotherMemoryPool); OperatorContext operatorContext1 = createContexts(sqlTask1); SqlTask sqlTask2 = newSqlTask(); OperatorContext operatorContext2 = createContexts(sqlTask2); List<SqlTask> tasks = ImmutableList.of(sqlTask1, sqlTask2); MemoryRevokingScheduler scheduler = new MemoryRevokingScheduler(asList(memoryPool, anotherMemoryPool), () -> tasks, executor, 1.0, 1.0); allOperatorContexts = ImmutableSet.of(operatorContext1, operatorContext2); /* * sqlTask1 fills its pool */ operatorContext1.localRevocableMemoryContext().setBytes(12); requestMemoryRevoking(scheduler); assertMemoryRevokingRequestedFor(operatorContext1); /* * When sqlTask2 fills its pool */ operatorContext2.localRevocableMemoryContext().setBytes(12); requestMemoryRevoking(scheduler); /* * Then sqlTask2 should be asked to revoke its memory too */ assertMemoryRevokingRequestedFor(operatorContext1, operatorContext2); }
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(); }
private OperatorContext createContexts(SqlTask sqlTask) { TaskContext taskContext = sqlTask.getQueryContext().addTaskContext(new TaskStateMachine(new TaskId("q", 1, 1), executor), session, false, false, OptionalInt.empty()); PipelineContext pipelineContext = taskContext.addPipelineContext(0, false, false, false); DriverContext driverContext = pipelineContext.addDriverContext(); OperatorContext operatorContext = driverContext.addOperatorContext(1, new PlanNodeId("na"), "na"); return operatorContext; }
MemoryPool systemPool = new MemoryPool(new MemoryPoolId("testSystem"), new DataSize(10, MEGABYTE)); QueryContext queryContext = new QueryContext(new QueryId("query"), new DataSize(10, MEGABYTE), pool, systemPool, localQueryRunner.getExecutor()); LocalQueryRunner.MaterializedOutputFactory outputFactory = new LocalQueryRunner.MaterializedOutputFactory(); TaskContext taskContext = createTaskContext(queryContext, localQueryRunner.getExecutor(), session, new DataSize(0, BYTE));
@Override public synchronized void updateMemoryPoolAssignments(MemoryPoolAssignmentsRequest assignments) { if (coordinatorId != null && coordinatorId.equals(assignments.getCoordinatorId()) && assignments.getVersion() <= currentMemoryPoolAssignmentVersion) { return; } currentMemoryPoolAssignmentVersion = assignments.getVersion(); if (coordinatorId != null && !coordinatorId.equals(assignments.getCoordinatorId())) { log.warn("Switching coordinator affinity from " + coordinatorId + " to " + assignments.getCoordinatorId()); } coordinatorId = assignments.getCoordinatorId(); for (MemoryPoolAssignment assignment : assignments.getAssignments()) { queryContexts.getUnchecked(assignment.getQueryId()).setMemoryPool(localMemoryManager.getPool(assignment.getPoolId())); } }
SqlTask sqlTask2 = newSqlTask(); TaskContext taskContext1 = sqlTask1.getQueryContext().addTaskContext(new TaskStateMachine(new TaskId("q1", 1, 1), executor), session, false, false, OptionalInt.empty()); PipelineContext pipelineContext11 = taskContext1.addPipelineContext(0, false, false, false); DriverContext driverContext111 = pipelineContext11.addDriverContext(); OperatorContext operatorContext3 = driverContext112.addOperatorContext(3, new PlanNodeId("na"), "na"); TaskContext taskContext2 = sqlTask2.getQueryContext().addTaskContext(new TaskStateMachine(new TaskId("q2", 1, 1), executor), session, false, false, OptionalInt.empty()); PipelineContext pipelineContext21 = taskContext2.addPipelineContext(1, false, false, false); DriverContext driverContext211 = pipelineContext21.addDriverContext();
public SqlTaskExecution create(Session session, QueryContext queryContext, TaskStateMachine taskStateMachine, SharedBuffer sharedBuffer, PlanFragment fragment, List<TaskSource> sources) { boolean verboseStats = getVerboseStats(session); TaskContext taskContext = queryContext.addTaskContext( taskStateMachine, session, requireNonNull(operatorPreAllocatedMemory, "operatorPreAllocatedMemory is null"), verboseStats, cpuTimerEnabled); return createSqlTaskExecution( taskStateMachine, taskContext, sharedBuffer, fragment, sources, planner, taskExecutor, taskNotificationExecutor, queryMonitor); }
public static TaskContext createTaskContext(QueryContext queryContext, Executor executor, Session session, DataSize preallocated) { return queryContext.addTaskContext( new TaskStateMachine(new TaskId("query", "stage", "task"), checkNotSameThreadExecutor(executor, "executor is null")), session, preallocated, true, true); } }