1+ /*
2+ * Copyright 2016-2018 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license.
3+ */
4+
5+ package kotlinx.coroutines.log4j
6+
7+ import kotlinx.coroutines.ThreadContextElement
8+ import org.apache.logging.log4j.ThreadContext
9+ import kotlin.coroutines.AbstractCoroutineContextElement
10+ import kotlin.coroutines.CoroutineContext
11+
12+ /* *
13+ * Creates a new, immutable, [DiagnosticContext].
14+ *
15+ * See [DiagnosticContext] for usage instructions.
16+ *
17+ * If not modifying the [ThreadContext], this method is preferred over [MutableDiagnosticContext] as it performs fewer
18+ * unnecessary allocations.
19+ */
20+ public fun immutableDiagnosticContext (): DiagnosticContext = DiagnosticContext (ThreadContext .getImmutableContext())
21+
22+ /* *
23+ * Enables the use of Log4J 2's [ThreadContext] with coroutines.
24+ *
25+ * See [DiagnosticContext] for usage instructions.
26+ *
27+ * @param mappedContext Mapped diagnostic context to apply for the duration of the corresponding [CoroutineContext].
28+ */
29+ public class MutableDiagnosticContext private constructor(
30+ // Local reference so we can mutate the state
31+ private val mappedContext : MutableMap <String , String ?>
32+ ) : DiagnosticContext(mappedContext, ThreadContext .getImmutableContext()) {
33+
34+ /* *
35+ * Creates a new [MutableDiagnosticContext] populated with the current [ThreadContext].
36+ *
37+ * If not intending to modify the [ThreadContext], consider using [immutableDiagnosticContext] instead.
38+ * [immutableDiagnosticContext] is preferred in this case as it performs fewer unnecessary allocations.
39+ */
40+ public constructor () : this (ThreadContext .getContext())
41+
42+ /* *
43+ * Adds an entry to the Log4J context map.
44+ *
45+ * The entry will remain as part of the diagnostic context for the duration of the current coroutine context.
46+ *
47+ * This is the coroutine-compatible equivalent of [ThreadContext.put].
48+ *
49+ * @param key Key of the entry to add to the diagnostic context.
50+ * @param value Value of the entry to add to the diagnostic context.
51+ * @return This instance.
52+ */
53+ public fun put (key : String , value : String? ): MutableDiagnosticContext {
54+ mappedContext[key] = value
55+ return this
56+ }
57+
58+ /* *
59+ * Adds all entries to the Log4J context map.
60+ *
61+ * The entries will remain as part of the diagnostic context for the duration of the current coroutine context.
62+ *
63+ * This is the coroutine-compatible equivalent of [ThreadContext.putAll].
64+ *
65+ * @param from Entries to add to the diagnostic context.
66+ * @return This instance.
67+ */
68+ public fun putAll (from : Map <String , String ?>): MutableDiagnosticContext {
69+ mappedContext.putAll(from)
70+ return this
71+ }
72+ }
73+
74+ /* *
75+ * Enables the use of Log4J 2's [ThreadContext] with coroutines.
76+ *
77+ * # Example
78+ * The following example demonstrates usage of this class. All `assert`s pass. Note that only the mapped diagnostic
79+ * context is supported.
80+ *
81+ * ```kotlin
82+ * ThreadContext.put("kotlin", "rocks") // Put a value into the ThreadContext.
83+ * launch(immutableDiagnosticContext()) { // The contents of the ThreadContext are captured into the newly created CoroutineContext.
84+ * assert(ThreadContext.get("kotlin") == "rocks")
85+ *
86+ * withContext(MutableDiagnosticContext().put("kotlin", "is great") {
87+ * assert(ThreadContext.get("kotlin") == "is great")
88+ *
89+ * launch(Dispatchers.IO) {
90+ * assert(ThreadContext.get("kotlin") == "is great") // The diagnostic context is inherited by child CoroutineContexts.
91+ * }
92+ * }
93+ * assert(ThreadContext.get("kotlin") == "rocks") // The ThreadContext is reset when the CoroutineContext exits.
94+ * }
95+ * ```
96+ *
97+ * ## Combine with others
98+ * You may wish to combine this [ThreadContextElement] with other [CoroutineContext]s.
99+ *
100+ * ```kotlin
101+ * launch(Dispatchers.IO + immutableDiagnosticContext()) { ... }
102+ * ```
103+ *
104+ * # CloseableThreadContext
105+ * [org.apache.logging.log4j.CloseableThreadContext] is useful for automatically cleaning up the [ThreadContext] after a
106+ * block of code. The structured concurrency provided by coroutines offers the same functionality.
107+ *
108+ * In the following example, the modifications to the [ThreadContext] are cleaned up when the coroutine exits.
109+ *
110+ * ```kotlin
111+ * ThreadContext.put("kotlin", "rocks")
112+ *
113+ * withContext(MutableDiagnosticContext().put("kotlin", "is awesome") {
114+ * assert(ThreadContext.get("kotlin") == "is awesome")
115+ * }
116+ * assert(ThreadContext.get("kotlin") == "rocks")
117+ * ```
118+ *
119+ * @param mappedContextBefore Mapped diagnostic context to apply for the duration of the corresponding [CoroutineContext].
120+ * @param mappedContextAfter Mapped diagnostic context to restore when the corresponding [CoroutineContext] exits.
121+ */
122+ public open class DiagnosticContext internal constructor(
123+ private val mappedContextBefore : Map <String , String ?>,
124+ private val mappedContextAfter : Map <String , String ?> = mappedContextBefore
125+ ) : ThreadContextElement<DiagnosticContext>, AbstractCoroutineContextElement(Key ) {
126+
127+ /* *
128+ * Key of [DiagnosticContext] in [CoroutineContext].
129+ */
130+ public companion object Key : CoroutineContext.Key<DiagnosticContext>
131+
132+ /* * @suppress */
133+ final override fun updateThreadContext (context : CoroutineContext ): DiagnosticContext {
134+ setCurrent(mappedContextBefore)
135+ return this
136+ }
137+
138+ /* * @suppress */
139+ final override fun restoreThreadContext (context : CoroutineContext , oldState : DiagnosticContext ) {
140+ setCurrent(oldState.mappedContextAfter)
141+ }
142+
143+ private fun setCurrent (map : Map <String , String ?>) {
144+ /*
145+ * The logic here varies significantly from how CloseableThreadContext works. CloseableThreadContext has the
146+ * luxury of assuming that it is appending new state to the existing state of the current thread. We cannot make
147+ * this assumption. It is very realistic for us to be restoring a context to a thread that has loads of state
148+ * that we are not at all interested in, due to the Log4J ThreadContext being implemented as a ThreadLocal.
149+ *
150+ * So, to make sure that the ThreadLocal belonging to the Thread servicing this CoroutineContext is has the
151+ * correct state, we first clear everything existing, and then apply the desired state.
152+ */
153+ ThreadContext .clearMap()
154+ ThreadContext .putAll(map)
155+ }
156+ }
0 commit comments