From 654f50c7ac57c87f367533d3df843060f6c98b02 Mon Sep 17 00:00:00 2001
From: Przemyslaw Gomulka <przemyslaw.gomulka@elastic.co>
Date: Wed, 23 Jun 2021 15:01:14 +0200
Subject: [PATCH] RestController not using thread context directly from thread
 pool (#74293)

At the moment thread context is passed via dispatchRequest but in some
places thread context is fetched directly from thread pool
This is not a problem in production, because thread pool is initialized
with the same thread context as the one passed to dispatchRequest via
AbstractHttpServerTransport.
It might be harder to understand though and might cause problems in
testing in smaller units.
---
 .../main/java/org/elasticsearch/rest/RestController.java   | 7 ++++---
 .../org/elasticsearch/test/rest/RestActionTestCase.java    | 7 ++++---
 2 files changed, 8 insertions(+), 6 deletions(-)

diff --git a/server/src/main/java/org/elasticsearch/rest/RestController.java b/server/src/main/java/org/elasticsearch/rest/RestController.java
index 6dc768355297b..63bd4f3640b32 100644
--- a/server/src/main/java/org/elasticsearch/rest/RestController.java
+++ b/server/src/main/java/org/elasticsearch/rest/RestController.java
@@ -226,7 +226,9 @@ public void dispatchBadRequest(final RestChannel channel, final ThreadContext th
         }
     }
 
-    private void dispatchRequest(RestRequest request, RestChannel channel, RestHandler handler) throws Exception {
+    private void dispatchRequest(RestRequest request, RestChannel channel, RestHandler handler,
+                                 ThreadContext threadContext)
+        throws Exception {
         final int contentLength = request.contentLength();
         if (contentLength > 0) {
             final XContentType xContentType = request.getXContentType();
@@ -254,7 +256,6 @@ private void dispatchRequest(RestRequest request, RestChannel channel, RestHandl
                 request.ensureSafeBuffers();
             }
 
-            final ThreadContext threadContext = client.threadPool().getThreadContext();
             if (handler.allowSystemIndexAccessByDefault() == false) {
                 // The ELASTIC_PRODUCT_ORIGIN_HTTP_HEADER indicates that the request is coming from an Elastic product and
                 // therefore we should allow a subset of external system index access.
@@ -352,7 +353,7 @@ private void tryAllHandlers(final RestRequest request, final RestChannel channel
                         return;
                     }
                 } else {
-                    dispatchRequest(request, channel, handler);
+                    dispatchRequest(request, channel, handler, threadContext);
                     return;
                 }
             }
diff --git a/test/framework/src/main/java/org/elasticsearch/test/rest/RestActionTestCase.java b/test/framework/src/main/java/org/elasticsearch/test/rest/RestActionTestCase.java
index dd83e832cbfb1..a4d759242dcde 100644
--- a/test/framework/src/main/java/org/elasticsearch/test/rest/RestActionTestCase.java
+++ b/test/framework/src/main/java/org/elasticsearch/test/rest/RestActionTestCase.java
@@ -12,7 +12,6 @@
 import org.elasticsearch.action.ActionRequest;
 import org.elasticsearch.action.ActionResponse;
 import org.elasticsearch.action.ActionType;
-import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.common.util.concurrent.ThreadContext;
 import org.elasticsearch.indices.breaker.NoneCircuitBreakerService;
 import org.elasticsearch.rest.RestController;
@@ -65,8 +64,10 @@ protected RestController controller() {
      */
     protected void dispatchRequest(RestRequest request) {
         FakeRestChannel channel = new FakeRestChannel(request, false, 1);
-        ThreadContext threadContext = new ThreadContext(Settings.EMPTY);
-        controller.dispatchRequest(request, channel, threadContext);
+        ThreadContext threadContext = verifyingClient.threadPool().getThreadContext();
+        try(ThreadContext.StoredContext ignore = threadContext.stashContext()) {
+            controller.dispatchRequest(request, channel, threadContext);
+        }
     }
 
     /**