aboutsummaryrefslogtreecommitdiff
path: root/kamon-core/src/main/scala/kamon/TraceContextSwap.scala
blob: 4b5b66a9efa4c34c241b59c29158b1d1c18b27d8 (plain) (blame)
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
package kamon

import org.slf4j.MDC

/**
 *  Provides support for making a TraceContext available as ThreadLocal and cleanning up afterwards.
 */
trait TraceContextSwap {

  def withContext[A](ctx: Option[TraceContext], body: => A): A = withContext(ctx, body, body)

  def withContext[A](ctx: Option[TraceContext], primary: => A, fallback: => A): A = {
    ctx match {
      case Some(context) => {
        MDC.put("uow", context.userContext.get.asInstanceOf[String])
        Tracer.set(context)
        val bodyResult = primary
        Tracer.clear
        MDC.remove("uow")

        bodyResult
      }
      case None => fallback
    }

  }

}

object TraceContextSwap extends TraceContextSwap