@BeforeClass private void setUp() throws Exception { session = testSessionBuilder() .setSystemProperty("task_concurrency", "1") .setCatalog("tpch") .setSchema("tiny") .setClientInfo("{\"clientVersion\":\"testVersion\"}") .build(); queryRunner = new DistributedQueryRunner(session, 1); queryRunner.installPlugin(new TpchPlugin()); queryRunner.installPlugin(new TestingEventListenerPlugin(generatedEvents)); queryRunner.installPlugin(new ResourceGroupManagerPlugin()); queryRunner.createCatalog("tpch", "tpch", ImmutableMap.of("tpch.splits-per-node", Integer.toString(SPLITS_PER_NODE))); queryRunner.getCoordinator().getResourceGroupManager().get() .setConfigurationManager("file", ImmutableMap.of("resource-groups.config-file", getResourceFilePath("resource_groups_config_simple.json"))); }
public static DistributedQueryRunner createQueryRunner(Session session) throws Exception { DistributedQueryRunner queryRunner = DistributedQueryRunner.builder(session) .setNodeCount(2) .build(); try { queryRunner.installPlugin(new TpchPlugin()); queryRunner.createCatalog("tpch", "tpch"); return queryRunner; } catch (Exception e) { queryRunner.close(); throw e; } } }
public static DistributedQueryRunner createQueryRunner(Session session, Map<String, String> properties) throws Exception { DistributedQueryRunner queryRunner = new DistributedQueryRunner(session, 2, properties); try { queryRunner.installPlugin(new TpchPlugin()); queryRunner.createCatalog("tpch", "tpch"); return queryRunner; } catch (Exception e) { queryRunner.close(); throw e; } } }
@Override public DistributedQueryRunner build() throws Exception { DistributedQueryRunner queryRunner = buildWithoutCatalogs(); try { queryRunner.createCatalog("tpch", "tpch"); return queryRunner; } catch (Exception e) { queryRunner.close(); throw e; } }
private void assertResourceGroup(DistributedQueryRunner queryRunner, Session session, String query, ResourceGroupId expectedResourceGroup) throws InterruptedException { QueryId queryId = createQuery(queryRunner, session, query); waitForQueryState(queryRunner, queryId, ImmutableSet.of(RUNNING, FINISHED)); Optional<ResourceGroupId> resourceGroupId = queryRunner.getCoordinator().getQueryManager().getFullQueryInfo(queryId).getResourceGroupId(); assertTrue(resourceGroupId.isPresent(), "Query should have a resource group"); assertEquals(resourceGroupId.get(), expectedResourceGroup, format("Expected: '%s' resource group, found: %s", expectedResourceGroup, resourceGroupId.get())); }
@Test(timeOut = 240_000) public void testResourceGroupManagerWithTwoDashboardQueriesRequestedAtTheSameTime() throws Exception { try (DistributedQueryRunner queryRunner = createQueryRunner()) { queryRunner.installPlugin(new ResourceGroupManagerPlugin()); queryRunner.getCoordinator().getResourceGroupManager().get().setConfigurationManager("file", ImmutableMap.of("resource-groups.config-file", getResourceFilePath("resource_groups_config_dashboard.json"))); QueryId firstDashboardQuery = createDashboardQuery(queryRunner); QueryId secondDashboardQuery = createDashboardQuery(queryRunner); ImmutableSet<QueryState> queuedOrRunning = ImmutableSet.of(QUEUED, RUNNING); waitForQueryState(queryRunner, firstDashboardQuery, queuedOrRunning); waitForQueryState(queryRunner, secondDashboardQuery, queuedOrRunning); } }
private TableMetadata getTableMetadata(String catalog, String schema, String tableName) { Session session = getSession(); Metadata metadata = ((DistributedQueryRunner) getQueryRunner()).getCoordinator().getMetadata(); return transaction(getQueryRunner().getTransactionManager(), getQueryRunner().getAccessControl()) .readOnly() .execute(session, transactionSession -> { Optional<TableHandle> tableHandle = metadata.getTableHandle(transactionSession, new QualifiedObjectName(catalog, schema, tableName)); assertTrue(tableHandle.isPresent()); return metadata.getTableMetadata(transactionSession, tableHandle.get()); }); }
@Test(timeOut = 240_000) public void testSelectorResourceEstimateBasedSelection() throws Exception queryRunner.installPlugin(new ResourceGroupManagerPlugin()); queryRunner.getCoordinator().getResourceGroupManager().get() .setConfigurationManager("file", ImmutableMap.of("resource-groups.config-file", getResourceFilePath("resource_groups_resource_estimate_based_config.json"))); Optional.of(Duration.valueOf("4m")), Optional.empty(), Optional.of(DataSize.valueOf("400MB")))), LONG_LASTING_QUERY,
for (TestingPrestoServer server : queryRunner.getServers()) { for (MemoryPool pool : server.getLocalMemoryManager().getPools()) { assertTrue(pool.tryReserve(fakeQueryId, "test", pool.getMaxBytes())); queryFutures.add(executor.submit(() -> queryRunner.execute("SELECT COUNT(*), clerk FROM orders GROUP BY clerk"))); ClusterMemoryManager memoryManager = queryRunner.getCoordinator().getClusterMemoryManager(); ClusterMemoryPool reservedPool; while ((reservedPool = memoryManager.getPools().get(RESERVED_POOL)) == null) { assertNotNull(generalPool); List<BasicQueryInfo> currentQueryInfos = queryRunner.getCoordinator().getQueryManager().getQueries(); for (BasicQueryInfo info : currentQueryInfos) { assertFalse(info.getState().isDone()); currentQueryInfos = queryRunner.getCoordinator().getQueryManager().getQueries(); for (TestingPrestoServer server : queryRunner.getServers()) { Optional<MemoryPool> reserved = server.getLocalMemoryManager().getReservedPool(); assertTrue(reserved.isPresent()); for (BasicQueryInfo info : queryRunner.getCoordinator().getQueryManager().getQueries()) { for (TestingPrestoServer worker : queryRunner.getServers()) {
private void testNoLeak(@Language("SQL") String query) throws Exception { Map<String, String> properties = ImmutableMap.<String, String>builder() .put("task.verbose-stats", "true") .build(); try (DistributedQueryRunner queryRunner = createQueryRunner(TINY_SESSION, properties)) { executor.submit(() -> queryRunner.execute(query)).get(); for (BasicQueryInfo info : queryRunner.getCoordinator().getQueryManager().getQueries()) { assertEquals(info.getState(), FINISHED); } // Make sure we didn't leak any memory on the workers for (TestingPrestoServer worker : queryRunner.getServers()) { Optional<MemoryPool> reserved = worker.getLocalMemoryManager().getReservedPool(); assertTrue(reserved.isPresent()); assertEquals(reserved.get().getMaxBytes(), reserved.get().getFreeBytes()); MemoryPool general = worker.getLocalMemoryManager().getGeneralPool(); assertEquals(general.getMaxBytes(), general.getFreeBytes()); } } }
@Test(timeOut = SHUTDOWN_TIMEOUT_MILLIS) public void testShutdown() throws Exception Map<String, String> properties = ImmutableMap.<String, String>builder() .put("node-scheduler.include-coordinator", "false") .put("shutdown.grace-period", "10s") List<ListenableFuture<?>> queryFutures = new ArrayList<>(); for (int i = 0; i < 5; i++) { queryFutures.add(executor.submit(() -> queryRunner.execute("SELECT COUNT(*), clerk FROM orders GROUP BY clerk"))); TestingPrestoServer worker = queryRunner.getServers() .stream() .filter(server -> !server.isCoordinator()) .findFirst() .get(); TaskManager taskManager = worker.getTaskManager(); worker.getGracefulShutdownHandler().requestShutdown(); List<BasicQueryInfo> queryInfos = queryRunner.getCoordinator().getQueryManager().getQueries(); for (BasicQueryInfo info : queryInfos) { assertEquals(info.getState(), FINISHED);
public static DistributedQueryRunner createQueryRunner(String dbConfigUrl, H2ResourceGroupsDao dao, String environment) throws Exception { DistributedQueryRunner queryRunner = DistributedQueryRunner .builder(testSessionBuilder().setCatalog("tpch").setSchema("tiny").build()) .setNodeCount(2) .setEnvironment(environment) .build(); try { Plugin h2ResourceGroupManagerPlugin = new H2ResourceGroupManagerPlugin(); queryRunner.installPlugin(h2ResourceGroupManagerPlugin); queryRunner.getCoordinator().getResourceGroupManager().get() .setConfigurationManager(CONFIGURATION_MANAGER_TYPE, ImmutableMap.of("resource-groups.config-db-url", dbConfigUrl, "node.environment", environment)); queryRunner.installPlugin(new TpchPlugin()); queryRunner.createCatalog("tpch", "tpch"); setup(queryRunner, dao, environment); return queryRunner; } catch (Exception e) { queryRunner.close(); throw e; } }
private void testRejection() throws Exception { try (DistributedQueryRunner queryRunner = createQueryRunner()) { queryRunner.installPlugin(new ResourceGroupManagerPlugin()); queryRunner.getCoordinator().getResourceGroupManager().get().setConfigurationManager("file", ImmutableMap.of("resource-groups.config-file", getResourceFilePath("resource_groups_config_dashboard.json"))); QueryId queryId = createQuery(queryRunner, newRejectionSession(), LONG_LASTING_QUERY); waitForQueryState(queryRunner, queryId, FAILED); QueryManager queryManager = queryRunner.getCoordinator().getQueryManager(); assertEquals(queryManager.getQueryInfo(queryId).getErrorCode(), QUERY_REJECTED.toErrorCode()); } }
@Test(timeOut = 60_000) public void testRunningQuery() throws Exception { queryRunner.execute("SELECT COUNT(*), clerk FROM orders GROUP BY clerk"); while (true) { ResourceGroupInfo global = queryRunner.getCoordinator().getResourceGroupManager().get().getResourceGroupInfo(new ResourceGroupId(new ResourceGroupId("global"), "bi-user")); if (global.getSoftMemoryLimit().toBytes() > 0) { break; } TimeUnit.SECONDS.sleep(2); } }
@Test public void testNonLeafGroup() throws Exception { Session session = testSessionBuilder() .setCatalog("tpch") .setSchema("sf100000") .setSource("non-leaf") .build(); QueryManager queryManager = queryRunner.getCoordinator().getQueryManager(); InternalResourceGroupManager manager = queryRunner.getCoordinator().getResourceGroupManager().get(); DbResourceGroupConfigurationManager dbConfigurationManager = (DbResourceGroupConfigurationManager) manager.getConfigurationManager(); int originalSize = getSelectors(queryRunner).size(); // Add a selector for a non leaf group dao.insertSelector(3, 100, "user.*", "(?i).*non-leaf.*", null, null, null); dbConfigurationManager.load(); while (getSelectors(queryRunner).size() != originalSize + 1) { MILLISECONDS.sleep(500); } // Submit query with side effect of creating resource groups QueryId firstDashboardQuery = createQuery(queryRunner, dashboardSession(), LONG_LASTING_QUERY); waitForQueryState(queryRunner, firstDashboardQuery, RUNNING); cancelQuery(queryRunner, firstDashboardQuery); waitForQueryState(queryRunner, firstDashboardQuery, FAILED); // Submit a query to a non-leaf resource group QueryId invalidResourceGroupQuery = createQuery(queryRunner, session, LONG_LASTING_QUERY); waitForQueryState(queryRunner, invalidResourceGroupQuery, FAILED); assertEquals(queryRunner.getQueryInfo(invalidResourceGroupQuery).getErrorCode(), INVALID_RESOURCE_GROUP.toErrorCode()); }
public static DistributedQueryRunner createQueryRunner(Iterable<TpchTable<?>> tables, Map<String, String> extraProperties, String security, Map<String, String> extraHiveProperties) throws Exception assertEquals(DateTimeZone.getDefault(), TIME_ZONE, "Timezone not configured correctly. Add -Duser.timezone=America/Bahia_Banderas to your JVM arguments"); DistributedQueryRunner queryRunner = new DistributedQueryRunner(createSession(), 4, extraProperties); queryRunner.installPlugin(new TpchPlugin()); queryRunner.createCatalog("tpch", "tpch"); File baseDir = queryRunner.getCoordinator().getBaseDataDir().resolve("hive_data").toFile(); metastore.createDatabase(createDatabaseMetastoreObject(TPCH_SCHEMA)); metastore.createDatabase(createDatabaseMetastoreObject(TPCH_BUCKETED_SCHEMA)); queryRunner.installPlugin(new HivePlugin(HIVE_CATALOG, Optional.of(metastore))); Map<String, String> hiveProperties = ImmutableMap.<String, String>builder() .putAll(extraHiveProperties) .put("hive.time-zone", TIME_ZONE.getID()) .put("hive.compression-codec", "NONE") // so that the file is splittable .build(); queryRunner.createCatalog(HIVE_CATALOG, HIVE_CATALOG, hiveProperties); queryRunner.createCatalog(HIVE_BUCKETED_CATALOG, HIVE_CATALOG, hiveBucketedProperties); queryRunner.close(); throw e;
@Test public void testPathToRoot() throws Exception { try (DistributedQueryRunner queryRunner = TpchQueryRunnerBuilder.builder().build()) { queryRunner.installPlugin(new ResourceGroupManagerPlugin()); InternalResourceGroupManager<?> manager = getResourceGroupManager(queryRunner); manager.setConfigurationManager("file", ImmutableMap.of( "resource-groups.config-file", getResourceFilePath("resource_groups_config_dashboard.json"))); queryRunner.execute(testSessionBuilder().setCatalog("tpch").setSchema("tiny").setSource("dashboard-foo").build(), "SELECT COUNT(*), clerk FROM orders GROUP BY clerk"); List<ResourceGroupInfo> path = manager.getPathToRoot(new ResourceGroupId(new ResourceGroupId(new ResourceGroupId("global"), "user-user"), "dashboard-user")); assertEquals(path.size(), 3); assertTrue(path.get(1).getSubGroups() != null); assertEquals(path.get(2).getId(), new ResourceGroupId("global")); assertEquals(path.get(2).getHardConcurrencyLimit(), 100); assertEquals(path.get(2).getRunningQueries(), null); } }
public static List<ResourceGroupSelector> getSelectors(DistributedQueryRunner queryRunner) { try { return ((DbResourceGroupConfigurationManager) queryRunner.getCoordinator().getResourceGroupManager().get().getConfigurationManager()).getSelectors(); } catch (PrestoException e) { if (e.getErrorCode() == CONFIGURATION_INVALID.toErrorCode()) { return ImmutableList.of(); } throw e; } } }
@Test public void testSufficientInitialWorkerNodes() throws Exception { try (DistributedQueryRunner queryRunner = TpchQueryRunnerBuilder.builder() .setSingleCoordinatorProperty("query-manager.initialization-required-workers", "4") .setNodeCount(4) .build()) { queryRunner.execute("SELECT 1"); assertEquals(queryRunner.getCoordinator().refreshNodes().getActiveNodes().size(), 4); // Query should still be allowed to run if active workers drop down below the minimum required nodes queryRunner.getServers().get(0).close(); assertEquals(queryRunner.getCoordinator().refreshNodes().getActiveNodes().size(), 3); queryRunner.execute("SELECT 1"); } }
private HiveInsertTableHandle getHiveInsertTableHandle(Session session, String tableName) { Metadata metadata = ((DistributedQueryRunner) getQueryRunner()).getCoordinator().getMetadata(); return transaction(getQueryRunner().getTransactionManager(), getQueryRunner().getAccessControl()) .execute(session, transactionSession -> { QualifiedObjectName objectName = new QualifiedObjectName(catalog, TPCH_SCHEMA, tableName); Optional<TableHandle> handle = metadata.getTableHandle(transactionSession, objectName); InsertTableHandle insertTableHandle = metadata.beginInsert(transactionSession, handle.get()); HiveInsertTableHandle hiveInsertTableHandle = (HiveInsertTableHandle) insertTableHandle.getConnectorHandle(); metadata.finishInsert(transactionSession, insertTableHandle, ImmutableList.of(), ImmutableList.of()); return hiveInsertTableHandle; }); }