Examples with MockManagedQueryExecution io.prestosql.execution.MockManagedQueryExecution used on opensource projects

Example 1 with MockManagedQueryExecution

use of io.prestosql.execution.MockManagedQueryExecution in project hetu-core by openlookeng.

the class TestStateCacheStore method testQueryCachedStates.

@Test
public void testQueryCachedStates() {
    Map<String, SharedQueryState> queryStates = new HashMap<>();
    MockManagedQueryExecution query1 = new MockManagedQueryExecution(0);
    MockManagedQueryExecution query2 = new MockManagedQueryExecution(0);
    queryStates.put(query1.toString(), getSharedQueryState(query1));
    queryStates.put(query2.toString(), getSharedQueryState(query2));
    StateCacheStore.get().setCachedStates(StateStoreConstants.QUERY_STATE_COLLECTION_NAME, queryStates);
    Map resultQueryStates = StateCacheStore.get().getCachedStates(StateStoreConstants.QUERY_STATE_COLLECTION_NAME);
    assertQueryStatesEquals(resultQueryStates, queryStates);
}

Example 2 with MockManagedQueryExecution

use of io.prestosql.execution.MockManagedQueryExecution in project hetu-core by openlookeng.

the class TestStateCacheStore method testResourceGroupCachedStates.

@Test
public void testResourceGroupCachedStates() {
    MockManagedQueryExecution query1 = new MockManagedQueryExecution(0);
    MockManagedQueryExecution query2 = new MockManagedQueryExecution(0);
    MockManagedQueryExecution query3 = new MockManagedQueryExecution(0);
    ResourceGroupId group1 = new ResourceGroupId("group1");
    ResourceGroupId group2 = new ResourceGroupId("group2");
    ResourceGroupId group3 = new ResourceGroupId("group3");
    SharedResourceGroupState resourceGroupState1 = new SharedResourceGroupState(group1);
    resourceGroupState1.addQueryState(getSharedQueryState(query1));
    SharedResourceGroupState resourceGroupState2 = new SharedResourceGroupState(group2);
    resourceGroupState2.addQueryState(getSharedQueryState(query2));
    SharedResourceGroupState resourceGroupState3 = new SharedResourceGroupState(group3);
    resourceGroupState3.addQueryState(getSharedQueryState(query3));
    Map<ResourceGroupId, SharedResourceGroupState> resourceGroupStates = ImmutableMap.of(group1, resourceGroupState1, group2, resourceGroupState2, group3, resourceGroupState3);
    StateCacheStore.get().setCachedStates(StateStoreConstants.RESOURCE_GROUP_STATE_COLLECTION_NAME, resourceGroupStates);
    Map<ResourceGroupId, SharedResourceGroupState> result = StateCacheStore.get().getCachedStates(StateStoreConstants.RESOURCE_GROUP_STATE_COLLECTION_NAME);
    assertResourceGroupStateEquals(result, resourceGroupStates);
}

Example 3 with MockManagedQueryExecution

use of io.prestosql.execution.MockManagedQueryExecution in project hetu-core by openlookeng.

the class TestDistributedResourceGroup method testSubgroupMemoryLimit.

@Test
public void testSubgroupMemoryLimit() {
    synchronized (lock) {
        DistributedResourceGroup root = new DistributedResourceGroup(Optional.empty(), "root", (group, export) -> {
        }, directExecutor(), statestore);
        resourceGroupBasicSetUp(root, TEN_BYTE, 3, 4);
        DistributedResourceGroup subgroup = root.getOrCreateSubGroup("subgroup");
        resourceGroupBasicSetUp(subgroup, ONE_BYTE, 3, 4);
        MockManagedQueryExecution query1 = new MockManagedQueryExecution(2);
        subgroup.run(query1);
        // Process the group to refresh stats
        root.processQueuedQueries();
        assertEquals(query1.getState(), RUNNING);
        query1.setResourceGroupId(subgroup.getId());
        updateQueryStateCache(query1);
        MockManagedQueryExecution query2 = new MockManagedQueryExecution(0);
        subgroup.run(query2);
        assertEquals(query2.getState(), QUEUED);
        query2.setResourceGroupId(subgroup.getId());
        updateQueryStateCache(query2);
        MockManagedQueryExecution query3 = new MockManagedQueryExecution(0);
        subgroup.run(query3);
        assertEquals(query3.getState(), QUEUED);
        query3.setResourceGroupId(subgroup.getId());
        updateQueryStateCache(query3);
        query1.complete();
        updateQueryStateCache(query1);
        root.processQueuedQueries();
        assertEquals(query2.getState(), RUNNING);
        assertEquals(query3.getState(), RUNNING);
        StateCacheStore.get().resetCachedStates();
    }
}

Example 4 with MockManagedQueryExecution

use of io.prestosql.execution.MockManagedQueryExecution in project hetu-core by openlookeng.

the class TestDistributedResourceGroup method testQueryKillMemoryFinishPercent.

@Test
public void testQueryKillMemoryFinishPercent() {
    synchronized (lock) {
        DistributedResourceGroup root = new DistributedResourceGroup(Optional.empty(), "root", (group, export) -> {
        }, directExecutor(), statestore);
        resourceGroupBasicSetUp(root, new DataSize(130, BYTE), 3, 4);
        root.setKillPolicy(KillPolicy.HIGH_MEMORY_QUERIES);
        Set<MockManagedQueryExecution> queries = new HashSet<>();
        // query1 running in remote
        MockManagedQueryExecution query1 = new MockManagedQueryExecution(2, "q1");
        query1.startWaitingForResources();
        query1.setResourceGroupId(root.getId());
        updateQueryStateCache(query1);
        // Process the group to refresh stats
        root.run(query1);
        assertEquals(query1.getState(), RUNNING);
        MockManagedQueryExecution query2 = new MockManagedQueryExecution(100, "q2", 20.0);
        query2.setResourceGroupId(root.getId());
        root.run(query2);
        MockManagedQueryExecution query3 = new MockManagedQueryExecution(95, "q3", 10.0);
        query3.setResourceGroupId(root.getId());
        root.run(query3);
        MockManagedQueryExecution query4 = new MockManagedQueryExecution(120, "q4", 5.0);
        query4.setResourceGroupId(root.getId());
        root.run(query4);
        updateQueryStateCache(query2);
        updateQueryStateCache(query3);
        updateQueryStateCache(query4);
        assertEquals(query3.getState(), RUNNING);
        root.processQueuedQueries();
        assertEquals(query1.getState(), RUNNING);
        assertEquals(query2.getState(), RUNNING);
        assertEquals(query3.getState(), FAILED);
        assertEquals(query4.getState(), FAILED);
        query1.complete();
        query2.complete();
        StateCacheStore.get().resetCachedStates();
    }
}

Example 5 with MockManagedQueryExecution

use of io.prestosql.execution.MockManagedQueryExecution in project hetu-core by openlookeng.

the class TestDistributedResourceGroup method testSoftReservedMemory.

@Test
public void testSoftReservedMemory() {
    synchronized (lock) {
        DistributedResourceGroup root = new DistributedResourceGroup(Optional.empty(), "root", (group, export) -> {
        }, directExecutor(), statestore);
        resourceGroupBasicSetUp(root, TEN_BYTE, 10, 10);
        // three subgroup with setSoftReservedMemory = 1 byte
        DistributedResourceGroup group1 = root.getOrCreateSubGroup("1");
        resourceGroupBasicSetUp(group1, FIVE_BYTE, 3, 4);
        group1.setSoftReservedMemory(ONE_BYTE);
        DistributedResourceGroup group2 = root.getOrCreateSubGroup("2");
        resourceGroupBasicSetUp(group2, FIVE_BYTE, 3, 4);
        group2.setSoftReservedMemory(ONE_BYTE);
        DistributedResourceGroup group3 = root.getOrCreateSubGroup("3");
        resourceGroupBasicSetUp(group3, FIVE_BYTE, 3, 4);
        group3.setSoftReservedMemory(ONE_BYTE);
        MockManagedQueryExecution query1 = new MockManagedQueryExecution(5);
        query1.setResourceGroupId(group1.getId());
        query1.startWaitingForResources();
        updateQueryStateCache(query1);
        MockManagedQueryExecution query2 = new MockManagedQueryExecution(4);
        query2.setResourceGroupId(group2.getId());
        query2.startWaitingForResources();
        updateQueryStateCache(query2);
        // group1 memory usage = 5M, group2 memory usage = 4M, group3 memory usage = 0
        assertEquals(query1.getState(), RUNNING);
        assertEquals(query2.getState(), RUNNING);
        root.processQueuedQueries();
        MockManagedQueryExecution extraQuery = new MockManagedQueryExecution(1);
        group2.run(extraQuery);
        assertEquals(extraQuery.getState(), QUEUED);
        MockManagedQueryExecution query3 = new MockManagedQueryExecution(1);
        group3.run(query3);
        assertEquals(query3.getState(), RUNNING);
        StateCacheStore.get().resetCachedStates();
    }
}