1 files changed, 35 insertions, 0 deletions
diff --git a/src/main/scala/xyz/driver/pdsuicommon/concurrent/MdcExecutionContext.scala b/src/main/scala/xyz/driver/pdsuicommon/concurrent/MdcExecutionContext.scala
new file mode 100644
index 0000000..3dee8ea
--- /dev/null
+++ b/src/main/scala/xyz/driver/pdsuicommon/concurrent/MdcExecutionContext.scala
@@ -0,0 +1,35 @@
+package xyz.driver.pdsuicommon.concurrent
+
+import org.slf4j.MDC
+
+import scala.concurrent.{ExecutionContext, ExecutionContextExecutor}
+
+object MdcExecutionContext {
+  def from(orig: ExecutionContext): ExecutionContext = new MdcExecutionContext(orig)
+}
+
+class MdcExecutionContext(orig: ExecutionContext) extends ExecutionContextExecutor {
+
+  def execute(runnable: Runnable): Unit = {
+    val parentMdcContext = MDC.getCopyOfContextMap
+
+    orig.execute(new Runnable {
+      def run(): Unit = {
+        val saveMdcContext = MDC.getCopyOfContextMap
+        setContextMap(parentMdcContext)
+
+        try {
+          runnable.run()
+        } finally {
+          setContextMap(saveMdcContext)
+        }
+      }
+    })
+  }
+
+  private[this] def setContextMap(context: java.util.Map[String, String]): Unit =
+    Option(context).fold(MDC.clear())(MDC.setContextMap)
+
+  def reportFailure(t: Throwable): Unit = orig.reportFailure(t)
+
+}