Examples with AsynchronousSearchPersistenceContext org.opensearch.search.asynchronous.context.persistence.AsynchronousSearchPersistenceContext used on opensource projects

Example 1 with AsynchronousSearchPersistenceContext

use of org.opensearch.search.asynchronous.context.persistence.AsynchronousSearchPersistenceContext 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);
        }));
    }
}

Example 2 with AsynchronousSearchPersistenceContext

use of org.opensearch.search.asynchronous.context.persistence.AsynchronousSearchPersistenceContext 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));
    }
}

Example 3 with AsynchronousSearchPersistenceContext

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

the class AsynchronousSearchServiceUpdateContextTests method testUpdateContextWhenContextCloseAndKeepOnCompletionTrue.

public void testUpdateContextWhenContextCloseAndKeepOnCompletionTrue() throws InterruptedException, IOException {
    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 = timeValueHours(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);
        long oldExpirationTimeMillis = asActiveContext.getExpirationTimeMillis();
        simulateIsAlive = false;
        simulateTimedOut = false;
        CountDownLatch updateLatch = new CountDownLatch(1);
        when(mockStore.getContext(asContextId)).thenReturn(Optional.of(asActiveContext));
        asService.updateKeepAliveAndGetContext(asActiveContext.getAsynchronousSearchId(), keepAlive, asActiveContext.getContextId(), user1, new LatchedActionListener<>(wrap(r -> {
            assertTrue(r instanceof AsynchronousSearchPersistenceContext);
            // assert active context expiration time is not updated
            assertEquals(asActiveContext.getExpirationTimeMillis(), oldExpirationTimeMillis);
        }, e -> {
            fail("expected successful update got " + e.getMessage());
        }), updateLatch));
        updateLatch.await();
        mockClusterService.stop();
    } finally {
        ThreadPool.terminate(testThreadPool, 30, TimeUnit.SECONDS);
    }
}

Example 4 with AsynchronousSearchPersistenceContext

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

the class AsynchronousSearchServiceUpdateContextTests method testUpdateClosedContext.

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