Skip to content

Commit

Permalink
Fix and enable Flaky Test TestQueues#testQueuedQueryInteraction
Browse files Browse the repository at this point in the history
  • Loading branch information
imjalpreet authored and tdcmeehan committed Apr 16, 2024
1 parent 06edc73 commit ef197ad
Showing 1 changed file with 25 additions and 25 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@
import com.facebook.airlift.json.JsonObjectMapperProvider;
import com.facebook.presto.Session;
import com.facebook.presto.dispatcher.DispatchManager;
import com.facebook.presto.plugin.blackhole.BlackHolePlugin;
import com.facebook.presto.resourceGroups.ResourceGroupManagerPlugin;
import com.facebook.presto.server.BasicQueryInfo;
import com.facebook.presto.spi.QueryId;
Expand All @@ -27,7 +28,6 @@
import com.facebook.presto.spi.session.ResourceEstimates;
import com.facebook.presto.sql.Serialization;
import com.facebook.presto.tests.DistributedQueryRunner;
import com.facebook.presto.tests.tpch.TpchQueryRunnerBuilder;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.module.SimpleModule;
import com.google.common.collect.ImmutableList;
Expand Down Expand Up @@ -56,8 +56,8 @@
import static com.facebook.presto.execution.QueryState.RUNNING;
import static com.facebook.presto.execution.TestQueryRunnerUtil.cancelQuery;
import static com.facebook.presto.execution.TestQueryRunnerUtil.createQuery;
import static com.facebook.presto.execution.TestQueryRunnerUtil.createQueryRunner;
import static com.facebook.presto.execution.TestQueryRunnerUtil.waitForQueryState;
import static com.facebook.presto.execution.resourceGroups.db.H2TestUtil.getSimpleQueryRunner;
import static com.facebook.presto.metadata.FunctionAndTypeManager.createTestFunctionAndTypeManager;
import static com.facebook.presto.spi.StandardErrorCode.ADMINISTRATIVELY_KILLED;
import static com.facebook.presto.spi.StandardErrorCode.ADMINISTRATIVELY_PREEMPTED;
Expand All @@ -76,7 +76,7 @@
@Test(singleThreaded = true)
public class TestQueues
{
public static final String LONG_LASTING_QUERY = "SELECT COUNT(*) FROM lineitem";
public static final String LONG_LASTING_QUERY = "SELECT COUNT(*) FROM blackhole.default.dummy";

private DistributedQueryRunner queryRunner;
private ObjectMapper objectMapper;
Expand All @@ -86,7 +86,12 @@ public class TestQueues
public void setup()
throws Exception
{
queryRunner = createQueryRunner();
queryRunner = getSimpleQueryRunner();
queryRunner.installPlugin(new BlackHolePlugin());
queryRunner.createCatalog("blackhole", "blackhole");
queryRunner.execute(
"CREATE TABLE blackhole.default.dummy (col BIGINT) " +
"WITH (split_count = 1, rows_per_page = 1, pages_per_split = 1, page_processing_delay = '10m')");
client = new JettyHttpClient();
objectMapper = new JsonObjectMapperProvider().get();
objectMapper.registerModule(new SimpleModule()
Expand Down Expand Up @@ -318,20 +323,17 @@ public void testSelectorResourceEstimateBasedSelection()
public void testQueryTypeBasedSelection()
throws Exception
{
try (DistributedQueryRunner queryRunner = TpchQueryRunnerBuilder.builder().build()) {
queryRunner.installPlugin(new ResourceGroupManagerPlugin());
queryRunner.getCoordinator().getResourceGroupManager().get()
.forceSetConfigurationManager("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", "control"));
}
queryRunner.installPlugin(new ResourceGroupManagerPlugin());
queryRunner.getCoordinator().getResourceGroupManager().get()
.forceSetConfigurationManager("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", "control"));
}

// This test is flaky: https://github.com/prestodb/presto/issues/19633
@Test(timeOut = 60_000, enabled = false)
@Test(timeOut = 240_000)
public void testQueuedQueryInteraction()
throws Exception
{
Expand Down Expand Up @@ -385,15 +387,13 @@ private void assertResourceGroup(DistributedQueryRunner queryRunner, Session ses
private void testRejection()
throws Exception
{
try (DistributedQueryRunner queryRunner = createQueryRunner()) {
queryRunner.installPlugin(new ResourceGroupManagerPlugin());
queryRunner.getCoordinator().getResourceGroupManager().get().forceSetConfigurationManager("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);
DispatchManager dispatchManager = queryRunner.getCoordinator().getDispatchManager();
assertEquals(dispatchManager.getQueryInfo(queryId).getErrorCode(), MISSING_RESOURCE_GROUP_SELECTOR.toErrorCode());
}
queryRunner.installPlugin(new ResourceGroupManagerPlugin());
queryRunner.getCoordinator().getResourceGroupManager().get().forceSetConfigurationManager("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);
DispatchManager dispatchManager = queryRunner.getCoordinator().getDispatchManager();
assertEquals(dispatchManager.getQueryInfo(queryId).getErrorCode(), MISSING_RESOURCE_GROUP_SELECTOR.toErrorCode());
}

private QueryId createDashboardQuery(DistributedQueryRunner queryRunner)
Expand Down

0 comments on commit ef197ad

Please sign in to comment.