private static InternalResourceGroupManager<?> getResourceGroupManager(DistributedQueryRunner queryRunner) { return queryRunner.getCoordinator().getResourceGroupManager() .orElseThrow(() -> new IllegalArgumentException("no resource group manager")); } }
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(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(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); } }
@Test(timeOut = 240_000) public void testClientTagsBasedSelection() 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_client_tags_based_config.json"))); assertResourceGroup(queryRunner, newSessionWithTags(ImmutableSet.of("a")), LONG_LASTING_QUERY, createResourceGroupId("global", "a", "default")); assertResourceGroup(queryRunner, newSessionWithTags(ImmutableSet.of("b")), LONG_LASTING_QUERY, createResourceGroupId("global", "b")); assertResourceGroup(queryRunner, newSessionWithTags(ImmutableSet.of("a", "c")), LONG_LASTING_QUERY, createResourceGroupId("global", "a", "c")); } }
@Test(timeOut = 60_000) public void testSelectorPriority() throws Exception { InternalResourceGroupManager manager = queryRunner.getCoordinator().getResourceGroupManager().get(); QueryManager queryManager = queryRunner.getCoordinator().getQueryManager(); DbResourceGroupConfigurationManager dbConfigurationManager = (DbResourceGroupConfigurationManager) manager.getConfigurationManager(); QueryId firstQuery = createQuery(queryRunner, dashboardSession(), LONG_LASTING_QUERY); waitForQueryState(queryRunner, firstQuery, RUNNING); Optional<ResourceGroupId> resourceGroup = queryManager.getFullQueryInfo(firstQuery).getResourceGroupId(); assertTrue(resourceGroup.isPresent()); assertEquals(resourceGroup.get().toString(), "global.user-user.dashboard-user"); // create a new resource group that rejects all queries submitted to it dao.insertResourceGroup(8, "reject-all-queries", "1MB", 0, 0, 0, null, null, null, null, null, 3L, TEST_ENVIRONMENT); // add a new selector that has a higher priority than the existing dashboard selector and that routes queries to the "reject-all-queries" resource group dao.insertSelector(8, 200, "user.*", "(?i).*dashboard.*", null, null, null); // reload the configuration dbConfigurationManager.load(); QueryId secondQuery = createQuery(queryRunner, dashboardSession(), LONG_LASTING_QUERY); waitForQueryState(queryRunner, secondQuery, FAILED); resourceGroup = queryManager.getFullQueryInfo(secondQuery).getResourceGroupId(); assertTrue(resourceGroup.isPresent()); assertEquals(resourceGroup.get(), createResourceGroupId("global", "user-user", "reject-all-queries")); }
@Test(timeOut = 240_000) public void testResourceGroupManagerWithTooManyQueriesScheduled() 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); waitForQueryState(queryRunner, firstDashboardQuery, RUNNING); QueryId secondDashboardQuery = createDashboardQuery(queryRunner); waitForQueryState(queryRunner, secondDashboardQuery, QUEUED); QueryId thirdDashboardQuery = createDashboardQuery(queryRunner); waitForQueryState(queryRunner, thirdDashboardQuery, FAILED); } }
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()); } }
@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"))); }
queryRunner.getCoordinator().getResourceGroupManager().get() .setConfigurationManager("file", ImmutableMap.of("resource-groups.config-file", getResourceFilePath("resource_groups_resource_estimate_based_config.json")));
@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()); }
@Test(timeOut = 240_000) public void testQueryTypeBasedSelection() throws Exception { try (DistributedQueryRunner queryRunner = TpchQueryRunnerBuilder.builder().build()) { queryRunner.installPlugin(new ResourceGroupManagerPlugin()); queryRunner.getCoordinator().getResourceGroupManager().get() .setConfigurationManager("file", ImmutableMap.of("resource-groups.config-file", getResourceFilePath("resource_groups_query_type_based_config.json"))); assertResourceGroup(queryRunner, newAdhocSession(), LONG_LASTING_QUERY, createResourceGroupId("global", "select")); assertResourceGroup(queryRunner, newAdhocSession(), "SHOW TABLES", createResourceGroupId("global", "describe")); assertResourceGroup(queryRunner, newAdhocSession(), "EXPLAIN " + LONG_LASTING_QUERY, createResourceGroupId("global", "explain")); assertResourceGroup(queryRunner, newAdhocSession(), "DESCRIBE lineitem", createResourceGroupId("global", "describe")); assertResourceGroup(queryRunner, newAdhocSession(), "RESET SESSION " + HASH_PARTITION_COUNT, createResourceGroupId("global", "data_definition")); } }
@Test(timeOut = 240_000) public void testResourceGroupManager() 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"))); // submit first "dashboard" query QueryId firstDashboardQuery = createDashboardQuery(queryRunner); // wait for the first "dashboard" query to start waitForQueryState(queryRunner, firstDashboardQuery, RUNNING); // submit second "dashboard" query QueryId secondDashboardQuery = createDashboardQuery(queryRunner); // wait for the second "dashboard" query to be queued ("dashboard.${USER}" queue strategy only allows one "dashboard" query to be accepted for execution) waitForQueryState(queryRunner, secondDashboardQuery, QUEUED); // submit first non "dashboard" query QueryId firstNonDashboardQuery = createAdHocQuery(queryRunner); // wait for the first non "dashboard" query to start waitForQueryState(queryRunner, firstNonDashboardQuery, RUNNING); // submit second non "dashboard" query QueryId secondNonDashboardQuery = createAdHocQuery(queryRunner); // wait for the second non "dashboard" query to start waitForQueryState(queryRunner, secondNonDashboardQuery, RUNNING); // cancel first "dashboard" query, second "dashboard" query and second non "dashboard" query should start running cancelQuery(queryRunner, firstDashboardQuery); waitForQueryState(queryRunner, firstDashboardQuery, FAILED); waitForQueryState(queryRunner, secondDashboardQuery, RUNNING); } }
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; } }
@Test(timeOut = 90_000) public void testTooManyQueries() throws Exception { QueryId firstDashboardQuery = createQuery(queryRunner, dashboardSession(), LONG_LASTING_QUERY); waitForQueryState(queryRunner, firstDashboardQuery, RUNNING); QueryId secondDashboardQuery = createQuery(queryRunner, dashboardSession(), LONG_LASTING_QUERY); waitForQueryState(queryRunner, secondDashboardQuery, QUEUED); QueryId thirdDashboardQuery = createQuery(queryRunner, dashboardSession(), LONG_LASTING_QUERY); waitForQueryState(queryRunner, thirdDashboardQuery, FAILED); // Allow one more query to run and resubmit third query dao.updateResourceGroup(3, "user-${USER}", "1MB", 3, 4, 4, null, null, null, null, null, 1L, TEST_ENVIRONMENT); dao.updateResourceGroup(5, "dashboard-${USER}", "1MB", 1, 2, 2, null, null, null, null, null, 3L, TEST_ENVIRONMENT); InternalResourceGroupManager manager = queryRunner.getCoordinator().getResourceGroupManager().get(); DbResourceGroupConfigurationManager dbConfigurationManager = (DbResourceGroupConfigurationManager) manager.getConfigurationManager(); // Trigger reload to make the test more deterministic dbConfigurationManager.load(); waitForQueryState(queryRunner, secondDashboardQuery, RUNNING); thirdDashboardQuery = createQuery(queryRunner, dashboardSession(), LONG_LASTING_QUERY); waitForQueryState(queryRunner, thirdDashboardQuery, QUEUED); // Lower running queries in dashboard resource groups and reload the config dao.updateResourceGroup(5, "dashboard-${USER}", "1MB", 1, 1, 1, null, null, null, null, null, 3L, TEST_ENVIRONMENT); dbConfigurationManager.load(); // Cancel query and verify that third query is still queued cancelQuery(queryRunner, firstDashboardQuery); waitForQueryState(queryRunner, firstDashboardQuery, FAILED); MILLISECONDS.sleep(2000); waitForQueryState(queryRunner, thirdDashboardQuery, QUEUED); }
queryRunner.getCoordinator().getResourceGroupManager().get().setConfigurationManager("file", ImmutableMap.of("resource-groups.config-file", getResourceFilePath("resource_groups_config_soft_limits.json")));
private static InternalResourceGroupManager<?> getResourceGroupManager(DistributedQueryRunner queryRunner) { return queryRunner.getCoordinator().getResourceGroupManager() .orElseThrow(() -> new IllegalArgumentException("no resource group manager")); } }
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(timeOut = 240_000) public void testClientTagsBasedSelection() 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_client_tags_based_config.json"))); assertResourceGroup(queryRunner, newSessionWithTags(ImmutableSet.of("a")), LONG_LASTING_QUERY, createResourceGroupId("global", "a", "default")); assertResourceGroup(queryRunner, newSessionWithTags(ImmutableSet.of("b")), LONG_LASTING_QUERY, createResourceGroupId("global", "b")); assertResourceGroup(queryRunner, newSessionWithTags(ImmutableSet.of("a", "c")), LONG_LASTING_QUERY, createResourceGroupId("global", "a", "c")); } }
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()); } }