public static SqlTask createSqlTask( TaskId taskId, URI location, String nodeId, QueryContext queryContext, SqlTaskExecutionFactory sqlTaskExecutionFactory, ExecutorService taskNotificationExecutor, Function<SqlTask, ?> onDone, DataSize maxBufferSize, CounterStat failedTasks) { SqlTask sqlTask = new SqlTask(taskId, location, nodeId, queryContext, sqlTaskExecutionFactory, taskNotificationExecutor, maxBufferSize); sqlTask.initialize(onDone, failedTasks); return sqlTask; }
@Override public TaskInfo abortTaskResults(TaskId taskId, OutputBufferId bufferId) { requireNonNull(taskId, "taskId is null"); requireNonNull(bufferId, "bufferId is null"); return tasks.getUnchecked(taskId).abortTaskResults(bufferId); }
@Override public TaskInfo cancelTask(TaskId taskId) { requireNonNull(taskId, "taskId is null"); return tasks.getUnchecked(taskId).cancel(); }
@Override public TaskInfo getTaskInfo(TaskId taskId) { requireNonNull(taskId, "taskId is null"); SqlTask sqlTask = tasks.getUnchecked(taskId); sqlTask.recordHeartbeat(); return sqlTask.getTaskInfo(); }
@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); }
@Test public void testSimpleQuery() throws Exception { SqlTask sqlTask = createInitialTask(); TaskInfo taskInfo = sqlTask.updateTask(TEST_SESSION, Optional.of(PLAN_FRAGMENT), ImmutableList.of(new TaskSource(TABLE_SCAN_NODE_ID, ImmutableSet.of(SPLIT), true)), createInitialEmptyOutputBuffers(PARTITIONED).withBuffer(OUT, 0).withNoMoreBufferIds(), OptionalInt.empty()); assertEquals(taskInfo.getTaskStatus().getState(), TaskState.RUNNING); taskInfo = sqlTask.getTaskInfo(); assertEquals(taskInfo.getTaskStatus().getState(), TaskState.RUNNING); BufferResult results = sqlTask.getTaskResults(OUT, 0, new DataSize(1, 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 = sqlTask.getTaskResults(OUT, results.getToken() + results.getSerializedPages().size(), new DataSize(1, MEGABYTE)).get(); } assertEquals(results.getSerializedPages().size(), 0); // complete the task by calling abort on it TaskInfo info = sqlTask.abortTaskResults(OUT); assertEquals(info.getOutputBuffers().getState(), BufferState.FINISHED); taskInfo = sqlTask.getTaskInfo(taskInfo.getTaskStatus().getState()).get(1, SECONDS); assertEquals(taskInfo.getTaskStatus().getState(), TaskState.FINISHED); taskInfo = sqlTask.getTaskInfo(); assertEquals(taskInfo.getTaskStatus().getState(), TaskState.FINISHED); }
@Test public void testBufferCloseOnCancel() throws Exception { SqlTask sqlTask = createInitialTask(); updateTask(sqlTask, EMPTY_SOURCES, createInitialEmptyOutputBuffers(PARTITIONED).withBuffer(OUT, 0).withNoMoreBufferIds()); ListenableFuture<BufferResult> bufferResult = sqlTask.getTaskResults(OUT, 0, new DataSize(1, MEGABYTE)); assertFalse(bufferResult.isDone()); sqlTask.cancel(); assertEquals(sqlTask.getTaskInfo().getTaskStatus().getState(), TaskState.CANCELED); // buffer future will complete.. the event is async so wait a bit for event to propagate bufferResult.get(1, SECONDS); bufferResult = sqlTask.getTaskResults(OUT, 0, new DataSize(1, MEGABYTE)); assertTrue(bufferResult.isDone()); assertTrue(bufferResult.get().isBufferComplete()); }
@Test public void testAbort() throws Exception { SqlTask sqlTask = createInitialTask(); TaskInfo taskInfo = sqlTask.updateTask(TEST_SESSION, Optional.of(PLAN_FRAGMENT), ImmutableList.of(new TaskSource(TABLE_SCAN_NODE_ID, ImmutableSet.of(SPLIT), true)), createInitialEmptyOutputBuffers(PARTITIONED).withBuffer(OUT, 0).withNoMoreBufferIds(), OptionalInt.empty()); assertEquals(taskInfo.getTaskStatus().getState(), TaskState.RUNNING); taskInfo = sqlTask.getTaskInfo(); assertEquals(taskInfo.getTaskStatus().getState(), TaskState.RUNNING); sqlTask.abortTaskResults(OUT); taskInfo = sqlTask.getTaskInfo(taskInfo.getTaskStatus().getState()).get(1, SECONDS); assertEquals(taskInfo.getTaskStatus().getState(), TaskState.FINISHED); taskInfo = sqlTask.getTaskInfo(); assertEquals(taskInfo.getTaskStatus().getState(), TaskState.FINISHED); }
@Test public void testCancel() { SqlTask sqlTask = createInitialTask(); TaskInfo taskInfo = sqlTask.updateTask(TEST_SESSION, Optional.of(PLAN_FRAGMENT), ImmutableList.of(), createInitialEmptyOutputBuffers(PARTITIONED) .withBuffer(OUT, 0) .withNoMoreBufferIds(), OptionalInt.empty()); assertEquals(taskInfo.getTaskStatus().getState(), TaskState.RUNNING); assertNull(taskInfo.getStats().getEndTime()); taskInfo = sqlTask.getTaskInfo(); assertEquals(taskInfo.getTaskStatus().getState(), TaskState.RUNNING); assertNull(taskInfo.getStats().getEndTime()); taskInfo = sqlTask.cancel(); assertEquals(taskInfo.getTaskStatus().getState(), TaskState.CANCELED); assertNotNull(taskInfo.getStats().getEndTime()); taskInfo = sqlTask.getTaskInfo(); assertEquals(taskInfo.getTaskStatus().getState(), TaskState.CANCELED); assertNotNull(taskInfo.getStats().getEndTime()); }
@Test public void testBufferNotCloseOnFail() throws Exception { SqlTask sqlTask = createInitialTask(); updateTask(sqlTask, EMPTY_SOURCES, createInitialEmptyOutputBuffers(PARTITIONED).withBuffer(OUT, 0).withNoMoreBufferIds()); ListenableFuture<BufferResult> bufferResult = sqlTask.getTaskResults(OUT, 0, new DataSize(1, MEGABYTE)); assertFalse(bufferResult.isDone()); TaskState taskState = sqlTask.getTaskInfo().getTaskStatus().getState(); sqlTask.failed(new Exception("test")); assertEquals(sqlTask.getTaskInfo(taskState).get(1, SECONDS).getTaskStatus().getState(), TaskState.FAILED); // buffer will not be closed by fail event. event is async so wait a bit for event to fire try { assertTrue(bufferResult.get(1, SECONDS).isBufferComplete()); fail("expected TimeoutException"); } catch (TimeoutException expected) { // expected } assertFalse(sqlTask.getTaskResults(OUT, 0, new DataSize(1, MEGABYTE)).isDone()); }
public void failAbandonedTasks() { DateTime now = DateTime.now(); DateTime oldestAllowedHeartbeat = now.minus(clientTimeout.toMillis()); for (SqlTask sqlTask : tasks.asMap().values()) { try { TaskInfo taskInfo = sqlTask.getTaskInfo(); TaskStatus taskStatus = taskInfo.getTaskStatus(); if (taskStatus.getState().isDone()) { continue; } DateTime lastHeartbeat = taskInfo.getLastHeartbeat(); if (lastHeartbeat != null && lastHeartbeat.isBefore(oldestAllowedHeartbeat)) { log.info("Failing abandoned task %s", taskStatus.getTaskId()); sqlTask.failed(new PrestoException(ABANDONED_TASK, format("Task %s has not been accessed since %s: currentTime %s", taskStatus.getTaskId(), lastHeartbeat, now))); } } catch (RuntimeException e) { log.warn(e, "Error while inspecting age of task %s", sqlTask.getTaskId()); } } }
@Test public void testBufferCloseOnFinish() throws Exception { SqlTask sqlTask = createInitialTask(); OutputBuffers outputBuffers = INITIAL_EMPTY_OUTPUT_BUFFERS.withBuffer(OUT, 0).withNoMoreBufferIds(); updateTask(sqlTask, EMPTY_SOURCES, outputBuffers); CompletableFuture<BufferResult> bufferResult = sqlTask.getTaskResults(OUT, 0, new DataSize(1, MEGABYTE)); assertFalse(bufferResult.isDone()); // close the sources (no splits will ever be added) updateTask(sqlTask, ImmutableList.of(new TaskSource(TABLE_SCAN_NODE_ID, ImmutableSet.<ScheduledSplit>of(), true)), outputBuffers); // finish the task by calling abort on it TaskInfo taskInfo = sqlTask.abortTaskResults(OUT); // buffer will be closed by cancel event (wait for event to fire) assertTrue(bufferResult.get(1, SECONDS).isBufferComplete()); assertEquals(sqlTask.getTaskInfo().getOutputBuffers().getState(), BufferState.FINISHED); taskInfo = sqlTask.getTaskInfo(taskInfo.getState()).get(1, SECONDS); assertEquals(taskInfo.getState(), TaskState.FINISHED); // verify the buffer is closed bufferResult = sqlTask.getTaskResults(OUT, 0, new DataSize(1, MEGABYTE)); assertTrue(bufferResult.isDone()); assertTrue(bufferResult.get().isBufferComplete()); }
public TaskInfo cancel() { taskStateMachine.cancel(); return getTaskInfo(); }
@Test public void testEmptyQuery() { SqlTask sqlTask = createInitialTask(); TaskInfo taskInfo = sqlTask.updateTask(TEST_SESSION, Optional.of(PLAN_FRAGMENT), ImmutableList.of(), createInitialEmptyOutputBuffers(PARTITIONED) .withNoMoreBufferIds(), OptionalInt.empty()); assertEquals(taskInfo.getTaskStatus().getState(), TaskState.RUNNING); taskInfo = sqlTask.getTaskInfo(); assertEquals(taskInfo.getTaskStatus().getState(), TaskState.RUNNING); taskInfo = sqlTask.updateTask(TEST_SESSION, Optional.of(PLAN_FRAGMENT), ImmutableList.of(new TaskSource(TABLE_SCAN_NODE_ID, ImmutableSet.of(), true)), createInitialEmptyOutputBuffers(PARTITIONED) .withNoMoreBufferIds(), OptionalInt.empty()); assertEquals(taskInfo.getTaskStatus().getState(), TaskState.FINISHED); taskInfo = sqlTask.getTaskInfo(); assertEquals(taskInfo.getTaskStatus().getState(), TaskState.FINISHED); }
failed(e); throw e; failed(e); return getTaskInfo();
@Test public void testBufferCloseOnFinish() throws Exception { SqlTask sqlTask = createInitialTask(); OutputBuffers outputBuffers = createInitialEmptyOutputBuffers(PARTITIONED).withBuffer(OUT, 0).withNoMoreBufferIds(); updateTask(sqlTask, EMPTY_SOURCES, outputBuffers); ListenableFuture<BufferResult> bufferResult = sqlTask.getTaskResults(OUT, 0, new DataSize(1, MEGABYTE)); assertFalse(bufferResult.isDone()); // close the sources (no splits will ever be added) updateTask(sqlTask, ImmutableList.of(new TaskSource(TABLE_SCAN_NODE_ID, ImmutableSet.of(), true)), outputBuffers); // finish the task by calling abort on it sqlTask.abortTaskResults(OUT); // buffer will be closed by cancel event (wait for event to fire) bufferResult.get(1, SECONDS); // verify the buffer is closed bufferResult = sqlTask.getTaskResults(OUT, 0, new DataSize(1, MEGABYTE)); assertTrue(bufferResult.isDone()); assertTrue(bufferResult.get().isBufferComplete()); }
@Override @PreDestroy public void close() { boolean taskCanceled = false; for (SqlTask task : tasks.asMap().values()) { if (task.getTaskStatus().getState().isDone()) { continue; } task.failed(new PrestoException(SERVER_SHUTTING_DOWN, format("Server is shutting down. Task %s has been canceled", task.getTaskId()))); taskCanceled = true; } if (taskCanceled) { try { TimeUnit.SECONDS.sleep(5); } catch (InterruptedException e) { Thread.currentThread().interrupt(); } } taskNotificationExecutor.shutdownNow(); }
@Override public SqlTask load(TaskId taskId) throws Exception { return new SqlTask( taskId, locationFactory.createLocalTaskLocation(taskId), queryContexts.getUnchecked(taskId.getQueryId()), sqlTaskExecutionFactory, taskNotificationExecutor, sqlTask -> { finishedTaskStats.merge(sqlTask.getIoStats()); return null; }, maxBufferSize ); } });
private void updateStats() { SqlTaskIoStats tempIoStats = new SqlTaskIoStats(); tempIoStats.merge(finishedTaskStats); // there is a race here between task completion, which merges stats into // finishedTaskStats, and getting the stats from the task. Since we have // already merged the final stats, we could miss the stats from this task // which would result in an under-count, but we will not get an over-count. tasks.asMap().values().stream() .filter(task -> !task.getTaskInfo().getState().isDone()) .forEach(task -> tempIoStats.merge(task.getIoStats())); cachedStats.resetTo(tempIoStats); }
@Override public ListenableFuture<BufferResult> getTaskResults(TaskId taskId, OutputBufferId bufferId, long startingSequenceId, DataSize maxSize) { requireNonNull(taskId, "taskId is null"); requireNonNull(bufferId, "bufferId is null"); checkArgument(startingSequenceId >= 0, "startingSequenceId is negative"); requireNonNull(maxSize, "maxSize is null"); return tasks.getUnchecked(taskId).getTaskResults(bufferId, startingSequenceId, maxSize); }