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

Example 1 with StateStore

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

the class TestDynamicFiltersCollector method TestCollectingGlobalDynamicFilters.

@Test
public void TestCollectingGlobalDynamicFilters() throws InterruptedException {
    final QueryId queryId = new QueryId("test_query");
    final String filterId = "1";
    final String columnName = "column";
    final TestingColumnHandle columnHandle = new TestingColumnHandle(columnName);
    final Set<String> valueSet = ImmutableSet.of("1", "2", "3");
    TaskContext taskContext = mock(TaskContext.class);
    Session session = testSessionBuilder().setQueryId(queryId).setSystemProperty(ENABLE_DYNAMIC_FILTERING, "true").setSystemProperty(DYNAMIC_FILTERING_DATA_TYPE, "HASHSET").build();
    when(taskContext.getSession()).thenReturn(session);
    // set up state store and merged dynamic filters map
    Map mockMap = new HashMap<>();
    StateStoreProvider stateStoreProvider = mock(StateStoreProvider.class);
    StateStore stateStore = mock(StateStore.class);
    StateMap stateMap = new MockStateMap<>("test-map", mockMap);
    when(stateStoreProvider.getStateStore()).thenReturn(stateStore);
    when(stateStore.getStateCollection(any())).thenReturn(stateMap);
    when(stateStore.createStateMap(any())).thenReturn(stateMap);
    when(stateStore.getOrCreateStateCollection(any(), any())).thenReturn(stateMap);
    // set up state store listener and dynamic filter cache
    StateStoreListenerManager stateStoreListenerManager = new StateStoreListenerManager(stateStoreProvider);
    DynamicFilterCacheManager dynamicFilterCacheManager = new DynamicFilterCacheManager();
    stateStoreListenerManager.addStateStoreListener(new DynamicFilterListener(dynamicFilterCacheManager), MERGED_DYNAMIC_FILTERS);
    LocalDynamicFiltersCollector collector = new LocalDynamicFiltersCollector(taskContext, Optional.empty(), dynamicFilterCacheManager);
    TableScanNode tableScan = mock(TableScanNode.class);
    when(tableScan.getAssignments()).thenReturn(ImmutableMap.of(new Symbol(columnName), columnHandle));
    List<DynamicFilters.Descriptor> dynamicFilterDescriptors = ImmutableList.of(new DynamicFilters.Descriptor(filterId, new VariableReferenceExpression(columnName, BIGINT)));
    collector.initContext(ImmutableList.of(dynamicFilterDescriptors), SymbolUtils.toLayOut(tableScan.getOutputSymbols()));
    assertTrue(collector.getDynamicFilters(tableScan).isEmpty(), "there should be no dynamic filter available");
    // put some values in state store as a new dynamic filter
    // and wait for the listener to process the event
    stateMap.put(createKey(DynamicFilterUtils.FILTERPREFIX, filterId, queryId.getId()), valueSet);
    TimeUnit.MILLISECONDS.sleep(100);
    // get available dynamic filter and verify it
    List<Map<ColumnHandle, DynamicFilter>> dynamicFilters = collector.getDynamicFilters(tableScan);
    assertEquals(dynamicFilters.size(), 1, "there should be a new dynamic filter");
    assertEquals(dynamicFilters.size(), 1);
    DynamicFilter dynamicFilter = dynamicFilters.get(0).get(columnHandle);
    assertTrue(dynamicFilter instanceof HashSetDynamicFilter, "new dynamic filter should be hashset");
    assertEquals(dynamicFilter.getSize(), valueSet.size(), "new dynamic filter should have correct size");
    for (String value : valueSet) {
        assertTrue(dynamicFilter.contains(value), "new dynamic filter should contain correct values");
    }
    // clean up when task finishes
    collector.removeDynamicFilter(true);
    DynamicFilter cachedFilter = dynamicFilterCacheManager.getDynamicFilter(DynamicFilterCacheManager.createCacheKey(filterId, queryId.getId()));
    assertNull(cachedFilter, "cached dynamic filter should have been removed");
}

Example 2 with StateStore

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

the class TestDynamicFilterUtil method setupMockStateStore.

public static StateStore setupMockStateStore(Map mergeMap, Map<String, String> dfTypeMap, Set<String> tasks, Set partial, String queryId, String filterId) {
    StateMap mockMergeMap = mock(StateMap.class);
    StateMap mockDFTypeMap = mock(StateMap.class);
    StateSet mockPartialSet = mock(StateSet.class);
    StateSet mockTasksSet = mock(StateSet.class);
    StateStore stateStore = mock(StateStore.class);
    when(mockMergeMap.put(anyString(), any())).thenAnswer(i -> mergeMap.put(i.getArguments()[0], i.getArguments()[1]));
    when(mockDFTypeMap.put(anyString(), anyString())).thenAnswer(i -> dfTypeMap.put((String) i.getArguments()[0], (String) i.getArguments()[1]));
    when(mockTasksSet.add(anyString())).thenAnswer(i -> tasks.add((String) i.getArguments()[0]));
    when(mockPartialSet.add(any())).thenAnswer(i -> partial.add(i.getArguments()[0]));
    when(mockMergeMap.get(anyString())).thenAnswer(i -> mergeMap.get(i.getArguments()[0]));
    when(mockDFTypeMap.get(anyString())).thenAnswer(i -> dfTypeMap.get(i.getArguments()[0]));
    when(mockMergeMap.getAll()).thenReturn(mergeMap);
    when(mockDFTypeMap.getAll()).thenReturn(dfTypeMap);
    when(mockPartialSet.getAll()).thenReturn(partial);
    when(mockPartialSet.size()).thenAnswer(i -> partial.size());
    when(mockTasksSet.size()).thenAnswer(i -> tasks.size());
    when(stateStore.getStateCollection(DynamicFilterUtils.MERGED_DYNAMIC_FILTERS)).thenReturn(mockMergeMap);
    when(stateStore.createStateCollection(DynamicFilterUtils.MERGED_DYNAMIC_FILTERS, StateCollection.Type.MAP)).thenReturn(mockMergeMap);
    when(stateStore.getOrCreateStateCollection(DynamicFilterUtils.MERGED_DYNAMIC_FILTERS, StateCollection.Type.MAP)).thenReturn(mockMergeMap);
    when(stateStore.getStateCollection(DynamicFilterUtils.createKey(DynamicFilterUtils.TASKSPREFIX, filterId, queryId))).thenReturn(mockTasksSet);
    when(stateStore.createStateCollection(DynamicFilterUtils.createKey(DynamicFilterUtils.TASKSPREFIX, filterId, queryId), StateCollection.Type.SET)).thenReturn(mockTasksSet);
    when(stateStore.getStateCollection(DynamicFilterUtils.createKey(DynamicFilterUtils.PARTIALPREFIX, filterId, queryId))).thenReturn(mockPartialSet);
    when(stateStore.createStateCollection(DynamicFilterUtils.createKey(DynamicFilterUtils.PARTIALPREFIX, filterId, queryId), StateCollection.Type.SET)).thenReturn(mockPartialSet);
    // In statestore, set and map are destroyed and set to null after query finishes, however, in the UT we just assume the set and map to be empty.
    doAnswer(i -> {
        tasks.clear();
        return null;
    }).when(mockTasksSet).destroy();
    doAnswer(i -> {
        partial.clear();
        return null;
    }).when(mockPartialSet).destroy();
    return stateStore;
}

Example 3 with StateStore

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

the class TestStateStoreLauncherAndProvider method testRegisterDiscoveryService.

@Test(timeOut = 5000, expectedExceptions = ThreadTimeoutException.class)
public void testRegisterDiscoveryService() throws Exception {
    String failurehost = "failurehost";
    String otherhost = "otherhost";
    String localHostName = "localhost";
    int port = 8888;
    URI uri = new URI("http://" + localHostName + ":" + port);
    MockStateMap discoveryServiceMap = new MockStateMap(DISCOVERY_SERVICE, new HashMap<>());
    // Mock
    StateStore stateStore = mock(StateStore.class);
    Lock lock = mock(ReentrantLock.class);
    InternalCommunicationConfig internalCommunicationConfig = mock(InternalCommunicationConfig.class);
    HttpServerInfo httpServerInfo = mock(HttpServerInfo.class);
    when(httpServerInfo.getHttpUri()).thenReturn(uri);
    when(internalCommunicationConfig.isHttpsRequired()).thenReturn(false);
    when(stateStore.getStateCollection(DISCOVERY_SERVICE)).thenReturn(discoveryServiceMap);
    when(stateStore.getLock(DISCOVERY_SERVICE_LOCK)).thenReturn(lock);
    EmbeddedStateStoreLauncher launcher = new EmbeddedStateStoreLauncher(new SeedStoreManager(new FileSystemClientManager()), internalCommunicationConfig, httpServerInfo, new HetuConfig());
    launcher.setStateStore(stateStore);
    when(lock.tryLock(DISCOVERY_REGISTRY_LOCK_TIMEOUT, TimeUnit.MILLISECONDS)).thenReturn(true);
    // discoveryServiceMap is empty, so the current coordinator can get the lock and register itself(register=true)
    discoveryServiceMap.clear();
    assertTrue(launcher.registerDiscoveryService(failurehost));
    assertEquals(discoveryServiceMap.size(), 1);
    assertTrue(discoveryServiceMap.getAll().keySet().contains(localHostName));
    // discoveryServiceMap contains the failure host, so the current coordinator can get the lock and register itself(register=true)
    discoveryServiceMap.clear();
    discoveryServiceMap.put(failurehost, String.valueOf(port));
    assertTrue(launcher.registerDiscoveryService(failurehost));
    assertEquals(discoveryServiceMap.size(), 1);
    assertTrue(discoveryServiceMap.getAll().keySet().contains(localHostName));
    // discoveryServiceMap is already updated by other coordinator(otherhosts)
    // the current coordinator can grab the lock but will not register itself(register=false)
    discoveryServiceMap.clear();
    discoveryServiceMap.put(otherhost, String.valueOf(port));
    assertFalse(launcher.registerDiscoveryService(failurehost));
    assertEquals(discoveryServiceMap.size(), 1);
    assertFalse(discoveryServiceMap.containsKey(localHostName));
    when(lock.tryLock(DISCOVERY_REGISTRY_LOCK_TIMEOUT, TimeUnit.MILLISECONDS)).thenReturn(false);
    // discoveryServiceMap is already updated by other coordinator(otherhosts)
    // the current coordinator cannot grab the lock and not register itself
    discoveryServiceMap.clear();
    discoveryServiceMap.put(otherhost, String.valueOf(port));
    assertFalse(launcher.registerDiscoveryService(failurehost));
    assertEquals(discoveryServiceMap.size(), 1);
    assertFalse(discoveryServiceMap.containsKey(localHostName));
    // discoveryServiceMap contains failure host.
    // The current coordinator cannot get the lock and retry will cause timeout exception
    discoveryServiceMap.clear();
    discoveryServiceMap.put(failurehost, String.valueOf(port));
    launcher.registerDiscoveryService(failurehost);
}

Example 4 with StateStore

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

the class HetuServiceInventory method getServiceDescriptors.

@Override
public Iterable<ServiceDescriptor> getServiceDescriptors(String type) {
    if ("discovery".equals(type)) {
        if (hetuConfig.isMultipleCoordinatorEnabled()) {
            try {
                StateStore stateStore = stateStoreProvider.getStateStore();
                if (stateStore == null) {
                    throw new PrestoException(GENERIC_INTERNAL_ERROR, "State store has not been loaded yet");
                }
                Map.Entry<String, String> entry = ((StateMap<String, String>) stateStore.getStateCollection(StateStoreConstants.DISCOVERY_SERVICE_COLLECTION_NAME)).getAll().entrySet().stream().findFirst().get();
                ImmutableMap.Builder properties = new ImmutableMap.Builder();
                if (internalCommunicationConfig != null && internalCommunicationConfig.isHttpsRequired()) {
                    properties.put("https", "https://" + entry.getKey() + ":" + entry.getValue());
                } else {
                    properties.put("http", "http://" + entry.getKey() + ":" + entry.getValue());
                }
                return ImmutableList.of(new ServiceDescriptor(UUID.randomUUID(), nodeId, "discovery", null, null, ServiceState.RUNNING, properties.build()));
            } catch (Exception e) {
                logServerError("Select service from state store failed: " + e);
                // return an empty list here.
                return ImmutableList.of();
            }
        } else if (seedStoreManager != null && seedStoreManager.isSeedStoreOnYarnEnabled()) {
            try {
                boolean httpsRequired = (internalCommunicationConfig != null && internalCommunicationConfig.isHttpsRequired());
                String location = seedStoreManager.getLatestSeedLocation(SeedStoreSubType.ON_YARN, httpsRequired);
                if (location == null || location.isEmpty()) {
                    throw new PrestoException(GENERIC_INTERNAL_ERROR, "Seed store has not been initialized yet");
                }
                ImmutableMap.Builder properties = new ImmutableMap.Builder();
                if (httpsRequired) {
                    properties.put("https", location);
                } else {
                    properties.put("http", location);
                }
                return ImmutableList.of(new ServiceDescriptor(UUID.randomUUID(), nodeId, "discovery", null, null, ServiceState.RUNNING, properties.build()));
            } catch (Exception e) {
                logServerError("Select service from seed store failed: " + e);
                // return an empty list here.
                return ImmutableList.of();
            }
        }
    }
    return super.getServiceDescriptors(type);
}

Example 5 with StateStore

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