@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 setThreadName = 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);
}
}
}
}
}
@Override
public MaterializedResult execute(Session session, @Language("SQL") String sql) {
lock.readLock().lock();
try {
MaterializedOutputFactory outputFactory = new MaterializedOutputFactory();
TaskContext taskContext = createTaskContext(executor, session);
List<Driver> drivers = createDrivers(session, sql, outputFactory, taskContext);
boolean done = false;
while (!done) {
boolean processed = false;
for (Driver driver : drivers) {
if (!driver.isFinished()) {
driver.process();
processed = true;
}
}
done = !processed;
}
return outputFactory.getMaterializingOperator().getMaterializedResult();
} finally {
lock.readLock().unlock();
}
}
@Override
protected List<Driver> createDrivers(TaskContext taskContext) {
if (lookupSourceSupplier == null) {
OperatorFactory ordersTableScan =
createTableScanOperator(0, "orders", "orderkey", "totalprice");
HashBuilderOperatorFactory hashBuilder =
new HashBuilderOperatorFactory(1, ordersTableScan.getTypes(), Ints.asList(0), 1_500_000);
DriverContext driverContext = taskContext.addPipelineContext(false, false).addDriverContext();
Driver driver =
new DriverFactory(false, false, ordersTableScan, hashBuilder).createDriver(driverContext);
while (!driver.isFinished()) {
driver.process();
}
lookupSourceSupplier = hashBuilder.getLookupSourceSupplier();
}
OperatorFactory lineItemTableScan =
createTableScanOperator(0, "lineitem", "orderkey", "quantity");
OperatorFactory joinOperator =
LookupJoinOperators.innerJoin(
1, lookupSourceSupplier, lineItemTableScan.getTypes(), Ints.asList(0));
NullOutputOperatorFactory output = new NullOutputOperatorFactory(2, joinOperator.getTypes());
DriverFactory driverFactory =
new DriverFactory(true, true, lineItemTableScan, joinOperator, output);
DriverContext driverContext = taskContext.addPipelineContext(true, true).addDriverContext();
Driver driver = driverFactory.createDriver(driverContext);
return ImmutableList.of(driver);
}
@Test
public void testAbruptFinish() {
List<Type> types = ImmutableList.<Type>of(VARCHAR, BIGINT, BIGINT);
ValuesOperator source =
new ValuesOperator(
driverContext.addOperatorContext(0, new PlanNodeId("test"), "values"),
types,
rowPagesBuilder(types).addSequencePage(10, 20, 30, 40).build());
PageConsumerOperator sink = createSinkOperator(source);
Driver driver = new Driver(driverContext, source, sink);
assertSame(driver.getDriverContext(), driverContext);
assertFalse(driver.isFinished());
driver.close();
assertTrue(driver.isFinished());
// finish is only called in normal operations
assertFalse(source.isFinished());
assertFalse(sink.isFinished());
// close is always called (values operator doesn't have a closed state)
assertTrue(sink.isClosed());
}
@Test
public void testBrokenOperatorCloseWhileProcessing() throws Exception {
BrokenOperator brokenOperator =
new BrokenOperator(
driverContext.addOperatorContext(0, new PlanNodeId("test"), "source"), false);
final Driver driver =
new Driver(driverContext, brokenOperator, createSinkOperator(brokenOperator));
assertSame(driver.getDriverContext(), driverContext);
// block thread in operator processing
Future<Boolean> driverProcessFor =
executor.submit(
new Callable<Boolean>() {
@Override
public Boolean call() throws Exception {
return driver.processFor(new Duration(1, TimeUnit.MILLISECONDS)).isDone();
}
});
brokenOperator.waitForLocked();
driver.close();
assertTrue(driver.isFinished());
try {
driverProcessFor.get(1, TimeUnit.SECONDS);
fail("Expected InterruptedException");
} catch (ExecutionException e) {
checkArgument(
getRootCause(e) instanceof InterruptedException,
"Expected root cause exception to be an instance of InterruptedException");
}
}
@Test
public void testBrokenOperatorProcessWhileClosing() throws Exception {
BrokenOperator brokenOperator =
new BrokenOperator(
driverContext.addOperatorContext(0, new PlanNodeId("test"), "source"), true);
final Driver driver =
new Driver(driverContext, brokenOperator, createSinkOperator(brokenOperator));
assertSame(driver.getDriverContext(), driverContext);
// block thread in operator close
Future<Boolean> driverClose =
executor.submit(
new Callable<Boolean>() {
@Override
public Boolean call() throws Exception {
driver.close();
return true;
}
});
brokenOperator.waitForLocked();
assertTrue(driver.processFor(new Duration(1, TimeUnit.MILLISECONDS)).isDone());
assertTrue(driver.isFinished());
brokenOperator.unlock();
assertTrue(driverClose.get());
}
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;
}
@Override
public void close() {
Driver driver;
synchronized (this) {
closed = true;
driver = this.driver;
}
if (driver != null) {
driver.close();
}
}
@Override
public synchronized ListenableFuture<?> processFor(Duration duration) {
if (driver == null) {
driver = driverSplitRunnerFactory.createDriver(partitionedSplit);
}
return driver.processFor(duration);
}
public MaterializedResult execute(@Language("SQL") String sql) {
MaterializedOutputFactory outputFactory = new MaterializedOutputFactory();
List<Driver> drivers = createDrivers(sql, outputFactory);
boolean done = false;
while (!done) {
boolean processed = false;
for (Driver driver : drivers) {
if (!driver.isFinished()) {
driver.process();
processed = true;
}
}
done = !processed;
}
return outputFactory.getMaterializingOperator().getMaterializedResult();
}
@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 synchronized boolean isFinished() {
if (closed) {
return true;
}
if (driver == null) {
return false;
}
return driver.isFinished();
}
@Test
public void testAbruptFinish() {
ValuesOperator source =
new ValuesOperator(
driverContext.addOperatorContext(0, "values"),
rowPagesBuilder(SINGLE_VARBINARY, SINGLE_LONG, SINGLE_LONG)
.addSequencePage(10, 20, 30, 40)
.build());
MaterializingOperator sink = createSinkOperator(source);
Driver driver = new Driver(driverContext, source, sink);
assertSame(driver.getDriverContext(), driverContext);
assertFalse(driver.isFinished());
driver.close();
assertTrue(driver.isFinished());
assertTrue(sink.isFinished());
assertTrue(source.isFinished());
}
@Test
public void testNormalFinish() {
ValuesOperator source =
new ValuesOperator(
driverContext.addOperatorContext(0, "values"),
rowPagesBuilder(SINGLE_VARBINARY, SINGLE_LONG, SINGLE_LONG)
.addSequencePage(10, 20, 30, 40)
.build());
MaterializingOperator sink = createSinkOperator(source);
Driver driver = new Driver(driverContext, source, sink);
assertSame(driver.getDriverContext(), driverContext);
assertFalse(driver.isFinished());
ListenableFuture<?> blocked = driver.processFor(new Duration(1, TimeUnit.SECONDS));
assertTrue(blocked.isDone());
assertTrue(driver.isFinished());
assertTrue(sink.isFinished());
assertTrue(source.isFinished());
}
@Test
public void testNormalFinish() {
List<Type> types = ImmutableList.<Type>of(VARCHAR, BIGINT, BIGINT);
ValuesOperator source =
new ValuesOperator(
driverContext.addOperatorContext(0, new PlanNodeId("test"), "values"),
types,
rowPagesBuilder(types).addSequencePage(10, 20, 30, 40).build());
Operator sink = createSinkOperator(source);
Driver driver = new Driver(driverContext, source, sink);
assertSame(driver.getDriverContext(), driverContext);
assertFalse(driver.isFinished());
ListenableFuture<?> blocked = driver.processFor(new Duration(1, TimeUnit.SECONDS));
assertTrue(blocked.isDone());
assertTrue(driver.isFinished());
assertTrue(sink.isFinished());
assertTrue(source.isFinished());
}
@Test
public void testAddSourceFinish() {
PlanNodeId sourceId = new PlanNodeId("source");
TableScanOperator source =
new TableScanOperator(
driverContext.addOperatorContext(99, "values"),
sourceId,
new DataStreamProvider() {
@Override
public Operator createNewDataStream(
OperatorContext operatorContext, Split split, List<ColumnHandle> columns) {
return new ValuesOperator(
driverContext.addOperatorContext(0, "values"),
rowPagesBuilder(SINGLE_VARBINARY, SINGLE_LONG, SINGLE_LONG)
.addSequencePage(10, 20, 30, 40)
.build());
}
},
ImmutableList.of(SINGLE_VARBINARY, SINGLE_LONG, SINGLE_LONG),
ImmutableList.<ColumnHandle>of());
MaterializingOperator sink = createSinkOperator(source);
Driver driver = new Driver(driverContext, source, sink);
assertSame(driver.getDriverContext(), driverContext);
assertFalse(driver.isFinished());
// todo TableScanOperator should be blocked until split is set
assertTrue(driver.processFor(new Duration(1, TimeUnit.MILLISECONDS)).isDone());
assertFalse(driver.isFinished());
driver.updateSource(
new TaskSource(sourceId, ImmutableSet.of(new ScheduledSplit(0, new MockSplit())), true));
assertFalse(driver.isFinished());
assertTrue(driver.processFor(new Duration(1, TimeUnit.SECONDS)).isDone());
assertTrue(driver.isFinished());
assertTrue(sink.isFinished());
assertTrue(source.isFinished());
}
@Test
public void testAddSourceFinish() {
PlanNodeId sourceId = new PlanNodeId("source");
final List<Type> types = ImmutableList.<Type>of(VARCHAR, BIGINT, BIGINT);
TableScanOperator source =
new TableScanOperator(
driverContext.addOperatorContext(99, new PlanNodeId("test"), "values"),
sourceId,
new PageSourceProvider() {
@Override
public ConnectorPageSource createPageSource(
Session session, Split split, List<ColumnHandle> columns) {
return new FixedPageSource(
rowPagesBuilder(types).addSequencePage(10, 20, 30, 40).build());
}
},
types,
ImmutableList.<ColumnHandle>of());
PageConsumerOperator sink = createSinkOperator(source);
Driver driver = new Driver(driverContext, source, sink);
assertSame(driver.getDriverContext(), driverContext);
assertFalse(driver.isFinished());
assertFalse(driver.processFor(new Duration(1, TimeUnit.MILLISECONDS)).isDone());
assertFalse(driver.isFinished());
driver.updateSource(
new TaskSource(
sourceId, ImmutableSet.of(new ScheduledSplit(0, sourceId, newMockSplit())), true));
assertFalse(driver.isFinished());
assertTrue(driver.processFor(new Duration(1, TimeUnit.SECONDS)).isDone());
assertTrue(driver.isFinished());
assertTrue(sink.isFinished());
assertTrue(source.isFinished());
}
@Test
public void testBrokenOperatorAddSource() throws Exception {
PlanNodeId sourceId = new PlanNodeId("source");
final List<Type> types = ImmutableList.<Type>of(VARCHAR, BIGINT, BIGINT);
// create a table scan operator that does not block, which will cause the driver loop to busy
// wait
TableScanOperator source =
new NotBlockedTableScanOperator(
driverContext.addOperatorContext(99, new PlanNodeId("test"), "values"),
sourceId,
new PageSourceProvider() {
@Override
public ConnectorPageSource createPageSource(
Session session, Split split, List<ColumnHandle> columns) {
return new FixedPageSource(
rowPagesBuilder(types).addSequencePage(10, 20, 30, 40).build());
}
},
types,
ImmutableList.<ColumnHandle>of());
BrokenOperator brokenOperator =
new BrokenOperator(driverContext.addOperatorContext(0, new PlanNodeId("test"), "source"));
final Driver driver = new Driver(driverContext, source, brokenOperator);
// block thread in operator processing
Future<Boolean> driverProcessFor =
executor.submit(
new Callable<Boolean>() {
@Override
public Boolean call() throws Exception {
return driver.processFor(new Duration(1, TimeUnit.MILLISECONDS)).isDone();
}
});
brokenOperator.waitForLocked();
assertSame(driver.getDriverContext(), driverContext);
assertFalse(driver.isFinished());
// processFor always returns NOT_BLOCKED, because DriveLockResult was not acquired
assertTrue(driver.processFor(new Duration(1, TimeUnit.MILLISECONDS)).isDone());
assertFalse(driver.isFinished());
driver.updateSource(
new TaskSource(
sourceId, ImmutableSet.of(new ScheduledSplit(0, sourceId, newMockSplit())), true));
assertFalse(driver.isFinished());
// processFor always returns NOT_BLOCKED, because DriveLockResult was not acquired
assertTrue(driver.processFor(new Duration(1, TimeUnit.SECONDS)).isDone());
assertFalse(driver.isFinished());
driver.close();
assertTrue(driver.isFinished());
try {
driverProcessFor.get(1, TimeUnit.SECONDS);
fail("Expected InterruptedException");
} catch (ExecutionException e) {
checkArgument(
getRootCause(e) instanceof InterruptedException,
"Expected root cause exception to be an instance of InterruptedException");
}
}
private static LookupSourceSupplier buildHash(
boolean parallelBuild,
TaskContext taskContext,
List<Integer> hashChannels,
RowPagesBuilder buildPages) {
if (parallelBuild) {
ParallelHashBuilder parallelHashBuilder =
new ParallelHashBuilder(
buildPages.getTypes(),
hashChannels,
buildPages.getHashChannel(),
100,
PARTITION_COUNT);
// collect input data
DriverContext collectDriverContext =
taskContext.addPipelineContext(true, true).addDriverContext();
ValuesOperatorFactory valuesOperatorFactory =
new ValuesOperatorFactory(
0, new PlanNodeId("test"), buildPages.getTypes(), buildPages.build());
OperatorFactory collectOperatorFactory =
parallelHashBuilder.getCollectOperatorFactory(1, new PlanNodeId("test"));
Driver driver =
new Driver(
collectDriverContext,
valuesOperatorFactory.createOperator(collectDriverContext),
collectOperatorFactory.createOperator(collectDriverContext));
while (!driver.isFinished()) {
driver.process();
}
// build hash tables
PipelineContext buildPipeline = taskContext.addPipelineContext(true, true);
OperatorFactory buildOperatorFactory =
parallelHashBuilder.getBuildOperatorFactory(new PlanNodeId("test"));
for (int i = 0; i < PARTITION_COUNT; i++) {
DriverContext buildDriverContext = buildPipeline.addDriverContext();
Driver buildDriver =
new Driver(buildDriverContext, buildOperatorFactory.createOperator(buildDriverContext));
while (!buildDriver.isFinished()) {
buildDriver.process();
}
}
return parallelHashBuilder.getLookupSourceSupplier();
} else {
DriverContext driverContext = taskContext.addPipelineContext(true, true).addDriverContext();
ValuesOperatorFactory valuesOperatorFactory =
new ValuesOperatorFactory(
0, new PlanNodeId("test"), buildPages.getTypes(), buildPages.build());
HashBuilderOperatorFactory hashBuilderOperatorFactory =
new HashBuilderOperatorFactory(
1,
new PlanNodeId("test"),
buildPages.getTypes(),
hashChannels,
buildPages.getHashChannel(),
100);
Driver driver =
new Driver(
driverContext,
valuesOperatorFactory.createOperator(driverContext),
hashBuilderOperatorFactory.createOperator(driverContext));
while (!driver.isFinished()) {
driver.process();
}
return hashBuilderOperatorFactory.getLookupSourceSupplier();
}
}
public List<Driver> createDrivers(
Session session,
@Language("SQL") String sql,
OutputFactory outputFactory,
TaskContext taskContext) {
Statement statement = sqlParser.createStatement(sql);
assertFormattedSql(sqlParser, statement);
PlanNodeIdAllocator idAllocator = new PlanNodeIdAllocator();
FeaturesConfig featuresConfig =
new FeaturesConfig()
.setExperimentalSyntaxEnabled(true)
.setDistributedIndexJoinsEnabled(false)
.setOptimizeHashGeneration(true);
PlanOptimizersFactory planOptimizersFactory =
new PlanOptimizersFactory(metadata, sqlParser, indexManager, featuresConfig, true);
QueryExplainer queryExplainer =
new QueryExplainer(
planOptimizersFactory.get(),
metadata,
accessControl,
sqlParser,
dataDefinitionTask,
featuresConfig.isExperimentalSyntaxEnabled());
Analyzer analyzer =
new Analyzer(
session,
metadata,
sqlParser,
accessControl,
Optional.of(queryExplainer),
featuresConfig.isExperimentalSyntaxEnabled());
Analysis analysis = analyzer.analyze(statement);
Plan plan =
new LogicalPlanner(session, planOptimizersFactory.get(), idAllocator, metadata)
.plan(analysis);
if (printPlan) {
System.out.println(
PlanPrinter.textLogicalPlan(plan.getRoot(), plan.getTypes(), metadata, session));
}
SubPlan subplan = new PlanFragmenter().createSubPlans(plan);
if (!subplan.getChildren().isEmpty()) {
throw new AssertionError("Expected subplan to have no children");
}
LocalExecutionPlanner executionPlanner =
new LocalExecutionPlanner(
metadata,
sqlParser,
pageSourceManager,
indexManager,
pageSinkManager,
null,
compiler,
new IndexJoinLookupStats(),
new CompilerConfig()
.setInterpreterEnabled(false), // make sure tests fail if compiler breaks
new TaskManagerConfig().setTaskDefaultConcurrency(4));
// plan query
LocalExecutionPlan localExecutionPlan =
executionPlanner.plan(
session,
subplan.getFragment().getRoot(),
subplan.getFragment().getOutputLayout(),
plan.getTypes(),
subplan.getFragment().getDistribution(),
outputFactory);
// generate sources
List<TaskSource> sources = new ArrayList<>();
long sequenceId = 0;
for (TableScanNode tableScan : findTableScanNodes(subplan.getFragment().getRoot())) {
TableLayoutHandle layout = tableScan.getLayout().get();
SplitSource splitSource = splitManager.getSplits(session, layout);
ImmutableSet.Builder<ScheduledSplit> scheduledSplits = ImmutableSet.builder();
while (!splitSource.isFinished()) {
for (Split split : getFutureValue(splitSource.getNextBatch(1000))) {
scheduledSplits.add(new ScheduledSplit(sequenceId++, split));
}
}
sources.add(new TaskSource(tableScan.getId(), scheduledSplits.build(), true));
}
// create drivers
List<Driver> drivers = new ArrayList<>();
Map<PlanNodeId, Driver> driversBySource = new HashMap<>();
for (DriverFactory driverFactory : localExecutionPlan.getDriverFactories()) {
for (int i = 0; i < driverFactory.getDriverInstances(); i++) {
DriverContext driverContext =
taskContext
.addPipelineContext(driverFactory.isInputDriver(), driverFactory.isOutputDriver())
.addDriverContext();
Driver driver = driverFactory.createDriver(driverContext);
drivers.add(driver);
for (PlanNodeId sourceId : driver.getSourceIds()) {
driversBySource.put(sourceId, driver);
}
}
driverFactory.close();
}
// add sources to the drivers
for (TaskSource source : sources) {
for (Driver driver : driversBySource.values()) {
driver.updateSource(source);
}
}
return ImmutableList.copyOf(drivers);
}
@Override
public synchronized void close() {
if (driver != null) {
driver.close();
}
}
public List<Driver> createDrivers(
@Language("SQL") String sql, OutputFactory outputFactory, TaskContext taskContext) {
Statement statement = SqlParser.createStatement(sql);
if (printPlan) {
assertFormattedSql(statement);
}
PlanNodeIdAllocator idAllocator = new PlanNodeIdAllocator();
FeaturesConfig featuresConfig = new FeaturesConfig().setExperimentalSyntaxEnabled(true);
PlanOptimizersFactory planOptimizersFactory =
new PlanOptimizersFactory(metadata, splitManager, indexManager, featuresConfig);
QueryExplainer queryExplainer =
new QueryExplainer(
session,
planOptimizersFactory.get(),
metadata,
featuresConfig.isExperimentalSyntaxEnabled());
Analyzer analyzer =
new Analyzer(
session,
metadata,
Optional.of(queryExplainer),
featuresConfig.isExperimentalSyntaxEnabled());
Analysis analysis = analyzer.analyze(statement);
Plan plan =
new LogicalPlanner(session, planOptimizersFactory.get(), idAllocator, metadata)
.plan(analysis);
if (printPlan) {
System.out.println(PlanPrinter.textLogicalPlan(plan.getRoot(), plan.getTypes(), metadata));
}
SubPlan subplan =
new DistributedLogicalPlanner(session, metadata, idAllocator).createSubPlans(plan, true);
assertTrue(subplan.getChildren().isEmpty(), "Expected subplan to have no children");
LocalExecutionPlanner executionPlanner =
new LocalExecutionPlanner(
new NodeInfo(new NodeConfig().setEnvironment("test").setNodeId("test-node")),
metadata,
dataStreamProvider,
indexManager,
storageManager,
recordSinkManager,
null,
compiler);
// plan query
LocalExecutionPlan localExecutionPlan =
executionPlanner.plan(
session, subplan.getFragment().getRoot(), plan.getTypes(), outputFactory);
// generate sources
List<TaskSource> sources = new ArrayList<>();
long sequenceId = 0;
for (PlanNode sourceNode : subplan.getFragment().getSources()) {
if (sourceNode instanceof ValuesNode) {
continue;
}
TableScanNode tableScan = (TableScanNode) sourceNode;
SplitSource splitSource =
splitManager.getPartitionSplits(tableScan.getTable(), getPartitions(tableScan));
ImmutableSet.Builder<ScheduledSplit> scheduledSplits = ImmutableSet.builder();
while (!splitSource.isFinished()) {
try {
for (Split split : splitSource.getNextBatch(1000)) {
scheduledSplits.add(new ScheduledSplit(sequenceId++, split));
}
} catch (InterruptedException e) {
Thread.currentThread().interrupt();
throw Throwables.propagate(e);
}
}
sources.add(new TaskSource(tableScan.getId(), scheduledSplits.build(), true));
}
// create drivers
List<Driver> drivers = new ArrayList<>();
Map<PlanNodeId, Driver> driversBySource = new HashMap<>();
for (DriverFactory driverFactory : localExecutionPlan.getDriverFactories()) {
DriverContext driverContext =
taskContext
.addPipelineContext(driverFactory.isInputDriver(), driverFactory.isOutputDriver())
.addDriverContext();
Driver driver = driverFactory.createDriver(driverContext);
drivers.add(driver);
for (PlanNodeId sourceId : driver.getSourceIds()) {
driversBySource.put(sourceId, driver);
}
driverFactory.close();
}
// add sources to the drivers
for (TaskSource source : sources) {
for (Driver driver : driversBySource.values()) {
driver.updateSource(source);
}
}
return ImmutableList.copyOf(drivers);
}
@Test
public void testBrokenOperatorAddSource() throws Exception {
PlanNodeId sourceId = new PlanNodeId("source");
TableScanOperator source =
new TableScanOperator(
driverContext.addOperatorContext(99, "values"),
sourceId,
new DataStreamProvider() {
@Override
public Operator createNewDataStream(
OperatorContext operatorContext, Split split, List<ColumnHandle> columns) {
return new ValuesOperator(
driverContext.addOperatorContext(0, "values"),
rowPagesBuilder(SINGLE_VARBINARY, SINGLE_LONG, SINGLE_LONG)
.addSequencePage(10, 20, 30, 40)
.build());
}
},
ImmutableList.of(SINGLE_VARBINARY, SINGLE_LONG, SINGLE_LONG),
ImmutableList.<ColumnHandle>of());
BrokenOperator brokenOperator =
new BrokenOperator(driverContext.addOperatorContext(0, "source"));
final Driver driver = new Driver(driverContext, source, brokenOperator);
// block thread in operator processing
Future<Boolean> driverProcessFor =
executor.submit(
new Callable<Boolean>() {
@Override
public Boolean call() throws Exception {
return driver.processFor(new Duration(1, TimeUnit.MILLISECONDS)).isDone();
}
});
brokenOperator.waitForLocked();
assertSame(driver.getDriverContext(), driverContext);
assertFalse(driver.isFinished());
// todo TableScanOperator should be blocked until split is set
assertTrue(driver.processFor(new Duration(1, TimeUnit.MILLISECONDS)).isDone());
assertFalse(driver.isFinished());
driver.updateSource(
new TaskSource(sourceId, ImmutableSet.of(new ScheduledSplit(0, new MockSplit())), true));
assertFalse(driver.isFinished());
assertTrue(driver.processFor(new Duration(1, TimeUnit.SECONDS)).isDone());
assertFalse(driver.isFinished());
driver.close();
assertTrue(driver.isFinished());
try {
driverProcessFor.get(1, TimeUnit.SECONDS);
fail("Expected InterruptedException");
} catch (ExecutionException e) {
checkArgument(
getRootCause(e) instanceof InterruptedException,
"Expected root cause exception to be an instance of InterruptedException");
}
}
public synchronized DriverContext getDriverContext() {
if (driver == null) {
return null;
}
return driver.getDriverContext();
}
@Override
public synchronized boolean isFinished() {
return driver.isFinished();
}