Search in sources :

Example 1 with AsynchronousSearchProgressListener

use of org.opensearch.search.asynchronous.listener.AsynchronousSearchProgressListener in project asynchronous-search by opensearch-project.

the class AsynchronousSearchService method createAndStoreContext.

/**
 * Creates a new active asynchronous search for a newly submitted asynchronous search.
 *
 * @param request                 the SubmitAsynchronousSearchRequest
 * @param relativeStartTimeMillis the relative start time of the search in millis
 * @param user                    current user
 * @param reduceContextBuilder    the reference for the reduceContextBuilder
 * @return the AsynchronousSearchContext for the submitted request
 */
public AsynchronousSearchContext createAndStoreContext(SubmitAsynchronousSearchRequest request, long relativeStartTimeMillis, Supplier<InternalAggregation.ReduceContextBuilder> reduceContextBuilder, User user) {
    validateRequest(request);
    AsynchronousSearchContextId asynchronousSearchContextId = new AsynchronousSearchContextId(UUIDs.base64UUID(), idGenerator.incrementAndGet());
    contextEventListener.onNewContext(asynchronousSearchContextId);
    AsynchronousSearchProgressListener progressActionListener = new AsynchronousSearchProgressListener(relativeStartTimeMillis, (response) -> asynchronousSearchPostProcessor.processSearchResponse(response, asynchronousSearchContextId), (e) -> asynchronousSearchPostProcessor.processSearchFailure(e, asynchronousSearchContextId), threadPool.executor(AsynchronousSearchPlugin.OPEN_DISTRO_ASYNC_SEARCH_GENERIC_THREAD_POOL_NAME), threadPool::relativeTimeInMillis, reduceContextBuilder);
    AsynchronousSearchActiveContext asynchronousSearchContext = new AsynchronousSearchActiveContext(asynchronousSearchContextId, clusterService.localNode().getId(), request.getKeepAlive(), request.getKeepOnCompletion(), threadPool, currentTimeSupplier, progressActionListener, user, () -> persistSearchFailure);
    asynchronousSearchActiveStore.putContext(asynchronousSearchContextId, asynchronousSearchContext, contextEventListener::onContextRejected);
    contextEventListener.onContextInitialized(asynchronousSearchContextId);
    return asynchronousSearchContext;
}
Also used : AsynchronousSearchContextId(org.opensearch.search.asynchronous.context.AsynchronousSearchContextId) AsynchronousSearchActiveContext(org.opensearch.search.asynchronous.context.active.AsynchronousSearchActiveContext) AsynchronousSearchProgressListener(org.opensearch.search.asynchronous.listener.AsynchronousSearchProgressListener)

Example 2 with AsynchronousSearchProgressListener

use of org.opensearch.search.asynchronous.listener.AsynchronousSearchProgressListener 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 3 with AsynchronousSearchProgressListener

use of org.opensearch.search.asynchronous.listener.AsynchronousSearchProgressListener in project asynchronous-search by opensearch-project.

the class AsynchronousSearchActiveContextTests method testInitializeContext.

public void testInitializeContext() {
    TestThreadPool threadPool = null;
    try {
        int writeThreadPoolSize = randomIntBetween(1, 2);
        int writeThreadPoolQueueSize = randomIntBetween(1, 2);
        Settings settings = Settings.builder().put("thread_pool." + AsynchronousSearchPlugin.OPEN_DISTRO_ASYNC_SEARCH_GENERIC_THREAD_POOL_NAME + ".size", writeThreadPoolSize).put("thread_pool." + AsynchronousSearchPlugin.OPEN_DISTRO_ASYNC_SEARCH_GENERIC_THREAD_POOL_NAME + ".queue_size", writeThreadPoolQueueSize).build();
        final int availableProcessors = OpenSearchExecutors.allocatedProcessors(settings);
        ScalingExecutorBuilder scalingExecutorBuilder = new ScalingExecutorBuilder(AsynchronousSearchPlugin.OPEN_DISTRO_ASYNC_SEARCH_GENERIC_THREAD_POOL_NAME, 1, Math.min(2 * availableProcessors, Math.max(128, 512)), TimeValue.timeValueMinutes(30));
        threadPool = new TestThreadPool("Tests", settings, scalingExecutorBuilder);
        String node = UUID.randomUUID().toString();
        AsynchronousSearchProgressListener asProgressListener = mockAsynchronousSearchProgressListener(threadPool);
        AsynchronousSearchContextId asContextId = new AsynchronousSearchContextId(UUID.randomUUID().toString(), randomNonNegativeLong());
        boolean keepOnCompletion = randomBoolean();
        User user = TestClientUtils.randomUser();
        TimeValue keepAlive = TimeValue.timeValueDays(randomInt(100));
        AsynchronousSearchActiveContext context = new AsynchronousSearchActiveContext(asContextId, node, keepAlive, keepOnCompletion, threadPool, threadPool::absoluteTimeInMillis, asProgressListener, user, () -> true);
        assertEquals(AsynchronousSearchState.INIT, context.getAsynchronousSearchState());
        assertNull(context.getTask());
        assertNull(context.getAsynchronousSearchId());
        assertEquals(context.getAsynchronousSearchState(), AsynchronousSearchState.INIT);
    } finally {
        ThreadPool.terminate(threadPool, 30, TimeUnit.SECONDS);
    }
}
Also used : ScalingExecutorBuilder(org.opensearch.threadpool.ScalingExecutorBuilder) User(org.opensearch.commons.authuser.User) AsynchronousSearchContextId(org.opensearch.search.asynchronous.context.AsynchronousSearchContextId) Matchers.containsString(org.hamcrest.Matchers.containsString) TestThreadPool(org.opensearch.threadpool.TestThreadPool) Settings(org.opensearch.common.settings.Settings) TimeValue(org.opensearch.common.unit.TimeValue) AsynchronousSearchProgressListener(org.opensearch.search.asynchronous.listener.AsynchronousSearchProgressListener)

Example 4 with AsynchronousSearchProgressListener

use of org.opensearch.search.asynchronous.listener.AsynchronousSearchProgressListener in project asynchronous-search by opensearch-project.

the class TransportGetAsynchronousSearchAction method handleWaitForCompletion.

private void handleWaitForCompletion(AsynchronousSearchContext context, GetAsynchronousSearchRequest request, ActionListener<AsynchronousSearchResponse> listener) {
    // We wait for a response only if a wait for completion is non-null and the search execution is still in progress.
    if (context.isRunning() && request.getWaitForCompletionTimeout() != null) {
        logger.debug("Context is running for asynchronous search id [{}]", context.getAsynchronousSearchId());
        AsynchronousSearchProgressListener progressActionListener = context.getAsynchronousSearchProgressListener();
        assert progressActionListener != null : "progress listener cannot be null";
        PrioritizedActionListener<AsynchronousSearchResponse> wrappedListener = AsynchronousSearchTimeoutWrapper.wrapScheduledTimeout(threadPool, request.getWaitForCompletionTimeout(), OPEN_DISTRO_ASYNC_SEARCH_GENERIC_THREAD_POOL_NAME, listener, (actionListener) -> {
            progressActionListener.searchProgressActionListener().removeListener(actionListener);
            listener.onResponse(context.getAsynchronousSearchResponse());
        });
        progressActionListener.searchProgressActionListener().addOrExecuteListener(wrappedListener);
    } else {
        // we don't need to wait any further on search completion
        logger.debug("Context is not running for asynchronous search id [{}]", context.getAsynchronousSearchId());
        listener.onResponse(context.getAsynchronousSearchResponse());
    }
}
Also used : AsynchronousSearchResponse(org.opensearch.search.asynchronous.response.AsynchronousSearchResponse) AsynchronousSearchProgressListener(org.opensearch.search.asynchronous.listener.AsynchronousSearchProgressListener)

Example 5 with AsynchronousSearchProgressListener

use of org.opensearch.search.asynchronous.listener.AsynchronousSearchProgressListener in project asynchronous-search by opensearch-project.

the class AsynchronousSearchServiceFreeContextTests method testFreeContextPermitAcquisitionFailure.

public void testFreeContextPermitAcquisitionFailure() 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());
        MockAsynchronousSearchActiveContext asActiveContext = new MockAsynchronousSearchActiveContext(asContextId, discoveryNode.getId(), keepAlive, true, testThreadPool, testThreadPool::absoluteTimeInMillis, asProgressListener, user1);
        // bootstrap search
        AsynchronousSearchTask task = new AsynchronousSearchTask(randomNonNegativeLong(), "transport", SearchAction.NAME, TaskId.EMPTY_TASK_ID, emptyMap(), asActiveContext, null, (c) -> {
        });
        asActiveContext.setTask(task);
        simulateUncheckedException = true;
        persisted = false;
        when(mockStore.getContext(asContextId)).thenReturn(Optional.of(asActiveContext));
        CountDownLatch latch = new CountDownLatch(1);
        asService.freeContext(asActiveContext.getAsynchronousSearchId(), asActiveContext.getContextId(), null, new LatchedActionListener<>(wrap(r -> fail("Expected resource_not_found_exception. Got acknowledgement " + r), e -> {
            assertTrue(e.getClass().getName(), e instanceof ResourceNotFoundException);
        }), latch));
        latch.await();
        assertEquals(1, (int) mockClient.deleteCount);
        mockClusterService.stop();
    } finally {
        ThreadPool.terminate(testThreadPool, 30, TimeUnit.SECONDS);
    }
}
Also used : NamedWriteableRegistry(org.opensearch.common.io.stream.NamedWriteableRegistry) SubmitAsynchronousSearchRequest(org.opensearch.search.asynchronous.request.SubmitAsynchronousSearchRequest) SearchRequest(org.opensearch.action.search.SearchRequest) DiscoveryNode(org.opensearch.cluster.node.DiscoveryNode) TestClientUtils.randomUser(org.opensearch.search.asynchronous.utils.TestClientUtils.randomUser) User(org.opensearch.commons.authuser.User) SubmitAsynchronousSearchRequest(org.opensearch.search.asynchronous.request.SubmitAsynchronousSearchRequest) AsynchronousSearchContextId(org.opensearch.search.asynchronous.context.AsynchronousSearchContextId) TestThreadPool(org.opensearch.threadpool.TestThreadPool) ThreadPool(org.opensearch.threadpool.ThreadPool) InternalAsynchronousSearchStats(org.opensearch.search.asynchronous.stats.InternalAsynchronousSearchStats) TestThreadPool(org.opensearch.threadpool.TestThreadPool) CountDownLatch(java.util.concurrent.CountDownLatch) AsynchronousSearchActiveStore(org.opensearch.search.asynchronous.context.active.AsynchronousSearchActiveStore) AsynchronousSearchProgressListener(org.opensearch.search.asynchronous.listener.AsynchronousSearchProgressListener) AsynchronousSearchTestCase.mockAsynchronousSearchProgressListener(org.opensearch.search.asynchronous.commons.AsynchronousSearchTestCase.mockAsynchronousSearchProgressListener) AsynchronousSearchTask(org.opensearch.search.asynchronous.task.AsynchronousSearchTask) ClusterService(org.opensearch.cluster.service.ClusterService) ResourceNotFoundException(org.opensearch.ResourceNotFoundException) TimeValue(org.opensearch.common.unit.TimeValue)

Aggregations

AsynchronousSearchProgressListener (org.opensearch.search.asynchronous.listener.AsynchronousSearchProgressListener)22 TimeValue (org.opensearch.common.unit.TimeValue)19 AsynchronousSearchContextId (org.opensearch.search.asynchronous.context.AsynchronousSearchContextId)19 TestThreadPool (org.opensearch.threadpool.TestThreadPool)19 SearchRequest (org.opensearch.action.search.SearchRequest)15 User (org.opensearch.commons.authuser.User)15 SubmitAsynchronousSearchRequest (org.opensearch.search.asynchronous.request.SubmitAsynchronousSearchRequest)15 AsynchronousSearchTask (org.opensearch.search.asynchronous.task.AsynchronousSearchTask)15 CountDownLatch (java.util.concurrent.CountDownLatch)14 DiscoveryNode (org.opensearch.cluster.node.DiscoveryNode)14 ClusterService (org.opensearch.cluster.service.ClusterService)14 ThreadPool (org.opensearch.threadpool.ThreadPool)14 NamedWriteableRegistry (org.opensearch.common.io.stream.NamedWriteableRegistry)13 AsynchronousSearchActiveStore (org.opensearch.search.asynchronous.context.active.AsynchronousSearchActiveStore)13 AsynchronousSearchTestCase.mockAsynchronousSearchProgressListener (org.opensearch.search.asynchronous.commons.AsynchronousSearchTestCase.mockAsynchronousSearchProgressListener)12 InternalAsynchronousSearchStats (org.opensearch.search.asynchronous.stats.InternalAsynchronousSearchStats)12 TestClientUtils.randomUser (org.opensearch.search.asynchronous.utils.TestClientUtils.randomUser)11 AsynchronousSearchActiveContext (org.opensearch.search.asynchronous.context.active.AsynchronousSearchActiveContext)8 ResourceNotFoundException (org.opensearch.ResourceNotFoundException)6 Settings (org.opensearch.common.settings.Settings)6