@Override public String toString() { return toStringHelper(this) .add("taskId", taskStatus.getTaskId()) .add("state", taskStatus.getState()) .toString(); }
private void addTask(RemoteTask task) { if (remoteTasks.add(task)) { task.addStateChangeListener(taskStatus -> { if (taskStatus.getState().isDone()) { remoteTasks.remove(task); } }); // Check if task state is already done before adding the listener if (task.getTaskStatus().getState().isDone()) { remoteTasks.remove(task); } } }
public boolean isCompleteInfo() { return state.isDone() && tasks.stream().allMatch(taskInfo -> taskInfo.getTaskStatus().getState().isDone()); } }
private static Optional<TaskInfo> findFailedTask(StageInfo stageInfo) { for (StageInfo subStage : stageInfo.getSubStages()) { Optional<TaskInfo> task = findFailedTask(subStage); if (task.isPresent()) { return task; } } return stageInfo.getTasks().stream() .filter(taskInfo -> taskInfo.getTaskStatus().getState() == TaskState.FAILED) .findFirst(); }
private static boolean isDone(TaskInfo taskInfo) { return taskInfo.getTaskStatus().getState().isDone(); } }
public boolean isFinalStageInfo() { return state.isDone() && tasks.stream().allMatch(taskInfo -> taskInfo.getTaskStatus().getState().isDone()); }
private List<TaskInfo> getActiveTasks() { return sqlTaskManager.getAllTaskInfo() .stream() .filter(taskInfo -> !taskInfo.getTaskStatus().getState().isDone()) .collect(toImmutableList()); }
@Override public synchronized void setOutputBuffers(OutputBuffers newOutputBuffers) { if (getTaskStatus().getState().isDone()) { return; } if (newOutputBuffers.getVersion() > outputBuffers.get().getVersion()) { outputBuffers.set(newOutputBuffers); needsUpdate.set(true); scheduleUpdate(); } }
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.getTaskStatus().getState().isDone()) .forEach(task -> tempIoStats.merge(task.getIoStats())); cachedStats.resetTo(tempIoStats); }
@Override public synchronized void abort() { if (getTaskStatus().getState().isDone()) { return; } abort(failWith(getTaskStatus(), ABORTED, ImmutableList.of())); }
@Override public int getQueuedPartitionedSplitCount() { TaskStatus taskStatus = getTaskStatus(); if (taskStatus.getState().isDone()) { return 0; } return getPendingSourceSplitCount() + taskStatus.getQueuedPartitionedDrivers(); }
@Override public void onSuccess(JsonResponse<TaskInfo> result) { try { updateTaskInfo(result.getValue()); } finally { if (!getTaskInfo().getTaskStatus().getState().isDone()) { cleanUpLocally(); } } }
@Override public int getPartitionedSplitCount() { TaskStatus taskStatus = getTaskStatus(); if (taskStatus.getState().isDone()) { return 0; } return getPendingSourceSplitCount() + taskStatus.getQueuedPartitionedDrivers() + taskStatus.getRunningPartitionedDrivers(); }
/** * Move the task directly to the failed state if there was a failure in this task */ private void failTask(Throwable cause) { TaskStatus taskStatus = getTaskStatus(); if (!taskStatus.getState().isDone()) { log.debug(cause, "Remote task %s failed with %s", taskStatus.getSelf(), cause); } abort(failWith(getTaskStatus(), FAILED, ImmutableList.of(toFailure(cause)))); }
@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 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); } }
public TaskInfo summarize() { if (taskStatus.getState().isDone()) { return new TaskInfo(taskStatus, lastHeartbeat, outputBuffers.summarize(), noMoreSplits, stats.summarizeFinal(), needsPlan); } return new TaskInfo(taskStatus, lastHeartbeat, outputBuffers.summarize(), noMoreSplits, stats.summarize(), needsPlan); }
private synchronized void abort(TaskStatus status) { checkState(status.getState().isDone(), "cannot abort task with an incomplete status"); try (SetThreadName ignored = new SetThreadName("HttpRemoteTask-%s", taskId)) { taskStatusFetcher.updateTaskStatus(status); // send abort to task HttpUriBuilder uriBuilder = getHttpUriBuilder(getTaskStatus()); Request request = prepareDelete() .setUri(uriBuilder.build()) .build(); scheduleAsyncCleanupRequest(createCleanupBackoff(), request, "abort"); } }
@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()); }
@Override public synchronized void cancel() { try (SetThreadName ignored = new SetThreadName("HttpRemoteTask-%s", taskId)) { TaskStatus taskStatus = getTaskStatus(); if (taskStatus.getState().isDone()) { return; } // send cancel to task and ignore response HttpUriBuilder uriBuilder = getHttpUriBuilder(taskStatus).addParameter("abort", "false"); Request request = prepareDelete() .setUri(uriBuilder.build()) .build(); scheduleAsyncCleanupRequest(createCleanupBackoff(), request, "cancel"); } }