1 /*
2  * Copyright 2016-2020 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license.
3  */
4 
5 package kotlinx.coroutines.slf4j
6 
7 import kotlinx.coroutines.*
8 import org.slf4j.MDC
9 import kotlin.coroutines.AbstractCoroutineContextElement
10 import kotlin.coroutines.CoroutineContext
11 
12 /**
13  * The value of [MDC] context map.
14  * See [MDC.getCopyOfContextMap].
15  */
16 public typealias MDCContextMap = Map<String, String>?
17 
18 /**
19  * [MDC] context element for [CoroutineContext].
20  *
21  * Example:
22  *
23  * ```
24  * MDC.put("kotlin", "rocks") // Put a value into the MDC context
25  *
26  * launch(MDCContext()) {
27  *     logger.info { "..." }   // The MDC context contains the mapping here
28  * }
29  * ```
30  *
31  * Note that you cannot update MDC context from inside of the coroutine simply
32  * using [MDC.put]. These updates are going to be lost on the next suspension and
33  * reinstalled to the MDC context that was captured or explicitly specified in
34  * [contextMap] when this object was created on the next resumption.
35  * Use `withContext(MDCContext()) { ... }` to capture updated map of MDC keys and values
36  * for the specified block of code.
37  *
38  * @param contextMap the value of [MDC] context map.
39  * Default value is the copy of the current thread's context map that is acquired via
40  * [MDC.getCopyOfContextMap].
41  */
42 public class MDCContext(
43     /**
44      * The value of [MDC] context map.
45      */
46     public val contextMap: MDCContextMap = MDC.getCopyOfContextMap()
47 ) : ThreadContextElement<MDCContextMap>, AbstractCoroutineContextElement(Key) {
48     /**
49      * Key of [MDCContext] in [CoroutineContext].
50      */
51     public companion object Key : CoroutineContext.Key<MDCContext>
52 
53     /** @suppress */
updateThreadContextnull54     override fun updateThreadContext(context: CoroutineContext): MDCContextMap {
55         val oldState = MDC.getCopyOfContextMap()
56         setCurrent(contextMap)
57         return oldState
58     }
59 
60     /** @suppress */
restoreThreadContextnull61     override fun restoreThreadContext(context: CoroutineContext, oldState: MDCContextMap) {
62         setCurrent(oldState)
63     }
64 
setCurrentnull65     private fun setCurrent(contextMap: MDCContextMap) {
66         if (contextMap == null) {
67             MDC.clear()
68         } else {
69             MDC.setContextMap(contextMap)
70         }
71     }
72 }
73