Search in sources :

Example 1 with AsynchronousSearchContext

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

the class AsynchronousSearchService method findContext.

/**
 * Tries to find an {@linkplain AsynchronousSearchActiveContext}. If not found,
 * queries the {@linkplain AsynchronousSearchPersistenceService}  for a hit. If a response is found, it builds and returns an
 * {@linkplain AsynchronousSearchPersistenceContext}, else throws
 * {@linkplain ResourceNotFoundException}
 *
 * @param id                   The asynchronous search id
 * @param asynchronousSearchContextId the Async search context id
 * @param user                 current user
 * @param listener             to be invoked on finding an {@linkplain AsynchronousSearchContext}
 */
public void findContext(String id, AsynchronousSearchContextId asynchronousSearchContextId, User user, ActionListener<AsynchronousSearchContext> listener) {
    ActionListener<AsynchronousSearchContext> exceptionTranslationListener = getExceptionTranslationWrapper(id, listener);
    Optional<AsynchronousSearchActiveContext> optionalAsynchronousSearchActiveContext = asynchronousSearchActiveStore.getContext(asynchronousSearchContextId);
    // so most likely a CLOSED context is stale
    if (optionalAsynchronousSearchActiveContext.isPresent() && optionalAsynchronousSearchActiveContext.get().isAlive()) {
        logger.debug("Active context is present for asynchronous search ID [{}]", id);
        AsynchronousSearchActiveContext asynchronousSearchActiveContext = optionalAsynchronousSearchActiveContext.get();
        if (isUserValid(user, asynchronousSearchActiveContext.getUser()) == false) {
            logger.debug("Invalid user requesting GET active context for asynchronous search id {}", id);
            exceptionTranslationListener.onFailure(new OpenSearchSecurityException("User doesn't have necessary roles to access the asynchronous search with id " + id, RestStatus.FORBIDDEN));
        } else {
            exceptionTranslationListener.onResponse(asynchronousSearchActiveContext);
        }
    } else {
        logger.debug("Active context is not present for asynchronous search ID [{}]", id);
        persistenceService.getResponse(id, user, wrap((persistenceModel) -> exceptionTranslationListener.onResponse(new AsynchronousSearchPersistenceContext(id, asynchronousSearchContextId, persistenceModel, currentTimeSupplier, namedWriteableRegistry)), ex -> {
            logger.debug(() -> new ParameterizedMessage("Context not found for ID  in the system index {}", id), ex);
            exceptionTranslationListener.onFailure(ex);
        }));
    }
}
Also used : OpenSearchSecurityException(org.opensearch.OpenSearchSecurityException) InternalAsynchronousSearchStats(org.opensearch.search.asynchronous.stats.InternalAsynchronousSearchStats) SearchStartedEvent(org.opensearch.search.asynchronous.context.state.event.SearchStartedEvent) PERSISTING(org.opensearch.search.asynchronous.context.state.AsynchronousSearchState.PERSISTING) LongSupplier(java.util.function.LongSupplier) ActionListener.runAfter(org.opensearch.action.ActionListener.runAfter) AsynchronousSearchState(org.opensearch.search.asynchronous.context.state.AsynchronousSearchState) TimeoutException(java.util.concurrent.TimeoutException) SearchDeletedEvent(org.opensearch.search.asynchronous.context.state.event.SearchDeletedEvent) GroupedActionListener(org.opensearch.action.support.GroupedActionListener) SUCCEEDED(org.opensearch.search.asynchronous.context.state.AsynchronousSearchState.SUCCEEDED) Locale(java.util.Locale) Map(java.util.Map) ActionListener(org.opensearch.action.ActionListener) SearchSuccessfulEvent(org.opensearch.search.asynchronous.context.state.event.SearchSuccessfulEvent) CancelTasksResponse(org.opensearch.action.admin.cluster.node.tasks.cancel.CancelTasksResponse) EnumSet(java.util.EnumSet) AsynchronousSearchTransition(org.opensearch.search.asynchronous.context.state.AsynchronousSearchTransition) AsynchronousSearchContext(org.opensearch.search.asynchronous.context.AsynchronousSearchContext) Client(org.opensearch.client.Client) TimeValue(org.opensearch.common.unit.TimeValue) SearchTask(org.opensearch.action.search.SearchTask) AsynchronousSearchPostProcessor(org.opensearch.search.asynchronous.processor.AsynchronousSearchPostProcessor) ExceptionsHelper(org.opensearch.ExceptionsHelper) Set(java.util.Set) Settings(org.opensearch.common.settings.Settings) SearchResponsePersistedEvent(org.opensearch.search.asynchronous.context.state.event.SearchResponsePersistedEvent) RestStatus(org.opensearch.rest.RestStatus) AsynchronousSearchPersistenceContext(org.opensearch.search.asynchronous.context.persistence.AsynchronousSearchPersistenceContext) Collectors(java.util.stream.Collectors) CLOSED(org.opensearch.search.asynchronous.context.state.AsynchronousSearchState.CLOSED) Objects(java.util.Objects) AbstractLifecycleComponent(org.opensearch.common.component.AbstractLifecycleComponent) RUNNING(org.opensearch.search.asynchronous.context.state.AsynchronousSearchState.RUNNING) Logger(org.apache.logging.log4j.Logger) SubmitAsynchronousSearchRequest(org.opensearch.search.asynchronous.request.SubmitAsynchronousSearchRequest) OpenSearchTimeoutException(org.opensearch.OpenSearchTimeoutException) Optional(java.util.Optional) AsynchronousSearchExceptionUtils(org.opensearch.search.asynchronous.utils.AsynchronousSearchExceptionUtils) ActionListener.wrap(org.opensearch.action.ActionListener.wrap) InternalAggregation(org.opensearch.search.aggregations.InternalAggregation) AsynchronousSearchActiveContext(org.opensearch.search.asynchronous.context.active.AsynchronousSearchActiveContext) AsynchronousSearchContextPermits(org.opensearch.search.asynchronous.context.permits.AsynchronousSearchContextPermits) SearchAction(org.opensearch.action.search.SearchAction) LegacyOpendistroAsynchronousSearchSettings(org.opensearch.search.asynchronous.settings.LegacyOpendistroAsynchronousSearchSettings) PERSIST_SUCCEEDED(org.opensearch.search.asynchronous.context.state.AsynchronousSearchState.PERSIST_SUCCEEDED) ThreadPool(org.opensearch.threadpool.ThreadPool) AsynchronousSearchProgressListener(org.opensearch.search.asynchronous.listener.AsynchronousSearchProgressListener) AsynchronousSearchContextId(org.opensearch.search.asynchronous.context.AsynchronousSearchContextId) AsynchronousSearchPlugin(org.opensearch.search.asynchronous.plugin.AsynchronousSearchPlugin) Releasable(org.opensearch.common.lease.Releasable) ResourceNotFoundException(org.opensearch.ResourceNotFoundException) ClusterStateListener(org.opensearch.cluster.ClusterStateListener) ParameterizedMessage(org.apache.logging.log4j.message.ParameterizedMessage) AtomicReference(java.util.concurrent.atomic.AtomicReference) Supplier(java.util.function.Supplier) UserAuthUtils.isUserValid(org.opensearch.search.asynchronous.utils.UserAuthUtils.isUserValid) NamedWriteableRegistry(org.opensearch.common.io.stream.NamedWriteableRegistry) BeginPersistEvent(org.opensearch.search.asynchronous.context.state.event.BeginPersistEvent) AsynchronousSearchStateMachine(org.opensearch.search.asynchronous.context.state.AsynchronousSearchStateMachine) UUIDs(org.opensearch.common.UUIDs) AsynchronousSearchContextEventListener(org.opensearch.search.asynchronous.listener.AsynchronousSearchContextEventListener) SearchFailureEvent(org.opensearch.search.asynchronous.context.state.event.SearchFailureEvent) AsynchronousSearchStateMachineClosedException(org.opensearch.search.asynchronous.context.state.AsynchronousSearchStateMachineClosedException) Setting(org.opensearch.common.settings.Setting) TaskId(org.opensearch.tasks.TaskId) SearchResponsePersistFailedEvent(org.opensearch.search.asynchronous.context.state.event.SearchResponsePersistFailedEvent) OpenSearchSecurityException(org.opensearch.OpenSearchSecurityException) AtomicLong(java.util.concurrent.atomic.AtomicLong) FAILED(org.opensearch.search.asynchronous.context.state.AsynchronousSearchState.FAILED) INIT(org.opensearch.search.asynchronous.context.state.AsynchronousSearchState.INIT) User(org.opensearch.commons.authuser.User) CancelTasksRequest(org.opensearch.action.admin.cluster.node.tasks.cancel.CancelTasksRequest) ClusterService(org.opensearch.cluster.service.ClusterService) PERSIST_FAILED(org.opensearch.search.asynchronous.context.state.AsynchronousSearchState.PERSIST_FAILED) AsynchronousSearchActiveStore(org.opensearch.search.asynchronous.context.active.AsynchronousSearchActiveStore) AsynchronousSearchStats(org.opensearch.search.asynchronous.stats.AsynchronousSearchStats) LogManager(org.apache.logging.log4j.LogManager) Collections(java.util.Collections) ClusterChangedEvent(org.opensearch.cluster.ClusterChangedEvent) AsynchronousSearchPersistenceContext(org.opensearch.search.asynchronous.context.persistence.AsynchronousSearchPersistenceContext) AsynchronousSearchContext(org.opensearch.search.asynchronous.context.AsynchronousSearchContext) ParameterizedMessage(org.apache.logging.log4j.message.ParameterizedMessage) AsynchronousSearchActiveContext(org.opensearch.search.asynchronous.context.active.AsynchronousSearchActiveContext)

Example 2 with AsynchronousSearchContext

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

the class AsynchronousSearchService method updateKeepAliveAndGetContext.

/**
 * If an active context is found, a permit is acquired from
 * {@linkplain AsynchronousSearchContextPermits}
 * and on acquisition of permit, a check is performed to see if response has been persisted in system index. If true, we update
 * expiration in index. Else we update expiration field in {@linkplain AsynchronousSearchActiveContext}.
 *
 * @param id                   asynchronous search id
 * @param keepAlive            the new keep alive duration
 * @param asynchronousSearchContextId asynchronous search context id
 * @param user                 current user
 * @param listener             listener to invoke after updating expiration.
 */
public void updateKeepAliveAndGetContext(String id, TimeValue keepAlive, AsynchronousSearchContextId asynchronousSearchContextId, User user, ActionListener<AsynchronousSearchContext> listener) {
    ActionListener<AsynchronousSearchContext> exceptionTranslationWrapper = getExceptionTranslationWrapper(id, listener);
    validateKeepAlive(keepAlive);
    long requestedExpirationTime = currentTimeSupplier.getAsLong() + keepAlive.getMillis();
    // find an active context on this node if one exists
    Optional<AsynchronousSearchActiveContext> asynchronousSearchContextOptional = asynchronousSearchActiveStore.getContext(asynchronousSearchContextId);
    // for all other stages we don't really care much as those contexts are destined to be discarded
    if (asynchronousSearchContextOptional.isPresent()) {
        AsynchronousSearchActiveContext asynchronousSearchActiveContext = asynchronousSearchContextOptional.get();
        asynchronousSearchActiveContext.acquireContextPermitIfRequired(wrap(releasable -> {
            ActionListener<AsynchronousSearchContext> releasableActionListener = runAfter(exceptionTranslationWrapper, releasable::close);
            // At this point it's possible that the response would have been persisted to system index
            if (asynchronousSearchActiveContext.isAlive() == false && asynchronousSearchActiveContext.keepOnCompletion()) {
                logger.debug("Updating persistence store after state is PERSISTED asynchronous search id [{}] " + "for updating context", asynchronousSearchActiveContext.getAsynchronousSearchId());
                persistenceService.updateExpirationTime(id, requestedExpirationTime, user, wrap((actionResponse) -> releasableActionListener.onResponse(new AsynchronousSearchPersistenceContext(id, asynchronousSearchContextId, actionResponse, currentTimeSupplier, namedWriteableRegistry)), releasableActionListener::onFailure));
            } else {
                if (isUserValid(user, asynchronousSearchActiveContext.getUser())) {
                    logger.debug("Updating persistence store: NO as state is NOT PERSISTED yet asynchronous search id [{}] " + "for updating context", asynchronousSearchActiveContext.getAsynchronousSearchId());
                    asynchronousSearchActiveContext.setExpirationTimeMillis(requestedExpirationTime);
                    releasableActionListener.onResponse(asynchronousSearchActiveContext);
                } else {
                    releasableActionListener.onFailure(new OpenSearchSecurityException("User doesn't have necessary roles to access the " + "asynchronous search with id " + id, RestStatus.FORBIDDEN));
                }
            }
        }, exception -> {
            Throwable cause = ExceptionsHelper.unwrapCause(exception);
            if (cause instanceof TimeoutException) {
                // this should ideally not happen. This would mean we couldn't acquire permits within the timeout
                logger.debug(() -> new ParameterizedMessage("Failed to acquire permits for " + "asynchronous search id [{}] for updating context within timeout 5s", asynchronousSearchActiveContext.getAsynchronousSearchId()), exception);
                listener.onFailure(new OpenSearchTimeoutException(id));
            } else {
                // best effort we try an update the doc if one exists
                if (asynchronousSearchActiveContext.keepOnCompletion()) {
                    logger.debug("Updating persistence store after failing to acquire permits for asynchronous search id [{}] for " + "updating context with expiration time [{}]", asynchronousSearchActiveContext.getAsynchronousSearchId(), requestedExpirationTime);
                    persistenceService.updateExpirationTime(id, requestedExpirationTime, user, wrap((actionResponse) -> exceptionTranslationWrapper.onResponse(new AsynchronousSearchPersistenceContext(id, asynchronousSearchContextId, actionResponse, currentTimeSupplier, namedWriteableRegistry)), exceptionTranslationWrapper::onFailure));
                } else {
                    exceptionTranslationWrapper.onFailure(new ResourceNotFoundException(asynchronousSearchActiveContext.getAsynchronousSearchId()));
                }
            }
        }), TimeValue.timeValueSeconds(5), "update keep alive");
    } else {
        // try update the doc on the index assuming there exists one.
        logger.debug("Updating persistence store after active context evicted for asynchronous search id [{}] " + "for updating context", id);
        persistenceService.updateExpirationTime(id, requestedExpirationTime, user, wrap((actionResponse) -> exceptionTranslationWrapper.onResponse(new AsynchronousSearchPersistenceContext(id, asynchronousSearchContextId, actionResponse, currentTimeSupplier, namedWriteableRegistry)), exceptionTranslationWrapper::onFailure));
    }
}
Also used : InternalAsynchronousSearchStats(org.opensearch.search.asynchronous.stats.InternalAsynchronousSearchStats) SearchStartedEvent(org.opensearch.search.asynchronous.context.state.event.SearchStartedEvent) PERSISTING(org.opensearch.search.asynchronous.context.state.AsynchronousSearchState.PERSISTING) LongSupplier(java.util.function.LongSupplier) ActionListener.runAfter(org.opensearch.action.ActionListener.runAfter) AsynchronousSearchState(org.opensearch.search.asynchronous.context.state.AsynchronousSearchState) TimeoutException(java.util.concurrent.TimeoutException) SearchDeletedEvent(org.opensearch.search.asynchronous.context.state.event.SearchDeletedEvent) GroupedActionListener(org.opensearch.action.support.GroupedActionListener) SUCCEEDED(org.opensearch.search.asynchronous.context.state.AsynchronousSearchState.SUCCEEDED) Locale(java.util.Locale) Map(java.util.Map) ActionListener(org.opensearch.action.ActionListener) SearchSuccessfulEvent(org.opensearch.search.asynchronous.context.state.event.SearchSuccessfulEvent) CancelTasksResponse(org.opensearch.action.admin.cluster.node.tasks.cancel.CancelTasksResponse) EnumSet(java.util.EnumSet) AsynchronousSearchTransition(org.opensearch.search.asynchronous.context.state.AsynchronousSearchTransition) AsynchronousSearchContext(org.opensearch.search.asynchronous.context.AsynchronousSearchContext) Client(org.opensearch.client.Client) TimeValue(org.opensearch.common.unit.TimeValue) SearchTask(org.opensearch.action.search.SearchTask) AsynchronousSearchPostProcessor(org.opensearch.search.asynchronous.processor.AsynchronousSearchPostProcessor) ExceptionsHelper(org.opensearch.ExceptionsHelper) Set(java.util.Set) Settings(org.opensearch.common.settings.Settings) SearchResponsePersistedEvent(org.opensearch.search.asynchronous.context.state.event.SearchResponsePersistedEvent) RestStatus(org.opensearch.rest.RestStatus) AsynchronousSearchPersistenceContext(org.opensearch.search.asynchronous.context.persistence.AsynchronousSearchPersistenceContext) Collectors(java.util.stream.Collectors) CLOSED(org.opensearch.search.asynchronous.context.state.AsynchronousSearchState.CLOSED) Objects(java.util.Objects) AbstractLifecycleComponent(org.opensearch.common.component.AbstractLifecycleComponent) RUNNING(org.opensearch.search.asynchronous.context.state.AsynchronousSearchState.RUNNING) Logger(org.apache.logging.log4j.Logger) SubmitAsynchronousSearchRequest(org.opensearch.search.asynchronous.request.SubmitAsynchronousSearchRequest) OpenSearchTimeoutException(org.opensearch.OpenSearchTimeoutException) Optional(java.util.Optional) AsynchronousSearchExceptionUtils(org.opensearch.search.asynchronous.utils.AsynchronousSearchExceptionUtils) ActionListener.wrap(org.opensearch.action.ActionListener.wrap) InternalAggregation(org.opensearch.search.aggregations.InternalAggregation) AsynchronousSearchActiveContext(org.opensearch.search.asynchronous.context.active.AsynchronousSearchActiveContext) AsynchronousSearchContextPermits(org.opensearch.search.asynchronous.context.permits.AsynchronousSearchContextPermits) SearchAction(org.opensearch.action.search.SearchAction) LegacyOpendistroAsynchronousSearchSettings(org.opensearch.search.asynchronous.settings.LegacyOpendistroAsynchronousSearchSettings) PERSIST_SUCCEEDED(org.opensearch.search.asynchronous.context.state.AsynchronousSearchState.PERSIST_SUCCEEDED) ThreadPool(org.opensearch.threadpool.ThreadPool) AsynchronousSearchProgressListener(org.opensearch.search.asynchronous.listener.AsynchronousSearchProgressListener) AsynchronousSearchContextId(org.opensearch.search.asynchronous.context.AsynchronousSearchContextId) AsynchronousSearchPlugin(org.opensearch.search.asynchronous.plugin.AsynchronousSearchPlugin) Releasable(org.opensearch.common.lease.Releasable) ResourceNotFoundException(org.opensearch.ResourceNotFoundException) ClusterStateListener(org.opensearch.cluster.ClusterStateListener) ParameterizedMessage(org.apache.logging.log4j.message.ParameterizedMessage) AtomicReference(java.util.concurrent.atomic.AtomicReference) Supplier(java.util.function.Supplier) UserAuthUtils.isUserValid(org.opensearch.search.asynchronous.utils.UserAuthUtils.isUserValid) NamedWriteableRegistry(org.opensearch.common.io.stream.NamedWriteableRegistry) BeginPersistEvent(org.opensearch.search.asynchronous.context.state.event.BeginPersistEvent) AsynchronousSearchStateMachine(org.opensearch.search.asynchronous.context.state.AsynchronousSearchStateMachine) UUIDs(org.opensearch.common.UUIDs) AsynchronousSearchContextEventListener(org.opensearch.search.asynchronous.listener.AsynchronousSearchContextEventListener) SearchFailureEvent(org.opensearch.search.asynchronous.context.state.event.SearchFailureEvent) AsynchronousSearchStateMachineClosedException(org.opensearch.search.asynchronous.context.state.AsynchronousSearchStateMachineClosedException) Setting(org.opensearch.common.settings.Setting) TaskId(org.opensearch.tasks.TaskId) SearchResponsePersistFailedEvent(org.opensearch.search.asynchronous.context.state.event.SearchResponsePersistFailedEvent) OpenSearchSecurityException(org.opensearch.OpenSearchSecurityException) AtomicLong(java.util.concurrent.atomic.AtomicLong) FAILED(org.opensearch.search.asynchronous.context.state.AsynchronousSearchState.FAILED) INIT(org.opensearch.search.asynchronous.context.state.AsynchronousSearchState.INIT) User(org.opensearch.commons.authuser.User) CancelTasksRequest(org.opensearch.action.admin.cluster.node.tasks.cancel.CancelTasksRequest) ClusterService(org.opensearch.cluster.service.ClusterService) PERSIST_FAILED(org.opensearch.search.asynchronous.context.state.AsynchronousSearchState.PERSIST_FAILED) AsynchronousSearchActiveStore(org.opensearch.search.asynchronous.context.active.AsynchronousSearchActiveStore) AsynchronousSearchStats(org.opensearch.search.asynchronous.stats.AsynchronousSearchStats) LogManager(org.apache.logging.log4j.LogManager) Collections(java.util.Collections) ClusterChangedEvent(org.opensearch.cluster.ClusterChangedEvent) OpenSearchSecurityException(org.opensearch.OpenSearchSecurityException) OpenSearchTimeoutException(org.opensearch.OpenSearchTimeoutException) AsynchronousSearchPersistenceContext(org.opensearch.search.asynchronous.context.persistence.AsynchronousSearchPersistenceContext) AsynchronousSearchContext(org.opensearch.search.asynchronous.context.AsynchronousSearchContext) AsynchronousSearchActiveContext(org.opensearch.search.asynchronous.context.active.AsynchronousSearchActiveContext) GroupedActionListener(org.opensearch.action.support.GroupedActionListener) ActionListener(org.opensearch.action.ActionListener) ParameterizedMessage(org.apache.logging.log4j.message.ParameterizedMessage) ResourceNotFoundException(org.opensearch.ResourceNotFoundException) TimeoutException(java.util.concurrent.TimeoutException) OpenSearchTimeoutException(org.opensearch.OpenSearchTimeoutException)

Example 3 with AsynchronousSearchContext

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

the class AsynchronousSearchServiceTests method testFreeContext.

public void testFreeContext() 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);
        CountDownLatch latch = new CountDownLatch(1);
        asService.freeContext(context.getAsynchronousSearchId(), context.getContextId(), user1, new LatchedActionListener<>(wrap(Assert::assertTrue, e -> fail()), latch));
        latch.await();
    } finally {
        ThreadPool.terminate(testThreadPool, 200, TimeUnit.MILLISECONDS);
    }
}
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) User(org.opensearch.commons.authuser.User) SubmitAsynchronousSearchRequest(org.opensearch.search.asynchronous.request.SubmitAsynchronousSearchRequest) AsynchronousSearchContext(org.opensearch.search.asynchronous.context.AsynchronousSearchContext) 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) CountDownLatch(java.util.concurrent.CountDownLatch) AsynchronousSearchActiveStore(org.opensearch.search.asynchronous.context.active.AsynchronousSearchActiveStore) AsynchronousSearchTask(org.opensearch.search.asynchronous.task.AsynchronousSearchTask) ClusterService(org.opensearch.cluster.service.ClusterService) Assert(org.junit.Assert) TimeValue(org.opensearch.common.unit.TimeValue)

Example 4 with AsynchronousSearchContext

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

the class AsynchronousSearchServiceTests method testFindContext.

public void testFindContext() 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 = TestUtils.createClusterService(settings, 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();
        User user2 = 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) -> {
        });
        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);
        CountDownLatch findContextLatch = new CountDownLatch(3);
        ActionListener<AsynchronousSearchContext> expectedSuccessfulActive = new LatchedActionListener<>(wrap(r -> {
            assertTrue(r instanceof AsynchronousSearchActiveContext);
            assertEquals(r, context);
        }, e -> fail("Find context shouldn't have failed. " + e.getMessage())), findContextLatch);
        ActionListener<AsynchronousSearchContext> expectedSecurityException = new LatchedActionListener<>(wrap(r -> fail("Expecting security exception"), e -> assertTrue(e instanceof ResourceNotFoundException)), findContextLatch);
        asService.findContext(asActiveContext.getAsynchronousSearchId(), asActiveContext.getContextId(), user1, expectedSuccessfulActive);
        asService.findContext(asActiveContext.getAsynchronousSearchId(), asActiveContext.getContextId(), user2, expectedSecurityException);
        asService.findContext(asActiveContext.getAsynchronousSearchId(), asActiveContext.getContextId(), null, expectedSuccessfulActive);
        findContextLatch.await();
        AsynchronousSearchProgressListener asProgressListener = asActiveContext.getAsynchronousSearchProgressListener();
        boolean success = randomBoolean();
        if (success) {
            // successful search response
            asProgressListener.onResponse(getMockSearchResponse());
        } else {
            // exception occurred in search
            asProgressListener.onFailure(new RuntimeException("test"));
        }
        waitUntil(() -> asService.getAllActiveContexts().isEmpty());
        if (keepOnCompletion) {
            // persist to disk
            assertEquals(1, fakeClient.persistenceCount.intValue());
        } else {
            assertEquals(fakeClient.persistenceCount, Integer.valueOf(0));
            CountDownLatch freeContextLatch = new CountDownLatch(1);
            asService.findContext(context.getAsynchronousSearchId(), context.getContextId(), null, new LatchedActionListener<>(wrap(r -> fail("No context should have been found but found " + asService.getAllActiveContexts().size()), e -> assertTrue(e instanceof ResourceNotFoundException)), freeContextLatch));
            freeContextLatch.await();
        }
    } finally {
        ThreadPool.terminate(testThreadPool, 200, TimeUnit.MILLISECONDS);
    }
}
Also used : NamedWriteableRegistry(org.opensearch.common.io.stream.NamedWriteableRegistry) InternalAsynchronousSearchStats(org.opensearch.search.asynchronous.stats.InternalAsynchronousSearchStats) ExecutorBuilder(org.opensearch.threadpool.ExecutorBuilder) AsynchronousSearchState(org.opensearch.search.asynchronous.context.state.AsynchronousSearchState) TestThreadPool(org.opensearch.threadpool.TestThreadPool) Version(org.opensearch.Version) ActionRequest(org.opensearch.action.ActionRequest) DiscoveryNode(org.opensearch.cluster.node.DiscoveryNode) LatchedActionListener(org.opensearch.action.LatchedActionListener) InternalAggregations(org.opensearch.search.aggregations.InternalAggregations) ScalingExecutorBuilder(org.opensearch.threadpool.ScalingExecutorBuilder) ActionListener(org.opensearch.action.ActionListener) ActionResponse(org.opensearch.action.ActionResponse) ActionType(org.opensearch.action.ActionType) AsynchronousSearchContext(org.opensearch.search.asynchronous.context.AsynchronousSearchContext) TimeValue(org.opensearch.common.unit.TimeValue) SearchHit(org.opensearch.search.SearchHit) Collections.emptyList(java.util.Collections.emptyList) OpenSearchTestCase(org.opensearch.test.OpenSearchTestCase) NoOpClient(org.opensearch.test.client.NoOpClient) Set(java.util.Set) Settings(org.opensearch.common.settings.Settings) AsynchronousSearchTask(org.opensearch.search.asynchronous.task.AsynchronousSearchTask) DiscoveryNodeRole(org.opensearch.cluster.node.DiscoveryNodeRole) Collectors(java.util.stream.Collectors) CountDownLatch(java.util.concurrent.CountDownLatch) List(java.util.List) Stream(java.util.stream.Stream) SubmitAsynchronousSearchRequest(org.opensearch.search.asynchronous.request.SubmitAsynchronousSearchRequest) OpenSearchTimeoutException(org.opensearch.OpenSearchTimeoutException) IndexAction(org.opensearch.action.index.IndexAction) ActionListener.wrap(org.opensearch.action.ActionListener.wrap) Matchers.greaterThan(org.hamcrest.Matchers.greaterThan) AsynchronousSearchActiveContext(org.opensearch.search.asynchronous.context.active.AsynchronousSearchActiveContext) ClusterServiceUtils(org.opensearch.test.ClusterServiceUtils) SearchProfileShardResults(org.opensearch.search.profile.SearchProfileShardResults) SearchAction(org.opensearch.action.search.SearchAction) ThreadPool(org.opensearch.threadpool.ThreadPool) AsynchronousSearchProgressListener(org.opensearch.search.asynchronous.listener.AsynchronousSearchProgressListener) AsynchronousSearchContextId(org.opensearch.search.asynchronous.context.AsynchronousSearchContextId) AsynchronousSearchPlugin(org.opensearch.search.asynchronous.plugin.AsynchronousSearchPlugin) TimeValue.timeValueHours(org.opensearch.common.unit.TimeValue.timeValueHours) SearchHits(org.opensearch.search.SearchHits) ResourceNotFoundException(org.opensearch.ResourceNotFoundException) OpenSearchExecutors(org.opensearch.common.util.concurrent.OpenSearchExecutors) NamedWriteableRegistry(org.opensearch.common.io.stream.NamedWriteableRegistry) ArrayList(java.util.ArrayList) SearchRequest(org.opensearch.action.search.SearchRequest) SearchResponse(org.opensearch.action.search.SearchResponse) ClusterSettings(org.opensearch.common.settings.ClusterSettings) Before(org.junit.Before) Collections.emptyMap(java.util.Collections.emptyMap) Setting(org.opensearch.common.settings.Setting) TestUtils(org.opensearch.search.asynchronous.utils.TestUtils) TaskId(org.opensearch.tasks.TaskId) InternalSearchResponse(org.opensearch.search.internal.InternalSearchResponse) TotalHits(org.apache.lucene.search.TotalHits) TimeUnit(java.util.concurrent.TimeUnit) User(org.opensearch.commons.authuser.User) TestClientUtils(org.opensearch.search.asynchronous.utils.TestClientUtils) Suggest(org.opensearch.search.suggest.Suggest) ShardSearchFailure(org.opensearch.action.search.ShardSearchFailure) ClusterService(org.opensearch.cluster.service.ClusterService) AsynchronousSearchActiveStore(org.opensearch.search.asynchronous.context.active.AsynchronousSearchActiveStore) Assert(org.junit.Assert) Collections(java.util.Collections) SubmitAsynchronousSearchRequest(org.opensearch.search.asynchronous.request.SubmitAsynchronousSearchRequest) SearchRequest(org.opensearch.action.search.SearchRequest) DiscoveryNode(org.opensearch.cluster.node.DiscoveryNode) User(org.opensearch.commons.authuser.User) AsynchronousSearchContext(org.opensearch.search.asynchronous.context.AsynchronousSearchContext) TestThreadPool(org.opensearch.threadpool.TestThreadPool) ThreadPool(org.opensearch.threadpool.ThreadPool) TestThreadPool(org.opensearch.threadpool.TestThreadPool) LatchedActionListener(org.opensearch.action.LatchedActionListener) ResourceNotFoundException(org.opensearch.ResourceNotFoundException) TimeValue(org.opensearch.common.unit.TimeValue) SubmitAsynchronousSearchRequest(org.opensearch.search.asynchronous.request.SubmitAsynchronousSearchRequest) InternalAsynchronousSearchStats(org.opensearch.search.asynchronous.stats.InternalAsynchronousSearchStats) AsynchronousSearchActiveContext(org.opensearch.search.asynchronous.context.active.AsynchronousSearchActiveContext) CountDownLatch(java.util.concurrent.CountDownLatch) AsynchronousSearchActiveStore(org.opensearch.search.asynchronous.context.active.AsynchronousSearchActiveStore) AsynchronousSearchProgressListener(org.opensearch.search.asynchronous.listener.AsynchronousSearchProgressListener) AsynchronousSearchTask(org.opensearch.search.asynchronous.task.AsynchronousSearchTask) ClusterService(org.opensearch.cluster.service.ClusterService)

Example 5 with AsynchronousSearchContext

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

the class AsynchronousSearchServiceTests method testFindContextsToReap.

public void testFindContextsToReap() {
    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) {

            @Override
            public long absoluteTimeInMillis() {
                // simulate search has over run)
                return System.currentTimeMillis() - 24 * 3600 * 1000;
            }
        };
        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(asService.getContextsToReap().contains(context));
    } 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) User(org.opensearch.commons.authuser.User) SubmitAsynchronousSearchRequest(org.opensearch.search.asynchronous.request.SubmitAsynchronousSearchRequest) AsynchronousSearchContext(org.opensearch.search.asynchronous.context.AsynchronousSearchContext) TestThreadPool(org.opensearch.threadpool.TestThreadPool) ThreadPool(org.opensearch.threadpool.ThreadPool) InternalAsynchronousSearchStats(org.opensearch.search.asynchronous.stats.InternalAsynchronousSearchStats) TestThreadPool(org.opensearch.threadpool.TestThreadPool) AsynchronousSearchActiveStore(org.opensearch.search.asynchronous.context.active.AsynchronousSearchActiveStore) ClusterService(org.opensearch.cluster.service.ClusterService) TimeValue(org.opensearch.common.unit.TimeValue)

Aggregations

AsynchronousSearchContext (org.opensearch.search.asynchronous.context.AsynchronousSearchContext)16 SubmitAsynchronousSearchRequest (org.opensearch.search.asynchronous.request.SubmitAsynchronousSearchRequest)15 ClusterService (org.opensearch.cluster.service.ClusterService)14 NamedWriteableRegistry (org.opensearch.common.io.stream.NamedWriteableRegistry)14 TimeValue (org.opensearch.common.unit.TimeValue)14 AsynchronousSearchActiveContext (org.opensearch.search.asynchronous.context.active.AsynchronousSearchActiveContext)14 AsynchronousSearchActiveStore (org.opensearch.search.asynchronous.context.active.AsynchronousSearchActiveStore)14 InternalAsynchronousSearchStats (org.opensearch.search.asynchronous.stats.InternalAsynchronousSearchStats)14 ThreadPool (org.opensearch.threadpool.ThreadPool)14 SearchRequest (org.opensearch.action.search.SearchRequest)12 User (org.opensearch.commons.authuser.User)12 DiscoveryNode (org.opensearch.cluster.node.DiscoveryNode)11 TestThreadPool (org.opensearch.threadpool.TestThreadPool)11 CountDownLatch (java.util.concurrent.CountDownLatch)9 AsynchronousSearchTask (org.opensearch.search.asynchronous.task.AsynchronousSearchTask)9 ResourceNotFoundException (org.opensearch.ResourceNotFoundException)8 OpenSearchTimeoutException (org.opensearch.OpenSearchTimeoutException)6 AsynchronousSearchContextId (org.opensearch.search.asynchronous.context.AsynchronousSearchContextId)6 AsynchronousSearchProgressListener (org.opensearch.search.asynchronous.listener.AsynchronousSearchProgressListener)6 TaskId (org.opensearch.tasks.TaskId)6