Examples with DispatchQuery io.prestosql.dispatcher.DispatchQuery used on opensource projects

Example 1 with DispatchQuery

use of io.prestosql.dispatcher.DispatchQuery in project hetu-core by openlookeng.

the class TestStateUpdater method mockDispatchQueryData.

private DispatchQuery mockDispatchQueryData(boolean userError) {
    DispatchQuery dispatchQuery = Mockito.mock(LocalDispatchQuery.class);
    BasicQueryInfo basicQueryInfo = createBasicQueryInfo();
    when(dispatchQuery.getBasicQueryInfo()).then(new Returns(basicQueryInfo));
    when(dispatchQuery.getSession()).then(new Returns(TEST_SESSION));
    ErrorCode errorCode;
    if (!userError) {
        errorCode = CLUSTER_OUT_OF_MEMORY.toErrorCode();
    } else {
        errorCode = new ErrorCode(ERROR_CODE_VALUE_INDEX_TIME_NO_INVOCATION, ERROR_CODE_TEST, USER_ERROR);
    }
    Optional<ErrorCode> optionalErrorCode = Optional.of(errorCode);
    when(dispatchQuery.getErrorCode()).then(new Returns(optionalErrorCode));
    DataSize userDataSize = new DataSize(USER_DATA_SIZE, DataSize.Unit.BYTE);
    DataSize totalDataSize = new DataSize(TOTAL_DATA_SIZE, DataSize.Unit.BYTE);
    when(dispatchQuery.getUserMemoryReservation()).then(new Returns(userDataSize));
    when(dispatchQuery.getTotalMemoryReservation()).then(new Returns(totalDataSize));
    when(dispatchQuery.getTotalCpuTime()).then(new Returns(new Duration(ERROR_CODE_VALUE_INDEX_TIME_NO_INVOCATION, MILLISECONDS)));
    when(dispatchQuery.getExecutionStartTime()).then(new Returns(Optional.of(new DateTime(DateTimeZone.UTC))));
    return dispatchQuery;
}

Example 2 with DispatchQuery

use of io.prestosql.dispatcher.DispatchQuery in project hetu-core by openlookeng.

the class TestStateUpdater method testRegisterQuery.

@Test
public void testRegisterQuery() {
    StateStoreProvider stateStoreProvider = Mockito.mock(LocalStateStoreProvider.class);
    StateUpdater stateUpdater = new StateUpdater(stateStoreProvider, updateInterval);
    DispatchQuery dispatchQuery = Mockito.mock(DispatchQuery.class);
    stateUpdater.registerQuery(STATE_COLLECTION_QUERY, dispatchQuery);
}

Example 3 with DispatchQuery

use of io.prestosql.dispatcher.DispatchQuery 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 4 with DispatchQuery

use of io.prestosql.dispatcher.DispatchQuery in project hetu-core by openlookeng.

the class AutoVacuumScanner method startVacuum.

private void startVacuum(Catalog catalog, String vacuumTable, boolean isFull) {
    String catalogNameVacuumTable = catalog.getCatalogName() + "." + vacuumTable;
    if (vacuumInProgressMap.containsKey(catalogNameVacuumTable)) {
        log.debug("return Present in vacuumInProgressMap %s ", catalogNameVacuumTable);
        return;
    }
    long attempts = 0;
    QueryId queryId = dispatchManager.createQueryId();
    String slug = "x" + randomUUID().toString().toLowerCase(ENGLISH).replace("-", "");
    String vacuumQuery;
    if (isFull) {
        vacuumQuery = "vacuum table " + catalogNameVacuumTable + "  full";
    } else {
        vacuumQuery = "vacuum table " + catalogNameVacuumTable;
    }
    Session.SessionBuilder sessionBuilder = Session.builder(sessionPropertyManager).setQueryId(queryId).setIdentity(new Identity("openLooKeng", Optional.empty())).setSource("auto-vacuum");
    Session session = sessionBuilder.build();
    AutoVacuumSessionContext sessionContext = new AutoVacuumSessionContext(session);
    vacuumInProgressMap.put(catalogNameVacuumTable, System.currentTimeMillis());
    log.debug("Query.create queryId %s  catalogNameVacuumTable: %s ", queryId.toString(), catalogNameVacuumTable);
    ListenableFuture<?> lf = waitForDispatched(queryId, slug, sessionContext, vacuumQuery);
    Futures.addCallback(lf, new FutureCallback<Object>() {

        @Override
        public void onSuccess(@Nullable Object result) {
            try {
                DispatchQuery dispatchQuery = dispatchManager.getQuery(queryId);
                dispatchQuery.addStateChangeListener((state) -> {
                    Query query = getQuery(queryId, slug);
                    if ((null != query) && (!dispatchManager.getQueryInfo(queryId).getState().isDone())) {
                        query.waitForResults(attempts, Duration.valueOf("1s"), DataSize.valueOf("1MB"));
                    }
                    if (state.isDone()) {
                        log.debug("STATUS  %s QueryID %s Query %s", state.name(), queryId.toString(), vacuumQuery);
                        vacuumInProgressMap.remove(catalogNameVacuumTable);
                    }
                });
            } catch (Throwable e) {
                vacuumInProgressMap.remove(catalogNameVacuumTable);
                log.error("Filed to execute vacuum for table %s QueryID %s", catalogNameVacuumTable, queryId.toString(), e.getMessage());
            }
        }

        @Override
        public void onFailure(Throwable t) {
            vacuumInProgressMap.remove(catalogNameVacuumTable);
            log.error("Query %s request to start vacuum scan failed at queryId[%s]: %s ", vacuumQuery, queryId, t.getMessage());
        }
    }, directExecutor());
}

Example 5 with DispatchQuery

use of io.prestosql.dispatcher.DispatchQuery in project hetu-core by openlookeng.

the class TestStateUpdater method testStateChangeListenerOfRegisterQuery.

@Test
public void testStateChangeListenerOfRegisterQuery() throws JsonProcessingException {
    StateStoreProvider stateStoreProvider = Mockito.mock(LocalStateStoreProvider.class);
    StateUpdater stateUpdater = new StateUpdater(stateStoreProvider, updateInterval);
    DispatchQuery dispatchQuery = mockDispatchQueryData(true);
    updateStateChange(dispatchQuery);
    stateUpdater.registerQuery(STATE_COLLECTION_QUERY, dispatchQuery);
    when(stateStoreProvider.getStateStore()).then(new Returns(stateStore));
    when(stateStoreProvider.getStateStore().getStateCollection(any())).then(new Returns(Mockito.mock(StateMap.class)));
    when(stateStoreProvider.getStateStore().getStateCollection(any()).getType()).then(new Returns(StateCollection.Type.MAP));
    stateUpdater.updateStates();
    int numberOfCalls = mockingDetails(stateStoreProvider.getStateStore().getStateCollection(any())).getInvocations().size();
    assertNotEquals(numberOfCalls, ERROR_CODE_VALUE_INDEX_TIME_NO_INVOCATION);
}