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

`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.

server/src/main/java/org/elasticsearch/action/admin/cluster/remote/RemoteClusterNodesAction.java

@@ -115,8 +115,7 @@ public class RemoteClusterNodesAction {
         }
 
         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());

server/src/main/java/org/elasticsearch/action/bulk/IncrementalBulkService.java

@@ -141,8 +141,7 @@ public class IncrementalBulkService {
                     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 class IncrementalBulkService {
             } 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.

server/src/main/java/org/elasticsearch/cluster/coordination/CoordinationDiagnosticsService.java

@@ -30,7 +30,6 @@ import org.elasticsearch.common.io.stream.Writeable;
 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 class CoordinationDiagnosticsService implements ClusterStateListener {
          * 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();
             }
         }

server/src/main/java/org/elasticsearch/cluster/coordination/LagDetector.java

@@ -24,7 +24,6 @@ import org.elasticsearch.common.settings.Setting;
 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 class LagDetector {
                     @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(

server/src/main/java/org/elasticsearch/cluster/service/ClusterApplierService.java

@@ -334,8 +334,7 @@ public class ClusterApplierService extends AbstractLifecycleComponent implements
         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,

server/src/main/java/org/elasticsearch/common/util/concurrent/ThreadContext.java

@@ -101,7 +101,7 @@ public final class ThreadContext implements Writeable, TraceContext {
     }
 
     /**
-     * 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 final class ThreadContext implements Writeable, TraceContext {
         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 final class ThreadContext implements Writeable, TraceContext {
         };
     }
 
+    /**
+     * 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 @@ public final class ThreadContext implements Writeable, TraceContext {
         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();
             }
         }

server/src/main/java/org/elasticsearch/index/seqno/RetentionLeaseBackgroundSyncAction.java

@@ -23,7 +23,6 @@ import org.elasticsearch.cluster.service.ClusterService;
 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 @@ public class RetentionLeaseBackgroundSyncAction extends TransportReplicationActi
     }
 
     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));
         }
     }
 

server/src/main/java/org/elasticsearch/index/seqno/RetentionLeaseSyncAction.java

@@ -29,7 +29,6 @@ import org.elasticsearch.cluster.service.ClusterService;
 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 @@ public class RetentionLeaseSyncAction extends TransportWriteAction<
         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);
+                    }
+                }
+            );
         }
     }
 

server/src/main/java/org/elasticsearch/index/shard/RefreshListeners.java

@@ -147,8 +147,7 @@ public final class RefreshListeners implements ReferenceManager.RefreshListener,
             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);
                     }
                 };

server/src/main/java/org/elasticsearch/transport/InboundHandler.java

@@ -109,7 +109,8 @@ public class InboundHandler {
 
         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);

server/src/main/java/org/elasticsearch/transport/OutboundHandler.java

@@ -23,7 +23,6 @@ import org.elasticsearch.common.network.CloseableChannel;
 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 @@ final class OutboundHandler {
         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) {

server/src/main/java/org/elasticsearch/transport/RemoteClusterConnection.java

@@ -129,9 +129,8 @@ public final class RemoteClusterConnection implements Closeable {
             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

server/src/main/java/org/elasticsearch/transport/SniffConnectionStrategy.java

@@ -328,10 +328,8 @@ public class SniffConnectionStrategy extends RemoteConnectionStrategy {
                     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,

server/src/test/java/org/elasticsearch/common/util/concurrent/ThreadContextTests.java

@@ -1122,6 +1122,57 @@ public class ThreadContextTests extends ESTestCase {
         }
     }
 
+    public void testNewEmptyContext() {
+        final var threadContext = new ThreadContext(Settings.EMPTY);
+        final var header = randomBoolean() ? randomIdentifier() : randomFrom(HEADERS_TO_COPY);
+        threadContext.putHeader(header, randomIdentifier());
+
+        try (var ignored = threadContext.newEmptyContext()) {
+            assertTrue(threadContext.isDefaultContext());
+            assertNull(threadContext.getHeader(header));
+            assertTrue(threadContext.getHeaders().isEmpty());
+        }
+
+        assertNotNull(threadContext.getHeader(header));
+    }
+
+    public void testNewEmptySystemContext() {
+        final var threadContext = new ThreadContext(Settings.EMPTY);
+        final var header = randomBoolean() ? randomIdentifier() : randomFrom(HEADERS_TO_COPY);
+        threadContext.putHeader(header, randomIdentifier());
+
+        try (var ignored = threadContext.newEmptySystemContext()) {
+            assertTrue(threadContext.isSystemContext());
+            assertNull(threadContext.getHeader(header));
+            assertTrue(threadContext.getHeaders().isEmpty());
+        }
+
+        assertNotNull(threadContext.getHeader(header));
+    }
+
+    public void testRestoreExistingContext() {
+        final var threadContext = new ThreadContext(Settings.EMPTY);
+        final var header = randomIdentifier();
+        final var originalValue = randomIdentifier();
+        threadContext.putHeader(header, originalValue);
+        try (var originalContext = threadContext.newStoredContext()) {
+            assertEquals(originalValue, threadContext.getHeader(header));
+
+            try (var ignored1 = threadContext.newEmptyContext()) {
+                final var updatedValue = randomIdentifier();
+                threadContext.putHeader(header, updatedValue);
+
+                try (var ignored2 = threadContext.restoreExistingContext(originalContext)) {
+                    assertEquals(originalValue, threadContext.getHeader(header));
+                }
+
+                assertEquals(updatedValue, threadContext.getHeader(header));
+            }
+
+            assertEquals(originalValue, threadContext.getHeader(header));
+        }
+    }
+
     private String randomCase(String original) {
         int i = randomInt(original.length() - 1);
         StringBuilder sb = new StringBuilder(original);

x-pack/plugin/ccr/src/main/java/org/elasticsearch/xpack/ccr/CcrLicenseChecker.java

@@ -479,8 +479,7 @@ public class CcrLicenseChecker {
                 ActionListener<Response> listener
             ) {
                 final Supplier<ThreadContext.StoredContext> supplier = threadContext.newRestorableContext(false);
-                try (ThreadContext.StoredContext ignore = threadContext.stashContext()) {
-                    threadContext.markAsSystemContext();
+                try (var ignore = threadContext.newEmptySystemContext()) {
                     delegate.execute(connection, action, request, new ContextPreservingActionListener<>(supplier, listener));
                 }
             }

x-pack/plugin/ccr/src/main/java/org/elasticsearch/xpack/ccr/action/TransportUnfollowAction.java

@@ -195,9 +195,8 @@ public class TransportUnfollowAction extends AcknowledgedTransportMasterNodeActi
                     threadContext.newRestorableContext(true),
                     listener
                 );
-                try (ThreadContext.StoredContext ignore = threadPool.getThreadContext().stashContext()) {
+                try (var ignore = threadPool.getThreadContext().newEmptySystemContext()) {
                     // we have to execute under the system context so that if security is enabled the removal is authorized
-                    threadContext.markAsSystemContext();
                     CcrRetentionLeases.asyncRemoveRetentionLease(leaderShardId, retentionLeaseId, remoteClient, preservedListener);
                 }
             }

x-pack/plugin/ccr/src/main/java/org/elasticsearch/xpack/ccr/repository/CcrRepository.java

@@ -47,7 +47,6 @@ import org.elasticsearch.common.metrics.CounterMetric;
 import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.common.unit.ByteSizeValue;
 import org.elasticsearch.common.util.Maps;
-import org.elasticsearch.common.util.concurrent.ThreadContext;
 import org.elasticsearch.core.CheckedConsumer;
 import org.elasticsearch.core.IOUtils;
 import org.elasticsearch.core.Releasable;
@@ -447,10 +446,8 @@ public class CcrRepository extends AbstractLifecycleComponent implements Reposit
             // schedule renewals to run during the restore
             final Scheduler.Cancellable renewable = threadPool.scheduleWithFixedDelay(() -> {
                 logger.trace("{} background renewal of retention lease [{}] during restore", shardId, retentionLeaseId);
-                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 renewal is authorized
-                    threadContext.markAsSystemContext();
                     CcrRetentionLeases.asyncRenewRetentionLease(
                         leaderShardId,
                         retentionLeaseId,

x-pack/plugin/core/src/main/java/org/elasticsearch/license/RemoteClusterLicenseChecker.java

@@ -214,10 +214,8 @@ public final class RemoteClusterLicenseChecker {
             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();
-
             final XPackInfoRequest request = new XPackInfoRequest();
             request.setCategories(EnumSet.of(XPackInfoRequest.Category.LICENSE));
             try {

x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ClientHelper.java

@@ -267,7 +267,7 @@ public final class ClientHelper {
                 return supplier.get();
             }
         } else {
-            try (ThreadContext.StoredContext ignore = client.threadPool().getThreadContext().stashContext()) {
+            try (var ignore = client.threadPool().getThreadContext().stashContext()) {
                 client.threadPool().getThreadContext().copyHeaders(filteredHeaders.entrySet());
                 return supplier.get();
             }