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.SessionTestUtils.TEST_SESSION;
 import static com.facebook.presto.execution.StateMachine.StateChangeListener;
 import static com.facebook.presto.spi.type.VarcharType.VARCHAR;
 import static com.facebook.presto.util.Failures.toFailures;
 import static com.google.common.base.Preconditions.checkNotNull;
 import static io.airlift.units.DataSize.Unit.MEGABYTE;
 
 public class MockRemoteTaskFactory
         implements RemoteTaskFactory
 {
     private final Executor executor;
 
     MockRemoteTaskFactory(Executor executor)
     {
         this. = executor;
     }
 
     public RemoteTask createTableScanTask(Node newNodeList<Splitsplits)
     {
         TaskId taskId = new TaskId(new StageId("test""1"), "1");
         Symbol symbol = new Symbol("column");
         PlanNodeId tableScanNodeId = new PlanNodeId("test");
         PlanNodeId sourceId = new PlanNodeId("sourceId");
         PlanFragment testFragment = new PlanFragment(
                 new PlanFragmentId("test"),
                 new TableScanNode(
                         new PlanNodeId("test"),
                         new TableHandle("test"new TestingTableHandle()),
                         ImmutableList.of(symbol),
                         ImmutableMap.of(symbolnew ColumnHandle("test"new TestingColumnHandle("column"))),
                         null,
                         Optional.empty()),
                 ImmutableMap.<SymbolType>of(symbol),
                 ImmutableList.of(symbol),
                 ..,
                 tableScanNodeId,
                 ..,
                 ImmutableList.<Symbol>of(),
                 Optional.empty()
         );
 
         ImmutableMultimap.Builder<PlanNodeIdSplitinitialSplits = ImmutableMultimap.builder();
         for (Split sourceSplit : splits) {
             initialSplits.put(sourceIdsourceSplit);
         }
         return createRemoteTask(taskIdnewNodetestFragmentinitialSplits.build(), .);
    }
    @Override
            Session session,
            TaskId taskId,
            Node node,
            PlanFragment fragment,
            Multimap<PlanNodeIdSplitinitialSplits,
            OutputBuffers outputBuffers)
    {
        return new MockRemoteTask(taskIdfragmentnode.getNodeIdentifier(), initialSplits);
    }
    private class MockRemoteTask
            implements RemoteTask
    {
        private final AtomicLong nextTaskInfoVersion = new AtomicLong(.);
        private final URI location;
        private final TaskStateMachine taskStateMachine;
        private final TaskContext taskContext;
        private final SharedBuffer sharedBuffer;
        private final String nodeId;
        private final PlanFragment fragment;
        @GuardedBy("this")
        private final Set<PlanNodeIdnoMoreSplits = new HashSet<>();
        @GuardedBy("this")
        private final Multimap<PlanNodeIdSplitsplits = HashMultimap.create();
        public MockRemoteTask(TaskId taskId,
                PlanFragment fragment,
                String nodeId,
                Executor executor,
                Multimap<PlanNodeIdSplitinitialSplits)
        {
            this. = new TaskStateMachine(checkNotNull(taskId"taskId is null"), checkNotNull(executor"executor is null"));
            this. = new TaskContext(executornew DataSize(256, ), new DataSize(1, ), truetrue);
            this. = URI.create("fake://task/" + taskId);
            this. = new SharedBuffer(taskIdexecutorcheckNotNull(new DataSize(1, ..), "maxBufferSize is null"));
            this. = checkNotNull(fragment"fragment is null");
            this. = checkNotNull(nodeId"nodeId is null");
            .putAll(initialSplits);
        }
        @Override
        public String getNodeId()
        {
            return ;
        }
        @Override
        public TaskInfo getTaskInfo()
        {
            TaskState state = .getState();
            List<ExecutionFailureInfofailures = ImmutableList.of();
            if (state == .) {
                failures = toFailures(.getFailureCauses());
            }
            return new TaskInfo(
                    .getTaskId(),
                    Optional.empty(),
                    .getAndIncrement(),
                    state,
                    ,
                    DateTime.now(),
                    .getInfo(),
                    ImmutableSet.<PlanNodeId>of(),
                    .getTaskStats(),
                    failures);
        }
        @Override
        public void start()
        {
        }
        @Override
        public void addSplits(PlanNodeId sourceIdIterable<Splitsplits)
        {
            checkNotNull(splits"splits is null");
            for (Split split : splits) {
                this..put(sourceIdsplit);
            }
        }
        @Override
        public void noMoreSplits(PlanNodeId sourceId)
        {
            .add(sourceId);
            boolean allSourcesComplete = Stream.concat(Stream.of(.getPartitionedSourceNode()), .getRemoteSourceNodes().stream())
                    .filter(Objects::nonNull)
                    .map(PlanNode::getId)
                    .allMatch(::contains);
            if (allSourcesComplete) {
                .finished();
            }
        }
        @Override
        public void setOutputBuffers(OutputBuffers outputBuffers)
        {
            .setOutputBuffers(outputBuffers);
        }
        @Override
        public void addStateChangeListener(final StateChangeListener<TaskInfostateChangeListener)
        {
            {
                @Override
                public void stateChanged(TaskState newValue)
                {
                    stateChangeListener.stateChanged(getTaskInfo());
                }
            });
        }
        @Override
        public void cancel()
        {
            .cancel();
        }
        @Override
        public void abort()
        {
            .abort();
        }
        @Override
        public int getPartitionedSplitCount()
        {
            if (.getState().isDone()) {
                return 0;
            }
            return .size();
        }
        @Override
        public int getQueuedPartitionedSplitCount()
        {
            if (.getState().isDone()) {
                return 0;
            }
            return .size();
        }
        @Override
        public Duration waitForTaskToFinish(Duration maxWait)
                throws InterruptedException
        {
            while (true) {
                TaskState currentState = .getState();
                if (maxWait.toMillis() <= 1 || currentState.isDone()) {
                    return maxWait;
                }
                maxWait = .waitForStateChange(currentStatemaxWait);
            }
        }
    }
New to GrepCode? Check out our FAQ X