Preserve project-id when stashing thread context (#1755)

"stashing" the thread context is typically used to execute something
with a differenet set of privileges (access to system indices, or a
different user). In these cases we don't want to lose the project-id
from the context (just like we don't want to lose APM tracing
information).

This change preserves the project-id during stash.

Author: tvernum

server/src/main/java/org/elasticsearch/common/util/concurrent/ThreadContext.java

@@ -101,16 +101,16 @@ public ThreadContext(Settings settings) {
     }
 
     /**
-     * Removes the current context and resets a default context except for headers involved in task tracing. The removed context can be
-     * restored by closing the returned {@link StoredContext}.
+     * Removes the current context and resets a default context except for headers involved in task tracing and project awareness.
+     * The removed context can be restored by closing the returned {@link StoredContext}.
      * @return a stored context that will restore the current context to its state at the point this method was called
      */
     public StoredContext stashContext() {
         return stashContextPreservingRequestHeaders(Collections.emptySet());
     }
 
     /**
-     * Just like {@link #stashContext()} but preserves request headers specified via {@code requestHeaders},
+     * Just like {@link #stashContext()} but also preserves request headers specified via {@code requestHeaders},
      * if these exist in the context before stashing.
      */
     public StoredContext stashContextPreservingRequestHeaders(Set<String> requestHeaders) {

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

@@ -65,7 +65,8 @@ public class Task implements Traceable {
         X_OPAQUE_ID_HTTP_HEADER,
         TRACE_PARENT_HTTP_HEADER,
         TRACE_ID,
-        X_ELASTIC_PRODUCT_ORIGIN_HTTP_HEADER
+        X_ELASTIC_PRODUCT_ORIGIN_HTTP_HEADER,
+        X_ELASTIC_PROJECT_ID_HTTP_HEADER
     );
 
     private final long id;

server/src/test/java/org/elasticsearch/common/util/concurrent/ThreadContextTests.java

@@ -18,6 +18,7 @@
 import org.elasticsearch.tasks.Task;
 import org.elasticsearch.test.ESTestCase;
 import org.elasticsearch.test.MockLog;
+import org.hamcrest.Matcher;
 
 import java.io.IOException;
 import java.util.Arrays;
@@ -468,14 +469,25 @@ public void testStashWithOrigin() {
             threadContext.putHeader("foo", "bar");
         }
 
+        final Matcher<? super String> matchesProjectId;
+        if (randomBoolean()) {
+            final String projectId = randomUUID();
+            matchesProjectId = equalTo(projectId);
+            threadContext.putHeader(Task.X_ELASTIC_PROJECT_ID_HTTP_HEADER, projectId);
+        } else {
+            matchesProjectId = nullValue();
+        }
+
         assertNull(threadContext.getTransient(ThreadContext.ACTION_ORIGIN_TRANSIENT_NAME));
         try (ThreadContext.StoredContext storedContext = threadContext.stashWithOrigin(origin)) {
             assertEquals(origin, threadContext.getTransient(ThreadContext.ACTION_ORIGIN_TRANSIENT_NAME));
             assertNull(threadContext.getTransient("foo"));
             assertNull(threadContext.getTransient("bar"));
+            assertThat(threadContext.getHeader(Task.X_ELASTIC_PROJECT_ID_HTTP_HEADER), matchesProjectId);
         }
 
         assertNull(threadContext.getTransient(ThreadContext.ACTION_ORIGIN_TRANSIENT_NAME));
+        assertThat(threadContext.getHeader(Task.X_ELASTIC_PROJECT_ID_HTTP_HEADER), matchesProjectId);
 
         if (setOtherValues) {
             assertEquals("bar", threadContext.getTransient("foo"));