use of io.prestosql.execution.TestSqlTaskManager.MockLocationFactory in project hetu-core by openlookeng.
the class TestSqlStageExecution method testFinalStageInfoInternal.
private void testFinalStageInfoInternal() throws Exception {
NodeTaskMap nodeTaskMap = new NodeTaskMap(new FinalizerService());
StageId stageId = new StageId(new QueryId("query"), 0);
SqlStageExecution stage = createSqlStageExecution(stageId, new MockLocationFactory().createStageLocation(stageId), createExchangePlanFragment(), ImmutableMap.of(), new MockRemoteTaskFactory(executor, scheduledExecutor), TEST_SESSION, true, nodeTaskMap, executor, new NoOpFailureDetector(), new SplitSchedulerStats(), new DynamicFilterService(new LocalStateStoreProvider(new SeedStoreManager(new FileSystemClientManager()))), new QuerySnapshotManager(stageId.getQueryId(), NOOP_SNAPSHOT_UTILS, TEST_SESSION));
stage.setOutputBuffers(createInitialEmptyOutputBuffers(ARBITRARY));
// add listener that fetches stage info when the final status is available
SettableFuture<StageInfo> finalStageInfo = SettableFuture.create();
stage.addFinalStageInfoListener(finalStageInfo::set);
// in a background thread add a ton of tasks
CountDownLatch latch = new CountDownLatch(1000);
Future<?> addTasksTask = executor.submit(() -> {
try {
for (int i = 0; i < 1_000_000; i++) {
if (Thread.interrupted()) {
return;
}
InternalNode node = new InternalNode("source" + i, URI.create("http://10.0.0." + (i / 10_000) + ":" + (i % 10_000)), NodeVersion.UNKNOWN, false);
stage.scheduleTask(node, i, OptionalInt.empty());
latch.countDown();
}
} finally {
while (latch.getCount() > 0) {
latch.countDown();
}
}
});
// wait for some tasks to be created, and then abort the query
latch.await(1, MINUTES);
assertFalse(stage.getStageInfo().getTasks().isEmpty());
stage.abort();
// once the final stage info is available, verify that it is complete
StageInfo stageInfo = finalStageInfo.get(1, MINUTES);
assertFalse(stageInfo.getTasks().isEmpty());
assertTrue(stageInfo.isCompleteInfo());
assertSame(stage.getStageInfo(), stageInfo);
// cancel the background thread adding tasks
addTasksTask.cancel(true);
}
Aggregations