Preserve context in ResultDeduplicator & snapshotDeletionListeners (#84100)

Today the `ResultDeduplicator` may complete a collection of listeners in
contexts different from the ones in which they were submitted.
`snapshotDeletionListeners` has a similar problem. This commit makes
sure that the context is preserved in these listeners.

Backports #84038, #84089 and #84093 to 8.0 - they could not be
backported separately due to failures.

Author: DaveCTurner

docs/changelog/84038.yaml

@@ -0,0 +1,6 @@
+pr: 84038
+summary: Preserve context in `ResultDeduplicator`
+area: Infra/Core
+type: bug
+issues:
+ - 84036

docs/changelog/84089.yaml

@@ -0,0 +1,6 @@
+pr: 84089
+summary: Preserve context in `snapshotDeletionListeners`
+area: Snapshot/Restore
+type: bug
+issues:
+ - 84036

server/src/main/java/org/elasticsearch/action/ResultDeduplicator.java

@@ -8,7 +8,9 @@
 
 package org.elasticsearch.action;
 
+import org.elasticsearch.action.support.ContextPreservingActionListener;
 import org.elasticsearch.common.util.concurrent.ConcurrentCollections;
+import org.elasticsearch.common.util.concurrent.ThreadContext;
 
 import java.util.ArrayList;
 import java.util.List;
@@ -22,8 +24,14 @@
  */
 public final class ResultDeduplicator<T, R> {
 
+    private final ThreadContext threadContext;
     private final ConcurrentMap<T, CompositeListener> requests = ConcurrentCollections.newConcurrentMap();
 
+    public ResultDeduplicator(ThreadContext threadContext) {
+        assert threadContext != null;
+        this.threadContext = threadContext;
+    }
+
     /**
      * Ensures a given request not executed multiple times when another equal request is already in-flight.
      * If the request is not yet known to the deduplicator it will invoke the passed callback with an {@link ActionListener}
@@ -35,7 +43,8 @@ public final class ResultDeduplicator<T, R> {
      * @param callback Callback to be invoked with request and completion listener the first time the request is added to the deduplicator
      */
     public void executeOnce(T request, ActionListener<R> listener, BiConsumer<T, ActionListener<R>> callback) {
-        ActionListener<R> completionListener = requests.computeIfAbsent(request, CompositeListener::new).addListener(listener);
+        ActionListener<R> completionListener = requests.computeIfAbsent(request, CompositeListener::new)
+            .addListener(ContextPreservingActionListener.wrapPreservingContext(listener, threadContext));
         if (completionListener != null) {
             callback.accept(request, completionListener);
         }

server/src/main/java/org/elasticsearch/cluster/action/shard/ShardStateAction.java

@@ -80,7 +80,7 @@ public class ShardStateAction {
     private final ThreadPool threadPool;
 
     // we deduplicate these shard state requests in order to avoid sending duplicate failed/started shard requests for a shard
-    private final ResultDeduplicator<TransportRequest, Void> remoteShardStateUpdateDeduplicator = new ResultDeduplicator<>();
+    private final ResultDeduplicator<TransportRequest, Void> remoteShardStateUpdateDeduplicator;
 
     @Inject
     public ShardStateAction(
@@ -93,6 +93,7 @@ public ShardStateAction(
         this.transportService = transportService;
         this.clusterService = clusterService;
         this.threadPool = threadPool;
+        this.remoteShardStateUpdateDeduplicator = new ResultDeduplicator<>(threadPool.getThreadContext());
 
         transportService.registerRequestHandler(
             SHARD_STARTED_ACTION_NAME,

server/src/main/java/org/elasticsearch/repositories/blobstore/BlobStoreRepository.java

@@ -402,6 +402,7 @@ protected BlobStoreRepository(
         this.namedXContentRegistry = namedXContentRegistry;
         this.basePath = basePath;
         this.maxSnapshotCount = MAX_SNAPSHOTS_SETTING.get(metadata.settings());
+        this.repoDataDeduplicator = new ResultDeduplicator<>(threadPool.getThreadContext());
     }
 
     @Override
@@ -1864,7 +1865,7 @@ public void clusterStateProcessed(String source, ClusterState oldState, ClusterS
      * {@link #bestEffortConsistency} must be {@code false}, in which case we can assume that the {@link RepositoryData} loaded is
      * unique for a given value of {@link #metadata} at any point in time.
      */
-    private final ResultDeduplicator<RepositoryMetadata, RepositoryData> repoDataDeduplicator = new ResultDeduplicator<>();
+    private final ResultDeduplicator<RepositoryMetadata, RepositoryData> repoDataDeduplicator;
 
     private void doGetRepositoryData(ActionListener<RepositoryData> listener) {
         // Retry loading RepositoryData in a loop in case we run into concurrent modifications of the repository.

server/src/main/java/org/elasticsearch/snapshots/SnapshotShardsService.java

@@ -82,8 +82,7 @@ public class SnapshotShardsService extends AbstractLifecycleComponent implements
     private final Map<Snapshot, Map<ShardId, IndexShardSnapshotStatus>> shardSnapshots = new HashMap<>();
 
     // A map of snapshots to the shardIds that we already reported to the master as failed
-    private final ResultDeduplicator<UpdateIndexShardSnapshotStatusRequest, Void> remoteFailedRequestDeduplicator =
-        new ResultDeduplicator<>();
+    private final ResultDeduplicator<UpdateIndexShardSnapshotStatusRequest, Void> remoteFailedRequestDeduplicator;
 
     public SnapshotShardsService(
         Settings settings,
@@ -97,6 +96,7 @@ public SnapshotShardsService(
         this.transportService = transportService;
         this.clusterService = clusterService;
         this.threadPool = transportService.getThreadPool();
+        this.remoteFailedRequestDeduplicator = new ResultDeduplicator<>(threadPool.getThreadContext());
         if (DiscoveryNode.canContainData(settings)) {
             // this is only useful on the nodes that can hold data
             clusterService.addListener(this);

server/src/main/java/org/elasticsearch/snapshots/SnapshotsService.java

@@ -2286,7 +2286,8 @@ private static boolean isWritingToRepository(SnapshotsInProgress.Entry entry) {
     }
 
     private void addDeleteListener(String deleteUUID, ActionListener<Void> listener) {
-        snapshotDeletionListeners.computeIfAbsent(deleteUUID, k -> new CopyOnWriteArrayList<>()).add(listener);
+        snapshotDeletionListeners.computeIfAbsent(deleteUUID, k -> new CopyOnWriteArrayList<>())
+            .add(ContextPreservingActionListener.wrapPreservingContext(listener, threadPool.getThreadContext()));
     }
 
     /**

server/src/main/java/org/elasticsearch/tasks/TaskCancellationService.java

@@ -44,11 +44,12 @@ public class TaskCancellationService {
     private static final Logger logger = LogManager.getLogger(TaskCancellationService.class);
     private final TransportService transportService;
     private final TaskManager taskManager;
-    private final ResultDeduplicator<CancelRequest, Void> deduplicator = new ResultDeduplicator<>();
+    private final ResultDeduplicator<CancelRequest, Void> deduplicator;
 
     public TaskCancellationService(TransportService transportService) {
         this.transportService = transportService;
         this.taskManager = transportService.getTaskManager();
+        this.deduplicator = new ResultDeduplicator<>(transportService.getThreadPool().getThreadContext());
         transportService.registerRequestHandler(
             BAN_PARENT_ACTION_NAME,
             ThreadPool.Names.SAME,

server/src/test/java/org/elasticsearch/repositories/RepositoriesServiceTests.java

@@ -30,6 +30,7 @@
 import org.elasticsearch.common.component.LifecycleListener;
 import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.common.util.MockBigArrays;
+import org.elasticsearch.common.util.concurrent.ThreadContext;
 import org.elasticsearch.index.shard.ShardId;
 import org.elasticsearch.index.snapshots.IndexShardSnapshotStatus;
 import org.elasticsearch.index.store.Store;
@@ -61,7 +62,9 @@ public class RepositoriesServiceTests extends ESTestCase {
     @Override
     public void setUp() throws Exception {
         super.setUp();
+        ThreadContext threadContext = new ThreadContext(Settings.EMPTY);
         ThreadPool threadPool = mock(ThreadPool.class);
+        when(threadPool.getThreadContext()).thenReturn(threadContext);
         final TransportService transportService = new TransportService(
             Settings.EMPTY,
             mock(Transport.class),

server/src/test/java/org/elasticsearch/tasks/TaskManagerTests.java

@@ -46,6 +46,7 @@
 import static org.hamcrest.Matchers.everyItem;
 import static org.hamcrest.Matchers.in;
 import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
 
 public class TaskManagerTests extends ESTestCase {
     private ThreadPool threadPool;
@@ -76,7 +77,9 @@ public void testResultsServiceRetryTotalTime() {
     public void testTrackingChannelTask() throws Exception {
         final TaskManager taskManager = new TaskManager(Settings.EMPTY, threadPool, Set.of());
         Set<Task> cancelledTasks = ConcurrentCollections.newConcurrentSet();
-        taskManager.setTaskCancellationService(new TaskCancellationService(mock(TransportService.class)) {
+        final var transportServiceMock = mock(TransportService.class);
+        when(transportServiceMock.getThreadPool()).thenReturn(threadPool);
+        taskManager.setTaskCancellationService(new TaskCancellationService(transportServiceMock) {
             @Override
             void cancelTaskAndDescendants(CancellableTask task, String reason, boolean waitForCompletion, ActionListener<Void> listener) {
                 assertThat(reason, equalTo("channel was closed"));
@@ -124,7 +127,9 @@ void cancelTaskAndDescendants(CancellableTask task, String reason, boolean waitF
     public void testTrackingTaskAndCloseChannelConcurrently() throws Exception {
         final TaskManager taskManager = new TaskManager(Settings.EMPTY, threadPool, Set.of());
         Set<CancellableTask> cancelledTasks = ConcurrentCollections.newConcurrentSet();
-        taskManager.setTaskCancellationService(new TaskCancellationService(mock(TransportService.class)) {
+        final var transportServiceMock = mock(TransportService.class);
+        when(transportServiceMock.getThreadPool()).thenReturn(threadPool);
+        taskManager.setTaskCancellationService(new TaskCancellationService(transportServiceMock) {
             @Override
             void cancelTaskAndDescendants(CancellableTask task, String reason, boolean waitForCompletion, ActionListener<Void> listener) {
                 assertTrue("task [" + task + "] was cancelled already", cancelledTasks.add(task));
@@ -180,7 +185,9 @@ void cancelTaskAndDescendants(CancellableTask task, String reason, boolean waitF
 
     public void testRemoveBansOnChannelDisconnects() throws Exception {
         final TaskManager taskManager = new TaskManager(Settings.EMPTY, threadPool, Set.of());
-        taskManager.setTaskCancellationService(new TaskCancellationService(mock(TransportService.class)) {
+        final var transportServiceMock = mock(TransportService.class);
+        when(transportServiceMock.getThreadPool()).thenReturn(threadPool);
+        taskManager.setTaskCancellationService(new TaskCancellationService(transportServiceMock) {
             @Override
             void cancelTaskAndDescendants(CancellableTask task, String reason, boolean waitForCompletion, ActionListener<Void> listener) {}
         });