use of com.facebook.presto.operator.OutputFactory in project presto by prestodb.
the class TestMemoryPools method testBlocking.
@Test
public void testBlocking() throws Exception {
Session session = testSessionBuilder().setCatalog("tpch").setSchema("tiny").setSystemProperty("task_default_concurrency", "1").build();
LocalQueryRunner localQueryRunner = queryRunnerWithInitialTransaction(session);
// add tpch
localQueryRunner.createCatalog("tpch", new TpchConnectorFactory(1), ImmutableMap.of());
// reserve all the memory in the pool
MemoryPool pool = new MemoryPool(new MemoryPoolId("test"), new DataSize(10, MEGABYTE));
QueryId fakeQueryId = new QueryId("fake");
assertTrue(pool.tryReserve(fakeQueryId, TEN_MEGABYTES));
MemoryPool systemPool = new MemoryPool(new MemoryPoolId("testSystem"), new DataSize(10, MEGABYTE));
QueryContext queryContext = new QueryContext(new QueryId("query"), new DataSize(10, MEGABYTE), pool, systemPool, localQueryRunner.getExecutor());
// discard all output
OutputFactory outputFactory = new PageConsumerOutputFactory(types -> (page -> {
}));
TaskContext taskContext = createTaskContext(queryContext, localQueryRunner.getExecutor(), session);
List<Driver> drivers = localQueryRunner.createDrivers("SELECT COUNT(*) FROM orders JOIN lineitem USING (orderkey)", outputFactory, taskContext);
// run driver, until it blocks
while (!isWaitingForMemory(drivers)) {
for (Driver driver : drivers) {
driver.process();
}
}
// driver should be blocked waiting for memory
for (Driver driver : drivers) {
assertFalse(driver.isFinished());
}
assertTrue(pool.getFreeBytes() <= 0);
pool.free(fakeQueryId, TEN_MEGABYTES);
do {
assertFalse(isWaitingForMemory(drivers));
boolean progress = false;
for (Driver driver : drivers) {
ListenableFuture<?> blocked = driver.process();
progress = progress | blocked.isDone();
}
// query should not block
assertTrue(progress);
} while (!drivers.stream().allMatch(Driver::isFinished));
}
Aggregations