Examples with StateCollection io.prestosql.spi.statestore.StateCollection used on opensource projects

Example 1 with StateCollection

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

the class StateFetcher method fetchRunningQueryStates.

/**
 * Fetch states only related to running queries from state store to cache store
 * This is used when not all the states are needed for better states fetching performance
 *
 * @throws IOException exception when failed to deserialize states
 */
public void fetchRunningQueryStates(StateStore stateStore) throws IOException {
    synchronized (this) {
        long start = System.currentTimeMillis();
        LOG.debug("fetchStates starts at current time milliseconds: %s, at format HH:mm:ss:SSS:%s", start, new SimpleDateFormat("HH:mm:ss:SSS").format(new Date(start)));
        StateCollection cpuUsageCollection = stateStore.getStateCollection(CPU_USAGE_STATE_COLLECTION_NAME);
        StateCollection queryStateCollection = stateStore.getStateCollection(QUERY_STATE_COLLECTION_NAME);
        StateCacheStore.get().setCachedStates(CPU_USAGE_STATE_COLLECTION_NAME, ((StateMap) cpuUsageCollection).getAll());
        Map<String, String> states = ((StateMap<String, String>) queryStateCollection).getAll();
        StateCacheStore.get().setCachedStates(QUERY_STATE_COLLECTION_NAME, deserializeFetchedStates(states));
        long end = System.currentTimeMillis();
        LOG.debug("fetchStates ends at current time milliseconds: %s, at format HH:mm:ss:SSS:%s, total time use: %s", end, new SimpleDateFormat("HH:mm:ss:SSS").format(new Date(end)), end - start);
    }
}

Example 2 with StateCollection

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

the class StateFetcher method fetchAllStates.

/**
 * Fetch all states from state store to cache store
 *
 * @throws IOException exception when failed to deserialize states
 */
public void fetchAllStates() throws IOException {
    synchronized (this) {
        // State store hasn't been loaded yet
        if (stateStoreProvider.getStateStore() == null) {
            return;
        }
        long start = System.currentTimeMillis();
        LOG.debug("fetchStates starts at current time milliseconds: %s, at format HH:mm:ss:SSS:%s", start, new SimpleDateFormat("HH:mm:ss:SSS").format(new Date(start)));
        for (String stateCollectionName : stateCollections) {
            StateCollection stateCollection = stateStoreProvider.getStateStore().getStateCollection(stateCollectionName);
            if (stateCollection == null) {
                continue;
            }
            if (stateCollectionName.equals(CPU_USAGE_STATE_COLLECTION_NAME)) {
                StateCacheStore.get().setCachedStates(stateCollectionName, ((StateMap) stateCollection).getAll());
                continue;
            }
            if (stateCollection.getType() == StateCollection.Type.MAP) {
                Map<String, String> states = ((StateMap<String, String>) stateCollection).getAll();
                StateCacheStore.get().setCachedStates(stateCollectionName, deserializeFetchedStates(states));
            } else {
                LOG.warn("Unsupported state collection type: %s", stateCollection.getType());
            }
        }
        long end = System.currentTimeMillis();
        LOG.debug("fetchStates ends at current time milliseconds: %s, at format HH:mm:ss:SSS:%s, total time use: %s", end, new SimpleDateFormat("HH:mm:ss:SSS").format(new Date(end)), end - start);
    }
}

Example 3 with StateCollection

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

the class StateFetcher method handleExpiredQueryState.

private void handleExpiredQueryState(SharedQueryState state) {
    // State store hasn't been loaded yet
    final StateStore stateStore = stateStoreProvider.getStateStore();
    if (stateStore == null) {
        return;
    }
    Lock lock = null;
    boolean locked = false;
    try {
        lock = stateStore.getLock(HANDLE_EXPIRED_QUERY_LOCK_NAME);
        locked = lock.tryLock(DEFAULT_ACQUIRED_LOCK_TIME_MS, TimeUnit.MILLISECONDS);
        if (locked) {
            LOG.debug(String.format("EXPIRED!!! REMOVING... Id: %s, state: %s, uri: %s, query: %s", state.getBasicQueryInfo().getQueryId().getId(), state.getBasicQueryInfo().getState().toString(), state.getBasicQueryInfo().getSelf().toString(), state.getBasicQueryInfo().getQuery()));
            // remove expired query from oom
            StateCollection stateCollection = stateStore.getStateCollection(OOM_QUERY_STATE_COLLECTION_NAME);
            removeState(stateCollection, Optional.of(state.getBasicQueryInfo().getQueryId()), LOG);
            // update query to failed in QUERY_STATE_COLLECTION_NAME if exists
            stateCollection = stateStore.getStateCollection(QUERY_STATE_COLLECTION_NAME);
            StateCollection finishStateCollection = stateStore.getStateCollection(FINISHED_QUERY_STATE_COLLECTION_NAME);
            if (stateCollection != null && stateCollection.getType().equals(StateCollection.Type.MAP)) {
                String queryState = ((StateMap<String, String>) stateCollection).get(state.getBasicQueryInfo().getQueryId().getId());
                if (queryState != null) {
                    BasicQueryInfo oldQueryInfo = state.getBasicQueryInfo();
                    SharedQueryState newState = createExpiredState(oldQueryInfo, state);
                    String stateJson = MAPPER.writeValueAsString(newState);
                    ((StateMap) finishStateCollection).put(newState.getBasicQueryInfo().getQueryId().getId(), stateJson);
                    removeState(stateCollection, Optional.of(state.getBasicQueryInfo().getQueryId()), LOG);
                }
            }
        }
    } catch (Exception e) {
        LOG.error("Error handleExpiredQueryState: " + e.getMessage());
    } finally {
        if (locked) {
            lock.unlock();
        }
    }
}

Example 4 with StateCollection

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

the class StateUpdater method removeFromStateCollection.

private void removeFromStateCollection(StateStore stateStore, String stateCollectionName, ManagedQueryExecution query) {
    unregisterQuery(stateCollectionName, query);
    StateCollection stateCollection = stateStore.getStateCollection(stateCollectionName);
    removeState(stateCollection, Optional.of(query.getBasicQueryInfo().getQueryId()), LOG);
}

Example 5 with StateCollection

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

the class StateUpdater method updateStates.

/**
 * Update local queued query states to state store
 *
 * @throws JsonProcessingException exception when fail to serialize states to json
 */
public void updateStates() throws JsonProcessingException {
    // State store hasn't been loaded yet
    final StateStore stateStore = stateStoreProvider.getStateStore();
    if (stateStore == null) {
        return;
    }
    long start = System.currentTimeMillis();
    LOG.debug("UpdateStates starts at current time milliseconds: %s, at format HH:mm:ss:SSS:%s", start, new SimpleDateFormat("HH:mm:ss:SSS").format(new Date(start)));
    StateCollection finishedQueries = stateStore.getStateCollection(FINISHED_QUERY_STATE_COLLECTION_NAME);
    StateCollection queries = stateStore.getStateCollection(QUERY_STATE_COLLECTION_NAME);
    List<DispatchQuery> queriesToUnregister = new LinkedList<>();
    synchronized (registeredQueries) {
        for (DispatchQuery query : registeredQueries.get(QUERY_STATE_COLLECTION_NAME)) {
            SharedQueryState state = SharedQueryState.create(query);
            String stateJson = MAPPER.writeValueAsString(state);
            if (state.getBasicQueryInfo().getState() == QueryState.FINISHED || state.getBasicQueryInfo().getState() == QueryState.FAILED) {
                // No need to update states for finished queries
                // also move finished queries to finished-query state collection
                queriesToUnregister.add(query);
                ((StateMap) finishedQueries).put(state.getBasicQueryInfo().getQueryId().getId(), stateJson);
                continue;
            }
            ((StateMap) queries).put(state.getBasicQueryInfo().getQueryId().getId(), stateJson);
        }
    }
    for (DispatchQuery query : queriesToUnregister) {
        removeFromStateCollection(stateStore, QUERY_STATE_COLLECTION_NAME, query);
    }
    long end = System.currentTimeMillis();
    LOG.debug("updateStates ends at current time milliseconds: %s, at format HH:mm:ss:SSS:%s, total time use: %s", end, new SimpleDateFormat("HH:mm:ss:SSS").format(new Date(end)), end - start);
}