Examples with ContextPreservingActionListener org.opensearch.action.support.ContextPreservingActionListener used on opensource projects

Example 1 with ContextPreservingActionListener

use of org.opensearch.action.support.ContextPreservingActionListener in project OpenSearch by opensearch-project.

the class IndexShardOperationPermits method acquire.

private void acquire(final ActionListener<Releasable> onAcquired, final String executorOnDelay, final boolean forceExecution, final Object debugInfo, final StackTraceElement[] stackTrace) {
    if (closed) {
        onAcquired.onFailure(new IndexShardClosedException(shardId));
        return;
    }
    final Releasable releasable;
    try {
        synchronized (this) {
            if (queuedBlockOperations > 0) {
                final Supplier<StoredContext> contextSupplier = threadPool.getThreadContext().newRestorableContext(false);
                final ActionListener<Releasable> wrappedListener;
                if (executorOnDelay != null) {
                    wrappedListener = ActionListener.delegateFailure(new ContextPreservingActionListener<>(contextSupplier, onAcquired), (l, r) -> threadPool.executor(executorOnDelay).execute(new ActionRunnable<Releasable>(l) {

                        @Override
                        public boolean isForceExecution() {
                            return forceExecution;
                        }

                        @Override
                        protected void doRun() {
                            listener.onResponse(r);
                        }

                        @Override
                        public void onRejection(Exception e) {
                            IOUtils.closeWhileHandlingException(r);
                            super.onRejection(e);
                        }
                    }));
                } else {
                    wrappedListener = new ContextPreservingActionListener<>(contextSupplier, onAcquired);
                }
                delayedOperations.add(new DelayedOperation(wrappedListener, debugInfo, stackTrace));
                return;
            } else {
                releasable = acquire(debugInfo, stackTrace);
            }
        }
    } catch (final InterruptedException e) {
        onAcquired.onFailure(e);
        return;
    }
    // execute this outside the synchronized block!
    onAcquired.onResponse(releasable);
}

Example 2 with ContextPreservingActionListener

use of org.opensearch.action.support.ContextPreservingActionListener in project OpenSearch by opensearch-project.

the class RemoteClusterConnection method collectNodes.

/**
 * Collects all nodes on the connected cluster and returns / passes a nodeID to {@link DiscoveryNode} lookup function
 * that returns <code>null</code> if the node ID is not found.
 *
 * The requests to get cluster state on the connected cluster are made in the system context because logically
 * they are equivalent to checking a single detail in the local cluster state and should not require that the
 * user who made the request that is using this method in its implementation is authorized to view the entire
 * cluster state.
 */
void collectNodes(ActionListener<Function<String, DiscoveryNode>> listener) {
    Runnable runnable = () -> {
        final ThreadContext threadContext = threadPool.getThreadContext();
        final ContextPreservingActionListener<Function<String, DiscoveryNode>> contextPreservingActionListener = new ContextPreservingActionListener<>(threadContext.newRestorableContext(false), listener);
        try (ThreadContext.StoredContext ignore = threadContext.stashContext()) {
            // we stash any context here since this is an internal execution and should not leak any existing context information
            threadContext.markAsSystemContext();
            final ClusterStateRequest request = new ClusterStateRequest();
            request.clear();
            request.nodes(true);
            // run this on the node that gets the request it's as good as any other
            request.local(true);
            Transport.Connection connection = remoteConnectionManager.getAnyRemoteConnection();
            transportService.sendRequest(connection, ClusterStateAction.NAME, request, TransportRequestOptions.EMPTY, new TransportResponseHandler<ClusterStateResponse>() {

                @Override
                public ClusterStateResponse read(StreamInput in) throws IOException {
                    return new ClusterStateResponse(in);
                }

                @Override
                public void handleResponse(ClusterStateResponse response) {
                    DiscoveryNodes nodes = response.getState().nodes();
                    contextPreservingActionListener.onResponse(nodes::get);
                }

                @Override
                public void handleException(TransportException exp) {
                    contextPreservingActionListener.onFailure(exp);
                }

                @Override
                public String executor() {
                    return ThreadPool.Names.SAME;
                }
            });
        }
    };
    try {
        // just in case if we are not connected for some reason we try to connect and if we fail we have to notify the listener
        // this will cause some back pressure on the search end and eventually will cause rejections but that's fine
        // we can't proceed with a search on a cluster level.
        // in the future we might want to just skip the remote nodes in such a case but that can already be implemented on the
        // caller end since they provide the listener.
        ensureConnected(ActionListener.wrap((x) -> runnable.run(), listener::onFailure));
    } catch (Exception ex) {
        listener.onFailure(ex);
    }
}