Examples with AsynchronousSearchState org.opensearch.search.asynchronous.context.state.AsynchronousSearchState used on opensource projects

Example 1 with AsynchronousSearchState

use of org.opensearch.search.asynchronous.context.state.AsynchronousSearchState in project asynchronous-search by opensearch-project.

the class AsynchronousSearchService method initStateMachine.

private AsynchronousSearchStateMachine initStateMachine() {
    AsynchronousSearchStateMachine stateMachine = new AsynchronousSearchStateMachine(EnumSet.allOf(AsynchronousSearchState.class), INIT, contextEventListener);
    stateMachine.markTerminalStates(EnumSet.of(CLOSED));
    stateMachine.registerTransition(new AsynchronousSearchTransition<>(INIT, RUNNING, (s, e) -> ((AsynchronousSearchActiveContext) e.asynchronousSearchContext()).setTask(e.getSearchTask()), (contextId, listener) -> listener.onContextRunning(contextId), SearchStartedEvent.class));
    stateMachine.registerTransition(new AsynchronousSearchTransition<>(RUNNING, SUCCEEDED, (s, e) -> ((AsynchronousSearchActiveContext) e.asynchronousSearchContext()).processSearchResponse(e.getSearchResponse()), (contextId, listener) -> listener.onContextCompleted(contextId), SearchSuccessfulEvent.class));
    stateMachine.registerTransition(new AsynchronousSearchTransition<>(RUNNING, FAILED, (s, e) -> ((AsynchronousSearchActiveContext) e.asynchronousSearchContext()).processSearchFailure(e.getException()), (contextId, listener) -> listener.onContextFailed(contextId), SearchFailureEvent.class));
    stateMachine.registerTransition(new AsynchronousSearchTransition<>(SUCCEEDED, PERSISTING, (s, e) -> asynchronousSearchPostProcessor.persistResponse((AsynchronousSearchActiveContext) e.asynchronousSearchContext(), e.getAsynchronousSearchPersistenceModel()), (contextId, listener) -> {
    }, BeginPersistEvent.class));
    stateMachine.registerTransition(new AsynchronousSearchTransition<>(FAILED, PERSISTING, (s, e) -> asynchronousSearchPostProcessor.persistResponse((AsynchronousSearchActiveContext) e.asynchronousSearchContext(), e.getAsynchronousSearchPersistenceModel()), (contextId, listener) -> {
    }, BeginPersistEvent.class));
    stateMachine.registerTransition(new AsynchronousSearchTransition<>(PERSISTING, PERSIST_SUCCEEDED, (s, e) -> {
    }, (contextId, listener) -> listener.onContextPersisted(contextId), SearchResponsePersistedEvent.class));
    stateMachine.registerTransition(new AsynchronousSearchTransition<>(PERSISTING, PERSIST_FAILED, (s, e) -> {
    }, (contextId, listener) -> listener.onContextPersistFailed(contextId), SearchResponsePersistFailedEvent.class));
    stateMachine.registerTransition(new AsynchronousSearchTransition<>(RUNNING, CLOSED, (s, e) -> asynchronousSearchActiveStore.freeContext(e.asynchronousSearchContext().getContextId()), (contextId, listener) -> listener.onRunningContextDeleted(contextId), SearchDeletedEvent.class));
    for (AsynchronousSearchState state : EnumSet.of(PERSISTING, PERSIST_SUCCEEDED, PERSIST_FAILED, SUCCEEDED, FAILED, INIT)) {
        stateMachine.registerTransition(new AsynchronousSearchTransition<>(state, CLOSED, (s, e) -> asynchronousSearchActiveStore.freeContext(e.asynchronousSearchContext().getContextId()), (contextId, listener) -> listener.onContextDeleted(contextId), SearchDeletedEvent.class));
    }
    return stateMachine;
}

Example 2 with AsynchronousSearchState

use of org.opensearch.search.asynchronous.context.state.AsynchronousSearchState in project asynchronous-search by opensearch-project.

the class AsyncSearchBackwardsCompatibilityIT method testSubmitWithRetainedResponse.

public void testSubmitWithRetainedResponse(boolean shouldUseLegacyApi) throws IOException {
    SearchRequest searchRequest = new SearchRequest("test");
    searchRequest.source(new SearchSourceBuilder());
    SubmitAsynchronousSearchRequest submitAsynchronousSearchRequest = new SubmitAsynchronousSearchRequest(searchRequest);
    submitAsynchronousSearchRequest.keepOnCompletion(true);
    submitAsynchronousSearchRequest.waitForCompletionTimeout(TimeValue.timeValueMillis(randomLongBetween(1, 500)));
    AsynchronousSearchResponse submitResponse = executeSubmitAsynchronousSearch(submitAsynchronousSearchRequest, shouldUseLegacyApi);
    List<AsynchronousSearchState> legalStates = Arrays.asList(AsynchronousSearchState.RUNNING, AsynchronousSearchState.SUCCEEDED, AsynchronousSearchState.PERSIST_SUCCEEDED, AsynchronousSearchState.PERSISTING, AsynchronousSearchState.CLOSED, AsynchronousSearchState.STORE_RESIDENT);
    assertNotNull(submitResponse.getId());
    assertTrue(submitResponse.getState().name(), legalStates.contains(submitResponse.getState()));
    GetAsynchronousSearchRequest getAsynchronousSearchRequest = new GetAsynchronousSearchRequest(submitResponse.getId());
    AsynchronousSearchResponse getResponse;
    do {
        getResponse = getAssertedAsynchronousSearchResponse(submitResponse, getAsynchronousSearchRequest, shouldUseLegacyApi);
        if (getResponse.getState() == AsynchronousSearchState.RUNNING && getResponse.getSearchResponse() != null) {
            assertEquals(getResponse.getSearchResponse().getHits().getHits().length, 0);
        } else {
            assertNotNull(getResponse.getSearchResponse());
            assertNotEquals(getResponse.getSearchResponse().getTook(), -1L);
        }
    } while (AsynchronousSearchState.STORE_RESIDENT.equals(getResponse.getState()) == false);
    getResponse = getAssertedAsynchronousSearchResponse(submitResponse, getAsynchronousSearchRequest, shouldUseLegacyApi);
    assertNotNull(getResponse.getSearchResponse());
    assertEquals(AsynchronousSearchState.STORE_RESIDENT, getResponse.getState());
    assertHitCount(getResponse.getSearchResponse(), 5);
    executeDeleteAsynchronousSearch(new DeleteAsynchronousSearchRequest(submitResponse.getId()), shouldUseLegacyApi);
}

Example 3 with AsynchronousSearchState

use of org.opensearch.search.asynchronous.context.state.AsynchronousSearchState in project asynchronous-search by opensearch-project.

the class AsynchronousSearchResponse method innerFromXContent.

public static AsynchronousSearchResponse innerFromXContent(XContentParser parser) throws IOException {
    ensureExpectedToken(XContentParser.Token.FIELD_NAME, parser.currentToken(), parser);
    String id = null;
    AsynchronousSearchState status = null;
    long startTimeMillis = -1;
    long expirationTimeMillis = -1;
    SearchResponse searchResponse = null;
    OpenSearchException error = null;
    String currentFieldName = null;
    for (XContentParser.Token token = parser.nextToken(); token != XContentParser.Token.END_OBJECT; token = parser.nextToken()) {
        currentFieldName = parser.currentName();
        if (RESPONSE.match(currentFieldName, parser.getDeprecationHandler())) {
            if (token == XContentParser.Token.START_OBJECT) {
                ensureExpectedToken(XContentParser.Token.FIELD_NAME, parser.nextToken(), parser);
                searchResponse = SearchResponse.innerFromXContent(parser);
            }
        } else if (ERROR.match(currentFieldName, parser.getDeprecationHandler())) {
            parser.nextToken();
            error = OpenSearchException.fromXContent(parser);
        } else if (token.isValue()) {
            if (ID.match(currentFieldName, parser.getDeprecationHandler())) {
                id = parser.text();
            } else if (START_TIME_IN_MILLIS.match(currentFieldName, parser.getDeprecationHandler())) {
                startTimeMillis = parser.longValue();
            } else if (EXPIRATION_TIME_IN_MILLIS.match(currentFieldName, parser.getDeprecationHandler())) {
                expirationTimeMillis = parser.longValue();
            } else if (STATE.match(currentFieldName, parser.getDeprecationHandler())) {
                status = AsynchronousSearchState.valueOf(parser.text());
            } else {
                parser.skipChildren();
            }
        }
    }
    return new AsynchronousSearchResponse(id, status, startTimeMillis, expirationTimeMillis, searchResponse, error);
}

Example 4 with AsynchronousSearchState

use of org.opensearch.search.asynchronous.context.state.AsynchronousSearchState in project asynchronous-search by opensearch-project.

the class AsynchronousSearchStateMachineTests method doConcurrentStateMachineTrigger.

private <T extends Exception, R extends Exception> void doConcurrentStateMachineTrigger(AsynchronousSearchStateMachine asStateMachine, AsynchronousSearchContextEvent event, AsynchronousSearchState finalState, Class<T> throwable, Optional<Class<R>> terminalStateException) throws InterruptedException, BrokenBarrierException {
    int numThreads = 10;
    List<Thread> operationThreads = new ArrayList<>();
    AtomicInteger numTriggerSuccess = new AtomicInteger();
    CyclicBarrier barrier = new CyclicBarrier(numThreads + 1);
    for (int i = 0; i < numThreads; i++) {
        Thread thread = new Thread(() -> {
            try {
                AsynchronousSearchState state = asStateMachine.trigger(event);
                assertEquals(state, finalState);
                numTriggerSuccess.getAndIncrement();
            } catch (Exception e) {
                if (terminalStateException.isPresent()) {
                    assertTrue(terminalStateException.get().isInstance(e) || throwable.isInstance(e));
                } else {
                    assertTrue(throwable.isInstance(e));
                }
            } finally {
                try {
                    barrier.await();
                } catch (InterruptedException | BrokenBarrierException e) {
                    fail("stage advancement failure");
                }
            }
        });
        operationThreads.add(thread);
    }
    operationThreads.forEach(Thread::start);
    barrier.await();
    for (Thread t : operationThreads) {
        t.join();
    }
    assertEquals(1, numTriggerSuccess.get());
}

Example 5 with AsynchronousSearchState

use of org.opensearch.search.asynchronous.context.state.AsynchronousSearchState in project asynchronous-search by opensearch-project.

the class ApiParamsValidationIT method testSubmitSearchOnInvalidIndex.

public void testSubmitSearchOnInvalidIndex() throws IOException {
    SubmitAsynchronousSearchRequest submitAsynchronousSearchRequest = new SubmitAsynchronousSearchRequest(new SearchRequest("invalid-index"));
    AsynchronousSearchResponse submitResponse = executeSubmitAsynchronousSearch(submitAsynchronousSearchRequest);
    List<AsynchronousSearchState> legalStates = Arrays.asList(AsynchronousSearchState.FAILED, AsynchronousSearchState.CLOSED);
    assertNull(submitResponse.getSearchResponse());
    assertNotNull(submitResponse.getError());
    assertThat(submitResponse.getError().getDetailedMessage(), containsString("index_not_found"));
    assertTrue(legalStates.contains(submitResponse.getState()));
}