From 40c4bd562844beda9013fde5160143f3bc1f5322 Mon Sep 17 00:00:00 2001
From: Yannick Welsch <yannick@welsch.lu>
Date: Tue, 19 Jun 2018 19:03:31 +0200
Subject: [PATCH] Preserve response headers on cluster update task (#31421)

#31241 changed the cluster state update tasks to run under system context. The context wrapping
did not preserve response headers, though. This has led to a test failure on 6.x #31408, as the
deprecation warnings were not carried back anymore to the caller when creating an index. This
commit changes the restorable context supplier to preserve response headers.
---
 .../elasticsearch/cluster/service/MasterService.java   |  2 +-
 .../cluster/service/MasterServiceTests.java            | 10 ++++++++++
 2 files changed, 11 insertions(+), 1 deletion(-)

diff --git a/server/src/main/java/org/elasticsearch/cluster/service/MasterService.java b/server/src/main/java/org/elasticsearch/cluster/service/MasterService.java
index 2543be4811c..8927adfd434 100644
--- a/server/src/main/java/org/elasticsearch/cluster/service/MasterService.java
+++ b/server/src/main/java/org/elasticsearch/cluster/service/MasterService.java
@@ -730,7 +730,7 @@ public class MasterService extends AbstractLifecycleComponent {
             return;
         }
         final ThreadContext threadContext = threadPool.getThreadContext();
-        final Supplier<ThreadContext.StoredContext> supplier = threadContext.newRestorableContext(false);
+        final Supplier<ThreadContext.StoredContext> supplier = threadContext.newRestorableContext(true);
         try (ThreadContext.StoredContext ignore = threadContext.stashContext()) {
             threadContext.markAsSystemContext();
 
diff --git a/server/src/test/java/org/elasticsearch/cluster/service/MasterServiceTests.java b/server/src/test/java/org/elasticsearch/cluster/service/MasterServiceTests.java
index 20587d31f53..1ef548bd681 100644
--- a/server/src/test/java/org/elasticsearch/cluster/service/MasterServiceTests.java
+++ b/server/src/test/java/org/elasticsearch/cluster/service/MasterServiceTests.java
@@ -54,6 +54,7 @@ import org.junit.Before;
 import org.junit.BeforeClass;
 
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.HashSet;
@@ -177,6 +178,8 @@ public class MasterServiceTests extends ESTestCase {
 
         try (ThreadContext.StoredContext ignored = threadPool.getThreadContext().stashContext()) {
             final Map<String, String> expectedHeaders = Collections.singletonMap("test", "test");
+            final Map<String, List<String>> expectedResponseHeaders = Collections.singletonMap("testResponse",
+                Arrays.asList("testResponse"));
             threadPool.getThreadContext().putHeader(expectedHeaders);
 
             final TimeValue ackTimeout = randomBoolean() ? TimeValue.ZERO : TimeValue.timeValueMillis(randomInt(10000));
@@ -187,6 +190,8 @@ public class MasterServiceTests extends ESTestCase {
                 public ClusterState execute(ClusterState currentState) {
                     assertTrue(threadPool.getThreadContext().isSystemContext());
                     assertEquals(Collections.emptyMap(), threadPool.getThreadContext().getHeaders());
+                    threadPool.getThreadContext().addResponseHeader("testResponse", "testResponse");
+                    assertEquals(expectedResponseHeaders, threadPool.getThreadContext().getResponseHeaders());
 
                     if (randomBoolean()) {
                         return ClusterState.builder(currentState).build();
@@ -201,6 +206,7 @@ public class MasterServiceTests extends ESTestCase {
                 public void onFailure(String source, Exception e) {
                     assertFalse(threadPool.getThreadContext().isSystemContext());
                     assertEquals(expectedHeaders, threadPool.getThreadContext().getHeaders());
+                    assertEquals(expectedResponseHeaders, threadPool.getThreadContext().getResponseHeaders());
                     latch.countDown();
                 }
 
@@ -208,6 +214,7 @@ public class MasterServiceTests extends ESTestCase {
                 public void clusterStateProcessed(String source, ClusterState oldState, ClusterState newState) {
                     assertFalse(threadPool.getThreadContext().isSystemContext());
                     assertEquals(expectedHeaders, threadPool.getThreadContext().getHeaders());
+                    assertEquals(expectedResponseHeaders, threadPool.getThreadContext().getResponseHeaders());
                     latch.countDown();
                 }
 
@@ -229,6 +236,7 @@ public class MasterServiceTests extends ESTestCase {
                 public void onAllNodesAcked(@Nullable Exception e) {
                     assertFalse(threadPool.getThreadContext().isSystemContext());
                     assertEquals(expectedHeaders, threadPool.getThreadContext().getHeaders());
+                    assertEquals(expectedResponseHeaders, threadPool.getThreadContext().getResponseHeaders());
                     latch.countDown();
                 }
 
@@ -236,6 +244,7 @@ public class MasterServiceTests extends ESTestCase {
                 public void onAckTimeout() {
                     assertFalse(threadPool.getThreadContext().isSystemContext());
                     assertEquals(expectedHeaders, threadPool.getThreadContext().getHeaders());
+                    assertEquals(expectedResponseHeaders, threadPool.getThreadContext().getResponseHeaders());
                     latch.countDown();
                 }
 
@@ -243,6 +252,7 @@ public class MasterServiceTests extends ESTestCase {
 
             assertFalse(threadPool.getThreadContext().isSystemContext());
             assertEquals(expectedHeaders, threadPool.getThreadContext().getHeaders());
+            assertEquals(Collections.emptyMap(), threadPool.getThreadContext().getResponseHeaders());
         }
 
         latch.await();