From bd736a01b8b0954eddcaf07962c549fccf14f40d Mon Sep 17 00:00:00 2001
From: Eden Li <eden.li@gmail.com>
Date: Thu, 2 Oct 2014 17:01:01 -0700
Subject: [PATCH] Preserve the MDC context during callbacks.

The MDC context is captured when the Deferred is created. If another
thread begins the callback chain, its MDC thread locals will be set to
the captured values. Once the chain is complete, the thread's MDC
context will be restored.
---
 src/Deferred.java | 27 +++++++++++++++++++++++++++
 1 file changed, 27 insertions(+)

diff --git a/src/Deferred.java b/src/Deferred.java
index 9a02a52..6be7a56 100644
--- a/src/Deferred.java
+++ b/src/Deferred.java
@@ -28,9 +28,11 @@
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.concurrent.atomic.AtomicIntegerFieldUpdater;
+import java.util.Map;
 
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
+import org.slf4j.MDC;
 
 /**
  * A thread-safe implementation of a deferred result for easy asynchronous
@@ -570,6 +572,12 @@ public final class Deferred<T> {
   private static final AtomicIntegerFieldUpdater<Deferred> stateUpdater =
     AtomicIntegerFieldUpdater.newUpdater(Deferred.class, "state");
 
+  /**
+   * Captures the current context map of the current thread. This map is later
+   * reconstituted on the thread that begins the callback chain.
+   */
+  private final Map mdcContext = MDC.getCopyOfContextMap();
+
   /**
    * Atomically compares and swaps the state of this Deferred.
    * @param cmp The expected state to compare against.
@@ -1223,10 +1231,28 @@ public String toString() {
     }
   };
 
+  /**
+   * Sets the MDC context to the given {@code map}. It returns the original map
+   * that was in effect before the call took place. If {@code map} is null, the
+   * MDC context is cleared.
+   * @param map The context map to set.
+   * @return the original map that was in effect before this call took place.
+   */
+  private Map setMdc(Map map) {
+    Map old = MDC.getCopyOfContextMap();
+    if (map != null) {
+      MDC.setContextMap(map);
+    } else {
+      MDC.clear();
+    }
+    return old;
+  }
+
   /**
    * Executes all the callbacks in the current chain.
    */
   private void runCallbacks() {
+    Map old = setMdc(mdcContext);
     while (true) {
       Callback cb = null;
       Callback eb = null;
@@ -1263,6 +1289,7 @@ private void runCallbacks() {
       //          + "), result=" + result
       //          + " in " + (System.nanoTime() - start) / 1000 + "us");
     }
+    setMdc(old);
   }
 
   /**