Suppress response headers in AllocationActionMultiListener (#93777)

Today `AllocationActionMultiListener` allows response headers to pass
between listeners, but this should not be allowed. This commit stops
preserving response headers.

Closes #93773

Author: DaveCTurner

docs/changelog/93777.yaml

@@ -0,0 +1,6 @@
+pr: 93777
+summary: Suppress response headers in `AllocationActionMultiListener`
+area: Allocation
+type: bug
+issues:
+ - 93773

server/src/main/java/org/elasticsearch/cluster/routing/allocation/allocator/AllocationActionMultiListener.java

@@ -9,13 +9,12 @@
 package org.elasticsearch.cluster.routing.allocation.allocator;
 
 import org.elasticsearch.action.ActionListener;
+import org.elasticsearch.action.support.ContextPreservingActionListener;
 import org.elasticsearch.common.util.concurrent.ThreadContext;
 
 import java.util.ArrayList;
 import java.util.List;
 
-import static org.elasticsearch.action.support.ContextPreservingActionListener.wrapPreservingContext;
-
 /**
  * This event listener might be needed to delay execution of multiple distinct tasks until followup reroute is complete.
  */
@@ -30,7 +29,7 @@ public AllocationActionMultiListener(ThreadContext context) {
     }
 
     public ActionListener<T> delay(ActionListener<T> delegate) {
-        final var wrappedDelegate = wrapPreservingContext(delegate, context);
+        final var wrappedDelegate = new ContextPreservingActionListener<>(context.newRestorableContext(false), delegate);
         return new ActionListener<T>() {
             @Override
             public void onResponse(T response) {

server/src/test/java/org/elasticsearch/cluster/routing/allocation/allocator/AllocationActionMultiListenerTests.java

@@ -9,20 +9,22 @@
 package org.elasticsearch.cluster.routing.allocation.allocator;
 
 import org.elasticsearch.action.ActionListener;
+import org.elasticsearch.action.support.RefCountingRunnable;
 import org.elasticsearch.action.support.master.AcknowledgedResponse;
 import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.common.util.concurrent.ThreadContext;
-import org.elasticsearch.core.Tuple;
 import org.elasticsearch.test.ESTestCase;
 import org.elasticsearch.threadpool.TestThreadPool;
 import org.elasticsearch.threadpool.ThreadPool;
+import org.junit.Assert;
 
+import java.util.ArrayList;
 import java.util.List;
+import java.util.Set;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;
-import java.util.concurrent.atomic.AtomicReference;
 
 import static org.hamcrest.Matchers.equalTo;
 
@@ -121,45 +123,60 @@ private static void awaitQuietly(CountDownLatch latch) {
 
     public void testShouldExecuteWithCorrectContext() {
 
-        var context = new ThreadContext(Settings.EMPTY);
-        var listener = new AllocationActionMultiListener<Integer>(context);
-
-        context.putHeader("header", "root");
-        var r1 = new AtomicReference<String>();
-        var r2 = new AtomicReference<String>();
-        var l1 = listener.delay(
-            ActionListener.wrap(
-                response -> r1.set(context.getHeader("header")),
-                exception -> { throw new AssertionError("Should not fail in test"); }
-            )
-        );
-        var l2 = listener.delay(
-            ActionListener.wrap(
-                response -> r2.set(context.getHeader("header")),
-                exception -> { throw new AssertionError("Should not fail in test"); }
-            )
-        );
+        final var requestHeaderName = "header";
+        final var responseHeaderName = "responseHeader";
 
-        executeInRandomOrder(
-            context,
-            List.of(
-                new Tuple<>("clusterStateUpdate1", () -> l1.onResponse(1)),
-                new Tuple<>("clusterStateUpdate2", () -> l2.onResponse(2)),
-                new Tuple<>("reroute", () -> listener.reroute().onResponse(null))
-            )
-        );
+        final var expectedRequestHeader = randomAlphaOfLength(10);
+        final var expectedResponseHeader = randomAlphaOfLength(10);
 
-        assertThat(r1.get(), equalTo("root"));
-        assertThat(r2.get(), equalTo("root"));
-    }
+        var context = new ThreadContext(Settings.EMPTY);
+        var listener = new AllocationActionMultiListener<>(context);
+
+        context.putHeader(requestHeaderName, expectedRequestHeader);
+        context.addResponseHeader(responseHeaderName, expectedResponseHeader);
+
+        var isComplete = new AtomicBoolean();
+        try (var refs = new RefCountingRunnable(() -> assertTrue(isComplete.compareAndSet(false, true)))) {
+
+            List<Runnable> actions = new ArrayList<>();
+
+            for (int i = between(0, 5); i > 0; i--) {
+                var expectedVal = new Object();
+                var delayedListener = listener.delay(
+                    ActionListener.releaseAfter(ActionListener.wrap(Assert::fail).delegateFailure((l, val) -> {
+                        assertSame(expectedVal, val);
+                        assertEquals(expectedRequestHeader, context.getHeader(requestHeaderName));
+                        assertEquals(List.of(expectedResponseHeader), context.getResponseHeaders().get(responseHeaderName));
+                        context.addResponseHeader(responseHeaderName, randomAlphaOfLength(10));
+                    }), refs.acquire())
+                );
+                actions.add(() -> delayedListener.onResponse(expectedVal));
+            }
 
-    private static void executeInRandomOrder(ThreadContext context, List<Tuple<String, Runnable>> actions) {
-        for (var action : shuffledList(actions)) {
-            try (var ignored = context.stashContext()) {
-                context.putHeader("header", action.v1());
-                action.v2().run();
+            final var additionalResponseHeader = randomAlphaOfLength(10);
+            context.addResponseHeader(responseHeaderName, additionalResponseHeader);
+
+            actions.add(() -> listener.reroute().onResponse(null));
+
+            for (var action : shuffledList(actions)) {
+                try (var ignored = context.stashContext()) {
+                    final var localRequestHeader = randomAlphaOfLength(10);
+                    final var localResponseHeader = randomAlphaOfLength(10);
+                    context.putHeader(requestHeaderName, localRequestHeader);
+                    context.addResponseHeader(responseHeaderName, localResponseHeader);
+                    action.run();
+                    assertEquals(localRequestHeader, context.getHeader(requestHeaderName));
+                    assertEquals(List.of(localResponseHeader), context.getResponseHeaders().get(responseHeaderName));
+                }
             }
+
+            assertEquals(
+                Set.of(expectedResponseHeader, additionalResponseHeader),
+                Set.copyOf(context.getResponseHeaders().get(responseHeaderName))
+            );
         }
+
+        assertTrue(isComplete.get());
     }
 
     private static ThreadContext createEmptyThreadContext() {