MdcPropagatingExecutionContext

MdcPropagatingExecutionContext.scala

import java.util.{Map => JMap}

import com.twitter.inject.logging.MDCInitializer
import org.slf4j.MDC

import scala.concurrent.ExecutionContext

/**
  * An `ExecutionContext` which takes the calling thread's MDC context map and copies it into
  * the thread of where the runnable will be executed. As the `prepare` method of `ExecutionContext`
  * is deprecated, it is advised to capture thread local data when the `ExecutionContext` is created.
  * As a result of this, this `ExecutionContext` should not be stored as a `val` as it will not
  * update the MDC value on each run.
  */
object MdcPropagatingExecutionContext {

  def apply(underlying: ExecutionContext): ExecutionContext = new ExecutionContext {

    private val context: JMap[String, String] = MDC.getCopyOfContextMap

    override def execute(r: Runnable): Unit = {
      val previousContext = MDC.getCopyOfContextMap
      underlying.execute(() => {
        MDCInitializer.let {
          try {
            setMdcContext(context)
            r.run()
          } finally {
            setMdcContext(previousContext)
          }
        }
      })
    }

    private def setMdcContext(context: JMap[String, String]): Unit =
      if (context != null) MDC.setContextMap(context) else MDC.clear()

    override def reportFailure(t: Throwable): Unit = underlying.reportFailure(t)
  }

  object Implicits {

    /**
      * Wrapper around the global execution context. Needs to be a `def` so that each call to it uses the
      * current thread's MDC values.
      */
    implicit def global: ExecutionContext = MdcPropagatingExecutionContext(ExecutionContext.Implicits.global)
  }
}