Examples with SharedQueryState io.prestosql.statestore.SharedQueryState used on opensource projects

Example 11 with SharedQueryState

use of io.prestosql.statestore.SharedQueryState in project hetu-core by openlookeng.

the class TestDistributedResourceGroupUtils method testUpdateResourceGroupLastExecutionTime.

@Test
public void testUpdateResourceGroupLastExecutionTime() {
    synchronized (lock) {
        ResourceGroupId root = new ResourceGroupId("root");
        // set up queryStates
        Map<String, SharedQueryState> queryStates = new HashMap<>();
        MockManagedQueryExecution query1 = new MockManagedQueryExecution(0);
        query1.setResourceGroupId(root);
        // query1 running
        query1.startWaitingForResources();
        SharedQueryState queryState1 = getSharedQueryState(query1);
        queryStates.put(query1.toString(), queryState1);
        Optional<DateTime> queryStartTime1 = queryState1.getExecutionStartTime();
        StateCacheStore.get().setCachedStates(StateStoreConstants.QUERY_STATE_COLLECTION_NAME, queryStates);
        // map QueryStates to ResourceGroups and update ResourceGroup LastExecutionTime
        DistributedResourceGroupUtils.mapCachedStates();
        Map resourceGroupStates = StateCacheStore.get().getCachedStates(StateStoreConstants.RESOURCE_GROUP_STATE_COLLECTION_NAME);
        // lastExecutionTime should be updated to query1's start execution time
        assertTrue(resourceGroupStates.containsKey(root));
        SharedResourceGroupState resourceGroupState = (SharedResourceGroupState) resourceGroupStates.get(root);
        assertEquals(queryStartTime1, resourceGroupState.getLastExecutionTime());
        // add new queries: query2 (finished) and query3(queued)
        MockManagedQueryExecution query2 = new MockManagedQueryExecution(0);
        query2.setResourceGroupId(root);
        query2.complete();
        SharedQueryState queryState2 = getSharedQueryState(query2);
        MockManagedQueryExecution query3 = new MockManagedQueryExecution(0);
        query3.setResourceGroupId(root);
        SharedQueryState queryState3 = getSharedQueryState(query3);
        queryStates.put(query2.toString(), queryState2);
        queryStates.put(query3.toString(), queryState3);
        Optional<DateTime> queryStartTime2 = queryState2.getExecutionStartTime();
        StateCacheStore.get().setCachedStates(StateStoreConstants.QUERY_STATE_COLLECTION_NAME, queryStates);
        DistributedResourceGroupUtils.mapCachedStates();
        // check ResourceGroupState lastExecutionTime
        resourceGroupState = (SharedResourceGroupState) StateCacheStore.get().getCachedStates(StateStoreConstants.RESOURCE_GROUP_STATE_COLLECTION_NAME).get(root);
        // lastExecutionTime should be updated to query2's start execution time
        assertEquals(queryStartTime2, resourceGroupState.getLastExecutionTime());
        StateCacheStore.get().resetCachedStates();
    }
}

Example 12 with SharedQueryState

use of io.prestosql.statestore.SharedQueryState in project hetu-core by openlookeng.

the class DistributedResourceGroupUtils method mapQueryStatesToResourceGroups.

/**
 * Map cached query states to resource groups
 */
private static void mapQueryStatesToResourceGroups() {
    Map<String, SharedQueryState> queryStates = StateCacheStore.get().getCachedStates(StateStoreConstants.QUERY_STATE_COLLECTION_NAME);
    if (queryStates == null) {
        return;
    }
    Map<ResourceGroupId, SharedResourceGroupState> resourceGroupStates = new HashMap<>();
    for (SharedQueryState state : queryStates.values()) {
        Optional<ResourceGroupId> resourceGroupId = state.getBasicQueryInfo().getResourceGroupId();
        if (!resourceGroupId.isPresent()) {
            continue;
        }
        if (!resourceGroupStates.containsKey(resourceGroupId.get())) {
            resourceGroupStates.put(resourceGroupId.get(), new SharedResourceGroupState(resourceGroupId.get()));
        }
        SharedResourceGroupState resourceGroupState = resourceGroupStates.get(resourceGroupId.get());
        resourceGroupState.addQueryState(state);
        updateResourceGroupLastExecutionTime(state, resourceGroupState);
        logQueryState(state);
    }
    StateCacheStore.get().setCachedStates(StateStoreConstants.RESOURCE_GROUP_STATE_COLLECTION_NAME, resourceGroupStates);
}