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.server;
 
 
 
 import java.net.URI;
 import java.util.List;
 
 import static com.facebook.presto.block.BlockAssertions.createStringsBlock;
 import static com.facebook.presto.util.Failures.toFailures;
 import static com.google.common.base.Preconditions.checkNotNull;
 import static com.google.common.base.Preconditions.checkState;
 import static io.airlift.http.client.HttpUriBuilder.uriBuilderFrom;
 import static io.airlift.units.DataSize.Unit.MEGABYTE;
 
 public class MockTaskManager
         implements TaskManager
 {
     private final Executor executor = Executors.newCachedThreadPool(Threads.daemonThreadsNamed("test-%d"));
 
     private final HttpServerInfo httpServerInfo;
     private final DataSize maxBufferSize;
     private final int initialPages;
 
     private final ConcurrentMap<TaskIdMockTasktasks = new ConcurrentHashMap<>();
 
     @Inject
     public MockTaskManager(HttpServerInfo httpServerInfo)
     {
         this(httpServerInfonew DataSize(100, .), 12);
     }
 
     public MockTaskManager(HttpServerInfo httpServerInfoDataSize maxBufferSizeint initialPages)
     {
         checkNotNull(httpServerInfo"httpServerInfo is null");
         Preconditions.checkArgument(maxBufferSize.toBytes() > 0, "pageBufferMax must be at least 1");
         Preconditions.checkArgument(initialPages >= 0, "initialPages is negative");
         Preconditions.checkArgument(initialPages <= maxBufferSize.toBytes(), "initialPages is greater than maxBufferSize");
         this. = httpServerInfo;
         this. = maxBufferSize;
         this. = initialPages;
     }
 
     @Override
     public synchronized List<TaskInfogetAllTaskInfo(boolean full)
     {
         ImmutableList.Builder<TaskInfobuilder = ImmutableList.builder();
         for (MockTask task : .values()) {
             builder.add(task.getTaskInfo());
         }
         return builder.build();
     }
 
     @Override
    public void waitForStateChange(TaskId taskIdTaskState currentStateDuration maxWait)
            throws InterruptedException
    {
    }
    @Override
    public synchronized TaskInfo getTaskInfo(TaskId taskIdboolean full)
    {
        checkNotNull(taskId"taskId is null");
        MockTask task = .get(taskId);
        if (task == null) {
            throw new NoSuchElementException();
        }
        return task.getTaskInfo();
    }
    @Override
    public synchronized TaskInfo updateTask(Session sessionTaskId taskIdPlanFragment ignoredList<TaskSourcesourcesOutputBuffers outputBuffers)
    {
        checkNotNull(session"session is null");
        checkNotNull(taskId"taskId is null");
        checkNotNull(sources"sources is null");
        checkNotNull(outputBuffers"outputBuffers is null");
        MockTask task = .get(taskId);
        if (task == null) {
            task = new MockTask(session,
                    taskId,
                    uriBuilderFrom(.getHttpUri()).appendPath("v1/task").appendPath(taskId.toString()).build(),
                    outputBuffers,
                    ,
                    ,
                    
            );
            .put(taskIdtask);
        }
        task.addOutputBuffers(outputBuffers);
        return task.getTaskInfo();
    }
    @Override
    public BufferResult getTaskResults(TaskId taskIdString outputIdlong startingSequenceIdDataSize maxSizeDuration maxWaitTime)
            throws InterruptedException
    {
        checkNotNull(taskId"taskId is null");
        checkNotNull(outputId"outputId is null");
        MockTask task;
        synchronized (this) {
            task = .get(taskId);
        }
        if (task == null) {
            throw new NoSuchElementException();
        }
        return task.getResults(outputIdstartingSequenceIdmaxSizemaxWaitTime);
    }
    @Override
    public synchronized TaskInfo abortTaskResults(TaskId taskIdString outputId)
    {
        checkNotNull(taskId"taskId is null");
        checkNotNull(outputId"outputId is null");
        MockTask task = .get(taskId);
        if (task == null) {
            throw new NoSuchElementException();
        }
        task.abortResults(outputId);
        return task.getTaskInfo();
    }
    @Override
    public synchronized TaskInfo cancelTask(TaskId taskId)
    {
        checkNotNull(taskId"taskId is null");
        MockTask task = .get(taskId);
        if (task == null) {
            return null;
        }
        task.cancel();
        return task.getTaskInfo();
    }
    public static class MockTask
    {
        private final AtomicLong nextTaskInfoVersion = new AtomicLong(.);
        private final URI location;
        private final TaskStateMachine taskStateMachine;
        private final TaskContext taskContext;
        private final SharedBuffer sharedBuffer;
        public MockTask(Session session,
                TaskId taskId,
                URI location,
                OutputBuffers outputBuffers,
                DataSize maxBufferSize,
                int initialPages,
                Executor executor)
        {
            this. = new TaskStateMachine(checkNotNull(taskId"taskId is null"), checkNotNull(executor"executor is null"));
            this. = new TaskContext(executorsessionnew DataSize(256, ), new DataSize(1, ), true);
            this. = checkNotNull(location"location is null");
            this. = new SharedBuffer(taskIdexecutorcheckNotNull(maxBufferSize"maxBufferSize is null"), outputBuffers);
            List<Stringdata = ImmutableList.of("apple""banana""cherry""date");
            // load initial pages
            for (int i = 0; i < initialPagesi++) {
                checkState(.enqueue(new Page(createStringsBlock(Iterables.concat(Collections.nCopies(i + 1, data))))).isDone(), "Unable to add page to buffer");
            }
            .finish();
        }
        public void abortResults(String outputId)
        {
            .abort(outputId);
        }
        public void addOutputBuffers(OutputBuffers outputBuffers)
        {
            .setOutputBuffers(outputBuffers);
        }
        public void cancel()
        {
            .cancel();
        }
        public BufferResult getResults(String outputIdlong startingSequenceIdDataSize maxSizeDuration maxWaitTime)
                throws InterruptedException
        {
            return .get(outputIdstartingSequenceIdmaxSizemaxWaitTime);
        }
        public TaskInfo getTaskInfo()
        {
            TaskState state = .getState();
            List<ExecutionFailureInfofailures = ImmutableList.of();
            if (state == .) {
                failures = toFailures(.getFailureCauses());
            }
            return new TaskInfo(
                    .getTaskId(),
                    .getAndIncrement(),
                    state,
                    ,
                    DateTime.now(),
                    .getInfo(),
                    ImmutableSet.<PlanNodeId>of(),
                    .getTaskStats(),
                    failures);
        }
    }
New to GrepCode? Check out our FAQ X