Correct context for CancellableSOCache listener (#83021)

Today the `CancellableSingleObjectCache` completes its listeners in the
thread context of the `get()` call that actually computes the value
which will be the correct context only if no batching took place. With
this commit we make sure to complete each listener in the context in
which it was passed to the corresponding `get()` call.

Author: DaveCTurner

docs/changelog/83021.yaml

@@ -0,0 +1,5 @@
+pr: 83021
+summary: Correct context for CancellableSOCache listener
+area: Stats
+type: bug
+issues: []

server/src/main/java/org/elasticsearch/action/admin/cluster/stats/TransportClusterStatsAction.java

@@ -29,6 +29,7 @@
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.common.util.CancellableSingleObjectCache;
+import org.elasticsearch.common.util.concurrent.ThreadContext;
 import org.elasticsearch.index.IndexService;
 import org.elasticsearch.index.engine.CommitStats;
 import org.elasticsearch.index.seqno.RetentionLeaseStats;
@@ -69,8 +70,8 @@ public class TransportClusterStatsAction extends TransportNodesAction<
     private final NodeService nodeService;
     private final IndicesService indicesService;
 
-    private final MetadataStatsCache<MappingStats> mappingStatsCache = new MetadataStatsCache<>(MappingStats::of);
-    private final MetadataStatsCache<AnalysisStats> analysisStatsCache = new MetadataStatsCache<>(AnalysisStats::of);
+    private final MetadataStatsCache<MappingStats> mappingStatsCache;
+    private final MetadataStatsCache<AnalysisStats> analysisStatsCache;
 
     @Inject
     public TransportClusterStatsAction(
@@ -95,6 +96,8 @@ public TransportClusterStatsAction(
         );
         this.nodeService = nodeService;
         this.indicesService = indicesService;
+        this.mappingStatsCache = new MetadataStatsCache<>(threadPool.getThreadContext(), MappingStats::of);
+        this.analysisStatsCache = new MetadataStatsCache<>(threadPool.getThreadContext(), AnalysisStats::of);
     }
 
     @Override
@@ -258,7 +261,8 @@ public void writeTo(StreamOutput out) throws IOException {
     private static class MetadataStatsCache<T> extends CancellableSingleObjectCache<Metadata, Long, T> {
         private final BiFunction<Metadata, Runnable, T> function;
 
-        MetadataStatsCache(BiFunction<Metadata, Runnable, T> function) {
+        MetadataStatsCache(ThreadContext threadContext, BiFunction<Metadata, Runnable, T> function) {
+            super(threadContext);
             this.function = function;
         }
 

server/src/main/java/org/elasticsearch/common/util/CancellableSingleObjectCache.java

@@ -9,7 +9,9 @@
 package org.elasticsearch.common.util;
 
 import org.elasticsearch.action.ActionListener;
+import org.elasticsearch.action.support.ContextPreservingActionListener;
 import org.elasticsearch.action.support.ListenableActionFuture;
+import org.elasticsearch.common.util.concurrent.ThreadContext;
 import org.elasticsearch.core.AbstractRefCounted;
 import org.elasticsearch.core.Nullable;
 import org.elasticsearch.tasks.TaskCancelledException;
@@ -41,8 +43,14 @@
  */
 public abstract class CancellableSingleObjectCache<Input, Key, Value> {
 
+    private final ThreadContext threadContext;
+
     private final AtomicReference<CachedItem> currentCachedItemRef = new AtomicReference<>();
 
+    protected CancellableSingleObjectCache(ThreadContext threadContext) {
+        this.threadContext = threadContext;
+    }
+
     /**
      * Compute a new value for the cache.
      * <p>
@@ -220,7 +228,7 @@ boolean addListener(ActionListener<Value> listener, BooleanSupplier isCancelled)
                     ActionListener.completeWith(listener, () -> future.actionGet(0L));
                 } else {
                     // Refresh is still pending; it's not cancelled because there are still references.
-                    future.addListener(listener);
+                    future.addListener(ContextPreservingActionListener.wrapPreservingContext(listener, threadContext));
                     final AtomicBoolean released = new AtomicBoolean();
                     cancellationChecks.add(() -> {
                         if (released.get() == false && isCancelled.getAsBoolean() && released.compareAndSet(false, true)) {

server/src/test/java/org/elasticsearch/common/util/CancellableSingleObjectCacheTests.java

@@ -12,7 +12,9 @@
 import org.elasticsearch.action.ActionListener;
 import org.elasticsearch.action.StepListener;
 import org.elasticsearch.action.support.PlainActionFuture;
+import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.common.util.concurrent.ConcurrentCollections;
+import org.elasticsearch.common.util.concurrent.ThreadContext;
 import org.elasticsearch.tasks.TaskCancelledException;
 import org.elasticsearch.test.ESTestCase;
 import org.elasticsearch.threadpool.TestThreadPool;
@@ -193,7 +195,8 @@ public void testExceptionCompletesListenersButIsNotCached() {
     public void testConcurrentRefreshesAndCancellation() throws InterruptedException {
         final ThreadPool threadPool = new TestThreadPool("test");
         try {
-            final CancellableSingleObjectCache<String, String, Integer> testCache = new CancellableSingleObjectCache<>() {
+            final ThreadContext threadContext = threadPool.getThreadContext();
+            final CancellableSingleObjectCache<String, String, Integer> testCache = new CancellableSingleObjectCache<>(threadContext) {
                 @Override
                 protected void refresh(
                     String s,
@@ -219,6 +222,7 @@ protected String getKey(String s) {
             final CountDownLatch startLatch = new CountDownLatch(1);
             final CountDownLatch finishLatch = new CountDownLatch(count);
             final BlockingQueue<Runnable> queue = ConcurrentCollections.newBlockingQueue();
+            final String contextHeader = "test-context-header";
 
             for (int i = 0; i < count; i++) {
                 final boolean cancel = randomBoolean();
@@ -233,11 +237,14 @@ protected String getKey(String s) {
                     final StepListener<Integer> stepListener = new StepListener<>();
                     final AtomicBoolean isComplete = new AtomicBoolean();
                     final AtomicBoolean isCancelled = new AtomicBoolean();
-                    testCache.get(
-                        input,
-                        isCancelled::get,
-                        ActionListener.runBefore(stepListener, () -> assertTrue(isComplete.compareAndSet(false, true)))
-                    );
+                    try (ThreadContext.StoredContext ignored = threadContext.stashContext()) {
+                        final String contextValue = randomAlphaOfLength(10);
+                        threadContext.putHeader(contextHeader, contextValue);
+                        testCache.get(input, isCancelled::get, ActionListener.runBefore(stepListener, () -> {
+                            assertTrue(isComplete.compareAndSet(false, true));
+                            assertThat(threadContext.getHeader(contextHeader), equalTo(contextValue));
+                        }));
+                    }
 
                     final Runnable next = queue.poll();
                     if (next != null) {
@@ -277,7 +284,9 @@ protected String getKey(String s) {
     public void testConcurrentRefreshesWithFreshnessCheck() throws InterruptedException {
         final ThreadPool threadPool = new TestThreadPool("test");
         try {
-            final CancellableSingleObjectCache<String, String, Integer> testCache = new CancellableSingleObjectCache<>() {
+            final CancellableSingleObjectCache<String, String, Integer> testCache = new CancellableSingleObjectCache<>(
+                threadPool.getThreadContext()
+            ) {
                 @Override
                 protected void refresh(
                     String s,
@@ -380,7 +389,7 @@ public void run() {
             }
         };
 
-        final CancellableSingleObjectCache<String, String, Integer> testCache = new CancellableSingleObjectCache<>() {
+        final CancellableSingleObjectCache<String, String, Integer> testCache = new CancellableSingleObjectCache<>(testThreadContext) {
             @Override
             protected void refresh(
                 String s,
@@ -424,10 +433,16 @@ protected String getKey(String s) {
         expectThrows(TaskCancelledException.class, () -> cancelledFuture.actionGet(0L));
     }
 
+    private static final ThreadContext testThreadContext = new ThreadContext(Settings.EMPTY);
+
     private static class TestCache extends CancellableSingleObjectCache<String, String, Integer> {
 
         private final LinkedList<StepListener<Function<String, Integer>>> pendingRefreshes = new LinkedList<>();
 
+        private TestCache() {
+            super(testThreadContext);
+        }
+
         @Override
         protected void refresh(
             String input,