Search in sources :

Example 1 with OutputFactory

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));
}
Also used : LocalQueryRunner(com.facebook.presto.testing.LocalQueryRunner) PageConsumerOutputFactory(com.facebook.presto.testing.PageConsumerOperator.PageConsumerOutputFactory) TaskContext(com.facebook.presto.operator.TaskContext) ListenableFuture(com.google.common.util.concurrent.ListenableFuture) TestingTaskContext.createTaskContext(com.facebook.presto.testing.TestingTaskContext.createTaskContext) ImmutableMap(com.google.common.collect.ImmutableMap) Session(com.facebook.presto.Session) MEGABYTE(io.airlift.units.DataSize.Unit.MEGABYTE) TpchConnectorFactory(com.facebook.presto.tpch.TpchConnectorFactory) OutputFactory(com.facebook.presto.operator.OutputFactory) Test(org.testng.annotations.Test) Driver(com.facebook.presto.operator.Driver) TestingSession.testSessionBuilder(com.facebook.presto.testing.TestingSession.testSessionBuilder) MemoryPoolId(com.facebook.presto.spi.memory.MemoryPoolId) DataSize(io.airlift.units.DataSize) List(java.util.List) QueryId(com.facebook.presto.spi.QueryId) Assert.assertTrue(org.testng.Assert.assertTrue) LocalQueryRunner.queryRunnerWithInitialTransaction(com.facebook.presto.testing.LocalQueryRunner.queryRunnerWithInitialTransaction) OperatorContext(com.facebook.presto.operator.OperatorContext) Assert.assertFalse(org.testng.Assert.assertFalse) TpchConnectorFactory(com.facebook.presto.tpch.TpchConnectorFactory) TaskContext(com.facebook.presto.operator.TaskContext) TestingTaskContext.createTaskContext(com.facebook.presto.testing.TestingTaskContext.createTaskContext) PageConsumerOutputFactory(com.facebook.presto.testing.PageConsumerOperator.PageConsumerOutputFactory) QueryId(com.facebook.presto.spi.QueryId) Driver(com.facebook.presto.operator.Driver) LocalQueryRunner(com.facebook.presto.testing.LocalQueryRunner) DataSize(io.airlift.units.DataSize) PageConsumerOutputFactory(com.facebook.presto.testing.PageConsumerOperator.PageConsumerOutputFactory) OutputFactory(com.facebook.presto.operator.OutputFactory) MemoryPoolId(com.facebook.presto.spi.memory.MemoryPoolId) Session(com.facebook.presto.Session) Test(org.testng.annotations.Test)

Aggregations

Session (com.facebook.presto.Session)1 Driver (com.facebook.presto.operator.Driver)1 OperatorContext (com.facebook.presto.operator.OperatorContext)1 OutputFactory (com.facebook.presto.operator.OutputFactory)1 TaskContext (com.facebook.presto.operator.TaskContext)1 QueryId (com.facebook.presto.spi.QueryId)1 MemoryPoolId (com.facebook.presto.spi.memory.MemoryPoolId)1 LocalQueryRunner (com.facebook.presto.testing.LocalQueryRunner)1 LocalQueryRunner.queryRunnerWithInitialTransaction (com.facebook.presto.testing.LocalQueryRunner.queryRunnerWithInitialTransaction)1 PageConsumerOutputFactory (com.facebook.presto.testing.PageConsumerOperator.PageConsumerOutputFactory)1 TestingSession.testSessionBuilder (com.facebook.presto.testing.TestingSession.testSessionBuilder)1 TestingTaskContext.createTaskContext (com.facebook.presto.testing.TestingTaskContext.createTaskContext)1 TpchConnectorFactory (com.facebook.presto.tpch.TpchConnectorFactory)1 ImmutableMap (com.google.common.collect.ImmutableMap)1 ListenableFuture (com.google.common.util.concurrent.ListenableFuture)1 DataSize (io.airlift.units.DataSize)1 MEGABYTE (io.airlift.units.DataSize.Unit.MEGABYTE)1 List (java.util.List)1 Assert.assertFalse (org.testng.Assert.assertFalse)1 Assert.assertTrue (org.testng.Assert.assertTrue)1