@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 @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(); }
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 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 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()); }
@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()); }
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()); } } }
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(); }
public TaskInfo abort() { taskStateMachine.abort(); 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()); }
@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();
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; }
taskId -> createSqlTask( taskId, locationFactory.createLocalTaskLocation(taskId), taskNotificationExecutor, sqlTask -> { finishedTaskStats.merge(sqlTask.getIoStats()); return null; },
private SqlTask newSqlTask() { TaskId taskId = new TaskId("query", 0, idGeneator.incrementAndGet()); URI location = URI.create("fake://task/" + taskId); return createSqlTask( taskId, location, "fake", new DefaultQueryContext(new QueryId("query"), new DataSize(1, MEGABYTE), new DataSize(2, MEGABYTE), memoryPool, new TestingGcMonitor(), executor, scheduledExecutor, new DataSize(1, GIGABYTE), spillSpaceTracker), sqlTaskExecutionFactory, executor, Functions.identity(), new DataSize(32, MEGABYTE), new CounterStat()); } }