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

Example 16 with AsynchronousSearchActiveContext

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

the class AsynchronousSearchServiceUpdateContextTests method testUpdateActiveContextInvalidUser.

public void testUpdateActiveContextInvalidUser() 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 = true;
        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);
        CountDownLatch latch = new CountDownLatch(1);
        User differenteUser = randomUser();
        asService.updateKeepAliveAndGetContext(context.getAsynchronousSearchId(), timeValueHours(9), context.getContextId(), differenteUser, new LatchedActionListener<>(wrap(r -> fail("expected security exception Users must be different, " + "actual user " + user1 + " random user " + differenteUser), e -> {
            assertTrue(e instanceof ResourceNotFoundException);
        }), latch));
        latch.await();
        mockClusterService.stop();
    } finally {
        ThreadPool.terminate(testThreadPool, 30, TimeUnit.SECONDS);
    }
}

Example 17 with AsynchronousSearchActiveContext

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

the class AsynchronousSearchServiceUpdateContextTests method testUpdateRunningContextValidUser.

public void testUpdateRunningContextValidUser() 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 = true;
        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);
        CountDownLatch latch = new CountDownLatch(1);
        Long oldExpiration = context.getExpirationTimeMillis();
        asService.updateKeepAliveAndGetContext(context.getAsynchronousSearchId(), timeValueHours(10), context.getContextId(), user1, new LatchedActionListener<>(wrap(r -> assertThat(context.getExpirationTimeMillis(), greaterThan(oldExpiration)), e -> {
            fail("Expected successful update but got failure " + e.getMessage());
        }), latch));
        latch.await();
        mockClusterService.stop();
    } finally {
        ThreadPool.terminate(testThreadPool, 30, TimeUnit.SECONDS);
    }
}

Example 18 with AsynchronousSearchActiveContext

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

the class AsynchronousSearchServiceFreeContextTests method testFreeActiveContextKeepOnCompletionFalse.

public void testFreeActiveContextKeepOnCompletionFalse() 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 mockStore = mock(AsynchronousSearchActiveStore.class);
        AsynchronousSearchPersistenceService persistenceService = new AsynchronousSearchPersistenceService(mockClient, mockClusterService, testThreadPool);
        AsynchronousSearchService asService = new AsynchronousSearchService(persistenceService, mockStore, mockClient, mockClusterService, testThreadPool, new InternalAsynchronousSearchStats(), new NamedWriteableRegistry(emptyList()));
        TimeValue keepAlive = timeValueDays(9);
        boolean keepOnCompletion = true;
        User user1 = randomBoolean() ? randomUser() : null;
        SearchRequest searchRequest = new SearchRequest();
        SubmitAsynchronousSearchRequest submitAsynchronousSearchRequest = new SubmitAsynchronousSearchRequest(searchRequest);
        submitAsynchronousSearchRequest.keepOnCompletion(keepOnCompletion);
        submitAsynchronousSearchRequest.keepAlive(keepAlive);
        AsynchronousSearchProgressListener asProgressListener = mockAsynchronousSearchProgressListener(testThreadPool);
        AsynchronousSearchContextId asContextId = new AsynchronousSearchContextId(UUID.randomUUID().toString(), randomNonNegativeLong());
        AsynchronousSearchActiveContext asActiveContext = new AsynchronousSearchActiveContext(asContextId, discoveryNode.getId(), keepAlive, false, testThreadPool, testThreadPool::absoluteTimeInMillis, asProgressListener, user1, () -> true);
        // bootstrap search
        AsynchronousSearchTask task = new AsynchronousSearchTask(randomNonNegativeLong(), "transport", SearchAction.NAME, TaskId.EMPTY_TASK_ID, emptyMap(), asActiveContext, null, (c) -> {
        }) {

            @Override
            public boolean isCancelled() {
                return true;
            }
        };
        asActiveContext.setTask(task);
        asActiveContext.setState(AsynchronousSearchState.RUNNING);
        when(mockStore.getContext(any())).thenReturn(Optional.of(asActiveContext));
        persisted = false;
        // task cancellation fails
        // user
        User user = randomBoolean() ? null : user1;
        CountDownLatch latch = new CountDownLatch(1);
        asService.freeContext(asActiveContext.getAsynchronousSearchId(), asActiveContext.getContextId(), user, new LatchedActionListener<>(ActionListener.wrap(Assert::assertTrue, e -> {
            fail("active context should have been deleted");
        }), latch));
        latch.await();
        CountDownLatch latch1 = new CountDownLatch(1);
        asService.freeContext(asActiveContext.getAsynchronousSearchId(), asActiveContext.getContextId(), user, new LatchedActionListener<>(ActionListener.wrap(r -> {
            fail("Expected resource_not_found_exception");
        }, e -> {
            assertTrue(e.getClass().getName(), e instanceof ResourceNotFoundException);
        }), latch1));
        latch1.await();
        mockClusterService.stop();
    } finally {
        ThreadPool.terminate(testThreadPool, 30, TimeUnit.SECONDS);
    }
}

Example 19 with AsynchronousSearchActiveContext

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

the class AsynchronousSearchServiceTests method testUpdateExpirationTimesOutBlockedOnPersistence.

public void testUpdateExpirationTimesOutBlockedOnPersistence() 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 = true;
        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);
        // 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(asActiveContext.getTask(), task);
        assertEquals(asActiveContext.getStartTimeMillis(), task.getStartTime());
        long originalExpirationTimeMillis = asActiveContext.getExpirationTimeMillis();
        assertEquals(originalExpirationTimeMillis, task.getStartTime() + keepAlive.millis());
        assertEquals(asActiveContext.getAsynchronousSearchState(), AsynchronousSearchState.RUNNING);
        blockPersistence = true;
        context.getAsynchronousSearchProgressListener().onResponse(getMockSearchResponse());
        CountDownLatch updateLatch = new CountDownLatch(1);
        TimeValue newKeepAlive = timeValueHours(10);
        fakeClient.awaitBlock();
        asService.updateKeepAliveAndGetContext(asActiveContext.getAsynchronousSearchId(), newKeepAlive, asActiveContext.getContextId(), null, new LatchedActionListener<>(wrap(r -> fail("expected update req to timeout"), e -> assertTrue("expected timeout got " + e.getClass(), e instanceof OpenSearchTimeoutException)), updateLatch));
        updateLatch.await();
        fakeClient.releaseBlock();
    } finally {
        ThreadPool.terminate(testThreadPool, 200, TimeUnit.MILLISECONDS);
    }
}

Example 20 with AsynchronousSearchActiveContext

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

the class AsynchronousSearchServiceTests method testUpdateExpirationOnRunningSearch.

public void testUpdateExpirationOnRunningSearch() 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 = 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);
        // 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(asActiveContext.getTask(), task);
        assertEquals(asActiveContext.getStartTimeMillis(), task.getStartTime());
        long originalExpirationTimeMillis = asActiveContext.getExpirationTimeMillis();
        assertEquals(originalExpirationTimeMillis, task.getStartTime() + keepAlive.millis());
        assertEquals(asActiveContext.getAsynchronousSearchState(), AsynchronousSearchState.RUNNING);
        CountDownLatch findContextLatch = new CountDownLatch(1);
        asService.findContext(asActiveContext.getAsynchronousSearchId(), asActiveContext.getContextId(), null, new LatchedActionListener<>(wrap(r -> {
            assertTrue(r instanceof AsynchronousSearchActiveContext);
            assertEquals(r, context);
        }, e -> fail("Find context shouldn't have failed")), findContextLatch));
        findContextLatch.await();
        CountDownLatch updateLatch = new CountDownLatch(1);
        TimeValue newKeepAlive = timeValueHours(10);
        asService.updateKeepAliveAndGetContext(asActiveContext.getAsynchronousSearchId(), newKeepAlive, asActiveContext.getContextId(), null, new LatchedActionListener<>(wrap(r -> {
            assertTrue(r instanceof AsynchronousSearchActiveContext);
            assertThat(r.getExpirationTimeMillis(), greaterThan(originalExpirationTimeMillis));
        }, e -> fail()), updateLatch));
        updateLatch.await();
    } finally {
        ThreadPool.terminate(testThreadPool, 200, TimeUnit.MILLISECONDS);
    }
}