Correct context for CancellableSOCache listener (#83021) (#83169)

* 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. * Fix changelog YAML * Spotless
elastic · Jan 26, 2022 · 3b494e4 · 3b494e4
1 parent 8668691
commit 3b494e4
Show file tree

Hide file tree

Showing 4 changed files with 44 additions and 15 deletions.
diff --git a/docs/changelog/83021.yaml b/docs/changelog/83021.yaml
@@ -0,0 +1,5 @@
+pr: 83021
+summary: Correct context for CancellableSOCache listener
+area: Features/Stats
+type: bug
+issues: []
diff --git a/...c/main/java/org/elasticsearch/action/admin/cluster/stats/TransportClusterStatsAction.java b/...c/main/java/org/elasticsearch/action/admin/cluster/stats/TransportClusterStatsAction.java
@@ -30,6 +30,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
@@ -264,7 +267,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;
         }
 

diff --git a/server/src/main/java/org/elasticsearch/common/util/CancellableSingleObjectCache.java b/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)) {

diff --git a/server/src/test/java/org/elasticsearch/common/util/CancellableSingleObjectCacheTests.java b/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,10 +195,11 @@ public void testExceptionCompletesListenersButIsNotCached() {
     public void testConcurrentRefreshesAndCancellation() throws InterruptedException {
         final ThreadPool threadPool = new TestThreadPool("test");
         try {
+            final ThreadContext threadContext = threadPool.getThreadContext();
             final CancellableSingleObjectCache<String, String, Integer> testCache = new CancellableSingleObjectCache<
                 String,
                 String,
-                Integer>() {
+                Integer>(threadContext) {
                 @Override
                 protected void refresh(
                     String s,
@@ -222,6 +225,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();
@@ -236,11 +240,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) {
@@ -283,7 +290,7 @@ public void testConcurrentRefreshesWithFreshnessCheck() throws InterruptedExcept
             final CancellableSingleObjectCache<String, String, Integer> testCache = new CancellableSingleObjectCache<
                 String,
                 String,
-                Integer>() {
+                Integer>(threadPool.getThreadContext()) {
                 @Override
                 protected void refresh(
                     String s,
@@ -386,10 +393,9 @@ public void run() {
             }
         };
 
-        final CancellableSingleObjectCache<String, String, Integer> testCache = new CancellableSingleObjectCache<
-            String,
-            String,
-            Integer>() {
+        final CancellableSingleObjectCache<String, String, Integer> testCache = new CancellableSingleObjectCache<String, String, Integer>(
+            testThreadContext
+        ) {
             @Override
             protected void refresh(
                 String s,
@@ -433,10 +439,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,