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

Example 6 with SharedQueryState

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

the class TestDistributedResourceGroup method updateQueryStateCache.

// update cached queryStates for a single query
private static void updateQueryStateCache(MockManagedQueryExecution query) {
    Map<String, SharedQueryState> queryStates = StateCacheStore.get().getCachedStates(StateStoreConstants.QUERY_STATE_COLLECTION_NAME);
    if (queryStates == null) {
        queryStates = new HashMap<>();
    }
    String queryKey = query.toString();
    queryStates.put(queryKey, getSharedQueryState(query));
    StateCacheStore.get().setCachedStates(StateStoreConstants.QUERY_STATE_COLLECTION_NAME, queryStates);
    DistributedResourceGroupUtils.mapCachedStates();
}

Example 7 with SharedQueryState

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

the class SqlQueryManager method killBlockingQuery.

/**
 * Kill query when cluster is in OOM state
 */
private void killBlockingQuery() {
    if (isMultiCoordinatorEnabled()) {
        return;
    }
    List<QueryExecution> localRunningQueries = queryTracker.getAllQueries().stream().filter(query -> query.getState() == RUNNING).collect(toImmutableList());
    Map<String, SharedQueryState> queries = StateCacheStore.get().getCachedStates(StateStoreConstants.OOM_QUERY_STATE_COLLECTION_NAME);
    if (queries != null) {
        for (SharedQueryState query : queries.values()) {
            for (QueryExecution localQuery : localRunningQueries) {
                if (query.getBasicQueryInfo().getQueryId().equals(localQuery.getQueryId())) {
                    memoryManager.killLocalQuery(localQuery);
                }
            }
        }
    }
}

Example 8 with SharedQueryState

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

the class DistributedResourceGroup method internalStartNext.

/**
 * Check if there is any queued query in any eligible sub-group that can be started
 *
 * @return if any query has been started in sub-groups
 */
protected boolean internalStartNext() {
    checkState(Thread.holdsLock(root), "Must hold lock to find next query");
    synchronized (root) {
        if (!canRunMore()) {
            return false;
        }
        // Only start the query if it exists locally
        Optional<SharedResourceGroupState> resourceGroupState = getSharedResourceGroupState();
        PriorityQueue<SharedQueryState> globalQueuedQueries = resourceGroupState.isPresent() ? resourceGroupState.get().getQueuedQueries() : new PriorityQueue<>();
        if (!globalQueuedQueries.isEmpty() && !queuedQueries.isEmpty()) {
            // Get queued query with longest queued time from state cache store.
            // Remove it if local queued queries contains it.
            SharedQueryState nextQuery = globalQueuedQueries.peek();
            for (ManagedQueryExecution localQuery : queuedQueries) {
                if (nextQuery.getBasicQueryInfo().getQueryId().equals(localQuery.getBasicQueryInfo().getQueryId())) {
                    Lock lock = stateStore.getLock(id.toString());
                    boolean locked = false;
                    try {
                        locked = lock.tryLock(MILLISECONDS_PER_SECOND, TimeUnit.MILLISECONDS);
                        if (locked) {
                            // Get the most recent cached state store status and check canRunMore again
                            // Avoid the race condition that state store is updated by other process
                            // Make sure queued query start is synchronized
                            DistributedResourceGroupUtils.mapCachedStates();
                            if (canRunMore()) {
                                queuedQueries.remove(localQuery);
                                startInBackground(localQuery);
                                return true;
                            }
                        }
                        return false;
                    } catch (InterruptedException | RuntimeException e) {
                        return false;
                    } finally {
                        if (locked) {
                            lock.unlock();
                        }
                    }
                }
            }
        }
        // Try to find least recently used eligible group
        DistributedResourceGroup chosenGroup = findLeastRecentlyExecutedSubgroup();
        if (chosenGroup == null) {
            return false;
        }
        // Try to start queued query in the group
        boolean started = chosenGroup.internalStartNext();
        long currentTime = System.currentTimeMillis();
        if (lastStartMillis != 0) {
            timeBetweenStartsSec.update(Math.max(0, (currentTime - lastStartMillis) / MILLISECONDS_PER_SECOND));
        }
        lastStartMillis = currentTime;
        return started;
    }
}

Example 9 with SharedQueryState

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

the class TestDistributedResourceGroup method updateStateCacheStore.

// update cached queryStates for a set of queries
private static void updateStateCacheStore(Set<MockManagedQueryExecution> queries) {
    Map<String, SharedQueryState> queryStates = new HashMap<>();
    for (Iterator<MockManagedQueryExecution> iterator = queries.iterator(); iterator.hasNext(); ) {
        MockManagedQueryExecution query = iterator.next();
        queryStates.put(query.toString(), getSharedQueryState(query));
    }
    StateCacheStore.get().setCachedStates(StateStoreConstants.QUERY_STATE_COLLECTION_NAME, queryStates);
    DistributedResourceGroupUtils.mapCachedStates();
}

Example 10 with SharedQueryState

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

the class TestDistributedResourceGroupUtils method testMapCachedStates.

@Test
public void testMapCachedStates() {
    synchronized (lock) {
        ResourceGroupId root = new ResourceGroupId("root");
        // set up queryStates
        MockManagedQueryExecution query1 = new MockManagedQueryExecution(0);
        MockManagedQueryExecution query2 = new MockManagedQueryExecution(0);
        MockManagedQueryExecution query3 = new MockManagedQueryExecution(0);
        query1.setResourceGroupId(root);
        query2.setResourceGroupId(root);
        query3.setResourceGroupId(root);
        // query1 running, query2 queued, query3 completed
        query1.startWaitingForResources();
        query3.complete();
        SharedQueryState queryState1 = getSharedQueryState(query1);
        SharedQueryState queryState2 = getSharedQueryState(query2);
        SharedQueryState queryState3 = getSharedQueryState(query3);
        Map<String, SharedQueryState> queryStates = ImmutableMap.of(query1.toString(), queryState1, query2.toString(), queryState2, query3.toString(), queryState3);
        StateCacheStore.get().setCachedStates(StateStoreConstants.QUERY_STATE_COLLECTION_NAME, queryStates);
        // set up cpu usage states
        Map<String, Long> cpuUsageStates = ImmutableMap.of(root.toString(), 1000L);
        StateCacheStore.get().setCachedStates(StateStoreConstants.CPU_USAGE_STATE_COLLECTION_NAME, cpuUsageStates);
        // map QueryStates and cpuUsageStates to ResourceGroups
        DistributedResourceGroupUtils.mapCachedStates();
        Map resourceGroupStates = StateCacheStore.get().getCachedStates(StateStoreConstants.RESOURCE_GROUP_STATE_COLLECTION_NAME);
        // check ResourceGroupState
        assertTrue(resourceGroupStates.containsKey(root));
        SharedResourceGroupState resourceGroupState = (SharedResourceGroupState) resourceGroupStates.get(root);
        assertEquals(resourceGroupState.getRunningQueries().size(), 1);
        assertTrue(resourceGroupState.getRunningQueries().contains(queryState1));
        assertEquals(resourceGroupState.getQueuedQueries().size(), 1);
        assertEquals(resourceGroupState.getQueuedQueries().peek(), queryState2);
        assertEquals(resourceGroupState.getFinishedQueries().size(), 1);
        assertEquals(resourceGroupState.getFinishedQueries().peek(), queryState3);
        StateCacheStore.get().resetCachedStates();
    }
}