1 files changed, 45 insertions, 5 deletions

diff --git a/core/src/test/java/org/elasticsearch/common/util/concurrent/EsExecutorsTests.java b/core/src/test/java/org/elasticsearch/common/util/concurrent/EsExecutorsTests.java
index b59c8dd1cb..fabfc6a94e 100644
--- a/core/src/test/java/org/elasticsearch/common/util/concurrent/EsExecutorsTests.java
+++ b/core/src/test/java/org/elasticsearch/common/util/concurrent/EsExecutorsTests.java
@@ -20,6 +20,7 @@
 package org.elasticsearch.common.util.concurrent;
 
 import org.elasticsearch.ExceptionsHelper;
+import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.test.ESTestCase;
 import org.hamcrest.Matcher;
 
@@ -38,12 +39,13 @@ import static org.hamcrest.Matchers.lessThan;
  */
 public class EsExecutorsTests extends ESTestCase {
 
+    private final ThreadContext threadContext = new ThreadContext(Settings.EMPTY);
     private TimeUnit randomTimeUnit() {
         return TimeUnit.values()[between(0, TimeUnit.values().length - 1)];
     }
 
     public void testFixedForcedExecution() throws Exception {
-        EsThreadPoolExecutor executor = EsExecutors.newFixed(getTestName(), 1, 1, EsExecutors.daemonThreadFactory("test"));
+        EsThreadPoolExecutor executor = EsExecutors.newFixed(getTestName(), 1, 1, EsExecutors.daemonThreadFactory("test"), threadContext);
         final CountDownLatch wait = new CountDownLatch(1);
 
         final CountDownLatch exec1Wait = new CountDownLatch(1);
@@ -105,7 +107,7 @@ public class EsExecutorsTests extends ESTestCase {
     }
 
     public void testFixedRejected() throws Exception {
-        EsThreadPoolExecutor executor = EsExecutors.newFixed(getTestName(), 1, 1, EsExecutors.daemonThreadFactory("test"));
+        EsThreadPoolExecutor executor = EsExecutors.newFixed(getTestName(), 1, 1, EsExecutors.daemonThreadFactory("test"), threadContext);
         final CountDownLatch wait = new CountDownLatch(1);
 
         final CountDownLatch exec1Wait = new CountDownLatch(1);
@@ -163,7 +165,7 @@ public class EsExecutorsTests extends ESTestCase {
         final int max = between(min + 1, 6);
         final ThreadBarrier barrier = new ThreadBarrier(max + 1);
 
-        ThreadPoolExecutor pool = EsExecutors.newScaling(getTestName(), min, max, between(1, 100), randomTimeUnit(), EsExecutors.daemonThreadFactory("test"));
+        ThreadPoolExecutor pool = EsExecutors.newScaling(getTestName(), min, max, between(1, 100), randomTimeUnit(), EsExecutors.daemonThreadFactory("test"), threadContext);
         assertThat("Min property", pool.getCorePoolSize(), equalTo(min));
         assertThat("Max property", pool.getMaximumPoolSize(), equalTo(max));
 
@@ -199,7 +201,7 @@ public class EsExecutorsTests extends ESTestCase {
         final int max = between(min + 1, 6);
         final ThreadBarrier barrier = new ThreadBarrier(max + 1);
 
-        final ThreadPoolExecutor pool = EsExecutors.newScaling(getTestName(), min, max, between(1, 100), TimeUnit.MILLISECONDS, EsExecutors.daemonThreadFactory("test"));
+        final ThreadPoolExecutor pool = EsExecutors.newScaling(getTestName(), min, max, between(1, 100), TimeUnit.MILLISECONDS, EsExecutors.daemonThreadFactory("test"), threadContext);
         assertThat("Min property", pool.getCorePoolSize(), equalTo(min));
         assertThat("Max property", pool.getMaximumPoolSize(), equalTo(max));
 
@@ -242,7 +244,7 @@ public class EsExecutorsTests extends ESTestCase {
         int queue = between(0, 100);
         int actions = queue + pool;
         final CountDownLatch latch = new CountDownLatch(1);
-        EsThreadPoolExecutor executor = EsExecutors.newFixed(getTestName(), pool, queue, EsExecutors.daemonThreadFactory("dummy"));
+        EsThreadPoolExecutor executor = EsExecutors.newFixed(getTestName(), pool, queue, EsExecutors.daemonThreadFactory("dummy"), threadContext);
         try {
             for (int i = 0; i < actions; i++) {
                 executor.execute(new Runnable() {
@@ -321,4 +323,42 @@ public class EsExecutorsTests extends ESTestCase {
             assertThat(message, containsString("completed tasks = " + actions));
         }
     }
+
+    public void testInheritContext() throws InterruptedException {
+        int pool = between(1, 10);
+        int queue = between(0, 100);
+        final CountDownLatch latch = new CountDownLatch(1);
+        final CountDownLatch executed = new CountDownLatch(1);
+
+        threadContext.putHeader("foo", "bar");
+        final Integer one = new Integer(1);
+        threadContext.putTransient("foo", one);
+        EsThreadPoolExecutor executor = EsExecutors.newFixed(getTestName(), pool, queue, EsExecutors.daemonThreadFactory("dummy"), threadContext);
+        try {
+            executor.execute(new Runnable() {
+                @Override
+                public void run() {
+                    try {
+                        latch.await();
+                    } catch (InterruptedException e) {
+                        fail();
+                    }
+                    assertEquals(threadContext.getHeader("foo"), "bar");
+                    assertSame(threadContext.getTransient("foo"), one);
+                    assertNull(threadContext.getHeader("bar"));
+                    assertNull(threadContext.getTransient("bar"));
+                    executed.countDown();
+                }
+            });
+            threadContext.putTransient("bar", "boom");
+            threadContext.putHeader("bar", "boom");
+            latch.countDown();
+            executed.await();
+
+        } finally {
+            latch.countDown();
+            terminate(executor);
+        }
+
+    }
 }