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 setThreadName = new SetThreadName("Task-%s", taskId)) {
// record driver is finished
remainingDrivers.decrementAndGet();
checkTaskCompletion();
queryMonitor.splitCompletionEvent(
taskId, splitRunner.getDriverContext().getDriverStats());
}
}
@Override
public void onFailure(Throwable cause) {
try (SetThreadName setThreadName = new SetThreadName("Task-%s", taskId)) {
taskStateMachine.failed(cause);
// record driver is finished
remainingDrivers.decrementAndGet();
DriverContext driverContext = splitRunner.getDriverContext();
DriverStats driverStats;
if (driverContext != null) {
driverStats = driverContext.getDriverStats();
} else {
// split runner did not start successfully
driverStats = new DriverStats();
}
// fire failed event with cause
queryMonitor.splitFailedEvent(taskId, driverStats, cause);
}
}
},
notificationExecutor);
}
}
public TestSqlTask() {
taskExecutor = new TaskExecutor(8);
taskExecutor.start();
taskNotificationExecutor = newScheduledThreadPool(5, threadsNamed("task-notification-%d"));
LocalExecutionPlanner planner = createTestingPlanner();
sqlTaskExecutionFactory =
new SqlTaskExecutionFactory(
taskNotificationExecutor,
taskExecutor,
planner,
new QueryMonitor(
new ObjectMapperProvider().get(), new NullEventClient(), new NodeInfo("test")),
new TaskManagerConfig());
}
private SqlTaskExecution(
Session 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 setThreadName = 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;
}
}
@Override
public void stateChanged(TaskState taskState) {
if (taskState.isDone()) {
taskExecutor.removeTask(taskHandle);
}
}
private SqlTaskExecution(
TaskStateMachine taskStateMachine,
TaskContext taskContext,
SharedBuffer sharedBuffer,
PlanFragment fragment,
LocalExecutionPlanner planner,
TaskExecutor taskExecutor,
QueryMonitor queryMonitor,
Executor notificationExecutor) {
this.taskStateMachine = checkNotNull(taskStateMachine, "taskStateMachine is null");
this.taskId = taskStateMachine.getTaskId();
this.taskContext = checkNotNull(taskContext, "taskContext is null");
this.sharedBuffer = checkNotNull(sharedBuffer, "sharedBuffer is null");
this.taskExecutor = checkNotNull(taskExecutor, "driverExecutor is null");
this.notificationExecutor = checkNotNull(notificationExecutor, "notificationExecutor is null");
this.queryMonitor = checkNotNull(queryMonitor, "queryMonitor is null");
try (SetThreadName ignored = new SetThreadName("Task-%s", taskId)) {
List<DriverFactory> driverFactories;
try {
OutputFactory outputOperatorFactory;
if (fragment.getOutputPartitioning() == OutputPartitioning.NONE) {
outputOperatorFactory = new TaskOutputFactory(sharedBuffer);
} else if (fragment.getOutputPartitioning() == OutputPartitioning.HASH) {
outputOperatorFactory = new PartitionedOutputFactory(sharedBuffer);
} else {
throw new PrestoException(
NOT_SUPPORTED,
format("OutputPartitioning %s is not supported", fragment.getOutputPartitioning()));
}
LocalExecutionPlan localExecutionPlan =
planner.plan(
taskContext.getSession(),
fragment.getRoot(),
fragment.getOutputLayout(),
fragment.getSymbols(),
fragment.getDistribution(),
outputOperatorFactory);
driverFactories = localExecutionPlan.getDriverFactories();
} catch (Throwable e) {
// planning failed
taskStateMachine.failed(e);
throw Throwables.propagate(e);
}
// 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;
// don't register the task if it is already completed (most likely failed during planning
// above)
if (!taskStateMachine.getState().isDone()) {
taskHandle = taskExecutor.addTask(taskId);
taskStateMachine.addStateChangeListener(
new RemoveTaskHandleWhenDone(taskExecutor, taskHandle));
} else {
taskHandle = null;
}
sharedBuffer.addStateChangeListener(
new CheckTaskCompletionOnBufferFinish(SqlTaskExecution.this));
}
}
@AfterClass
public void destroy() throws Exception {
taskExecutor.stop();
taskNotificationExecutor.shutdownNow();
}
@BeforeMethod
public void setUp() throws Exception {
DualMetadata dualMetadata = new DualMetadata();
TableHandle tableHandle =
dualMetadata.getTableHandle(new SchemaTableName("default", DualMetadata.NAME));
assertNotNull(tableHandle, "tableHandle is null");
ColumnHandle columnHandle = dualMetadata.getColumnHandle(tableHandle, DualMetadata.COLUMN_NAME);
assertNotNull(columnHandle, "columnHandle is null");
Symbol symbol = new Symbol(DualMetadata.COLUMN_NAME);
MetadataManager metadata = new MetadataManager(new FeaturesConfig());
metadata.addInternalSchemaMetadata(MetadataManager.INTERNAL_CONNECTOR_ID, dualMetadata);
DualSplitManager dualSplitManager = new DualSplitManager(new InMemoryNodeManager());
PartitionResult partitionResult =
dualSplitManager.getPartitions(tableHandle, TupleDomain.all());
SplitSource splitSource =
dualSplitManager.getPartitionSplits(tableHandle, partitionResult.getPartitions());
split = Iterables.getOnlyElement(splitSource.getNextBatch(1));
assertTrue(splitSource.isFinished());
LocalExecutionPlanner planner =
new LocalExecutionPlanner(
new NodeInfo("test"),
metadata,
new DataStreamManager(new DualDataStreamProvider()),
new MockLocalStorageManager(new File("target/temp")),
new RecordSinkManager(),
new MockExchangeClientSupplier(),
new ExpressionCompiler(metadata));
taskExecutor = new TaskExecutor(8);
taskExecutor.start();
tableScanNodeId = new PlanNodeId("tableScan");
PlanFragment testFragment =
new PlanFragment(
new PlanFragmentId("fragment"),
new TableScanNode(
tableScanNodeId,
tableHandle,
ImmutableList.of(symbol),
ImmutableMap.of(symbol, columnHandle),
null,
Optional.<GeneratedPartitions>absent()),
ImmutableMap.of(symbol, Type.VARCHAR),
PlanDistribution.SOURCE,
tableScanNodeId,
OutputPartitioning.NONE,
ImmutableList.<Symbol>of());
TaskId taskId = new TaskId("query", "stage", "task");
Session session = new Session("user", "test", "default", "default", "test", "test");
taskNotificationExecutor = Executors.newCachedThreadPool(threadsNamed("task-notification-%d"));
outputBuffers = OutputBuffers.INITIAL_EMPTY_OUTPUT_BUFFERS;
taskExecution =
SqlTaskExecution.createSqlTaskExecution(
session,
taskId,
URI.create("fake://task/" + taskId),
testFragment,
ImmutableList.<TaskSource>of(),
outputBuffers,
planner,
new DataSize(32, Unit.MEGABYTE),
taskExecutor,
taskNotificationExecutor,
new DataSize(256, Unit.MEGABYTE),
new DataSize(8, Unit.MEGABYTE),
new QueryMonitor(
new ObjectMapperProvider().get(), new NullEventClient(), new NodeInfo("test")),
false);
}
@AfterMethod
public void tearDown() throws Exception {
taskExecutor.stop();
taskNotificationExecutor.shutdownNow();
}