Examples with TestingPrestoServer com.facebook.presto.server.testing.TestingPrestoServer used on opensource projects

Example 36 with TestingPrestoServer

use of com.facebook.presto.server.testing.TestingPrestoServer in project presto by prestodb.

the class TestProcedureCall method setUp.

@BeforeClass
public void setUp() {
    TestingPrestoServer coordinator = ((DistributedQueryRunner) getQueryRunner()).getCoordinator();
    tester = coordinator.getProcedureTester();
    // register procedures in the bogus testing catalog
    ProcedureRegistry procedureRegistry = coordinator.getMetadata().getProcedureRegistry();
    TestingProcedures procedures = new TestingProcedures(coordinator.getProcedureTester());
    procedureRegistry.addProcedures(new ConnectorId(TESTING_CATALOG), procedures.getProcedures(PROCEDURE_SCHEMA));
    session = testSessionBuilder().setCatalog(TESTING_CATALOG).setSchema(PROCEDURE_SCHEMA).build();
}

Example 37 with TestingPrestoServer

use of com.facebook.presto.server.testing.TestingPrestoServer in project presto by prestodb.

the class TestResourceGroupPerQueryLimitEnforcement method setUp.

@BeforeClass
public void setUp() throws Exception {
    queryRunner = TpchQueryRunnerBuilder.builder().build();
    TestingPrestoServer server = queryRunner.getCoordinator();
    server.getResourceGroupManager().get().addConfigurationManagerFactory(new FileResourceGroupConfigurationManagerFactory());
    server.getResourceGroupManager().get().setConfigurationManager("file", ImmutableMap.of("resource-groups.config-file", getResourceFilePath("resource_groups_config_simple.json")));
}

Example 38 with TestingPrestoServer

use of com.facebook.presto.server.testing.TestingPrestoServer in project presto by prestodb.

the class TestDistributedQueryResource method setup.

@BeforeMethod
public void setup() throws Exception {
    client = new JettyHttpClient();
    DistributedQueryRunner runner = createQueryRunner(ImmutableMap.of("query.client.timeout", "20s"), 2);
    coordinator1 = runner.getCoordinator(0);
    coordinator2 = runner.getCoordinator(1);
    Optional<TestingPrestoServer> resourceManager = runner.getResourceManager();
    checkState(resourceManager.isPresent(), "resource manager not present");
    this.resourceManager = resourceManager.get();
    coordinator1.getResourceGroupManager().get().addConfigurationManagerFactory(new FileResourceGroupConfigurationManagerFactory());
    coordinator1.getResourceGroupManager().get().setConfigurationManager("file", ImmutableMap.of("resource-groups.config-file", getResourceFilePath("resource_groups_config_simple.json")));
    coordinator2.getResourceGroupManager().get().addConfigurationManagerFactory(new FileResourceGroupConfigurationManagerFactory());
    coordinator2.getResourceGroupManager().get().setConfigurationManager("file", ImmutableMap.of("resource-groups.config-file", getResourceFilePath("resource_groups_config_simple.json")));
}

Example 39 with TestingPrestoServer

use of com.facebook.presto.server.testing.TestingPrestoServer in project presto by prestodb.

the class TestMemoryManager method testOutOfMemoryKillerMultiCoordinator.

@Test(timeOut = 300_000, expectedExceptions = ExecutionException.class, expectedExceptionsMessageRegExp = ".*Query killed because the cluster is out of memory. Please try again in a few minutes.")
public void testOutOfMemoryKillerMultiCoordinator() throws Exception {
    Map<String, String> properties = ImmutableMap.<String, String>builder().put("task.verbose-stats", "true").put("query.low-memory-killer.delay", "5s").put("query.low-memory-killer.policy", "total-reservation").put("resource-manager.memory-pool-fetch-interval", "10ms").put("resource-manager.query-heartbeat-interval", "10ms").build();
    try (DistributedQueryRunner queryRunner = createQueryRunner(properties, 2)) {
        // Reserve all the memory
        QueryId fakeQueryId = new QueryId("fake");
        for (TestingPrestoServer server : queryRunner.getServers()) {
            for (MemoryPool pool : server.getLocalMemoryManager().getPools()) {
                assertTrue(pool.tryReserve(fakeQueryId, "test", pool.getMaxBytes()));
            }
        }
        List<Future<?>> queryFutures = new ArrayList<>();
        for (int i = 0; i < 2; i++) {
            int coordinator = i;
            Thread.sleep(500);
            queryFutures.add(executor.submit(() -> {
                queryRunner.execute(coordinator, "SELECT COUNT(*), clerk FROM orders GROUP BY clerk");
            }));
        }
        // Wait for one of the queries to die
        waitForQueryToBeKilled(queryRunner);
        // Release the memory in the reserved pool
        for (TestingPrestoServer server : queryRunner.getServers()) {
            Optional<MemoryPool> reserved = server.getLocalMemoryManager().getReservedPool();
            assertTrue(reserved.isPresent());
            // Free up the entire pool
            reserved.get().free(fakeQueryId, "test", reserved.get().getMaxBytes());
            assertTrue(reserved.get().getFreeBytes() > 0);
        }
        for (Future<?> query : queryFutures) {
            query.get();
        }
    }
}

Example 40 with TestingPrestoServer

use of com.facebook.presto.server.testing.TestingPrestoServer in project presto by prestodb.

the class TestMemoryManager method testClusterPoolsMultiCoordinator.

@Test(timeOut = 60_000)
public void testClusterPoolsMultiCoordinator() throws Exception {
    Map<String, String> properties = ImmutableMap.<String, String>builder().put("task.verbose-stats", "true").put("resource-manager.memory-pool-fetch-interval", "10ms").put("resource-manager.query-heartbeat-interval", "10ms").put("resource-manager.node-status-timeout", "5s").build();
    try (DistributedQueryRunner queryRunner = createQueryRunner(properties, ImmutableMap.of(), properties, 2)) {
        // Reserve all the memory
        QueryId fakeQueryId = new QueryId("fake");
        for (TestingPrestoServer server : queryRunner.getCoordinatorWorkers()) {
            for (MemoryPool pool : server.getLocalMemoryManager().getPools()) {
                assertTrue(pool.tryReserve(fakeQueryId, "test", pool.getMaxBytes()));
            }
        }
        List<Future<?>> queryFutures = new ArrayList<>();
        for (int i = 0; i < 2; i++) {
            int coordinator = i;
            queryFutures.add(executor.submit(() -> queryRunner.execute(coordinator, "SELECT COUNT(*), clerk FROM orders GROUP BY clerk")));
        }
        ClusterMemoryManager memoryManager = queryRunner.getCoordinator(0).getClusterMemoryManager();
        ClusterMemoryPoolInfo reservedPool;
        while ((reservedPool = memoryManager.getClusterInfo(RESERVED_POOL)) == null) {
            MILLISECONDS.sleep(10);
        }
        ClusterMemoryPoolInfo generalPool = memoryManager.getClusterInfo(GENERAL_POOL);
        assertNotNull(generalPool);
        // Wait for the queries to start running and get assigned to the expected pools
        while (generalPool.getAssignedQueries() != 1 || reservedPool.getAssignedQueries() != 1 || generalPool.getBlockedNodes() != 3 || reservedPool.getBlockedNodes() != 3) {
            generalPool = memoryManager.getClusterInfo(GENERAL_POOL);
            reservedPool = memoryManager.getClusterInfo(RESERVED_POOL);
            MILLISECONDS.sleep(10);
        }
        // Make sure the queries are blocked
        List<BasicQueryInfo> currentQueryInfos;
        do {
            currentQueryInfos = queryRunner.getCoordinators().stream().map(TestingPrestoServer::getQueryManager).map(QueryManager::getQueries).flatMap(Collection::stream).collect(toImmutableList());
            MILLISECONDS.sleep(10);
        } while (currentQueryInfos.size() != 2);
        for (BasicQueryInfo info : currentQueryInfos) {
            assertFalse(info.getState().isDone());
        }
        // Check that the pool information propagated to the query objects
        assertNotEquals(currentQueryInfos.get(0).getMemoryPool(), currentQueryInfos.get(1).getMemoryPool());
        while (!currentQueryInfos.stream().allMatch(TestMemoryManager::isBlockedWaitingForMemory)) {
            MILLISECONDS.sleep(10);
            currentQueryInfos = queryRunner.getCoordinators().stream().map(TestingPrestoServer::getQueryManager).map(QueryManager::getQueries).flatMap(Collection::stream).collect(toImmutableList());
            for (BasicQueryInfo info : currentQueryInfos) {
                assertFalse(info.getState().isDone());
            }
        }
        // Release the memory in the reserved pool
        for (TestingPrestoServer server : queryRunner.getCoordinatorWorkers()) {
            Optional<MemoryPool> reserved = server.getLocalMemoryManager().getReservedPool();
            assertTrue(reserved.isPresent());
            // Free up the entire pool
            reserved.get().free(fakeQueryId, "test", reserved.get().getMaxBytes());
            assertTrue(reserved.get().getFreeBytes() > 0);
        }
        // This also checks that the query in the general pool is successfully moved to the reserved pool.
        for (Future<?> query : queryFutures) {
            query.get();
        }
        queryRunner.getCoordinators().stream().map(TestingPrestoServer::getQueryManager).map(QueryManager::getQueries).flatMap(Collection::stream).forEach(info -> assertEquals(info.getState(), FINISHED));
        // Make sure we didn't leak any memory on the workers
        for (TestingPrestoServer worker : queryRunner.getCoordinatorWorkers()) {
            Optional<MemoryPool> reserved = worker.getLocalMemoryManager().getReservedPool();
            assertTrue(reserved.isPresent());
            assertEquals(reserved.get().getMaxBytes(), reserved.get().getFreeBytes());
            MemoryPool general = worker.getLocalMemoryManager().getGeneralPool();
            // Free up the memory we reserved earlier
            general.free(fakeQueryId, "test", general.getMaxBytes());
            assertEquals(general.getMaxBytes(), general.getFreeBytes());
        }
    }
}