package com.mdc
import org.slf4j.MDC

import scala.concurrent.{ExecutionContext, ExecutionContextExecutor}

object MDCHttpExecutionContext {
  def fromThread(delegate: ExecutionContext): ExecutionContextExecutor = new MDCHttpExecutionContext(delegate)
}

class MDCHttpExecutionContext(delegate: ExecutionContext) extends ExecutionContextExecutor{

  override def execute(runnable: Runnable) = delegate.execute(new Runnable {
    val mdcContext = MDC.getCopyOfContextMap
    override def run(): Unit = {
      val oldMDCContext = MDC.getCopyOfContextMap
      setContextMap(mdcContext)
      try {
        runnable.run()
      } finally {
        setContextMap(oldMDCContext)
      }
    }
  })

  private[this] def setContextMap(context: java.util.Map[String, String]) = {
    if(context == null){
      MDC.clear()
    } else {
      MDC.setContextMap(context)
    }
  }

  override def reportFailure(cause: Throwable) = delegate.reportFailure(cause)
}