Examples with MockQueryExecution com.facebook.presto.execution.MockQueryExecution used on opensource projects

Example 11 with MockQueryExecution

use of com.facebook.presto.execution.MockQueryExecution in project presto by prestodb.

the class TestResourceGroups method testPriorityScheduling.

@Test(timeOut = 10_000)
public void testPriorityScheduling() {
    RootInternalResourceGroup root = new RootInternalResourceGroup("root", (group, export) -> {
    }, directExecutor());
    root.setSoftMemoryLimit(new DataSize(1, MEGABYTE));
    root.setMaxQueuedQueries(100);
    // Start with zero capacity, so that nothing starts running until we've added all the queries
    root.setMaxRunningQueries(0);
    root.setSchedulingPolicy(QUERY_PRIORITY);
    InternalResourceGroup group1 = root.getOrCreateSubGroup("1");
    group1.setSoftMemoryLimit(new DataSize(1, MEGABYTE));
    group1.setMaxQueuedQueries(100);
    group1.setMaxRunningQueries(1);
    InternalResourceGroup group2 = root.getOrCreateSubGroup("2");
    group2.setSoftMemoryLimit(new DataSize(1, MEGABYTE));
    group2.setMaxQueuedQueries(100);
    group2.setMaxRunningQueries(1);
    SortedMap<Integer, MockQueryExecution> queries = new TreeMap<>();
    Random random = new Random();
    for (int i = 0; i < 100; i++) {
        int priority;
        do {
            priority = random.nextInt(1_000_000) + 1;
        } while (queries.containsKey(priority));
        MockQueryExecution query = new MockQueryExecution(0, "query_id", priority);
        if (random.nextBoolean()) {
            group1.run(query);
        } else {
            group2.run(query);
        }
        queries.put(priority, query);
    }
    root.setMaxRunningQueries(1);
    List<MockQueryExecution> orderedQueries = new ArrayList<>(queries.values());
    reverse(orderedQueries);
    for (MockQueryExecution query : orderedQueries) {
        root.processQueuedQueries();
        assertEquals(query.getState(), RUNNING);
        query.complete();
    }
}