@Test public void testRemoveOldTasks() throws Exception { try (SqlTaskManager sqlTaskManager = createSqlTaskManager(new TaskManagerConfig().setInfoMaxAge(new Duration(5, TimeUnit.MILLISECONDS)))) { TaskId taskId = TASK_ID; TaskInfo taskInfo = createTask(sqlTaskManager, taskId, createInitialEmptyOutputBuffers(PARTITIONED).withBuffer(OUT, 0).withNoMoreBufferIds()); assertEquals(taskInfo.getTaskStatus().getState(), TaskState.RUNNING); taskInfo = sqlTaskManager.cancelTask(taskId); assertEquals(taskInfo.getTaskStatus().getState(), TaskState.CANCELED); taskInfo = sqlTaskManager.getTaskInfo(taskId); assertEquals(taskInfo.getTaskStatus().getState(), TaskState.CANCELED); Thread.sleep(100); sqlTaskManager.removeOldTasks(); for (TaskInfo info : sqlTaskManager.getAllTaskInfo()) { assertNotEquals(info.getTaskStatus().getTaskId(), taskId); } } }
@Test public void testAbort() { try (SqlTaskManager sqlTaskManager = createSqlTaskManager(new TaskManagerConfig())) { TaskId taskId = TASK_ID; TaskInfo taskInfo = createTask(sqlTaskManager, taskId, createInitialEmptyOutputBuffers(PARTITIONED).withBuffer(OUT, 0).withNoMoreBufferIds()); assertEquals(taskInfo.getTaskStatus().getState(), TaskState.RUNNING); assertNull(taskInfo.getStats().getEndTime()); taskInfo = sqlTaskManager.getTaskInfo(taskId); assertEquals(taskInfo.getTaskStatus().getState(), TaskState.RUNNING); assertNull(taskInfo.getStats().getEndTime()); taskInfo = sqlTaskManager.abortTask(taskId); assertEquals(taskInfo.getTaskStatus().getState(), TaskState.ABORTED); assertNotNull(taskInfo.getStats().getEndTime()); taskInfo = sqlTaskManager.getTaskInfo(taskId); assertEquals(taskInfo.getTaskStatus().getState(), TaskState.ABORTED); assertNotNull(taskInfo.getStats().getEndTime()); } }
@Test public void testSimpleQuery() throws Exception { try (SqlTaskManager sqlTaskManager = createSqlTaskManager(new TaskManagerConfig())) { TaskId taskId = TASK_ID; TaskInfo taskInfo = createTask(sqlTaskManager, taskId, ImmutableSet.of(SPLIT), createInitialEmptyOutputBuffers(PARTITIONED).withBuffer(OUT, 0).withNoMoreBufferIds()); assertEquals(taskInfo.getTaskStatus().getState(), TaskState.RUNNING); taskInfo = sqlTaskManager.getTaskInfo(taskId); assertEquals(taskInfo.getTaskStatus().getState(), TaskState.RUNNING); BufferResult results = sqlTaskManager.getTaskResults(taskId, OUT, 0, new DataSize(1, Unit.MEGABYTE)).get(); assertEquals(results.isBufferComplete(), false); assertEquals(results.getSerializedPages().size(), 1); assertEquals(results.getSerializedPages().get(0).getPositionCount(), 1); for (boolean moreResults = true; moreResults; moreResults = !results.isBufferComplete()) { results = sqlTaskManager.getTaskResults(taskId, OUT, results.getToken() + results.getSerializedPages().size(), new DataSize(1, Unit.MEGABYTE)).get(); } assertEquals(results.isBufferComplete(), true); assertEquals(results.getSerializedPages().size(), 0); // complete the task by calling abort on it TaskInfo info = sqlTaskManager.abortTaskResults(taskId, OUT); assertEquals(info.getOutputBuffers().getState(), BufferState.FINISHED); taskInfo = sqlTaskManager.getTaskInfo(taskId, taskInfo.getTaskStatus().getState()).get(1, TimeUnit.SECONDS); assertEquals(taskInfo.getTaskStatus().getState(), TaskState.FINISHED); taskInfo = sqlTaskManager.getTaskInfo(taskId); assertEquals(taskInfo.getTaskStatus().getState(), TaskState.FINISHED); } }
@PostConstruct public void start() { taskManagementExecutor.scheduleWithFixedDelay(() -> { try { removeOldTasks(); } catch (Throwable e) { log.warn(e, "Error removing old tasks"); } try { failAbandonedTasks(); } catch (Throwable e) { log.warn(e, "Error canceling abandoned tasks"); } }, 200, 200, TimeUnit.MILLISECONDS); taskManagementExecutor.scheduleWithFixedDelay(() -> { try { updateStats(); } catch (Throwable e) { log.warn(e, "Error updating stats"); } }, 0, 1, TimeUnit.SECONDS); }
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()); }
@Test public void testAbortResults() throws Exception { try (SqlTaskManager sqlTaskManager = createSqlTaskManager(new TaskManagerConfig())) { TaskId taskId = TASK_ID; TaskInfo taskInfo = createTask(sqlTaskManager, taskId, ImmutableSet.of(SPLIT), createInitialEmptyOutputBuffers(PARTITIONED).withBuffer(OUT, 0).withNoMoreBufferIds()); assertEquals(taskInfo.getTaskStatus().getState(), TaskState.RUNNING); taskInfo = sqlTaskManager.getTaskInfo(taskId); assertEquals(taskInfo.getTaskStatus().getState(), TaskState.RUNNING); sqlTaskManager.abortTaskResults(taskId, OUT); taskInfo = sqlTaskManager.getTaskInfo(taskId, taskInfo.getTaskStatus().getState()).get(1, TimeUnit.SECONDS); assertEquals(taskInfo.getTaskStatus().getState(), TaskState.FINISHED); taskInfo = sqlTaskManager.getTaskInfo(taskId); assertEquals(taskInfo.getTaskStatus().getState(), TaskState.FINISHED); } }
@Test public void testSimpleQuery() throws Exception { try (SqlTaskManager sqlTaskManager = createSqlTaskManager(new TaskManagerConfig())) { TaskId taskId = TASK_ID; TaskInfo taskInfo = sqlTaskManager.updateTask(TEST_SESSION, taskId, Optional.of(PLAN_FRAGMENT), ImmutableList.of(new TaskSource(TABLE_SCAN_NODE_ID, ImmutableSet.of(SPLIT), true)), INITIAL_EMPTY_OUTPUT_BUFFERS.withBuffer(OUT, 0).withNoMoreBufferIds()); assertEquals(taskInfo.getState(), TaskState.RUNNING); taskInfo = sqlTaskManager.getTaskInfo(taskId); assertEquals(taskInfo.getState(), TaskState.RUNNING); BufferResult results = sqlTaskManager.getTaskResults(taskId, OUT, 0, new DataSize(1, Unit.MEGABYTE)).get(); assertEquals(results.isBufferComplete(), false); assertEquals(results.getPages().size(), 1); assertEquals(results.getPages().get(0).getPositionCount(), 1); results = sqlTaskManager.getTaskResults(taskId, OUT, results.getToken() + results.getPages().size(), new DataSize(1, Unit.MEGABYTE)).get(); assertEquals(results.isBufferComplete(), true); assertEquals(results.getPages().size(), 0); // complete the task by calling abort on it TaskInfo info = sqlTaskManager.abortTaskResults(taskId, OUT); assertEquals(info.getOutputBuffers().getState(), SharedBuffer.BufferState.FINISHED); taskInfo = sqlTaskManager.getTaskInfo(taskId, taskInfo.getState()).get(1, TimeUnit.SECONDS); assertEquals(taskInfo.getState(), TaskState.FINISHED); taskInfo = sqlTaskManager.getTaskInfo(taskId); assertEquals(taskInfo.getState(), TaskState.FINISHED); } }
@Test public void testAbortResults() throws Exception { try (SqlTaskManager sqlTaskManager = createSqlTaskManager(new TaskManagerConfig())) { TaskId taskId = TASK_ID; TaskInfo taskInfo = sqlTaskManager.updateTask(TEST_SESSION, taskId, Optional.of(PLAN_FRAGMENT), ImmutableList.of(new TaskSource(TABLE_SCAN_NODE_ID, ImmutableSet.of(SPLIT), true)), INITIAL_EMPTY_OUTPUT_BUFFERS.withBuffer(OUT, 0).withNoMoreBufferIds()); assertEquals(taskInfo.getState(), TaskState.RUNNING); taskInfo = sqlTaskManager.getTaskInfo(taskId); assertEquals(taskInfo.getState(), TaskState.RUNNING); sqlTaskManager.abortTaskResults(taskId, OUT); taskInfo = sqlTaskManager.getTaskInfo(taskId, taskInfo.getState()).get(1, TimeUnit.SECONDS); assertEquals(taskInfo.getState(), TaskState.FINISHED); taskInfo = sqlTaskManager.getTaskInfo(taskId); assertEquals(taskInfo.getState(), TaskState.FINISHED); } }
@Test public void testCancel() throws Exception { try (SqlTaskManager sqlTaskManager = createSqlTaskManager(new TaskManagerConfig())) { TaskId taskId = TASK_ID; TaskInfo taskInfo = sqlTaskManager.updateTask(TEST_SESSION, taskId, Optional.of(PLAN_FRAGMENT), ImmutableList.<TaskSource>of(), INITIAL_EMPTY_OUTPUT_BUFFERS); assertEquals(taskInfo.getState(), TaskState.RUNNING); assertNull(taskInfo.getStats().getEndTime()); taskInfo = sqlTaskManager.getTaskInfo(taskId); assertEquals(taskInfo.getState(), TaskState.RUNNING); assertNull(taskInfo.getStats().getEndTime()); taskInfo = sqlTaskManager.cancelTask(taskId); assertEquals(taskInfo.getState(), TaskState.CANCELED); assertNotNull(taskInfo.getStats().getEndTime()); taskInfo = sqlTaskManager.getTaskInfo(taskId); assertEquals(taskInfo.getState(), TaskState.CANCELED); assertNotNull(taskInfo.getStats().getEndTime()); } }
@Test public void testAbort() throws Exception { try (SqlTaskManager sqlTaskManager = createSqlTaskManager(new TaskManagerConfig())) { TaskId taskId = TASK_ID; TaskInfo taskInfo = sqlTaskManager.updateTask(TEST_SESSION, taskId, Optional.of(PLAN_FRAGMENT), ImmutableList.<TaskSource>of(), INITIAL_EMPTY_OUTPUT_BUFFERS); assertEquals(taskInfo.getState(), TaskState.RUNNING); assertNull(taskInfo.getStats().getEndTime()); taskInfo = sqlTaskManager.getTaskInfo(taskId); assertEquals(taskInfo.getState(), TaskState.RUNNING); assertNull(taskInfo.getStats().getEndTime()); taskInfo = sqlTaskManager.abortTask(taskId); assertEquals(taskInfo.getState(), TaskState.ABORTED); assertNotNull(taskInfo.getStats().getEndTime()); taskInfo = sqlTaskManager.getTaskInfo(taskId); assertEquals(taskInfo.getState(), TaskState.ABORTED); assertNotNull(taskInfo.getStats().getEndTime()); } }
@Test public void testCancel() { try (SqlTaskManager sqlTaskManager = createSqlTaskManager(new TaskManagerConfig())) { TaskId taskId = TASK_ID; TaskInfo taskInfo = createTask(sqlTaskManager, taskId, createInitialEmptyOutputBuffers(PARTITIONED).withBuffer(OUT, 0).withNoMoreBufferIds()); assertEquals(taskInfo.getTaskStatus().getState(), TaskState.RUNNING); assertNull(taskInfo.getStats().getEndTime()); taskInfo = sqlTaskManager.getTaskInfo(taskId); assertEquals(taskInfo.getTaskStatus().getState(), TaskState.RUNNING); assertNull(taskInfo.getStats().getEndTime()); taskInfo = sqlTaskManager.cancelTask(taskId); assertEquals(taskInfo.getTaskStatus().getState(), TaskState.CANCELED); assertNotNull(taskInfo.getStats().getEndTime()); taskInfo = sqlTaskManager.getTaskInfo(taskId); assertEquals(taskInfo.getTaskStatus().getState(), TaskState.CANCELED); assertNotNull(taskInfo.getStats().getEndTime()); } }
@Test public void testEmptyQuery() throws Exception { try (SqlTaskManager sqlTaskManager = createSqlTaskManager(new TaskManagerConfig())) { TaskId taskId = TASK_ID; TaskInfo taskInfo = sqlTaskManager.updateTask(TEST_SESSION, taskId, Optional.of(PLAN_FRAGMENT), ImmutableList.<TaskSource>of(), INITIAL_EMPTY_OUTPUT_BUFFERS); assertEquals(taskInfo.getState(), TaskState.RUNNING); taskInfo = sqlTaskManager.getTaskInfo(taskId); assertEquals(taskInfo.getState(), TaskState.RUNNING); taskInfo = sqlTaskManager.updateTask(TEST_SESSION, taskId, Optional.of(PLAN_FRAGMENT), ImmutableList.of(new TaskSource(TABLE_SCAN_NODE_ID, ImmutableSet.<ScheduledSplit>of(), true)), INITIAL_EMPTY_OUTPUT_BUFFERS.withNoMoreBufferIds()); assertEquals(taskInfo.getState(), TaskState.FINISHED); taskInfo = sqlTaskManager.getTaskInfo(taskId); assertEquals(taskInfo.getState(), TaskState.FINISHED); } }
public SqlTaskManager createSqlTaskManager(TaskManagerConfig config) { return new SqlTaskManager( createTestingPlanner(), new MockLocationFactory(), taskExecutor, createTestSplitMonitor(), new NodeInfo("test"), localMemoryManager, taskManagementExecutor, config, new NodeMemoryConfig(), localSpillManager, new NodeSpillConfig(), new TestingGcMonitor()); }
@Test(timeOut = 60_000L) public void testLegacyQueryContext() throws Exception { QueryManager queryManager = queryRunner.getCoordinator().getQueryManager(); QueryId queryId = queryManager.createQueryId(); queryManager.createQuery( queryId, new TestingSessionContext(TEST_SESSION), "SELECT * FROM lineitem") .get(); waitForQueryState(queryRunner, queryId, RUNNING); // cancel query queryManager.failQuery(queryId, new PrestoException(GENERIC_INTERNAL_ERROR, "mock exception")); // assert that LegacyQueryContext is used instead of the DefaultQueryContext SqlTaskManager taskManager = (SqlTaskManager) queryRunner.getServers().get(0).getTaskManager(); assertInstanceOf(taskManager.getQueryContext(queryId), LegacyQueryContext.class); } }
@Test public void testEmptyQuery() { try (SqlTaskManager sqlTaskManager = createSqlTaskManager(new TaskManagerConfig())) { TaskId taskId = TASK_ID; TaskInfo taskInfo = createTask(sqlTaskManager, taskId, createInitialEmptyOutputBuffers(PARTITIONED).withNoMoreBufferIds()); assertEquals(taskInfo.getTaskStatus().getState(), TaskState.RUNNING); taskInfo = sqlTaskManager.getTaskInfo(taskId); assertEquals(taskInfo.getTaskStatus().getState(), TaskState.RUNNING); taskInfo = createTask(sqlTaskManager, taskId, ImmutableSet.of(), createInitialEmptyOutputBuffers(PARTITIONED).withNoMoreBufferIds()); assertEquals(taskInfo.getTaskStatus().getState(), TaskState.FINISHED); taskInfo = sqlTaskManager.getTaskInfo(taskId); assertEquals(taskInfo.getTaskStatus().getState(), TaskState.FINISHED); } }
@PostConstruct public void start() { taskManagementExecutor.scheduleWithFixedDelay(() -> { try { removeOldTasks(); } catch (Throwable e) { log.warn(e, "Error removing old tasks"); } try { failAbandonedTasks(); } catch (Throwable e) { log.warn(e, "Error canceling abandoned tasks"); } }, 200, 200, TimeUnit.MILLISECONDS); taskManagementExecutor.scheduleWithFixedDelay(() -> { try { updateStats(); } catch (Throwable e) { log.warn(e, "Error updating stats"); } }, 0, 1, TimeUnit.SECONDS); }
public SqlTaskManager createSqlTaskManager(TaskManagerConfig config) { return new SqlTaskManager( createTestingPlanner(), new MockLocationFactory(), taskExecutor, new QueryMonitor(new ObjectMapperProvider().get(), new NullEventClient(), new NodeInfo("test"), new NodeVersion("testVersion")), new NodeInfo("test"), localMemoryManager, config, new MemoryManagerConfig()); }
@Test(timeOut = 60_000L) public void testLegacyQueryContext() throws Exception { QueryManager queryManager = queryRunner.getCoordinator().getQueryManager(); QueryId queryId = queryManager.createQueryId(); queryManager.createQuery( queryId, new TestingSessionContext(TEST_SESSION), "SELECT * FROM lineitem") .get(); waitForQueryState(queryRunner, queryId, RUNNING); // cancel query queryManager.failQuery(queryId, new PrestoException(GENERIC_INTERNAL_ERROR, "mock exception")); // assert that LegacyQueryContext is used instead of the DefaultQueryContext SqlTaskManager taskManager = (SqlTaskManager) queryRunner.getServers().get(0).getTaskManager(); assertInstanceOf(taskManager.getQueryContext(queryId), LegacyQueryContext.class); } }
@Test public void testRemoveOldTasks() throws Exception { try (SqlTaskManager sqlTaskManager = createSqlTaskManager(new TaskManagerConfig().setInfoMaxAge(new Duration(5, TimeUnit.MILLISECONDS)))) { TaskId taskId = TASK_ID; TaskInfo taskInfo = sqlTaskManager.updateTask(TEST_SESSION, taskId, Optional.of(PLAN_FRAGMENT), ImmutableList.<TaskSource>of(), INITIAL_EMPTY_OUTPUT_BUFFERS); assertEquals(taskInfo.getState(), TaskState.RUNNING); taskInfo = sqlTaskManager.cancelTask(taskId); assertEquals(taskInfo.getState(), TaskState.CANCELED); taskInfo = sqlTaskManager.getTaskInfo(taskId); assertEquals(taskInfo.getState(), TaskState.CANCELED); Thread.sleep(100); sqlTaskManager.removeOldTasks(); for (TaskInfo info : sqlTaskManager.getAllTaskInfo()) { assertNotEquals(info.getTaskId(), taskId); } } }