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

Example 6 with StateStore

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

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

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

Example 9 with StateStore

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

the class TestHazelcastClusterLifecycleListener method createStateStoreClient.

private StateStore createStateStoreClient(Set<Seed> seeds) {
    Map<String, String> config = new HashMap<>(0);
    config.put("hazelcast.discovery.mode", "tcp-ip");
    config.put("state-store.cluster", TEST_CLUSTER_NAME);
    config.put(DISCOVERY_PORT_CONFIG_NAME, PORT1);
    MockSeedStore mockSeedStore = new MockSeedStore();
    mockSeedStore.add(seeds);
    StateStoreFactory factory = new HazelcastStateStoreFactory();
    StateStore stateStore = factory.create("testHazelcast", mockSeedStore, config);
    stateStore.registerClusterFailureHandler(event -> isClusterShutdownEventCaptured = true);
    return stateStore;
}

Example 10 with StateStore

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

the class TestHazelcastClusterLifecycleListener method testClusterShutdown.

@Test
void testClusterShutdown() throws InterruptedException {
    final long sleep = 30000L;
    isClusterShutdownEventCaptured = false;
    final StateStore member1 = createStateStoreCluster(PORT1);
    final StateStore member2 = createStateStoreCluster(PORT2);
    HashSet<Seed> seeds = new HashSet<>(0);
    seeds.add(new MockSeed(MEMBER_1_ADDRESS));
    seeds.add(new MockSeed(MEMBER_2_ADDRESS));
    StateStore client = createStateStoreClient(seeds);
    assertFalse(isClusterShutdownEventCaptured);
    ((HazelcastStateStore) member1).shutdown();
    assertFalse(isClusterShutdownEventCaptured);
    ((HazelcastStateStore) member2).shutdown();
    Thread.sleep(sleep);
    assertTrue(isClusterShutdownEventCaptured);
}