Start line:  
End line:  

Snippet Preview

Snippet HTML Code

Stack Overflow Questions
  /*
   * Licensed under the Apache License, Version 2.0 (the "License");
   * you may not use this file except in compliance with the License.
   * You may obtain a copy of the License at
   *
   *     http://www.apache.org/licenses/LICENSE-2.0
   *
   * Unless required by applicable law or agreed to in writing, software
   * distributed under the License is distributed on an "AS IS" BASIS,
  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
 package com.facebook.presto.execution;
 
 
 
 import java.net.URI;
 import java.util.List;
 import java.util.Set;
 
 import static com.facebook.presto.util.Failures.toFailures;
 import static com.google.common.base.Preconditions.checkArgument;
 import static com.google.common.base.Preconditions.checkNotNull;
 
 public class SqlTask
 {
     private static final Logger log = Logger.get(SqlTask.class);
 
     private final TaskId taskId;
     private final URI location;
     private final TaskStateMachine taskStateMachine;
     private final SharedBuffer sharedBuffer;
 
 
     private final AtomicReference<DateTimelastHeartbeat = new AtomicReference<>(DateTime.now());
 
     private final AtomicReference<TaskHoldertaskHolderReference = new AtomicReference<>(new TaskHolder());
 
     public SqlTask(
             TaskId taskId,
             URI location,
             SqlTaskExecutionFactory sqlTaskExecutionFactory,
             ExecutorService taskNotificationExecutor,
             final Function<SqlTask, ?> onDone,
             DataSize maxBufferSize)
     {
         this. = checkNotNull(taskId"taskId is null");
         this. = checkNotNull(location"location is null");
         this. = checkNotNull(sqlTaskExecutionFactory"sqlTaskExecutionFactory is null");
         checkNotNull(taskNotificationExecutor"taskNotificationExecutor is null");
         checkNotNull(onDone"onDone is null");
         checkNotNull(maxBufferSize"maxBufferSize is null");
 
          = new SharedBuffer(taskIdtaskNotificationExecutormaxBufferSize);
          = new TaskStateMachine(taskIdtaskNotificationExecutor);
         {
             @Override
             public void stateChanged(TaskState taskState)
             {
                 if (!taskState.isDone()) {
                     return;
                 }
 
                 // store final task info
                 while (true) {
                     TaskHolder taskHolder = .get();
                     if (taskHolder.isFinished()) {
                         // another concurrent worker already set the final state
                         return;
                     }
 
                     if (.compareAndSet(taskHoldernew TaskHolder(createTaskInfo(taskHolder), taskHolder.getIoStats()))) {
                         break;
                     }
                }
                // make sure buffers are cleaned up
                if (taskState == . || taskState == .) {
                    // don't close buffers for a failed query
                    // closed buffers signal to upstream tasks that everything finished cleanly
                    .fail();
                }
                else {
                    .destroy();
                }
                try {
                    onDone.apply(SqlTask.this);
                }
                catch (Exception e) {
                    .warn(e"Error running task cleanup callback %s"SqlTask.this.);
                }
            }
        });
    }
    public SqlTaskIoStats getIoStats()
    {
        return .get().getIoStats();
    }
    public TaskId getTaskId()
    {
        return .getTaskId();
    }
    public TaskInfo getTaskInfo()
    {
        .set(DateTime.now());
        try (SetThreadName ignored = new SetThreadName("Task-%s")) {
            return createTaskInfo(.get());
        }
    }
    private TaskInfo createTaskInfo(TaskHolder taskHolder)
    {
        // Always return a new TaskInfo with a larger version number;
        // otherwise a client will not accept the update
        long versionNumber = .getAndIncrement();
        TaskState state = .getState();
        List<ExecutionFailureInfofailures = ImmutableList.of();
        if (state == .) {
            failures = toFailures(.getFailureCauses());
        }
        TaskStats taskStats;
        Set<PlanNodeIdnoMoreSplits;
        TaskInfo finalTaskInfo = taskHolder.getFinalTaskInfo();
        if (finalTaskInfo != null) {
            taskStats = finalTaskInfo.getStats();
            noMoreSplits = finalTaskInfo.getNoMoreSplits();
        }
        else {
            SqlTaskExecution taskExecution = taskHolder.getTaskExecution();
            if (taskExecution != null) {
                taskStats = taskExecution.getTaskContext().getTaskStats();
                noMoreSplits = taskExecution.getNoMoreSplits();
            }
            else {
                // if the task completed without creation, set end time
                DateTime endTime = state.isDone() ? DateTime.now() : null;
                taskStats = new TaskStats(.getCreatedTime(), endTime);
                noMoreSplits = ImmutableSet.of();
            }
        }
        return new TaskInfo(
                .getTaskId(),
                versionNumber,
                state,
                ,
                .get(),
                .getInfo(),
                noMoreSplits,
                taskStats,
                failures);
    }
    public ListenableFuture<TaskInfogetTaskInfo(TaskState callersCurrentState)
    {
        checkNotNull(callersCurrentState"callersCurrentState is null");
        .set(DateTime.now());
        // If the caller's current state is already done, just return the current
        // state of this task as it will either be done or possibly still running
        // (due to a bug in the caller), since we can not transition from a done
        // state.
        if (callersCurrentState.isDone()) {
            return Futures.immediateFuture(getTaskInfo());
        }
        ListenableFuture<TaskStatefutureTaskState = .getStateChange(callersCurrentState);
        return Futures.transform(futureTaskState, (TaskState input) -> getTaskInfo());
    }
    public TaskInfo updateTask(Session sessionPlanFragment fragmentList<TaskSourcesourcesOutputBuffers outputBuffers)
    {
        // assure the task execution is only created once
        SqlTaskExecution taskExecution;
        synchronized (this) {
            // is task already complete?
            TaskHolder taskHolder = .get();
            if (taskHolder.isFinished()) {
                return taskHolder.getFinalTaskInfo();
            }
            taskExecution = taskHolder.getTaskExecution();
            if (taskExecution == null) {
                try {
                    taskExecution = .create(sessionfragmentsources);
                    .compareAndSet(taskHoldernew TaskHolder(taskExecution));
                }
                catch (RuntimeException e) {
                    failed(e);
                }
            }
        }
        .set(DateTime.now());
        if (taskExecution != null) {
            // addSources checks for task completion, so update the buffers first and the task might complete earlier
            .setOutputBuffers(outputBuffers);
            taskExecution.addSources(sources);
        }
        return getTaskInfo();
    }
    public ListenableFuture<BufferResultgetTaskResults(TaskId outputNamelong startingSequenceIdDataSize maxSize)
    {
        checkNotNull(outputName"outputName is null");
        checkArgument(maxSize.toBytes() > 0, "maxSize must be at least 1 byte");
        .set(DateTime.now());
        return .get(outputNamestartingSequenceIdmaxSize);
    }
    public TaskInfo abortTaskResults(TaskId outputId)
    {
        checkNotNull(outputId"outputId is null");
        .set(DateTime.now());
        .debug("Aborting task %s output %s"outputId);
        .abort(outputId);
        return getTaskInfo();
    }
    public void failed(Throwable cause)
    {
        checkNotNull(cause"cause is null");
        .failed(cause);
    }
    public TaskInfo cancel()
    {
        .set(DateTime.now());
        .cancel();
        return getTaskInfo();
    }
    public TaskInfo abort()
    {
        .set(DateTime.now());
        .abort();
        return getTaskInfo();
    }
    @Override
    public String toString()
    {
        return .toString();
    }
    private static final class TaskHolder
    {
        private final SqlTaskExecution taskExecution;
        private final TaskInfo finalTaskInfo;
        private final SqlTaskIoStats finalIoStats;
        private TaskHolder()
        {
            this. = null;
            this. = null;
            this. = null;
        }
        private TaskHolder(SqlTaskExecution taskExecution)
        {
            this. = checkNotNull(taskExecution"taskExecution is null");
            this. = null;
            this. = null;
        }
        private TaskHolder(TaskInfo finalTaskInfoSqlTaskIoStats finalIoStats)
        {
            this. = null;
            this. = checkNotNull(finalTaskInfo"finalTaskInfo is null");
            this. = checkNotNull(finalIoStats"finalIoStats is null");
        }
        public boolean isFinished()
        {
            return  != null;
        }
        @Nullable
        public SqlTaskExecution getTaskExecution()
        {
            return ;
        }
        @Nullable
        public TaskInfo getFinalTaskInfo()
        {
            return ;
        }
        public SqlTaskIoStats getIoStats()
        {
            // if we are finished, return the final IoStats
            if ( != null) {
                return ;
            }
            // if we haven't started yet, return an empty IoStats
            if ( == null) {
                return new SqlTaskIoStats();
            }
            // get IoStats from the current task execution
            TaskContext taskContext = .getTaskContext();
            return new SqlTaskIoStats(taskContext.getInputDataSize(), taskContext.getInputPositions(), taskContext.getOutputDataSize(), taskContext.getOutputPositions());
        }
    }
New to GrepCode? Check out our FAQ X