Search in sources :

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();
    }
}
Also used : RootInternalResourceGroup(com.facebook.presto.execution.resourceGroups.InternalResourceGroup.RootInternalResourceGroup) RootInternalResourceGroup(com.facebook.presto.execution.resourceGroups.InternalResourceGroup.RootInternalResourceGroup) Random(java.util.Random) DataSize(io.airlift.units.DataSize) ArrayList(java.util.ArrayList) TreeMap(java.util.TreeMap) MockQueryExecution(com.facebook.presto.execution.MockQueryExecution) Test(org.testng.annotations.Test)

Aggregations

MockQueryExecution (com.facebook.presto.execution.MockQueryExecution)11 RootInternalResourceGroup (com.facebook.presto.execution.resourceGroups.InternalResourceGroup.RootInternalResourceGroup)10 DataSize (io.airlift.units.DataSize)10 Test (org.testng.annotations.Test)10 Duration (io.airlift.units.Duration)2 ResourceGroupInfo (com.facebook.presto.spi.resourceGroups.ResourceGroupInfo)1 ArrayList (java.util.ArrayList)1 HashSet (java.util.HashSet)1 Random (java.util.Random)1 TreeMap (java.util.TreeMap)1 BinomialDistribution (org.apache.commons.math3.distribution.BinomialDistribution)1