Examples with MemoryPoolAssignmentsRequest com.facebook.presto.memory.MemoryPoolAssignmentsRequest used on opensource projects

Example 1 with MemoryPoolAssignmentsRequest

use of com.facebook.presto.memory.MemoryPoolAssignmentsRequest in project presto by prestodb.

the class TestSqlTaskManager method testMultipleCoordinatorAssignments.

@Test
public void testMultipleCoordinatorAssignments() {
    try (SqlTaskManager sqlTaskManager = createSqlTaskManager(new TaskManagerConfig())) {
        TaskId taskId = TASK_ID;
        TaskInfo taskInfo = createTask(sqlTaskManager, taskId, createInitialEmptyOutputBuffers(PARTITIONED).withNoMoreBufferIds());
        assertEquals(taskInfo.getTaskStatus().getState(), TaskState.RUNNING);
        sqlTaskManager.updateMemoryPoolAssignments(new MemoryPoolAssignmentsRequest("coordinator1", 1, ImmutableList.of(new MemoryPoolAssignment(taskId.getQueryId(), RESERVED_POOL))));
        assertEquals(sqlTaskManager.getQueryContext(taskId.getQueryId()).getMemoryPool().getId(), RESERVED_POOL);
        // Update to reserved pool
        sqlTaskManager.updateMemoryPoolAssignments(new MemoryPoolAssignmentsRequest("coordinator1", 2, ImmutableList.of(new MemoryPoolAssignment(taskId.getQueryId(), GENERAL_POOL))));
        assertEquals(sqlTaskManager.getQueryContext(taskId.getQueryId()).getMemoryPool().getId(), GENERAL_POOL);
        // Prior versions are rejected
        sqlTaskManager.updateMemoryPoolAssignments(new MemoryPoolAssignmentsRequest("coordinator1", 1, ImmutableList.of(new MemoryPoolAssignment(taskId.getQueryId(), RESERVED_POOL))));
        assertEquals(sqlTaskManager.getQueryContext(taskId.getQueryId()).getMemoryPool().getId(), GENERAL_POOL);
        // Different coordinators are not rejected
        sqlTaskManager.updateMemoryPoolAssignments(new MemoryPoolAssignmentsRequest("coordinator2", 1, ImmutableList.of(new MemoryPoolAssignment(taskId.getQueryId(), RESERVED_POOL))));
        assertEquals(sqlTaskManager.getQueryContext(taskId.getQueryId()).getMemoryPool().getId(), RESERVED_POOL);
    }
}