src/main/scala/xyz/driver/pdsuicommon/concurrent/MdcExecutionContext.scala


1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
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)

}