private QueryContext createQueryContext( QueryId queryId, LocalMemoryManager localMemoryManager, NodeMemoryConfig nodeMemoryConfig, LocalSpillManager localSpillManager, GcMonitor gcMonitor, DataSize maxQueryUserMemoryPerNode, DataSize maxQueryTotalMemoryPerNode, DataSize maxQuerySpillPerNode) { return new QueryContext( queryId, maxQueryUserMemoryPerNode, maxQueryTotalMemoryPerNode, localMemoryManager.getGeneralPool(), gcMonitor, taskNotificationExecutor, driverYieldExecutor, maxQuerySpillPerNode, localSpillManager.getSpillSpaceTracker()); }
private static TaskContext createTaskContext(QueryContext queryContext, Session session, TaskStateMachine taskStateMachine) { return queryContext.addTaskContext( taskStateMachine, session, true, true, OptionalInt.empty()); }
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(); }
SpillSpaceTracker spillSpaceTracker = new SpillSpaceTracker(new DataSize(1, GIGABYTE)); TaskContext taskContext = new QueryContext( new QueryId("test"), new DataSize(256, MEGABYTE), new DataSize(256, MEGABYTE), spillSpaceTracker) .addTaskContext(new TaskStateMachine(new TaskId("query", 0, 0), localQueryRunner.getExecutor()), session, false,
QueryContext queryContext = new QueryContext( new QueryId("query"), new DataSize(10, BYTE), queryContext.getQueryMemoryContext().initializeLocalMemoryContexts("test"); LocalMemoryContext userMemoryContext = queryContext.getQueryMemoryContext().localUserMemoryContext(); LocalMemoryContext revocableMemoryContext = queryContext.getQueryMemoryContext().localRevocableMemoryContext(); assertTrue(userMemoryContext.setBytes(3).isDone()); assertTrue(revocableMemoryContext.setBytes(5).isDone()); queryContext.setMemoryPool(reservedPool);
@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())); } } }
@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(); }
@Override public R visitQueryContext(QueryContext queryContext, C visitContext) { return mergeResults(queryContext.acceptChildren(this, visitContext)); }
private synchronized ListenableFuture<?> updateSystemMemory(String allocationTag, long delta) { // We call memoryPool.getQueryMemoryReservation(queryId) instead of calling queryMemoryContext.getUserMemory() to // calculate the total memory size. // // Calling the latter can result in a deadlock: // * A thread doing a user allocation will acquire locks in this order: // 1. monitor of queryMemoryContext.userAggregateMemoryContext // 2. monitor of this (QueryContext) // * The current thread doing a system allocation will acquire locks in this order: // 1. monitor of this (QueryContext) // 2. monitor of queryMemoryContext.userAggregateMemoryContext // Deadlock is possible for concurrent user and system allocations when updateSystemMemory()/updateUserMemory // calls queryMemoryContext.getUserMemory()/queryMemoryContext.getSystemMemory(), respectively. For concurrent // allocations of the same type (e.g., tryUpdateUserMemory/updateUserMemory) it is not possible as they share // the same RootAggregatedMemoryContext instance, and one of the threads will be blocked on the monitor of that // RootAggregatedMemoryContext instance even before calling the QueryContext methods (the monitors of // RootAggregatedMemoryContext instance and this will be acquired in the same order). long totalMemory = memoryPool.getQueryMemoryReservation(queryId); if (delta >= 0) { enforceTotalMemoryLimit(totalMemory, delta, maxTotalMemory); return memoryPool.reserve(queryId, allocationTag, delta); } memoryPool.free(queryId, allocationTag, -delta); return NOT_BLOCKED; }
@BeforeMethod public void setUpTest() { memoryPool = new MemoryPool(new MemoryPoolId("test"), memoryPoolSize); queryContext = new QueryContext( new QueryId("test_query"), queryMaxMemory, queryMaxTotalMemory, memoryPool, new TestingGcMonitor(), notificationExecutor, yieldExecutor, queryMaxSpillSize, spillSpaceTracker); taskContext = queryContext.addTaskContext( new TaskStateMachine(new TaskId("query", 0, 0), notificationExecutor), testSessionBuilder().build(), true, true, OptionalInt.empty()); pipelineContext = taskContext.addPipelineContext(0, true, true, false); driverContext = pipelineContext.addDriverContext(); operatorContext = driverContext.addOperatorContext(1, new PlanNodeId("a"), "test-operator"); }
@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())); } } }
@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(); }
@Override public R visitQueryContext(QueryContext queryContext, C visitContext) { return mergeResults(queryContext.acceptChildren(this, visitContext)); }
private TaskContext newTestingTaskContext(ScheduledExecutorService taskNotificationExecutor, ScheduledExecutorService driverYieldExecutor, TaskStateMachine taskStateMachine) { QueryContext queryContext = new QueryContext( new QueryId("queryid"), new DataSize(1, MEGABYTE), new DataSize(2, MEGABYTE), new MemoryPool(new MemoryPoolId("test"), new DataSize(1, GIGABYTE)), new TestingGcMonitor(), taskNotificationExecutor, driverYieldExecutor, new DataSize(1, MEGABYTE), new SpillSpaceTracker(new DataSize(1, GIGABYTE))); return queryContext.addTaskContext(taskStateMachine, TEST_SESSION, false, false, OptionalInt.empty()); }
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()); }
public TaskContext build() { MemoryPool memoryPool = new MemoryPool(new MemoryPoolId("test"), memoryPoolSize); SpillSpaceTracker spillSpaceTracker = new SpillSpaceTracker(maxSpillSize); QueryContext queryContext = new QueryContext( queryId, queryMaxMemory, queryMaxTotalMemory, memoryPool, GC_MONITOR, notificationExecutor, yieldExecutor, queryMaxSpillSize, spillSpaceTracker); return createTaskContext(queryContext, session, taskStateMachine); } }
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(); }
@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); }
QueryContext queryContext = new QueryContext(taskId.getQueryId(), new DataSize(1, MEGABYTE), new DataSize(2, MEGABYTE), new DataSize(1, MEGABYTE), spillSpaceTracker); this.taskContext = queryContext.addTaskContext(taskStateMachine, TEST_SESSION, true, true, totalPartitions);
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()); }