Examples with ThreadContext org.opensearch.common.util.concurrent.ThreadContext used on opensource projects

Example 86 with ThreadContext

use of org.opensearch.common.util.concurrent.ThreadContext in project OpenSearch by opensearch-project.

the class ContextPreservingActionListenerTests method testOriginalContextIsPreservedAfterOnResponse.

public void testOriginalContextIsPreservedAfterOnResponse() throws IOException {
    ThreadContext threadContext = new ThreadContext(Settings.EMPTY);
    final boolean nonEmptyContext = randomBoolean();
    if (nonEmptyContext) {
        threadContext.putHeader("not empty", "value");
    }
    final ContextPreservingActionListener<Void> actionListener;
    try (ThreadContext.StoredContext ignore = threadContext.stashContext()) {
        threadContext.putHeader("foo", "bar");
        final ActionListener<Void> delegate = new ActionListener<Void>() {

            @Override
            public void onResponse(Void aVoid) {
                assertEquals("bar", threadContext.getHeader("foo"));
                assertNull(threadContext.getHeader("not empty"));
            }

            @Override
            public void onFailure(Exception e) {
                throw new RuntimeException("onFailure shouldn't be called", e);
            }
        };
        if (randomBoolean()) {
            actionListener = new ContextPreservingActionListener<>(threadContext.newRestorableContext(true), delegate);
        } else {
            actionListener = ContextPreservingActionListener.wrapPreservingContext(delegate, threadContext);
        }
    }
    assertNull(threadContext.getHeader("foo"));
    assertEquals(nonEmptyContext ? "value" : null, threadContext.getHeader("not empty"));
    actionListener.onResponse(null);
    assertNull(threadContext.getHeader("foo"));
    assertEquals(nonEmptyContext ? "value" : null, threadContext.getHeader("not empty"));
}

Example 87 with ThreadContext

use of org.opensearch.common.util.concurrent.ThreadContext in project OpenSearch by opensearch-project.

the class ContextPreservingActionListenerTests method testOriginalContextIsWhenListenerThrows.

public void testOriginalContextIsWhenListenerThrows() throws Exception {
    ThreadContext threadContext = new ThreadContext(Settings.EMPTY);
    final boolean nonEmptyContext = randomBoolean();
    if (nonEmptyContext) {
        threadContext.putHeader("not empty", "value");
    }
    final ContextPreservingActionListener<Void> actionListener;
    try (ThreadContext.StoredContext ignore = threadContext.stashContext()) {
        threadContext.putHeader("foo", "bar");
        final ActionListener<Void> delegate = new ActionListener<Void>() {

            @Override
            public void onResponse(Void aVoid) {
                assertEquals("bar", threadContext.getHeader("foo"));
                assertNull(threadContext.getHeader("not empty"));
                throw new RuntimeException("onResponse called");
            }

            @Override
            public void onFailure(Exception e) {
                assertEquals("bar", threadContext.getHeader("foo"));
                assertNull(threadContext.getHeader("not empty"));
                throw new RuntimeException("onFailure called");
            }
        };
        if (randomBoolean()) {
            actionListener = new ContextPreservingActionListener<>(threadContext.newRestorableContext(true), delegate);
        } else {
            actionListener = ContextPreservingActionListener.wrapPreservingContext(delegate, threadContext);
        }
    }
    assertNull(threadContext.getHeader("foo"));
    assertEquals(nonEmptyContext ? "value" : null, threadContext.getHeader("not empty"));
    RuntimeException e = expectThrows(RuntimeException.class, () -> actionListener.onResponse(null));
    assertEquals("onResponse called", e.getMessage());
    assertNull(threadContext.getHeader("foo"));
    assertEquals(nonEmptyContext ? "value" : null, threadContext.getHeader("not empty"));
    e = expectThrows(RuntimeException.class, () -> actionListener.onFailure(null));
    assertEquals("onFailure called", e.getMessage());
    assertNull(threadContext.getHeader("foo"));
    assertEquals(nonEmptyContext ? "value" : null, threadContext.getHeader("not empty"));
}

Example 88 with ThreadContext

use of org.opensearch.common.util.concurrent.ThreadContext in project OpenSearch by opensearch-project.

the class ContextPreservingActionListenerTests method testToStringIncludesDelegate.

public void testToStringIncludesDelegate() {
    ThreadContext threadContext = new ThreadContext(Settings.EMPTY);
    final ContextPreservingActionListener<Void> actionListener;
    try (ThreadContext.StoredContext ignore = threadContext.stashContext()) {
        final ActionListener<Void> delegate = new ActionListener<Void>() {

            @Override
            public void onResponse(Void aVoid) {
            }

            @Override
            public void onFailure(Exception e) {
            }

            @Override
            public String toString() {
                return "test delegate";
            }
        };
        actionListener = ContextPreservingActionListener.wrapPreservingContext(delegate, threadContext);
    }
    assertThat(actionListener.toString(), allOf(containsString("test delegate"), containsString("ContextPreservingActionListener")));
}

Example 89 with ThreadContext

use of org.opensearch.common.util.concurrent.ThreadContext in project OpenSearch by opensearch-project.

the class RetentionLeaseBackgroundSyncAction method backgroundSync.

final void backgroundSync(ShardId shardId, String primaryAllocationId, long primaryTerm, RetentionLeases retentionLeases) {
    final ThreadContext threadContext = threadPool.getThreadContext();
    try (ThreadContext.StoredContext ignore = threadContext.stashContext()) {
        // we have to execute under the system context so that if security is enabled the sync is authorized
        threadContext.markAsSystemContext();
        final Request request = new Request(shardId, retentionLeases);
        final ReplicationTask task = (ReplicationTask) taskManager.register("transport", "retention_lease_background_sync", request);
        transportService.sendChildRequest(clusterService.localNode(), transportPrimaryAction, new ConcreteShardRequest<>(request, primaryAllocationId, primaryTerm), task, transportOptions, new TransportResponseHandler<ReplicationResponse>() {

            @Override
            public ReplicationResponse read(StreamInput in) throws IOException {
                return newResponseInstance(in);
            }

            @Override
            public String executor() {
                return ThreadPool.Names.SAME;
            }

            @Override
            public void handleResponse(ReplicationResponse response) {
                task.setPhase("finished");
                taskManager.unregister(task);
            }

            @Override
            public void handleException(TransportException e) {
                task.setPhase("finished");
                taskManager.unregister(task);
                if (ExceptionsHelper.unwrap(e, NodeClosedException.class) != null) {
                    // node shutting down
                    return;
                }
                if (ExceptionsHelper.unwrap(e, IndexNotFoundException.class, AlreadyClosedException.class, IndexShardClosedException.class) != null) {
                    // the index was deleted or the shard is closed
                    return;
                }
                getLogger().warn(new ParameterizedMessage("{} retention lease background sync failed", shardId), e);
            }
        });
    }
}

Example 90 with ThreadContext

use of org.opensearch.common.util.concurrent.ThreadContext 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);
    }
}