Search in sources :

Example 1 with SearchTask

use of org.opensearch.action.search.SearchTask in project asynchronous-search by opensearch-project.

the class AsynchronousSearchRejectionIT method testSearchFailures.

public void testSearchFailures() throws Exception {
    for (int i = 0; i < 10; i++) {
        client().prepareIndex("test").setId(Integer.toString(i)).setSource("field", "1").get();
    }
    int numberOfAsyncOps = randomIntBetween(100, 200);
    final CountDownLatch latch = new CountDownLatch(numberOfAsyncOps);
    final CopyOnWriteArrayList<Object> responses = new CopyOnWriteArrayList<>();
    TestThreadPool threadPool = null;
    try {
        threadPool = new TestThreadPool(AsynchronousSearchProgressListenerIT.class.getName());
        for (int i = 0; i < numberOfAsyncOps; i++) {
            SearchRequest request = client().prepareSearch("test").setSearchType(SearchType.QUERY_THEN_FETCH).setQuery(QueryBuilders.matchQuery("field", "1")).request();
            AtomicReference<SearchResponse> responseRef = new AtomicReference<>();
            AtomicInteger reduceContextInvocation = new AtomicInteger();
            AsynchronousSearchProgressListener listener;
            SearchService service = internalCluster().getInstance(SearchService.class);
            InternalAggregation.ReduceContextBuilder reduceContextBuilder = service.aggReduceContextBuilder(request);
            AtomicReference<Exception> exceptionRef = new AtomicReference<>();
            Function<SearchResponse, AsynchronousSearchResponse> responseFunction = (r) -> null;
            Function<Exception, AsynchronousSearchResponse> failureFunction = (e) -> null;
            listener = new AsynchronousSearchProgressListener(threadPool.relativeTimeInMillis(), responseFunction, failureFunction, threadPool.generic(), threadPool::relativeTimeInMillis, () -> reduceContextBuilder) {

                @Override
                public void onResponse(SearchResponse searchResponse) {
                    assertTrue(responseRef.compareAndSet(null, searchResponse));
                    AsynchronousSearchAssertions.assertSearchResponses(searchResponse, this.partialResponse());
                    latch.countDown();
                }

                @Override
                public void onFailure(Exception exception) {
                    assertTrue(exceptionRef.compareAndSet(null, exception));
                    latch.countDown();
                }
            };
            client().execute(SearchAction.INSTANCE, new SearchRequest(request) {

                @Override
                public SearchTask createTask(long id, String type, String action, TaskId parentTaskId, Map<String, String> headers) {
                    SearchTask task = super.createTask(id, type, action, parentTaskId, headers);
                    task.setProgressListener(listener);
                    return task;
                }
            }, listener);
        }
        latch.await();
    } finally {
        ThreadPool.terminate(threadPool, 100, TimeUnit.MILLISECONDS);
    }
}
Also used : AsynchronousSearchResponse(org.opensearch.search.asynchronous.response.AsynchronousSearchResponse) SearchAction(org.opensearch.action.search.SearchAction) OpenSearchRejectedExecutionException(org.opensearch.common.util.concurrent.OpenSearchRejectedExecutionException) AsynchronousSearchAssertions(org.opensearch.search.asynchronous.utils.AsynchronousSearchAssertions) SubmitAsynchronousSearchAction(org.opensearch.search.asynchronous.action.SubmitAsynchronousSearchAction) AsynchronousSearchIntegTestCase(org.opensearch.search.asynchronous.commons.AsynchronousSearchIntegTestCase) ThreadPool(org.opensearch.threadpool.ThreadPool) AsynchronousSearchProgressListener(org.opensearch.search.asynchronous.listener.AsynchronousSearchProgressListener) TestThreadPool(org.opensearch.threadpool.TestThreadPool) AsynchronousSearchResponse(org.opensearch.search.asynchronous.response.AsynchronousSearchResponse) ResourceNotFoundException(org.opensearch.ResourceNotFoundException) AtomicReference(java.util.concurrent.atomic.AtomicReference) Function(java.util.function.Function) AcknowledgedResponse(org.opensearch.search.asynchronous.response.AcknowledgedResponse) LatchedActionListener(org.opensearch.action.LatchedActionListener) AtomicInteger(java.util.concurrent.atomic.AtomicInteger) Locale(java.util.Locale) Map(java.util.Map) SearchRequest(org.opensearch.action.search.SearchRequest) DeleteAsynchronousSearchAction(org.opensearch.search.asynchronous.action.DeleteAsynchronousSearchAction) ActionListener(org.opensearch.action.ActionListener) SearchResponse(org.opensearch.action.search.SearchResponse) AsynchronousSearchProgressListenerIT(org.opensearch.search.asynchronous.listener.AsynchronousSearchProgressListenerIT) SearchService(org.opensearch.search.SearchService) QueryBuilders(org.opensearch.index.query.QueryBuilders) TimeValue(org.opensearch.common.unit.TimeValue) SearchTask(org.opensearch.action.search.SearchTask) ExceptionsHelper(org.opensearch.ExceptionsHelper) SearchType(org.opensearch.action.search.SearchType) TaskId(org.opensearch.tasks.TaskId) Settings(org.opensearch.common.settings.Settings) DeleteAsynchronousSearchRequest(org.opensearch.search.asynchronous.request.DeleteAsynchronousSearchRequest) Matchers.startsWith(org.hamcrest.Matchers.startsWith) AsynchronousSearchService(org.opensearch.search.asynchronous.service.AsynchronousSearchService) TestLogging(org.opensearch.test.junit.annotations.TestLogging) TimeUnit(java.util.concurrent.TimeUnit) CountDownLatch(java.util.concurrent.CountDownLatch) SubmitAsynchronousSearchRequest(org.opensearch.search.asynchronous.request.SubmitAsynchronousSearchRequest) ShardSearchFailure(org.opensearch.action.search.ShardSearchFailure) Matchers.equalTo(org.hamcrest.Matchers.equalTo) OpenSearchTimeoutException(org.opensearch.OpenSearchTimeoutException) SearchPhaseExecutionException(org.opensearch.action.search.SearchPhaseExecutionException) InternalAggregation(org.opensearch.search.aggregations.InternalAggregation) OpenSearchIntegTestCase(org.opensearch.test.OpenSearchIntegTestCase) CopyOnWriteArrayList(java.util.concurrent.CopyOnWriteArrayList) SearchRequest(org.opensearch.action.search.SearchRequest) DeleteAsynchronousSearchRequest(org.opensearch.search.asynchronous.request.DeleteAsynchronousSearchRequest) SubmitAsynchronousSearchRequest(org.opensearch.search.asynchronous.request.SubmitAsynchronousSearchRequest) TaskId(org.opensearch.tasks.TaskId) TestThreadPool(org.opensearch.threadpool.TestThreadPool) SearchTask(org.opensearch.action.search.SearchTask) SearchService(org.opensearch.search.SearchService) AsynchronousSearchService(org.opensearch.search.asynchronous.service.AsynchronousSearchService) AtomicReference(java.util.concurrent.atomic.AtomicReference) CountDownLatch(java.util.concurrent.CountDownLatch) OpenSearchRejectedExecutionException(org.opensearch.common.util.concurrent.OpenSearchRejectedExecutionException) ResourceNotFoundException(org.opensearch.ResourceNotFoundException) OpenSearchTimeoutException(org.opensearch.OpenSearchTimeoutException) SearchPhaseExecutionException(org.opensearch.action.search.SearchPhaseExecutionException) AsynchronousSearchResponse(org.opensearch.search.asynchronous.response.AsynchronousSearchResponse) SearchResponse(org.opensearch.action.search.SearchResponse) AsynchronousSearchProgressListener(org.opensearch.search.asynchronous.listener.AsynchronousSearchProgressListener) InternalAggregation(org.opensearch.search.aggregations.InternalAggregation) AtomicInteger(java.util.concurrent.atomic.AtomicInteger) CopyOnWriteArrayList(java.util.concurrent.CopyOnWriteArrayList)

Example 2 with SearchTask

use of org.opensearch.action.search.SearchTask in project asynchronous-search by opensearch-project.

the class AsynchronousSearchPartialResponseIT method testCase.

private void testCase(Client client, SearchRequest request) throws Exception {
    AtomicReference<SearchResponse> responseRef = new AtomicReference<>();
    AtomicInteger reduceContextInvocation = new AtomicInteger();
    TestThreadPool threadPool = null;
    AsynchronousSearchProgressListener listener;
    try {
        threadPool = new TestThreadPool(AsynchronousSearchProgressListenerIT.class.getName());
        SearchService service = internalCluster().getInstance(SearchService.class);
        InternalAggregation.ReduceContextBuilder reduceContextBuilder = service.aggReduceContextBuilder(request);
        AtomicReference<Exception> exceptionRef = new AtomicReference<>();
        CountDownLatch latch = new CountDownLatch(1);
        Function<SearchResponse, AsynchronousSearchResponse> responseFunction = (r) -> null;
        Function<Exception, AsynchronousSearchResponse> failureFunction = (e) -> null;
        listener = new AsynchronousSearchProgressListener(threadPool.relativeTimeInMillis(), responseFunction, failureFunction, threadPool.generic(), threadPool::relativeTimeInMillis, () -> reduceContextBuilder) {

            @Override
            public void onResponse(SearchResponse searchResponse) {
                assertTrue(responseRef.compareAndSet(null, searchResponse));
                AsynchronousSearchAssertions.assertSearchResponses(responseRef.get(), this.partialResponse());
                latch.countDown();
            }

            @Override
            protected void onPartialReduce(List<SearchShard> shards, TotalHits totalHits, InternalAggregations aggs, int reducePhase) {
                super.onPartialReduce(shards, totalHits, aggs, reducePhase);
                Terms terms = this.partialResponse().getAggregations().get("keys");
                List<? extends Terms.Bucket> buckets = terms.getBuckets();
                assertThat(buckets.size(), lessThanOrEqualTo(aggregationSize));
                reduceContextInvocation.incrementAndGet();
            }

            @Override
            protected void onFinalReduce(List<SearchShard> shards, TotalHits totalHits, InternalAggregations aggs, int reducePhase) {
                super.onFinalReduce(shards, totalHits, aggs, reducePhase);
                Terms terms = this.partialResponse().getAggregations().get("keys");
                List<? extends Terms.Bucket> buckets = terms.getBuckets();
                assertThat(buckets.size(), equalTo(aggregationSize));
            }

            @Override
            public void onFailure(Exception exception) {
                assertTrue(exceptionRef.compareAndSet(null, exception));
                latch.countDown();
            }
        };
        client.execute(SearchAction.INSTANCE, new SearchRequest(request) {

            @Override
            public SearchTask createTask(long id, String type, String action, TaskId parentTaskId, Map<String, String> headers) {
                SearchTask task = super.createTask(id, type, action, parentTaskId, headers);
                task.setProgressListener(listener);
                return task;
            }
        }, listener);
        latch.await();
        Terms terms = responseRef.get().getAggregations().get("keys");
        assertThat(reduceContextInvocation.get(), equalTo(responseRef.get().getNumReducePhases() - 1));
        List<? extends Terms.Bucket> buckets = terms.getBuckets();
        assertThat(buckets.size(), equalTo(aggregationSize));
    } finally {
        ThreadPool.terminate(threadPool, 100, TimeUnit.MILLISECONDS);
    }
}
Also used : AsynchronousSearchResponse(org.opensearch.search.asynchronous.response.AsynchronousSearchResponse) IndexRequestBuilder(org.opensearch.action.index.IndexRequestBuilder) SearchAction(org.opensearch.action.search.SearchAction) AsynchronousSearchAssertions(org.opensearch.search.asynchronous.utils.AsynchronousSearchAssertions) XContentFactory.jsonBuilder(org.opensearch.common.xcontent.XContentFactory.jsonBuilder) BucketOrder(org.opensearch.search.aggregations.BucketOrder) Arrays(java.util.Arrays) ThreadPool(org.opensearch.threadpool.ThreadPool) TestThreadPool(org.opensearch.threadpool.TestThreadPool) OpenSearchAssertions.assertSearchResponse(org.opensearch.test.hamcrest.OpenSearchAssertions.assertSearchResponse) AsynchronousSearchResponse(org.opensearch.search.asynchronous.response.AsynchronousSearchResponse) AtomicReference(java.util.concurrent.atomic.AtomicReference) Function(java.util.function.Function) AggregationBuilders.terms(org.opensearch.search.aggregations.AggregationBuilders.terms) ArrayList(java.util.ArrayList) AtomicInteger(java.util.concurrent.atomic.AtomicInteger) QueryBuilders.matchAllQuery(org.opensearch.index.query.QueryBuilders.matchAllQuery) InternalAggregations(org.opensearch.search.aggregations.InternalAggregations) Map(java.util.Map) SearchRequest(org.opensearch.action.search.SearchRequest) SearchResponse(org.opensearch.action.search.SearchResponse) SearchService(org.opensearch.search.SearchService) OpenSearchAssertions.assertAcked(org.opensearch.test.hamcrest.OpenSearchAssertions.assertAcked) Client(org.opensearch.client.Client) SearchTask(org.opensearch.action.search.SearchTask) Matchers.lessThanOrEqualTo(org.hamcrest.Matchers.lessThanOrEqualTo) TaskId(org.opensearch.tasks.TaskId) Terms(org.opensearch.search.aggregations.bucket.terms.Terms) SearchShard(org.opensearch.action.search.SearchShard) TotalHits(org.apache.lucene.search.TotalHits) TimeUnit(java.util.concurrent.TimeUnit) CountDownLatch(java.util.concurrent.CountDownLatch) List(java.util.List) Matchers.equalTo(org.hamcrest.Matchers.equalTo) Matchers.is(org.hamcrest.Matchers.is) Aggregator(org.opensearch.search.aggregations.Aggregator) InternalAggregation(org.opensearch.search.aggregations.InternalAggregation) OpenSearchIntegTestCase(org.opensearch.test.OpenSearchIntegTestCase) TotalHits(org.apache.lucene.search.TotalHits) SearchRequest(org.opensearch.action.search.SearchRequest) TaskId(org.opensearch.tasks.TaskId) TestThreadPool(org.opensearch.threadpool.TestThreadPool) SearchTask(org.opensearch.action.search.SearchTask) SearchService(org.opensearch.search.SearchService) ArrayList(java.util.ArrayList) List(java.util.List) SearchShard(org.opensearch.action.search.SearchShard) Terms(org.opensearch.search.aggregations.bucket.terms.Terms) AtomicReference(java.util.concurrent.atomic.AtomicReference) CountDownLatch(java.util.concurrent.CountDownLatch) OpenSearchAssertions.assertSearchResponse(org.opensearch.test.hamcrest.OpenSearchAssertions.assertSearchResponse) AsynchronousSearchResponse(org.opensearch.search.asynchronous.response.AsynchronousSearchResponse) SearchResponse(org.opensearch.action.search.SearchResponse) InternalAggregation(org.opensearch.search.aggregations.InternalAggregation) InternalAggregations(org.opensearch.search.aggregations.InternalAggregations) AtomicInteger(java.util.concurrent.atomic.AtomicInteger)

Example 3 with SearchTask

use of org.opensearch.action.search.SearchTask 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);
    }
}
Also used : NamedWriteableRegistry(org.opensearch.common.io.stream.NamedWriteableRegistry) AsynchronousSearchResponse(org.opensearch.search.asynchronous.response.AsynchronousSearchResponse) SubmitAsynchronousSearchRequest(org.opensearch.search.asynchronous.request.SubmitAsynchronousSearchRequest) SearchRequest(org.opensearch.action.search.SearchRequest) DiscoveryNode(org.opensearch.cluster.node.DiscoveryNode) SubmitAsynchronousSearchRequest(org.opensearch.search.asynchronous.request.SubmitAsynchronousSearchRequest) TestThreadPool(org.opensearch.threadpool.TestThreadPool) ThreadPool(org.opensearch.threadpool.ThreadPool) InternalAsynchronousSearchStats(org.opensearch.search.asynchronous.stats.InternalAsynchronousSearchStats) TestThreadPool(org.opensearch.threadpool.TestThreadPool) AsynchronousSearchActiveContext(org.opensearch.search.asynchronous.context.active.AsynchronousSearchActiveContext) AsynchronousSearchActiveStore(org.opensearch.search.asynchronous.context.active.AsynchronousSearchActiveStore) AsynchronousSearchResponse(org.opensearch.search.asynchronous.response.AsynchronousSearchResponse) SearchResponse(org.opensearch.action.search.SearchResponse) InternalSearchResponse(org.opensearch.search.internal.InternalSearchResponse) AtomicBoolean(java.util.concurrent.atomic.AtomicBoolean) ClusterService(org.opensearch.cluster.service.ClusterService) SearchTask(org.opensearch.action.search.SearchTask) SearchStartedEvent(org.opensearch.search.asynchronous.context.state.event.SearchStartedEvent) AsynchronousSearchStateMachine(org.opensearch.search.asynchronous.context.state.AsynchronousSearchStateMachine) AsynchronousSearchPostProcessor(org.opensearch.search.asynchronous.processor.AsynchronousSearchPostProcessor)

Example 4 with SearchTask

use of org.opensearch.action.search.SearchTask 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);
    }
}
Also used : NamedWriteableRegistry(org.opensearch.common.io.stream.NamedWriteableRegistry) AsynchronousSearchResponse(org.opensearch.search.asynchronous.response.AsynchronousSearchResponse) SubmitAsynchronousSearchRequest(org.opensearch.search.asynchronous.request.SubmitAsynchronousSearchRequest) SearchRequest(org.opensearch.action.search.SearchRequest) DiscoveryNode(org.opensearch.cluster.node.DiscoveryNode) SubmitAsynchronousSearchRequest(org.opensearch.search.asynchronous.request.SubmitAsynchronousSearchRequest) SearchPhaseExecutionException(org.opensearch.action.search.SearchPhaseExecutionException) TestThreadPool(org.opensearch.threadpool.TestThreadPool) ThreadPool(org.opensearch.threadpool.ThreadPool) InternalAsynchronousSearchStats(org.opensearch.search.asynchronous.stats.InternalAsynchronousSearchStats) TestThreadPool(org.opensearch.threadpool.TestThreadPool) AsynchronousSearchActiveContext(org.opensearch.search.asynchronous.context.active.AsynchronousSearchActiveContext) AsynchronousSearchActiveStore(org.opensearch.search.asynchronous.context.active.AsynchronousSearchActiveStore) AtomicBoolean(java.util.concurrent.atomic.AtomicBoolean) ClusterService(org.opensearch.cluster.service.ClusterService) SearchTask(org.opensearch.action.search.SearchTask) SearchStartedEvent(org.opensearch.search.asynchronous.context.state.event.SearchStartedEvent) AsynchronousSearchStateMachine(org.opensearch.search.asynchronous.context.state.AsynchronousSearchStateMachine) AsynchronousSearchPostProcessor(org.opensearch.search.asynchronous.processor.AsynchronousSearchPostProcessor) ShardSearchFailure(org.opensearch.action.search.ShardSearchFailure) SearchDeletedEvent(org.opensearch.search.asynchronous.context.state.event.SearchDeletedEvent)

Example 5 with SearchTask

use of org.opensearch.action.search.SearchTask 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);
    }
}
Also used : NamedWriteableRegistry(org.opensearch.common.io.stream.NamedWriteableRegistry) AsynchronousSearchResponse(org.opensearch.search.asynchronous.response.AsynchronousSearchResponse) SubmitAsynchronousSearchRequest(org.opensearch.search.asynchronous.request.SubmitAsynchronousSearchRequest) SearchRequest(org.opensearch.action.search.SearchRequest) DiscoveryNode(org.opensearch.cluster.node.DiscoveryNode) SubmitAsynchronousSearchRequest(org.opensearch.search.asynchronous.request.SubmitAsynchronousSearchRequest) TestThreadPool(org.opensearch.threadpool.TestThreadPool) ThreadPool(org.opensearch.threadpool.ThreadPool) InternalAsynchronousSearchStats(org.opensearch.search.asynchronous.stats.InternalAsynchronousSearchStats) TestThreadPool(org.opensearch.threadpool.TestThreadPool) AsynchronousSearchActiveContext(org.opensearch.search.asynchronous.context.active.AsynchronousSearchActiveContext) AsynchronousSearchActiveStore(org.opensearch.search.asynchronous.context.active.AsynchronousSearchActiveStore) AsynchronousSearchResponse(org.opensearch.search.asynchronous.response.AsynchronousSearchResponse) SearchResponse(org.opensearch.action.search.SearchResponse) InternalSearchResponse(org.opensearch.search.internal.InternalSearchResponse) AtomicBoolean(java.util.concurrent.atomic.AtomicBoolean) ClusterService(org.opensearch.cluster.service.ClusterService) SearchTask(org.opensearch.action.search.SearchTask) SearchStartedEvent(org.opensearch.search.asynchronous.context.state.event.SearchStartedEvent) AsynchronousSearchStateMachine(org.opensearch.search.asynchronous.context.state.AsynchronousSearchStateMachine) AsynchronousSearchPostProcessor(org.opensearch.search.asynchronous.processor.AsynchronousSearchPostProcessor)

Aggregations

SearchRequest (org.opensearch.action.search.SearchRequest)9 SearchTask (org.opensearch.action.search.SearchTask)9 AsynchronousSearchResponse (org.opensearch.search.asynchronous.response.AsynchronousSearchResponse)9 TestThreadPool (org.opensearch.threadpool.TestThreadPool)9 ThreadPool (org.opensearch.threadpool.ThreadPool)9 SearchResponse (org.opensearch.action.search.SearchResponse)8 AtomicBoolean (java.util.concurrent.atomic.AtomicBoolean)6 SubmitAsynchronousSearchRequest (org.opensearch.search.asynchronous.request.SubmitAsynchronousSearchRequest)6 DiscoveryNode (org.opensearch.cluster.node.DiscoveryNode)5 ClusterService (org.opensearch.cluster.service.ClusterService)5 NamedWriteableRegistry (org.opensearch.common.io.stream.NamedWriteableRegistry)5 AsynchronousSearchActiveContext (org.opensearch.search.asynchronous.context.active.AsynchronousSearchActiveContext)5 AsynchronousSearchActiveStore (org.opensearch.search.asynchronous.context.active.AsynchronousSearchActiveStore)5 AsynchronousSearchStateMachine (org.opensearch.search.asynchronous.context.state.AsynchronousSearchStateMachine)5 SearchStartedEvent (org.opensearch.search.asynchronous.context.state.event.SearchStartedEvent)5 AsynchronousSearchPostProcessor (org.opensearch.search.asynchronous.processor.AsynchronousSearchPostProcessor)5 InternalAsynchronousSearchStats (org.opensearch.search.asynchronous.stats.InternalAsynchronousSearchStats)5 Map (java.util.Map)4 CountDownLatch (java.util.concurrent.CountDownLatch)4 TimeUnit (java.util.concurrent.TimeUnit)4