Examples with AsynchronousSearchActiveContext org.opensearch.search.asynchronous.context.active.AsynchronousSearchActiveContext used on opensource projects

Example 21 with AsynchronousSearchActiveContext

use of org.opensearch.search.asynchronous.context.active.AsynchronousSearchActiveContext in project asynchronous-search by opensearch-project.

the class AsynchronousSearchServiceTests method testFreeActiveContext.

public void testFreeActiveContext() throws InterruptedException {
    DiscoveryNode discoveryNode = new DiscoveryNode("node", OpenSearchTestCase.buildNewFakeTransportAddress(), emptyMap(), DiscoveryNodeRole.BUILT_IN_ROLES, Version.CURRENT);
    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()));
        TimeValue keepAlive = timeValueHours(9);
        boolean keepOnCompletion = randomBoolean();
        User user1 = TestClientUtils.randomUser();
        SearchRequest searchRequest = new SearchRequest();
        SubmitAsynchronousSearchRequest submitAsynchronousSearchRequest = new SubmitAsynchronousSearchRequest(searchRequest);
        submitAsynchronousSearchRequest.keepOnCompletion(keepOnCompletion);
        submitAsynchronousSearchRequest.keepAlive(keepAlive);
        AsynchronousSearchContext context = asService.createAndStoreContext(submitAsynchronousSearchRequest, System.currentTimeMillis(), () -> null, user1);
        assertTrue(context instanceof AsynchronousSearchActiveContext);
        AsynchronousSearchActiveContext asActiveContext = (AsynchronousSearchActiveContext) context;
        assertNull(asActiveContext.getTask());
        assertNull(asActiveContext.getAsynchronousSearchId());
        assertEquals(asActiveContext.getAsynchronousSearchState(), AsynchronousSearchState.INIT);
        assertEquals(asActiveContext.getUser(), user1);
        // bootstrap search
        AsynchronousSearchTask task = new AsynchronousSearchTask(randomNonNegativeLong(), "transport", SearchAction.NAME, TaskId.EMPTY_TASK_ID, emptyMap(), (AsynchronousSearchActiveContext) context, null, (c) -> {
        }) {

            @Override
            public boolean isCancelled() {
                return true;
            }
        };
        asService.bootstrapSearch(task, context.getContextId());
        assertEquals(asActiveContext.getTask(), task);
        assertEquals(asActiveContext.getStartTimeMillis(), task.getStartTime());
        assertEquals(asActiveContext.getExpirationTimeMillis(), task.getStartTime() + keepAlive.millis());
        assertEquals(asActiveContext.getAsynchronousSearchState(), AsynchronousSearchState.RUNNING);
        assertTrue(asService.freeActiveContext(asActiveContext));
        assertTrue(asService.getAllActiveContexts().isEmpty());
        CountDownLatch latch = new CountDownLatch(1);
        asService.freeContext(context.getAsynchronousSearchId(), context.getContextId(), user1, new LatchedActionListener<>(wrap(r -> fail(), e -> assertTrue(e instanceof ResourceNotFoundException)), latch));
        latch.await();
    } finally {
        ThreadPool.terminate(testThreadPool, 200, TimeUnit.MILLISECONDS);
    }
}

Example 22 with AsynchronousSearchActiveContext

use of org.opensearch.search.asynchronous.context.active.AsynchronousSearchActiveContext in project asynchronous-search by opensearch-project.

the class AsynchronousSearchServiceTests method testBootStrapOnClosedSearch.

public void testBootStrapOnClosedSearch() {
    DiscoveryNode discoveryNode = new DiscoveryNode("node", OpenSearchTestCase.buildNewFakeTransportAddress(), emptyMap(), DiscoveryNodeRole.BUILT_IN_ROLES, Version.CURRENT);
    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()));
        TimeValue keepAlive = timeValueHours(9);
        boolean keepOnCompletion = false;
        SearchRequest searchRequest = new SearchRequest();
        SubmitAsynchronousSearchRequest submitAsynchronousSearchRequest = new SubmitAsynchronousSearchRequest(searchRequest);
        submitAsynchronousSearchRequest.keepOnCompletion(keepOnCompletion);
        submitAsynchronousSearchRequest.keepAlive(keepAlive);
        AsynchronousSearchContext context = asService.createAndStoreContext(submitAsynchronousSearchRequest, System.currentTimeMillis(), () -> null, null);
        assertTrue(context instanceof AsynchronousSearchActiveContext);
        AsynchronousSearchActiveContext asActiveContext = (AsynchronousSearchActiveContext) context;
        assertNull(asActiveContext.getTask());
        assertNull(asActiveContext.getAsynchronousSearchId());
        assertEquals(asActiveContext.getAsynchronousSearchState(), AsynchronousSearchState.INIT);
        // close context
        asActiveContext.close();
        // bootstrap search
        AsynchronousSearchTask task = new AsynchronousSearchTask(randomNonNegativeLong(), "transport", SearchAction.NAME, TaskId.EMPTY_TASK_ID, emptyMap(), (AsynchronousSearchActiveContext) context, null, (c) -> {
        });
        asActiveContext.setState(AsynchronousSearchState.CLOSED);
        expectThrows(IllegalStateException.class, () -> asService.bootstrapSearch(task, context.getContextId()));
    } finally {
        ThreadPool.terminate(testThreadPool, 200, TimeUnit.MILLISECONDS);
    }
}

Example 23 with AsynchronousSearchActiveContext

use of org.opensearch.search.asynchronous.context.active.AsynchronousSearchActiveContext in project asynchronous-search by opensearch-project.

the class AsynchronousSearchServiceUpdateContextTests method testUpdateClosedContext.

public void testUpdateClosedContext() throws InterruptedException {
    DiscoveryNode discoveryNode = new DiscoveryNode("node", OpenSearchTestCase.buildNewFakeTransportAddress(), emptyMap(), DiscoveryNodeRole.BUILT_IN_ROLES, Version.CURRENT);
    ThreadPool testThreadPool = null;
    try {
        testThreadPool = new TestThreadPool(OPEN_DISTRO_ASYNC_SEARCH_GENERIC_THREAD_POOL_NAME, executorBuilder);
        ClusterService mockClusterService = getClusterService(discoveryNode, testThreadPool);
        MockClient mockClient = new MockClient(testThreadPool);
        AsynchronousSearchActiveStore asActiveStore = new AsynchronousSearchActiveStore(mockClusterService);
        AsynchronousSearchPersistenceService persistenceService = new AsynchronousSearchPersistenceService(mockClient, mockClusterService, testThreadPool);
        AsynchronousSearchService asService = new AsynchronousSearchService(persistenceService, asActiveStore, mockClient, mockClusterService, testThreadPool, new InternalAsynchronousSearchStats(), new NamedWriteableRegistry(emptyList()));
        TimeValue keepAlive = timeValueHours(9);
        boolean keepOnCompletion = randomBoolean();
        User user1 = randomUser();
        SearchRequest searchRequest = new SearchRequest();
        SubmitAsynchronousSearchRequest submitAsynchronousSearchRequest = new SubmitAsynchronousSearchRequest(searchRequest);
        submitAsynchronousSearchRequest.keepOnCompletion(keepOnCompletion);
        submitAsynchronousSearchRequest.keepAlive(keepAlive);
        AsynchronousSearchContext context = asService.createAndStoreContext(submitAsynchronousSearchRequest, System.currentTimeMillis(), () -> null, user1);
        assertTrue(context instanceof AsynchronousSearchActiveContext);
        AsynchronousSearchActiveContext activeContext = (AsynchronousSearchActiveContext) context;
        assertNull(activeContext.getTask());
        assertNull(activeContext.getAsynchronousSearchId());
        assertEquals(activeContext.getAsynchronousSearchState(), INIT);
        assertEquals(activeContext.getUser(), user1);
        // bootstrap search
        AsynchronousSearchTask task = new AsynchronousSearchTask(randomNonNegativeLong(), "transport", SearchAction.NAME, TaskId.EMPTY_TASK_ID, emptyMap(), (AsynchronousSearchActiveContext) context, null, (c) -> {
        });
        asService.bootstrapSearch(task, context.getContextId());
        assertEquals(activeContext.getTask(), task);
        assertEquals(activeContext.getStartTimeMillis(), task.getStartTime());
        assertEquals(activeContext.getExpirationTimeMillis(), task.getStartTime() + keepAlive.millis());
        assertEquals(activeContext.getAsynchronousSearchState(), RUNNING);
        context.setState(CLOSED);
        ((AsynchronousSearchActiveContext) context).close();
        CountDownLatch latch = new CountDownLatch(1);
        asService.updateKeepAliveAndGetContext(context.getAsynchronousSearchId(), timeValueHours(9), context.getContextId(), user1, new LatchedActionListener<>(wrap(r -> {
            if (keepOnCompletion) {
                assertTrue(r instanceof AsynchronousSearchPersistenceContext);
            } else {
                fail("expected resource not found exception, got result.");
            }
        }, e -> {
            if (keepOnCompletion) {
                fail("expected resource not found exception, got result");
            } else {
                assertTrue(e instanceof ResourceNotFoundException);
            }
        }), latch));
        latch.await();
        mockClusterService.stop();
    } finally {
        ThreadPool.terminate(testThreadPool, 30, TimeUnit.SECONDS);
    }
}

Example 24 with AsynchronousSearchActiveContext

use of org.opensearch.search.asynchronous.context.active.AsynchronousSearchActiveContext in project asynchronous-search by opensearch-project.

the class AsynchronousSearchPostProcessorTests method testProcessSearchResponseOnClosedContext.

public void testProcessSearchResponseOnClosedContext() 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);
        asStateMachine.trigger(new SearchStartedEvent(context, new SearchTask(0, "n/a", "n/a", () -> "test", null, Collections.emptyMap())));
        asStateMachine.trigger(new SearchDeletedEvent(context));
        SearchResponse mockSearchResponse = getMockSearchResponse();
        AsynchronousSearchResponse asResponse = postProcessor.processSearchResponse(mockSearchResponse, context.getContextId());
        assertNull(asResponse.getId());
        assertNull(asResponse.getError());
        assertEquals(-1L, asResponse.getExpirationTimeMillis());
        assertEquals(AsynchronousSearchState.SUCCEEDED, asResponse.getState());
        assertEquals(-1L, asResponse.getStartTimeMillis());
        AsynchronousSearchAssertions.assertSearchResponses(mockSearchResponse, asResponse.getSearchResponse());
        assertFalse(activeContextCleanUpConsumerInvocation.get());
        assertEquals(0, fakeClient.persistenceCount);
    } finally {
        ThreadPool.terminate(testThreadPool, 200, TimeUnit.MILLISECONDS);
    }
}

Example 25 with AsynchronousSearchActiveContext

use of org.opensearch.search.asynchronous.context.active.AsynchronousSearchActiveContext in project asynchronous-search by opensearch-project.

the class AsynchronousSearchPostProcessorTests method testProcessSearchResponsePersisted.

public void testProcessSearchResponsePersisted() 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.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());
        waitUntil(() -> context.getAsynchronousSearchState() == AsynchronousSearchState.CLOSED);
        assertEquals(AsynchronousSearchState.CLOSED, context.getAsynchronousSearchState());
        AsynchronousSearchAssertions.assertSearchResponses(mockSearchResponse, asResponse.getSearchResponse());
        assertFalse(activeContextCleanUpConsumerInvocation.get());
        assertEquals(1, fakeClient.persistenceCount);
    } catch (InterruptedException e) {
        fail("Test interrupted " + e.getMessage());
    } finally {
        ThreadPool.terminate(testThreadPool, 200, TimeUnit.MILLISECONDS);
    }
}