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

Example 1 with AsynchronousSearchStateMachine

use of org.opensearch.search.asynchronous.context.state.AsynchronousSearchStateMachine 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 AsynchronousSearchStateMachine

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

the class AsynchronousSearchPostProcessorTests method testProcessSearchResponseForExpiredContext.

public void testProcessSearchResponseForExpiredContext() throws AsynchronousSearchStateMachineClosedException {
    DiscoveryNode discoveryNode = new DiscoveryNode("node", OpenSearchTestCase.buildNewFakeTransportAddress(), Collections.emptyMap(), DiscoveryNodeRole.BUILT_IN_ROLES, Version.CURRENT);
    AtomicBoolean activeContextCleanUpConsumerInvocation = new AtomicBoolean();
    ThreadPool testThreadPool = null;
    try {
        testThreadPool = new TestThreadPool(AsynchronousSearchPlugin.OPEN_DISTRO_ASYNC_SEARCH_GENERIC_THREAD_POOL_NAME, executorBuilder);
        ClusterService mockClusterService = ClusterServiceUtils.createClusterService(testThreadPool, discoveryNode, clusterSettings);
        FakeClient fakeClient = new FakeClient(testThreadPool);
        AsynchronousSearchActiveStore asActiveStore = new AsynchronousSearchActiveStore(mockClusterService);
        AsynchronousSearchPersistenceService persistenceService = new AsynchronousSearchPersistenceService(fakeClient, mockClusterService, testThreadPool);
        AsynchronousSearchService asService = new AsynchronousSearchService(persistenceService, asActiveStore, fakeClient, mockClusterService, testThreadPool, new InternalAsynchronousSearchStats(), new NamedWriteableRegistry(emptyList()));
        AsynchronousSearchStateMachine asStateMachine = asService.getStateMachine();
        ClusterService clusterService = ClusterServiceUtils.createClusterService(testThreadPool, discoveryNode, clusterSettings);
        AsynchronousSearchPostProcessor postProcessor = new AsynchronousSearchPostProcessor(persistenceService, asActiveStore, asStateMachine, (context) -> activeContextCleanUpConsumerInvocation.compareAndSet(false, true), testThreadPool, clusterService);
        SubmitAsynchronousSearchRequest submitAsynchronousSearchRequest = new SubmitAsynchronousSearchRequest(new SearchRequest());
        submitAsynchronousSearchRequest.keepOnCompletion(true);
        submitAsynchronousSearchRequest.keepAlive(TimeValue.timeValueMillis(1));
        AsynchronousSearchActiveContext context = (AsynchronousSearchActiveContext) asService.createAndStoreContext(submitAsynchronousSearchRequest, System.currentTimeMillis(), () -> InternalAggregationTestCase.emptyReduceContextBuilder(), null);
        asStateMachine.trigger(new SearchStartedEvent(context, new SearchTask(0, "n/a", "n/a", () -> "test", null, Collections.emptyMap())));
        SearchResponse mockSearchResponse = getMockSearchResponse();
        waitUntil(() -> context.isExpired());
        AsynchronousSearchResponse asResponse = postProcessor.processSearchResponse(mockSearchResponse, context.getContextId());
        assertNotNull(asResponse.getId());
        assertNull(asResponse.getError());
        assertEquals(AsynchronousSearchState.SUCCEEDED, asResponse.getState());
        AsynchronousSearchAssertions.assertSearchResponses(mockSearchResponse, asResponse.getSearchResponse());
        assertTrue(activeContextCleanUpConsumerInvocation.get());
        assertEquals(0, fakeClient.persistenceCount);
    } catch (InterruptedException e) {
        fail("Interrupted exception" + e.getMessage());
    } finally {
        ThreadPool.terminate(testThreadPool, 200, TimeUnit.MILLISECONDS);
    }
}

Example 3 with AsynchronousSearchStateMachine

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

the class AsynchronousSearchPostProcessorTests method testProcessSearchFailureOnDeletedContext.

public void testProcessSearchFailureOnDeletedContext() throws AsynchronousSearchStateMachineClosedException {
    DiscoveryNode discoveryNode = new DiscoveryNode("node", OpenSearchTestCase.buildNewFakeTransportAddress(), Collections.emptyMap(), DiscoveryNodeRole.BUILT_IN_ROLES, Version.CURRENT);
    AtomicBoolean activeContextCleanUpConsumerInvocation = new AtomicBoolean();
    ThreadPool testThreadPool = null;
    try {
        testThreadPool = new TestThreadPool(this.getClass().getName(), executorBuilder);
        ClusterService mockClusterService = ClusterServiceUtils.createClusterService(testThreadPool, discoveryNode, clusterSettings);
        FakeClient fakeClient = new FakeClient(testThreadPool);
        AsynchronousSearchActiveStore asActiveStore = new AsynchronousSearchActiveStore(mockClusterService);
        AsynchronousSearchPersistenceService persistenceService = new AsynchronousSearchPersistenceService(fakeClient, mockClusterService, testThreadPool);
        AsynchronousSearchService asService = new AsynchronousSearchService(persistenceService, asActiveStore, fakeClient, mockClusterService, testThreadPool, new InternalAsynchronousSearchStats(), new NamedWriteableRegistry(emptyList()));
        AsynchronousSearchStateMachine asStateMachine = asService.getStateMachine();
        ClusterService clusterService = ClusterServiceUtils.createClusterService(testThreadPool, discoveryNode, clusterSettings);
        AsynchronousSearchPostProcessor postProcessor = new AsynchronousSearchPostProcessor(persistenceService, asActiveStore, asStateMachine, (context) -> activeContextCleanUpConsumerInvocation.compareAndSet(false, true), testThreadPool, clusterService);
        SubmitAsynchronousSearchRequest submitAsynchronousSearchRequest = new SubmitAsynchronousSearchRequest(new SearchRequest());
        submitAsynchronousSearchRequest.keepOnCompletion(true);
        submitAsynchronousSearchRequest.keepAlive(TimeValue.timeValueHours(1));
        AsynchronousSearchActiveContext context = (AsynchronousSearchActiveContext) asService.createAndStoreContext(submitAsynchronousSearchRequest, System.currentTimeMillis(), () -> InternalAggregationTestCase.emptyReduceContextBuilder(), null);
        ShardSearchFailure shardSearchFailure = new ShardSearchFailure(new RuntimeException("runtime-exception"));
        SearchPhaseExecutionException exception = new SearchPhaseExecutionException("phase", "msg", new NullPointerException(), new ShardSearchFailure[] { shardSearchFailure });
        asStateMachine.trigger(new SearchStartedEvent(context, new SearchTask(0, "n/a", "n/a", () -> "test", null, Collections.emptyMap())));
        asStateMachine.trigger(new SearchDeletedEvent(context));
        AsynchronousSearchResponse asResponse = postProcessor.processSearchFailure(exception, context.getContextId());
        assertNull(asResponse.getId());
        assertNull(asResponse.getSearchResponse());
        assertEquals(-1L, asResponse.getExpirationTimeMillis());
        assertEquals(-1L, asResponse.getStartTimeMillis());
        assertEquals(AsynchronousSearchState.FAILED, asResponse.getState());
        assertThat(asResponse.getError(), instanceOf(SearchPhaseExecutionException.class));
        assertFalse(activeContextCleanUpConsumerInvocation.get());
        assertEquals(0, fakeClient.persistenceCount);
    } finally {
        ThreadPool.terminate(testThreadPool, 200, TimeUnit.MILLISECONDS);
    }
}

Example 4 with AsynchronousSearchStateMachine

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

the class AsynchronousSearchPostProcessorTests method testProcessSearchResponseBeginPersistence.

public void testProcessSearchResponseBeginPersistence() throws AsynchronousSearchStateMachineClosedException, InterruptedException {
    DiscoveryNode discoveryNode = new DiscoveryNode("node", OpenSearchTestCase.buildNewFakeTransportAddress(), Collections.emptyMap(), DiscoveryNodeRole.BUILT_IN_ROLES, Version.CURRENT);
    AtomicBoolean activeContextCleanUpConsumerInvocation = new AtomicBoolean();
    ThreadPool testThreadPool = null;
    try {
        testThreadPool = new TestThreadPool(AsynchronousSearchPlugin.OPEN_DISTRO_ASYNC_SEARCH_GENERIC_THREAD_POOL_NAME, executorBuilder);
        ClusterService mockClusterService = ClusterServiceUtils.createClusterService(testThreadPool, discoveryNode, clusterSettings);
        FakeClient fakeClient = new FakeClient(testThreadPool);
        AsynchronousSearchActiveStore asActiveStore = new AsynchronousSearchActiveStore(mockClusterService);
        AsynchronousSearchPersistenceService persistenceService = new AsynchronousSearchPersistenceService(fakeClient, mockClusterService, testThreadPool);
        AsynchronousSearchService asService = new AsynchronousSearchService(persistenceService, asActiveStore, fakeClient, mockClusterService, testThreadPool, new InternalAsynchronousSearchStats(), new NamedWriteableRegistry(emptyList()));
        AsynchronousSearchStateMachine asStateMachine = asService.getStateMachine();
        ClusterService clusterService = ClusterServiceUtils.createClusterService(testThreadPool, discoveryNode, clusterSettings);
        AsynchronousSearchPostProcessor postProcessor = new AsynchronousSearchPostProcessor(persistenceService, asActiveStore, asStateMachine, (context) -> activeContextCleanUpConsumerInvocation.compareAndSet(false, true), testThreadPool, clusterService);
        SubmitAsynchronousSearchRequest submitAsynchronousSearchRequest = new SubmitAsynchronousSearchRequest(new SearchRequest());
        submitAsynchronousSearchRequest.keepOnCompletion(true);
        submitAsynchronousSearchRequest.keepAlive(TimeValue.timeValueHours(1));
        AsynchronousSearchActiveContext context = (AsynchronousSearchActiveContext) asService.createAndStoreContext(submitAsynchronousSearchRequest, System.currentTimeMillis(), () -> InternalAggregationTestCase.emptyReduceContextBuilder(), null);
        asStateMachine.trigger(new SearchStartedEvent(context, new SearchTask(0, "n/a", "n/a", () -> "test", null, Collections.emptyMap())));
        SearchResponse mockSearchResponse = getMockSearchResponse();
        AsynchronousSearchResponse asResponse = postProcessor.processSearchResponse(mockSearchResponse, context.getContextId());
        assertNotNull(asResponse.getId());
        assertNull(asResponse.getError());
        assertEquals(AsynchronousSearchState.PERSISTING, asResponse.getState());
        AsynchronousSearchAssertions.assertSearchResponses(mockSearchResponse, asResponse.getSearchResponse());
        assertFalse(activeContextCleanUpConsumerInvocation.get());
        waitUntil(() -> fakeClient.persistenceCount == 1);
        assertEquals(1, fakeClient.persistenceCount);
    } finally {
        ThreadPool.terminate(testThreadPool, 200, TimeUnit.MILLISECONDS);
    }
}

Example 5 with AsynchronousSearchStateMachine

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

the class AsynchronousSearchStateMachineTests method testStateMachine.

public void testStateMachine() throws InterruptedException, BrokenBarrierException {
    DiscoveryNode discoveryNode = new DiscoveryNode("node", OpenSearchTestCase.buildNewFakeTransportAddress(), Collections.emptyMap(), DiscoveryNodeRole.BUILT_IN_ROLES, Version.CURRENT);
    TestThreadPool threadPool = null;
    try {
        threadPool = new TestThreadPool("test", executorBuilder);
        ClusterService mockClusterService = createClusterService(settings, threadPool, discoveryNode, clusterSettings);
        FakeClient fakeClient = new FakeClient(threadPool);
        AsynchronousSearchActiveStore asActiveStore = new AsynchronousSearchActiveStore(mockClusterService);
        AsynchronousSearchPersistenceService persistenceService = new AsynchronousSearchPersistenceService(fakeClient, mockClusterService, threadPool);
        CustomContextListener customContextListener = new CustomContextListener();
        AsynchronousSearchService asService = new AsynchronousSearchService(persistenceService, asActiveStore, fakeClient, mockClusterService, threadPool, customContextListener, new NamedWriteableRegistry(emptyList()));
        AsynchronousSearchProgressListener asProgressListener = mockAsynchronousSearchProgressListener(threadPool);
        AsynchronousSearchContextId asContextId = new AsynchronousSearchContextId(UUID.randomUUID().toString(), randomNonNegativeLong());
        TimeValue keepAlive = TimeValue.timeValueDays(randomInt(100));
        AsynchronousSearchActiveContext context = new AsynchronousSearchActiveContext(asContextId, discoveryNode.getId(), keepAlive, true, threadPool, threadPool::absoluteTimeInMillis, asProgressListener, null, () -> true);
        assertNull(context.getTask());
        assertEquals(context.getAsynchronousSearchState(), INIT);
        AsynchronousSearchStateMachine stateMachine = asService.getStateMachine();
        AtomicInteger numCompleted = new AtomicInteger();
        AtomicInteger numFailure = new AtomicInteger();
        doConcurrentStateMachineTrigger(stateMachine, new SearchStartedEvent(context, new AsynchronousSearchTask(randomNonNegativeLong(), "transport", SearchAction.NAME, TaskId.EMPTY_TASK_ID, emptyMap(), context, null, (a) -> {
        })), RUNNING, IllegalStateException.class, Optional.empty());
        boolean success = randomBoolean();
        assertNotNull(context.getTask());
        if (randomBoolean()) {
            // delete running context
            doConcurrentStateMachineTrigger(stateMachine, new SearchDeletedEvent(context), CLOSED, AsynchronousSearchStateMachineClosedException.class, Optional.empty());
        } else {
            if (success) {
                doConcurrentStateMachineTrigger(stateMachine, new SearchFailureEvent(context, new RuntimeException("test")), FAILED, IllegalStateException.class, Optional.empty());
                numFailure.getAndIncrement();
            } else {
                // success or failure
                doConcurrentStateMachineTrigger(stateMachine, new SearchSuccessfulEvent(context, getMockSearchResponse()), SUCCEEDED, IllegalStateException.class, Optional.empty());
                numCompleted.getAndIncrement();
            }
            doConcurrentStateMachineTrigger(stateMachine, new BeginPersistEvent(context), PERSISTING, IllegalStateException.class, Optional.of(AsynchronousSearchStateMachineClosedException.class));
            waitUntil(() -> context.getAsynchronousSearchState().equals(CLOSED), 1, TimeUnit.MINUTES);
            assertTrue(context.getAsynchronousSearchState().toString() + " numFailure : " + numFailure.get() + " numSuccess : " + numCompleted.get(), context.getAsynchronousSearchState().equals(CLOSED));
            assertEquals(1, customContextListener.getPersistedCount() + customContextListener.getPersistFailedCount());
        }
        assertEquals(numCompleted.get(), customContextListener.getCompletedCount());
        assertEquals(numFailure.get(), customContextListener.getFailedCount());
        assertEquals("success:" + success, 0, customContextListener.getRunningCount());
        assertEquals(1, customContextListener.getDeletedCount());
    } finally {
        ThreadPool.terminate(threadPool, 100, TimeUnit.MILLISECONDS);
    }
}