Introduce utils for _really_ stashing the thread context (#114786)

DaveCTurner · web-flow · commit 6ec7a9ab23f8 · 2024-10-15T15:27:07.000+01:00
`ThreadContext#stashContext` does not yield a completely fresh context:
it preserves headers related to tracing the original request. That may
be appropriate in many situations, but sometimes we really do want to
detach processing entirely from the original task. This commit
introduces new utilities to do that.
diff --git a/server/src/main/java/org/elasticsearch/action/admin/cluster/remote/RemoteClusterNodesAction.java b/server/src/main/java/org/elasticsearch/action/admin/cluster/remote/RemoteClusterNodesAction.java
@@ -115,8 +115,7 @@ protected void doExecute(Task task, Request request, ActionListener<Response> li
         }
 
         private void executeWithSystemContext(Request request, ThreadContext threadContext, ActionListener<Response> listener) {
-            try (var ignore = threadContext.stashContext()) {
-                threadContext.markAsSystemContext();
+            try (var ignore = threadContext.newEmptySystemContext()) {
                 if (request.remoteClusterServer) {
                     final NodesInfoRequest nodesInfoRequest = new NodesInfoRequest().clear()
                         .addMetrics(NodesInfoMetrics.Metric.REMOTE_CLUSTER_SERVER.metricName());
diff --git a/server/src/main/java/org/elasticsearch/action/bulk/IncrementalBulkService.java b/server/src/main/java/org/elasticsearch/action/bulk/IncrementalBulkService.java
@@ -141,8 +141,7 @@ public void addItems(List<DocWriteRequest<?>> items, Releasable releasable, Runn
                     if (shouldBackOff()) {
                         final boolean isFirstRequest = incrementalRequestSubmitted == false;
                         incrementalRequestSubmitted = true;
-                        try (ThreadContext.StoredContext ignored = threadContext.stashContext()) {
-                            requestContext.restore();
+                        try (var ignored = threadContext.restoreExistingContext(requestContext)) {
                             final ArrayList<Releasable> toRelease = new ArrayList<>(releasables);
                             releasables.clear();
                             bulkInProgress = true;
@@ -188,8 +187,7 @@ public void lastItems(List<DocWriteRequest<?>> items, Releasable releasable, Act
             } else {
                 assert bulkRequest != null;
                 if (internalAddItems(items, releasable)) {
-                    try (ThreadContext.StoredContext ignored = threadContext.stashContext()) {
-                        requestContext.restore();
+                    try (var ignored = threadContext.restoreExistingContext(requestContext)) {
                         final ArrayList<Releasable> toRelease = new ArrayList<>(releasables);
                         releasables.clear();
                         // We do not need to set this back to false as this will be the last request.
diff --git a/server/src/main/java/org/elasticsearch/cluster/coordination/CoordinationDiagnosticsService.java b/server/src/main/java/org/elasticsearch/cluster/coordination/CoordinationDiagnosticsService.java
@@ -30,7 +30,6 @@
 import org.elasticsearch.common.settings.Setting;
 import org.elasticsearch.common.util.concurrent.EsRejectedExecutionException;
 import org.elasticsearch.common.util.concurrent.ListenableFuture;
-import org.elasticsearch.common.util.concurrent.ThreadContext;
 import org.elasticsearch.core.Nullable;
 import org.elasticsearch.core.Releasable;
 import org.elasticsearch.core.Releasables;
@@ -194,9 +193,7 @@ public void start() {
          * system context.
          */
         if (clusterService.localNode().isMasterNode() == false) {
-            final ThreadContext threadContext = transportService.getThreadPool().getThreadContext();
-            try (ThreadContext.StoredContext ignored = threadContext.stashContext()) {
-                threadContext.markAsSystemContext();
+            try (var ignored = transportService.getThreadPool().getThreadContext().newEmptySystemContext()) {
                 beginPollingRemoteMasterStabilityDiagnostic();
             }
         }
diff --git a/server/src/main/java/org/elasticsearch/cluster/coordination/LagDetector.java b/server/src/main/java/org/elasticsearch/cluster/coordination/LagDetector.java
@@ -24,7 +24,6 @@
 import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.common.util.concurrent.AbstractRunnable;
 import org.elasticsearch.common.util.concurrent.PrioritizedThrottledTaskRunner;
-import org.elasticsearch.common.util.concurrent.ThreadContext;
 import org.elasticsearch.core.Releasable;
 import org.elasticsearch.core.Releasables;
 import org.elasticsearch.core.TimeValue;
@@ -271,9 +270,7 @@ public void onFailure(Exception e) {
                     @Override
                     public void onResponse(Releasable releasable) {
                         boolean success = false;
-                        final ThreadContext threadContext = transportService.getThreadPool().getThreadContext();
-                        try (ThreadContext.StoredContext ignored = threadContext.stashContext()) {
-                            threadContext.markAsSystemContext();
+                        try (var ignored = transportService.getThreadPool().getThreadContext().newEmptySystemContext()) {
                             client.execute(
                                 TransportNodesHotThreadsAction.TYPE,
                                 new NodesHotThreadsRequest(
diff --git a/server/src/main/java/org/elasticsearch/cluster/service/ClusterApplierService.java b/server/src/main/java/org/elasticsearch/cluster/service/ClusterApplierService.java
@@ -334,8 +334,7 @@ private void submitStateUpdateTask(
         final ThreadContext threadContext = threadPool.getThreadContext();
         final Supplier<ThreadContext.StoredContext> storedContextSupplier = threadContext.newRestorableContext(true);
 
-        try (ThreadContext.StoredContext ignore = threadContext.stashContext()) {
-            threadContext.markAsSystemContext();
+        try (var ignore = threadContext.newEmptySystemContext()) {
             threadPoolExecutor.execute(
                 new UpdateTask(
                     priority,
diff --git a/server/src/main/java/org/elasticsearch/common/util/concurrent/ThreadContext.java b/server/src/main/java/org/elasticsearch/common/util/concurrent/ThreadContext.java
@@ -101,7 +101,7 @@ public ThreadContext(Settings settings) {
     }
 
     /**
-     * Removes the current context and resets a default context. The removed context can be
+     * Removes the current context and resets a default context except for headers involved in task tracing. The removed context can be
      * restored by closing the returned {@link StoredContext}.
      * @return a stored context that will restore the current context to its state at the point this method was called
      */
@@ -159,6 +159,28 @@ public StoredContext stashContextPreservingRequestHeaders(final String... reques
         return stashContextPreservingRequestHeaders(Set.of(requestHeaders));
     }
 
+    /**
+     * Removes the current context and replaces it with a completely empty default context, detaching execution entirely from the calling
+     * context. The calling context can be restored by closing the returned {@link StoredContext}. Similar to {@link #stashContext()} except
+     * that this method does not even preserve tracing-related headers.
+     */
+    public StoredContext newEmptyContext() {
+        final var callingContext = threadLocal.get();
+        threadLocal.set(DEFAULT_CONTEXT);
+        return storedOriginalContext(callingContext);
+    }
+
+    /**
+     * Removes the current context and replaces it with a completely empty system context, detaching execution entirely from the calling
+     * context. The calling context can be restored by closing the returned {@link StoredContext}. Similar to {@link #stashContext()} except
+     * that this method does not even preserve tracing-related headers.
+     */
+    public StoredContext newEmptySystemContext() {
+        final var callingContext = threadLocal.get();
+        threadLocal.set(DEFAULT_CONTEXT.setSystemContext());
+        return storedOriginalContext(callingContext);
+    }
+
     /**
      * When using a {@link org.elasticsearch.telemetry.tracing.Tracer} to capture activity in Elasticsearch, when a parent span is already
      * in progress, it is necessary to start a new context before beginning a child span. This method creates a context,
@@ -330,6 +352,16 @@ public StoredContext newStoredContextPreservingResponseHeaders() {
         };
     }
 
+    /**
+     * Capture the current context and then restore the given context, returning a {@link StoredContext} that reverts back to the current
+     * context again. Equivalent to using {@link #newStoredContext()} and then calling {@code existingContext.restore()}.
+     */
+    public StoredContext restoreExistingContext(StoredContext existingContext) {
+        final var originalContext = threadLocal.get();
+        existingContext.restore();
+        return storedOriginalContext(originalContext);
+    }
+
     /**
      * Just like {@link #stashContext()} but no default context is set.
      */
@@ -914,14 +946,13 @@ private class ContextPreservingRunnable implements WrappedRunnable {
         private final ThreadContext.StoredContext ctx;
 
         private ContextPreservingRunnable(Runnable in) {
-            ctx = newStoredContext();
+            this.ctx = newStoredContext();
             this.in = in;
         }
 
         @Override
         public void run() {
-            try (ThreadContext.StoredContext ignore = stashContext()) {
-                ctx.restore();
+            try (var ignore = restoreExistingContext(ctx)) {
                 in.run();
             }
         }
diff --git a/server/src/main/java/org/elasticsearch/index/seqno/RetentionLeaseBackgroundSyncAction.java b/server/src/main/java/org/elasticsearch/index/seqno/RetentionLeaseBackgroundSyncAction.java
@@ -23,7 +23,6 @@
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.common.settings.Settings;
-import org.elasticsearch.common.util.concurrent.ThreadContext;
 import org.elasticsearch.index.shard.IndexShard;
 import org.elasticsearch.index.shard.ShardId;
 import org.elasticsearch.indices.IndicesService;
@@ -95,14 +94,9 @@ protected void doExecute(Task task, Request request, ActionListener<ReplicationR
     }
 
     final void backgroundSync(ShardId shardId, String primaryAllocationId, long primaryTerm, RetentionLeases retentionLeases) {
-        final ThreadContext threadContext = threadPool.getThreadContext();
-        try (ThreadContext.StoredContext ignore = threadContext.stashContext()) {
+        try (var ignore = threadPool.getThreadContext().newEmptySystemContext()) {
             // 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);
-            try (var ignored = threadContext.newTraceContext()) {
-                sendRetentionLeaseSyncAction(shardId, primaryAllocationId, primaryTerm, request);
-            }
+            sendRetentionLeaseSyncAction(shardId, primaryAllocationId, primaryTerm, new Request(shardId, retentionLeases));
         }
     }
 
diff --git a/server/src/main/java/org/elasticsearch/index/seqno/RetentionLeaseSyncAction.java b/server/src/main/java/org/elasticsearch/index/seqno/RetentionLeaseSyncAction.java
@@ -29,7 +29,6 @@
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.common.settings.Settings;
-import org.elasticsearch.common.util.concurrent.ThreadContext;
 import org.elasticsearch.index.IndexNotFoundException;
 import org.elasticsearch.index.IndexingPressure;
 import org.elasticsearch.index.shard.IndexShard;
@@ -111,47 +110,43 @@ final void sync(
         RetentionLeases retentionLeases,
         ActionListener<ReplicationResponse> listener
     ) {
-        final ThreadContext threadContext = threadPool.getThreadContext();
-        try (ThreadContext.StoredContext ignore = threadContext.stashContext()) {
+        try (var ignore = threadPool.getThreadContext().newEmptySystemContext()) {
             // 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);
-            try (var ignored = threadContext.newTraceContext()) {
-                final ReplicationTask task = (ReplicationTask) taskManager.register("transport", "retention_lease_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 Executor executor() {
-                            return TransportResponseHandler.TRANSPORT_WORKER;
-                        }
-
-                        @Override
-                        public void handleResponse(ReplicationResponse response) {
-                            task.setPhase("finished");
-                            taskManager.unregister(task);
-                            listener.onResponse(response);
-                        }
-
-                        @Override
-                        public void handleException(TransportException e) {
-                            LOGGER.log(getExceptionLogLevel(e), () -> format("%s retention lease sync failed", shardId), e);
-                            task.setPhase("finished");
-                            taskManager.unregister(task);
-                            listener.onFailure(e);
-                        }
+            final ReplicationTask task = (ReplicationTask) taskManager.register("transport", "retention_lease_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 Executor executor() {
+                        return TransportResponseHandler.TRANSPORT_WORKER;
+                    }
+
+                    @Override
+                    public void handleResponse(ReplicationResponse response) {
+                        task.setPhase("finished");
+                        taskManager.unregister(task);
+                        listener.onResponse(response);
+                    }
+
+                    @Override
+                    public void handleException(TransportException e) {
+                        LOGGER.log(getExceptionLogLevel(e), () -> format("%s retention lease sync failed", shardId), e);
+                        task.setPhase("finished");
+                        taskManager.unregister(task);
+                        listener.onFailure(e);
+                    }
+                }
+            );
         }
     }
 
diff --git a/server/src/main/java/org/elasticsearch/index/shard/RefreshListeners.java b/server/src/main/java/org/elasticsearch/index/shard/RefreshListeners.java
@@ -147,8 +147,7 @@ public boolean addOrNotify(Translog.Location location, Consumer<Boolean> listene
             if (refreshForcers == 0 && roomForListener(maxRefreshes, listeners, checkpointRefreshListeners)) {
                 ThreadContext.StoredContext storedContext = threadContext.newStoredContextPreservingResponseHeaders();
                 Consumer<Boolean> contextPreservingListener = forced -> {
-                    try (ThreadContext.StoredContext ignore = threadContext.stashContext()) {
-                        storedContext.restore();
+                    try (var ignore = threadContext.restoreExistingContext(storedContext)) {
                         listener.accept(forced);
                     }
                 };
diff --git a/server/src/main/java/org/elasticsearch/transport/InboundHandler.java b/server/src/main/java/org/elasticsearch/transport/InboundHandler.java
@@ -109,7 +109,8 @@ private void messageReceived(TcpChannel channel, InboundMessage message, long st
 
         TransportResponseHandler<?> responseHandler = null;
         ThreadContext threadContext = threadPool.getThreadContext();
-        try (ThreadContext.StoredContext existing = threadContext.stashContext()) {
+        assert threadContext.isDefaultContext();
+        try (var ignored = threadContext.newStoredContext()) {
             // Place the context with the headers from the message
             threadContext.setHeaders(header.getHeaders());
             threadContext.putTransient("_remote_address", remoteAddress);
diff --git a/server/src/main/java/org/elasticsearch/transport/OutboundHandler.java b/server/src/main/java/org/elasticsearch/transport/OutboundHandler.java
@@ -23,7 +23,6 @@
 import org.elasticsearch.common.network.HandlingTimeTracker;
 import org.elasticsearch.common.recycler.Recycler;
 import org.elasticsearch.common.transport.NetworkExceptionHelper;
-import org.elasticsearch.common.util.concurrent.ThreadContext;
 import org.elasticsearch.core.Nullable;
 import org.elasticsearch.core.Releasable;
 import org.elasticsearch.core.Releasables;
@@ -237,7 +236,7 @@ private void internalSend(
         final long messageSize = reference.length();
         TransportLogger.logOutboundMessage(channel, reference);
         // stash thread context so that channel event loop is not polluted by thread context
-        try (ThreadContext.StoredContext existing = threadPool.getThreadContext().stashContext()) {
+        try (var ignored = threadPool.getThreadContext().newEmptyContext()) {
             channel.sendMessage(reference, new ActionListener<>() {
                 @Override
                 public void onResponse(Void v) {
diff --git a/server/src/main/java/org/elasticsearch/transport/RemoteClusterConnection.java b/server/src/main/java/org/elasticsearch/transport/RemoteClusterConnection.java
@@ -129,9 +129,8 @@ void collectNodes(ActionListener<Function<String, DiscoveryNode>> listener) {
             final ThreadContext threadContext = threadPool.getThreadContext();
             final ContextPreservingActionListener<Function<String, DiscoveryNode>> contextPreservingActionListener =
                 new ContextPreservingActionListener<>(threadContext.newRestorableContext(false), listener);
-            try (ThreadContext.StoredContext ignore = threadContext.stashContext()) {
+            try (var ignore = threadContext.newEmptySystemContext()) {
                 // we stash any context here since this is an internal execution and should not leak any existing context information
-                threadContext.markAsSystemContext();
                 Transport.Connection connection = remoteConnectionManager.getAnyRemoteConnection();
 
                 // Use different action to collect nodes information depending on the connection model
diff --git a/server/src/main/java/org/elasticsearch/transport/SniffConnectionStrategy.java b/server/src/main/java/org/elasticsearch/transport/SniffConnectionStrategy.java
@@ -328,10 +328,8 @@ private void collectRemoteNodes(Iterator<Supplier<DiscoveryNode>> seedNodesSuppl
                     sniffResponseHandler = new ClusterStateSniffResponseHandler(connection, listener, seedNodesSuppliers);
                 }
 
-                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();
+                try (var ignored = threadContext.newEmptySystemContext()) {
+                    // we stash any context here since this is an internal execution and should not leak any existing context information.
                     transportService.sendRequest(
                         connection,
                         action,
diff --git a/server/src/test/java/org/elasticsearch/common/util/concurrent/ThreadContextTests.java b/server/src/test/java/org/elasticsearch/common/util/concurrent/ThreadContextTests.java
diff --git a/x-pack/plugin/ccr/src/main/java/org/elasticsearch/xpack/ccr/CcrLicenseChecker.java b/x-pack/plugin/ccr/src/main/java/org/elasticsearch/xpack/ccr/CcrLicenseChecker.java
diff --git a/x-pack/plugin/ccr/src/main/java/org/elasticsearch/xpack/ccr/action/TransportUnfollowAction.java b/x-pack/plugin/ccr/src/main/java/org/elasticsearch/xpack/ccr/action/TransportUnfollowAction.java
diff --git a/x-pack/plugin/ccr/src/main/java/org/elasticsearch/xpack/ccr/repository/CcrRepository.java b/x-pack/plugin/ccr/src/main/java/org/elasticsearch/xpack/ccr/repository/CcrRepository.java
diff --git a/x-pack/plugin/core/src/main/java/org/elasticsearch/license/RemoteClusterLicenseChecker.java b/x-pack/plugin/core/src/main/java/org/elasticsearch/license/RemoteClusterLicenseChecker.java
diff --git a/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ClientHelper.java b/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ClientHelper.java

Original file line number	Diff line number	Diff line change
`@@ -115,8 +115,7 @@ protected void doExecute(Task task, Request request, ActionListener<Response> li`
`115`	`115`	`}`
`116`	`116`
`117`	`117`	`private void executeWithSystemContext(Request request, ThreadContext threadContext, ActionListener<Response> listener) {`
`118`		`- try (var ignore = threadContext.stashContext()) {`
`119`		`- threadContext.markAsSystemContext();`
	`118`	`+ try (var ignore = threadContext.newEmptySystemContext()) {`
`120`	`119`	`if (request.remoteClusterServer) {`
`121`	`120`	`final NodesInfoRequest nodesInfoRequest = new NodesInfoRequest().clear()`
`122`	`121`	`.addMetrics(NodesInfoMetrics.Metric.REMOTE_CLUSTER_SERVER.metricName());`