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

Example 6 with StateMap

use of io.prestosql.spi.statestore.StateMap 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 7 with StateMap

use of io.prestosql.spi.statestore.StateMap 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);
}

Example 8 with StateMap

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

the class SqlQueryManager method updateKillingQuery.

/**
 * Check if any query is getting killed and update state store
 *
 * @param runningQueries All the running query executions
 * @param queryStates All the query states
 */
private void updateKillingQuery(List<QueryExecution> runningQueries, Map<String, SharedQueryState> queryStates) {
    for (QueryExecution query : runningQueries) {
        SharedQueryExecution queryExecution = (SharedQueryExecution) query;
        if (((SharedQueryExecution) query).isGettingKilled()) {
            StateMap stateMap = (StateMap<String, String>) stateStoreProvider.getStateStore().getStateCollection(StateStoreConstants.OOM_QUERY_STATE_COLLECTION_NAME);
            if (stateMap != null) {
                ObjectMapper mapper = new ObjectMapperProvider().get();
                try {
                    String stateJson = mapper.writeValueAsString(queryStates.get(queryExecution.getQueryId().getId()));
                    stateMap.put(queryExecution.getQueryId().getId(), stateJson);
                } catch (JsonProcessingException e) {
                    log.warn("Query %s state serialization failed: %s", queryExecution.getQueryId().getId(), e.getMessage());
                }
            }
            break;
        }
    }
}

Example 9 with StateMap

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

the class SplitCacheStateUpdater method updateStateStore.

public void updateStateStore() {
    StateStore stateStore = provider.getStateStore();
    if (stateStore == null) {
        log.debug("State store not yet initialized.");
        return;
    }
    if (status.get() == InitializationStatus.INITIALIZING) {
        log.debug("Split cache map not yet initialized.");
        return;
    }
    // create state store collection
    if (stateStore.getStateCollection(StateStoreConstants.SPLIT_CACHE_METADATA_NAME) == null) {
        stateStore.createStateMap(StateStoreConstants.SPLIT_CACHE_METADATA_NAME, new SplitCacheStateStoreChangesListener(SplitCacheMap.getInstance(), mapper));
    }
    StateMap<String, String> stateMap = (StateMap<String, String>) stateStore.getStateCollection(StateStoreConstants.SPLIT_CACHE_METADATA_NAME);
    splitCacheMap.getAndClearDroppedCaches().forEach(stateMap::remove);
    splitCacheMap.tableCacheInfoMap().forEach((fqTableName, localCacheInfo) -> {
        try {
            // Async update works only when new cache predicates added and splits are scheduled or updated
            // It does not perform merge also in case if both local info and state store are updated.
            // The logic is very simple. Only update state store if local copy is more recent than the one in state store.
            TableCacheInfo stateStoreCacheInfo = stateMap.containsKey(fqTableName) ? mapper.readerFor(TableCacheInfo.class).readValue(stateMap.get(fqTableName)) : null;
            if (stateStoreCacheInfo == null || localCacheInfo.getLastUpdated().isAfter(stateStoreCacheInfo.getLastUpdated())) {
                log.info("Updating state store split cache map for table %s.", fqTableName);
                String json = mapper.writeValueAsString(localCacheInfo);
                stateMap.put(fqTableName, json);
            } else if (localCacheInfo.getLastUpdated().isBefore(stateStoreCacheInfo.getLastUpdated())) {
                log.debug("Not updating state store split cache map for table %s. Local copy is outdated. State store split cache map is more recent. Local split cache map should be updated.", fqTableName);
            }
        } catch (IOException e) {
            log.error(e, "Unable to update state store split cache map.");
        }
    });
}

Example 10 with StateMap

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

the class SqlQueryExecution method handleCrossRegionDynamicFilter.

private void handleCrossRegionDynamicFilter(PlanRoot plan) {
    if (!isCrossRegionDynamicFilterEnabled(getSession()) || plan == null) {
        return;
    }
    StateStore stateStore = stateStoreProvider.getStateStore();
    if (stateStore == null) {
        return;
    }
    String queryId = getSession().getQueryId().getId();
    log.debug("queryId=%s begin to find columnToColumnMapping.", queryId);
    PlanNode outputNode = plan.getRoot().getFragment().getRoot();
    Map<String, Set<String>> columnToSymbolMapping = new HashMap<>();
    if (outputNode != null && outputNode instanceof OutputNode) {
        List<String> queryColumnNames = ((OutputNode) outputNode).getColumnNames();
        List<Symbol> outputSymbols = outputNode.getOutputSymbols();
        Map<String, Set<String>> tmpMapping = new HashMap<>(outputSymbols.size());
        for (Symbol symbol : outputNode.getOutputSymbols()) {
            Set<String> sets = new HashSet();
            sets.add(symbol.getName());
            tmpMapping.put(symbol.getName(), sets);
        }
        for (PlanFragment fragment : plan.getRoot().getAllFragments()) {
            if ("0".equals(fragment.getId().toString())) {
                continue;
            }
            PlanNode sourceNode = fragment.getRoot();
            findMappingFromPlan(tmpMapping, sourceNode);
        }
        for (int i = 0; i < outputSymbols.size(); i++) {
            columnToSymbolMapping.put(queryColumnNames.get(i), tmpMapping.get(outputSymbols.get(i).getName()));
        }
    }
    // save mapping into stateStore
    StateMap<String, Object> mappingStateMap = (StateMap<String, Object>) stateStore.getOrCreateStateCollection(CROSS_REGION_DYNAMIC_FILTERS, StateCollection.Type.MAP);
    mappingStateMap.put(queryId + QUERY_COLUMN_NAME_TO_SYMBOL_MAPPING, columnToSymbolMapping);
    log.debug("queryId=%s, add columnToSymbolMapping into hazelcast success.", queryId + QUERY_COLUMN_NAME_TO_SYMBOL_MAPPING);
}