Java tutorial
/* * 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 com.facebook.presto.OutputBuffers; import com.facebook.presto.ScheduledSplit; import com.facebook.presto.TaskSource; import com.facebook.presto.event.query.QueryMonitor; import com.facebook.presto.execution.SharedBuffer.QueueState; import com.facebook.presto.execution.StateMachine.StateChangeListener; import com.facebook.presto.execution.TaskExecutor.TaskHandle; import com.facebook.presto.operator.Driver; import com.facebook.presto.operator.DriverContext; import com.facebook.presto.operator.DriverFactory; import com.facebook.presto.operator.DriverStats; import com.facebook.presto.operator.PipelineContext; import com.facebook.presto.operator.TaskContext; import com.facebook.presto.operator.TaskOutputOperator.TaskOutputFactory; import com.facebook.presto.spi.ConnectorSession; import com.facebook.presto.sql.planner.LocalExecutionPlanner; import com.facebook.presto.sql.planner.LocalExecutionPlanner.LocalExecutionPlan; import com.facebook.presto.sql.planner.PlanFragment; import com.facebook.presto.sql.planner.PlanFragment.PlanDistribution; import com.facebook.presto.sql.planner.plan.PlanNodeId; import com.facebook.presto.util.SetThreadName; import com.google.common.base.Objects; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableSet; import com.google.common.util.concurrent.FutureCallback; import com.google.common.util.concurrent.Futures; import com.google.common.util.concurrent.ListenableFuture; import io.airlift.units.DataSize; import io.airlift.units.Duration; import org.joda.time.DateTime; import javax.annotation.Nullable; import javax.annotation.concurrent.GuardedBy; import java.lang.ref.WeakReference; import java.net.URI; import java.util.ArrayList; import java.util.HashMap; import java.util.List; import java.util.Map; import java.util.Set; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ConcurrentMap; import java.util.concurrent.CopyOnWriteArrayList; import java.util.concurrent.Executor; import java.util.concurrent.ExecutorService; import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicInteger; import java.util.concurrent.atomic.AtomicLong; import java.util.concurrent.atomic.AtomicReference; import static com.facebook.presto.util.Failures.toFailures; import static com.google.common.base.Preconditions.checkArgument; import static com.google.common.base.Preconditions.checkNotNull; import static com.google.common.base.Preconditions.checkState; import static java.lang.Math.max; public class SqlTaskExecution implements TaskExecution { private final TaskId taskId; private final URI location; private final TaskExecutor taskExecutor; private final Executor notificationExecutor; private final TaskStateMachine taskStateMachine; private final TaskContext taskContext; private final SharedBuffer sharedBuffer; private final QueryMonitor queryMonitor; private final TaskHandle taskHandle; private final List<WeakReference<Driver>> drivers = new CopyOnWriteArrayList<>(); /** * Number of drivers that have been sent to the TaskExecutor that have not finished. */ private final AtomicInteger remainingDrivers = new AtomicInteger(); // guarded for update only @GuardedBy("this") private final ConcurrentMap<PlanNodeId, TaskSource> unpartitionedSources = new ConcurrentHashMap<>(); @GuardedBy("this") private long maxAcknowledgedSplit = Long.MIN_VALUE; private final AtomicReference<DateTime> lastHeartbeat = new AtomicReference<>(DateTime.now()); private final PlanNodeId partitionedSourceId; private final DriverSplitRunnerFactory partitionedDriverFactory; private final List<DriverSplitRunnerFactory> unpartitionedDriverFactories; private final AtomicLong nextTaskInfoVersion = new AtomicLong(TaskInfo.STARTING_VERSION); public static SqlTaskExecution createSqlTaskExecution(ConnectorSession session, TaskId taskId, URI location, PlanFragment fragment, List<TaskSource> sources, OutputBuffers outputBuffers, LocalExecutionPlanner planner, DataSize maxBufferSize, TaskExecutor taskExecutor, ExecutorService notificationExecutor, DataSize maxTaskMemoryUsage, DataSize operatorPreAllocatedMemory, QueryMonitor queryMonitor, boolean cpuTimerEnabled) { SqlTaskExecution task = new SqlTaskExecution(session, taskId, location, fragment, outputBuffers, planner, maxBufferSize, taskExecutor, maxTaskMemoryUsage, operatorPreAllocatedMemory, queryMonitor, notificationExecutor, cpuTimerEnabled); try (SetThreadName ignored = new SetThreadName("Task-%s", taskId)) { task.start(); task.addSources(sources); task.recordHeartbeat(); return task; } } private SqlTaskExecution(ConnectorSession session, TaskId taskId, URI location, PlanFragment fragment, OutputBuffers outputBuffers, LocalExecutionPlanner planner, DataSize maxBufferSize, TaskExecutor taskExecutor, DataSize maxTaskMemoryUsage, DataSize operatorPreAllocatedMemory, QueryMonitor queryMonitor, Executor notificationExecutor, boolean cpuTimerEnabled) { try (SetThreadName ignored = new SetThreadName("Task-%s", taskId)) { this.taskId = checkNotNull(taskId, "taskId is null"); this.location = checkNotNull(location, "location is null"); this.taskExecutor = checkNotNull(taskExecutor, "driverExecutor is null"); this.notificationExecutor = checkNotNull(notificationExecutor, "notificationExecutor is null"); this.taskStateMachine = new TaskStateMachine(taskId, notificationExecutor); taskStateMachine.addStateChangeListener(new StateChangeListener<TaskState>() { @Override public void stateChanged(TaskState taskState) { if (taskState.isDone()) { SqlTaskExecution.this.taskExecutor.removeTask(taskHandle); // make sure buffers are cleaned up if (taskState != TaskState.FAILED) { // don't close buffers for a failed query // closed buffers signal to upstream tasks that everything finished cleanly sharedBuffer.destroy(); } } } }); this.taskContext = new TaskContext(taskStateMachine, notificationExecutor, session, checkNotNull(maxTaskMemoryUsage, "maxTaskMemoryUsage is null"), checkNotNull(operatorPreAllocatedMemory, "operatorPreAllocatedMemory is null"), cpuTimerEnabled); this.sharedBuffer = new SharedBuffer(taskId, notificationExecutor, checkNotNull(maxBufferSize, "maxBufferSize is null"), outputBuffers); sharedBuffer.addStateChangeListener(new StateChangeListener<QueueState>() { @Override public void stateChanged(QueueState taskState) { if (taskState == QueueState.FINISHED) { checkTaskCompletion(); } } }); this.queryMonitor = checkNotNull(queryMonitor, "queryMonitor is null"); taskHandle = taskExecutor.addTask(taskId); LocalExecutionPlan localExecutionPlan = planner.plan(session, fragment.getRoot(), fragment.getSymbols(), new TaskOutputFactory(sharedBuffer)); List<DriverFactory> driverFactories = localExecutionPlan.getDriverFactories(); // index driver factories DriverSplitRunnerFactory partitionedDriverFactory = null; ImmutableList.Builder<DriverSplitRunnerFactory> unpartitionedDriverFactories = ImmutableList.builder(); for (DriverFactory driverFactory : driverFactories) { if (driverFactory.getSourceIds().contains(fragment.getPartitionedSource())) { checkState(partitionedDriverFactory == null, "multiple partitioned sources are not supported"); partitionedDriverFactory = new DriverSplitRunnerFactory(driverFactory); } else { unpartitionedDriverFactories.add(new DriverSplitRunnerFactory(driverFactory)); } } this.unpartitionedDriverFactories = unpartitionedDriverFactories.build(); if (fragment.getDistribution() == PlanDistribution.SOURCE) { checkArgument(partitionedDriverFactory != null, "Fragment is partitioned, but no partitioned driver found"); } this.partitionedSourceId = fragment.getPartitionedSource(); this.partitionedDriverFactory = partitionedDriverFactory; } } // // This code starts registers a callback with access to this class, and this // call back is access from another thread, so this code can not be placed in the constructor private void start() { // start unpartitioned drivers List<DriverSplitRunner> runners = new ArrayList<>(); for (DriverSplitRunnerFactory driverFactory : unpartitionedDriverFactories) { runners.add(driverFactory.createDriverRunner(null)); driverFactory.setNoMoreSplits(); } enqueueDrivers(true, runners); } @Override public TaskId getTaskId() { return taskId; } @Override public TaskContext getTaskContext() { return taskContext; } @Override public void waitForStateChange(TaskState currentState, Duration maxWait) throws InterruptedException { try (SetThreadName ignored = new SetThreadName("Task-%s", taskId)) { taskStateMachine.waitForStateChange(currentState, maxWait); } } @Override public TaskInfo getTaskInfo() { try (SetThreadName ignored = new SetThreadName("Task-%s", taskId)) { checkTaskCompletion(); TaskState state = taskStateMachine.getState(); List<ExecutionFailureInfo> failures = ImmutableList.of(); if (state == TaskState.FAILED) { failures = toFailures(taskStateMachine.getFailureCauses()); } return new TaskInfo(taskStateMachine.getTaskId(), nextTaskInfoVersion.getAndIncrement(), state, location, lastHeartbeat.get(), sharedBuffer.getInfo(), getNoMoreSplits(), taskContext.getTaskStats(), failures); } } @Override public void addSources(List<TaskSource> sources) { checkNotNull(sources, "sources is null"); checkState(!Thread.holdsLock(this), "Can not add sources while holding a lock on the %s", getClass().getSimpleName()); try (SetThreadName ignored = new SetThreadName("Task-%s", taskId)) { // update our record of sources and schedule drivers for new partitioned splits Map<PlanNodeId, TaskSource> updatedUnpartitionedSources = updateSources(sources); // tell existing drivers about the new splits; it is safe to update drivers // multiple times and out of order because sources contain full record of // the unpartitioned splits for (TaskSource source : updatedUnpartitionedSources.values()) { // tell all the existing drivers this source is finished for (WeakReference<Driver> driverReference : drivers) { Driver driver = driverReference.get(); // the driver can be GCed due to a failure or a limit if (driver != null) { driver.updateSource(source); } else { // remove the weak reference from the list to avoid a memory leak // NOTE: this is a concurrent safe operation on a CopyOnWriteArrayList drivers.remove(driverReference); } } } } } private synchronized Map<PlanNodeId, TaskSource> updateSources(List<TaskSource> sources) { Map<PlanNodeId, TaskSource> updatedUnpartitionedSources = new HashMap<>(); // don't update maxAcknowledgedSplit until the end because task sources may not // be in sorted order and if we updated early we could skip splits long newMaxAcknowledgedSplit = maxAcknowledgedSplit; for (TaskSource source : sources) { PlanNodeId sourceId = source.getPlanNodeId(); if (sourceId.equals(partitionedSourceId)) { // partitioned split ImmutableList.Builder<DriverSplitRunner> runners = ImmutableList.builder(); for (ScheduledSplit scheduledSplit : source.getSplits()) { // only add a split if we have not already scheduled it if (scheduledSplit.getSequenceId() > maxAcknowledgedSplit) { // create a new driver for the split runners.add(partitionedDriverFactory.createDriverRunner(scheduledSplit)); newMaxAcknowledgedSplit = max(scheduledSplit.getSequenceId(), newMaxAcknowledgedSplit); } } enqueueDrivers(false, runners.build()); if (source.isNoMoreSplits()) { partitionedDriverFactory.setNoMoreSplits(); } } else { // unpartitioned split // update newMaxAcknowledgedSplit for (ScheduledSplit scheduledSplit : source.getSplits()) { newMaxAcknowledgedSplit = max(scheduledSplit.getSequenceId(), newMaxAcknowledgedSplit); } // create new source TaskSource newSource; TaskSource currentSource = unpartitionedSources.get(sourceId); if (currentSource == null) { newSource = source; } else { newSource = currentSource.update(source); } // only record new source if something changed if (newSource != currentSource) { unpartitionedSources.put(sourceId, newSource); updatedUnpartitionedSources.put(sourceId, newSource); } } } maxAcknowledgedSplit = newMaxAcknowledgedSplit; return updatedUnpartitionedSources; } @Override public synchronized void addResultQueue(OutputBuffers outputBuffers) { checkNotNull(outputBuffers, "outputBuffers is null"); try (SetThreadName ignored = new SetThreadName("Task-%s", taskId)) { sharedBuffer.setOutputBuffers(outputBuffers); } } private synchronized void enqueueDrivers(boolean forceRunSplit, List<DriverSplitRunner> runners) { // schedule driver to be executed List<ListenableFuture<?>> finishedFutures = taskExecutor.enqueueSplits(taskHandle, forceRunSplit, runners); checkState(finishedFutures.size() == runners.size(), "Expected %s futures but got %s", runners.size(), finishedFutures.size()); // record new driver remainingDrivers.addAndGet(finishedFutures.size()); // when driver completes, update state and fire events for (int i = 0; i < finishedFutures.size(); i++) { ListenableFuture<?> finishedFuture = finishedFutures.get(i); final DriverSplitRunner splitRunner = runners.get(i); Futures.addCallback(finishedFuture, new FutureCallback<Object>() { @Override public void onSuccess(Object result) { try (SetThreadName ignored = new SetThreadName("Task-%s", taskId)) { // record driver is finished remainingDrivers.decrementAndGet(); checkTaskCompletion(); queryMonitor.splitCompletionEvent(taskId, getDriverStats()); } } @Override public void onFailure(Throwable cause) { try (SetThreadName ignored = new SetThreadName("Task-%s", taskId)) { taskStateMachine.failed(cause); // record driver is finished remainingDrivers.decrementAndGet(); // fire failed event with cause queryMonitor.splitFailedEvent(taskId, getDriverStats(), cause); } } private DriverStats getDriverStats() { DriverContext driverContext = splitRunner.getDriverContext(); DriverStats driverStats; if (driverContext != null) { driverStats = driverContext.getDriverStats(); } else { // split runner did not start successfully driverStats = new DriverStats(); } return driverStats; } }, notificationExecutor); } } private Set<PlanNodeId> getNoMoreSplits() { ImmutableSet.Builder<PlanNodeId> noMoreSplits = ImmutableSet.builder(); if (partitionedDriverFactory != null && partitionedDriverFactory.isNoMoreSplits()) { noMoreSplits.add(partitionedSourceId); } for (TaskSource taskSource : unpartitionedSources.values()) { if (taskSource.isNoMoreSplits()) { noMoreSplits.add(taskSource.getPlanNodeId()); } } return noMoreSplits.build(); } private synchronized void checkTaskCompletion() { if (taskStateMachine.getState().isDone()) { return; } // are there more partition splits expected? if (partitionedDriverFactory != null && !partitionedDriverFactory.isNoMoreSplits()) { return; } // do we still have running tasks? if (remainingDrivers.get() != 0) { return; } // no more output will be created sharedBuffer.finish(); // are there still pages in the output buffer if (!sharedBuffer.isFinished()) { return; } // Cool! All done! taskStateMachine.finished(); } @Override public void cancel() { try (SetThreadName ignored = new SetThreadName("Task-%s", taskId)) { taskStateMachine.cancel(); } } @Override public void fail(Throwable cause) { try (SetThreadName ignored = new SetThreadName("Task-%s", taskId)) { taskStateMachine.failed(cause); } } @Override public BufferResult getResults(String outputId, long startingSequenceId, DataSize maxSize, Duration maxWait) throws InterruptedException { checkNotNull(outputId, "outputId is null"); checkArgument(maxSize.toBytes() > 0, "maxSize must be at least 1 byte"); checkNotNull(maxWait, "maxWait is null"); checkState(!Thread.holdsLock(this), "Can not get result data while holding a lock on the %s", getClass().getSimpleName()); try (SetThreadName ignored = new SetThreadName("Task-%s", taskId)) { return sharedBuffer.get(outputId, startingSequenceId, maxSize, maxWait); } } @Override public void abortResults(String outputId) { try (SetThreadName ignored = new SetThreadName("Task-%s", taskId)) { sharedBuffer.abort(outputId); } } @Override public void recordHeartbeat() { this.lastHeartbeat.set(DateTime.now()); } @Override public String toString() { return Objects.toStringHelper(this).add("taskId", taskId).add("remainingDrivers", remainingDrivers) .add("unpartitionedSources", unpartitionedSources).toString(); } private class DriverSplitRunnerFactory { private final DriverFactory driverFactory; private final PipelineContext pipelineContext; private final AtomicInteger pendingCreation = new AtomicInteger(); private final AtomicBoolean noMoreSplits = new AtomicBoolean(); private DriverSplitRunnerFactory(DriverFactory driverFactory) { this.driverFactory = driverFactory; this.pipelineContext = taskContext.addPipelineContext(driverFactory.isInputDriver(), driverFactory.isOutputDriver()); } private DriverSplitRunner createDriverRunner(@Nullable ScheduledSplit partitionedSplit) { pendingCreation.incrementAndGet(); // create driver context immediately so the driver existence is recorded in the stats // the number of drivers is used to balance work across nodes DriverContext driverContext = pipelineContext.addDriverContext(); return new DriverSplitRunner(this, driverContext, partitionedSplit); } private Driver createDriver(DriverContext driverContext, @Nullable ScheduledSplit partitionedSplit) { Driver driver = driverFactory.createDriver(driverContext); // record driver so other threads add unpartitioned sources can see the driver // NOTE: this MUST be done before reading unpartitionedSources, so we see a consistent view of the unpartitioned sources drivers.add(new WeakReference<>(driver)); if (partitionedSplit != null) { // TableScanOperator requires partitioned split to be added before the first call to process driver.updateSource(new TaskSource(partitionedSourceId, ImmutableSet.of(partitionedSplit), true)); } // add unpartitioned sources for (TaskSource source : unpartitionedSources.values()) { driver.updateSource(source); } pendingCreation.decrementAndGet(); closeDriverFactoryIfFullyCreated(); return driver; } private boolean isNoMoreSplits() { return noMoreSplits.get(); } private void setNoMoreSplits() { noMoreSplits.set(true); closeDriverFactoryIfFullyCreated(); } private void closeDriverFactoryIfFullyCreated() { if (isNoMoreSplits() && pendingCreation.get() <= 0) { driverFactory.close(); } } } private static class DriverSplitRunner implements SplitRunner { private final DriverSplitRunnerFactory driverSplitRunnerFactory; private final DriverContext driverContext; @GuardedBy("this") private boolean closed; @Nullable private final ScheduledSplit partitionedSplit; @GuardedBy("this") private Driver driver; private DriverSplitRunner(DriverSplitRunnerFactory driverSplitRunnerFactory, DriverContext driverContext, @Nullable ScheduledSplit partitionedSplit) { this.driverSplitRunnerFactory = checkNotNull(driverSplitRunnerFactory, "driverFactory is null"); this.driverContext = checkNotNull(driverContext, "driverContext is null"); this.partitionedSplit = partitionedSplit; } public synchronized DriverContext getDriverContext() { if (driver == null) { return null; } return driver.getDriverContext(); } @Override public synchronized boolean isFinished() { if (closed) { return true; } if (driver == null) { return false; } return driver.isFinished(); } @Override public ListenableFuture<?> processFor(Duration duration) { Driver driver; synchronized (this) { // if close() was called before we get here, there's not point in even creating the driver if (closed) { return Futures.immediateFuture(null); } if (this.driver == null) { this.driver = driverSplitRunnerFactory.createDriver(driverContext, partitionedSplit); } driver = this.driver; } return driver.processFor(duration); } @Override public void close() { Driver driver; synchronized (this) { closed = true; driver = this.driver; } if (driver != null) { driver.close(); } } } }