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

Example 26 with AsynchronousSearchActiveContext

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

the class AsynchronousSearchService method freeContext.

/**
 * Attempts to find both an {@linkplain AsynchronousSearchActiveContext} and an {@linkplain AsynchronousSearchPersistenceContext}
 * and delete them. If at least one of the aforementioned objects are found and deleted successfully, the listener is invoked with
 * #true, else {@linkplain ResourceNotFoundException} is thrown.
 *
 * @param id                   asynchronous search id
 * @param asynchronousSearchContextId context id
 * @param user                 current user
 * @param listener             listener to invoke on deletion or failure to do so
 */
public void freeContext(String id, AsynchronousSearchContextId asynchronousSearchContextId, User user, ActionListener<Boolean> listener) {
    ActionListener<Boolean> exceptionTranslationWrapper = getExceptionTranslationWrapper(id, listener);
    Optional<AsynchronousSearchActiveContext> asynchronousSearchContextOptional = asynchronousSearchActiveStore.getContext(asynchronousSearchContextId);
    if (asynchronousSearchContextOptional.isPresent()) {
        logger.debug("Active context present for asynchronous search id [{}]", id);
        AsynchronousSearchActiveContext asynchronousSearchContext = asynchronousSearchContextOptional.get();
        if (isUserValid(user, asynchronousSearchContext.getUser())) {
            cancelAndFreeActiveAndPersistedContext(asynchronousSearchContext, exceptionTranslationWrapper, user);
        } else {
            exceptionTranslationWrapper.onFailure(new OpenSearchSecurityException("User doesn't have necessary roles to access the asynchronous search with id " + id, RestStatus.FORBIDDEN));
        }
    } else {
        logger.debug("Active context NOT present for asynchronous search [{}]", id);
        // asynchronous search context didn't exist so obviously we didn't delete
        // deleted persisted context if one exists. If not the listener returns acknowledged as false
        // we don't need to acquire lock if the in-memory context doesn't exist. For persistence context we have a distributed view
        // with the last writer wins policy
        logger.debug("Deleting asynchronous search [{}] from system index ", id);
        persistenceService.deleteResponse(id, user, exceptionTranslationWrapper);
    }
}

Example 27 with AsynchronousSearchActiveContext

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

the class AsynchronousSearchService method cancelAndFreeActiveAndPersistedContext.

// We are skipping user check in this while deleting from the persisted layer
// as we have already checked for user in the present active context.
private void cancelAndFreeActiveAndPersistedContext(AsynchronousSearchActiveContext asynchronousSearchContext, ActionListener<Boolean> listener, User user) {
    // if there are no context found to be cleaned up we throw a ResourceNotFoundException
    AtomicReference<Releasable> releasableReference = new AtomicReference<>(() -> {
    });
    ActionListener<Boolean> releasableListener = runAfter(listener, releasableReference.get()::close);
    GroupedActionListener<Boolean> groupedDeletionListener = new GroupedActionListener<>(wrap((responses) -> {
        if (responses.stream().anyMatch(r -> r)) {
            logger.debug("Free context for asynchronous search [{}] successful ", asynchronousSearchContext.getAsynchronousSearchId());
            releasableListener.onResponse(true);
        } else {
            logger.debug("Freeing context, asynchronous search [{}] not found ", asynchronousSearchContext.getAsynchronousSearchId());
            releasableListener.onFailure(new ResourceNotFoundException(asynchronousSearchContext.getAsynchronousSearchId()));
        }
    }, releasableListener::onFailure), 2);
    // We get a true or a ResourceNotFound from persistence layer. We want to translate it to either a true/false or any other exception
    // that should be surfaced up
    ActionListener<Boolean> translatedListener = wrap(groupedDeletionListener::onResponse, (ex) -> {
        if (ex instanceof ResourceNotFoundException) {
            groupedDeletionListener.onResponse(false);
        } else {
            logger.debug(() -> new ParameterizedMessage("Translating exception, received for asynchronous search [{}]", asynchronousSearchContext.getAsynchronousSearchId()), ex);
            groupedDeletionListener.onFailure(ex);
        }
    });
    String triggeredBy = user != null ? (" by user [" + user + "]") : "";
    String cancelTaskReason = "Delete asynchronous search [" + asynchronousSearchContext.getAsynchronousSearchId() + "] has been triggered" + triggeredBy + ". Attempting to cancel in-progress search task";
    // Intent of the lock here is to disallow ongoing migration to system index
    // as if that is underway we might end up creating a new document post a DELETE was executed
    asynchronousSearchContext.acquireContextPermitIfRequired(wrap(releasable -> {
        releasableReference.set(releasable);
        if (asynchronousSearchContext.keepOnCompletion()) {
            handleCancelTaskPermitAcquired(asynchronousSearchContext, groupedDeletionListener, cancelTaskReason);
            logger.debug("Deleting asynchronous search id [{}] from system index ", asynchronousSearchContext.getAsynchronousSearchId());
            persistenceService.deleteResponse(asynchronousSearchContext.getAsynchronousSearchId(), user, translatedListener);
        } else {
            // keep on completion is false. simply cancel task and clean up active context
            handleCancelTaskPermitAcquired(asynchronousSearchContext, wrap(r -> {
                if (r) {
                    releasableListener.onResponse(true);
                } else {
                    releasableListener.onFailure(new ResourceNotFoundException(asynchronousSearchContext.getAsynchronousSearchId()));
                }
            }, releasableListener::onFailure), cancelTaskReason);
        }
    }, 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", asynchronousSearchContext.getAsynchronousSearchId()), exception);
            listener.onFailure(new OpenSearchTimeoutException(asynchronousSearchContext.getAsynchronousSearchId()));
        } else {
            // best effort clean up with acknowledged as false
            if (asynchronousSearchContext.keepOnCompletion()) {
                handleCancelTaskPermitAcquisitionFailed(asynchronousSearchContext, groupedDeletionListener, cancelTaskReason, exception);
                logger.debug("Deleting asynchronous search id [{}] from system index ", asynchronousSearchContext.getAsynchronousSearchId());
                persistenceService.deleteResponse(asynchronousSearchContext.getAsynchronousSearchId(), user, translatedListener);
            } else {
                handleCancelTaskPermitAcquisitionFailed(asynchronousSearchContext, releasableListener, cancelTaskReason, exception);
            }
        }
    }), TimeValue.timeValueSeconds(5), "free context");
}

Example 28 with AsynchronousSearchActiveContext

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

the class AsynchronousSearchPostProcessor method processSearchFailure.

public AsynchronousSearchResponse processSearchFailure(Exception exception, AsynchronousSearchContextId asynchronousSearchContextId) {
    final Optional<AsynchronousSearchActiveContext> asynchronousSearchContextOptional = asynchronousSearchActiveStore.getContext(asynchronousSearchContextId);
    try {
        if (asynchronousSearchContextOptional.isPresent()) {
            AsynchronousSearchActiveContext asynchronousSearchContext = asynchronousSearchContextOptional.get();
            asynchronousSearchStateMachine.trigger(new SearchFailureEvent(asynchronousSearchContext, exception));
            handlePersist(asynchronousSearchContext);
            return asynchronousSearchContext.getAsynchronousSearchResponse();
        }
        // Best effort to return the response.
        return new AsynchronousSearchResponse(AsynchronousSearchState.FAILED, -1L, -1L, null, ExceptionsHelper.convertToOpenSearchException(exception));
    } catch (AsynchronousSearchStateMachineClosedException ex) {
        // Best effort to return the response.
        return new AsynchronousSearchResponse(AsynchronousSearchState.FAILED, -1L, -1L, null, ExceptionsHelper.convertToOpenSearchException(exception));
    }
}

Example 29 with AsynchronousSearchActiveContext

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

the class AsynchronousSearchPostProcessor method processSearchResponse.

public AsynchronousSearchResponse processSearchResponse(SearchResponse searchResponse, AsynchronousSearchContextId asynchronousSearchContextId) {
    final Optional<AsynchronousSearchActiveContext> asynchronousSearchContextOptional = asynchronousSearchActiveStore.getContext(asynchronousSearchContextId);
    try {
        if (asynchronousSearchContextOptional.isPresent()) {
            AsynchronousSearchActiveContext asynchronousSearchContext = asynchronousSearchContextOptional.get();
            asynchronousSearchStateMachine.trigger(new SearchSuccessfulEvent(asynchronousSearchContext, searchResponse));
            handlePersist(asynchronousSearchContext);
            return asynchronousSearchContext.getAsynchronousSearchResponse();
        }
        // Best effort to return the response.
        return new AsynchronousSearchResponse(AsynchronousSearchState.SUCCEEDED, -1L, -1L, searchResponse, null);
    } catch (AsynchronousSearchStateMachineClosedException ex) {
        // Best effort to return the response.
        return new AsynchronousSearchResponse(AsynchronousSearchState.SUCCEEDED, -1L, -1L, searchResponse, null);
    }
}

Example 30 with AsynchronousSearchActiveContext

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

the class TransportSubmitAsynchronousSearchAction method doExecute.

@Override
protected void doExecute(Task task, SubmitAsynchronousSearchRequest request, ActionListener<AsynchronousSearchResponse> listener) {
    AsynchronousSearchContext asynchronousSearchContext = null;
    String userStr = threadPool.getThreadContext().getTransient(ConfigConstants.OPENSEARCH_SECURITY_USER_INFO_THREAD_CONTEXT);
    User user = User.parse(userStr);
    try {
        final long relativeStartTimeInMillis = threadPool.relativeTimeInMillis();
        asynchronousSearchContext = asynchronousSearchService.createAndStoreContext(request, relativeStartTimeInMillis, () -> searchService.aggReduceContextBuilder(request.getSearchRequest()), user);
        assert asynchronousSearchContext.getAsynchronousSearchProgressListener() != null : "missing progress listener for an active context";
        AsynchronousSearchProgressListener progressListener = asynchronousSearchContext.getAsynchronousSearchProgressListener();
        // making it effectively final for usage in anonymous class.
        AsynchronousSearchContext context = asynchronousSearchContext;
        SearchRequest searchRequest = new SearchRequest(request.getSearchRequest()) {

            @Override
            public SearchTask createTask(long id, String type, String action, TaskId parentTaskId, Map<String, String> headers) {
                AsynchronousSearchTask asynchronousSearchTask = new AsynchronousSearchTask(id, type, AsynchronousSearchTask.NAME, parentTaskId, headers, (AsynchronousSearchActiveContext) context, request, asynchronousSearchService::onCancelledFreeActiveContext);
                asynchronousSearchService.bootstrapSearch(asynchronousSearchTask, context.getContextId());
                PrioritizedActionListener<AsynchronousSearchResponse> wrappedListener = AsynchronousSearchTimeoutWrapper.wrapScheduledTimeout(threadPool, request.getWaitForCompletionTimeout(), AsynchronousSearchPlugin.OPEN_DISTRO_ASYNC_SEARCH_GENERIC_THREAD_POOL_NAME, listener, (actionListener) -> {
                    progressListener.searchProgressActionListener().removeListener(actionListener);
                    listener.onResponse(context.getAsynchronousSearchResponse());
                });
                progressListener.searchProgressActionListener().addOrExecuteListener(wrappedListener);
                return asynchronousSearchTask;
            }
        };
        // set the parent task as the submit task for cancellation on connection close
        searchRequest.setParentTask(task.taskInfo(clusterService.localNode().getId(), false).getTaskId());
        transportSearchAction.execute(searchRequest, progressListener);
    } catch (Exception e) {
        logger.error(() -> new ParameterizedMessage("Failed to submit asynchronous search request [{}]", request), e);
        if (asynchronousSearchContext != null) {
            AsynchronousSearchActiveContext asynchronousSearchActiveContext = (AsynchronousSearchActiveContext) asynchronousSearchContext;
            asynchronousSearchService.freeContext(asynchronousSearchActiveContext.getAsynchronousSearchId(), asynchronousSearchActiveContext.getContextId(), user, ActionListener.wrap((r) -> {
                logger.debug(() -> new ParameterizedMessage("Successfully cleaned up context on submit asynchronous" + " search [{}] on failure", asynchronousSearchActiveContext.getAsynchronousSearchId()), e);
                listener.onFailure(e);
            }, (ex) -> {
                logger.debug(() -> new ParameterizedMessage("Failed to cleaned up context on submit asynchronous search" + " [{}] on failure", asynchronousSearchActiveContext.getAsynchronousSearchId()), ex);
                listener.onFailure(e);
            }));
        } else {
            listener.onFailure(e);
        }
    }
}