fixup! Add an integration with Log4J 2's ThreadContext

README.md

@@ -52,8 +52,8 @@ suspend fun main() = coroutineScope {
   * Android, JavaFX, and Swing.
 * [integration](integration/README.md) — modules that provide integration with various asynchronous callback- and future-based libraries:
   * JDK8 [CompletionStage.await], Guava [ListenableFuture.await], and Google Play Services [Task.await];
-  * SLF4J MDC integration via [MDCContext].
-  * Log4J 2 ThreadContext integration via [Log4JThreadContext]
+  * SLF4J MDC integration via [MDCContext];
+  * Log4J 2 ThreadContext integration via [Log4JThreadContext].
 
 ## Documentation
 

integration/kotlinx-coroutines-log4j/README.md

@@ -4,13 +4,17 @@ Integration with Log4J 2's [ThreadContext](https://logging.apache.org/log4j/2.x/
 
 ## Example
 
-Add [Log4JThreadContext] to the coroutine context so that the Log4J `ThreadContext` state is captured and passed into the coroutine.
+Add a [DiagnosticContext] to the coroutine context so that the Log4J `ThreadContext` state is set for the duration of
+coroutine context.
 
 ```kotlin
-ThreadContext.put("kotlin", "rocks")
+launch(MutableDiagnosticContext().put("kotlin", "rocks")) {
+  logger.info(...)  // The ThreadContext will contain the mapping here
+}
 
-launch(Log4JThreadContext()) {
-   logger.info(...)   // the ThreadContext will contain the mapping here
+// If not modifying the context state, use an immutable context for fewer allocations
+launch(immutableDiagnosticContext()) {
+  logger.info(...)
 }
 ```
 

integration/kotlinx-coroutines-log4j/build.gradle

@@ -1,6 +1,6 @@
 dependencies {
-    implementation 'org.apache.logging.log4j:log4j-api:2.13.0'
-    testImplementation 'org.apache.logging.log4j:log4j-core:2.13.0'
+    implementation 'org.apache.logging.log4j:log4j-api:2.13.1'
+    testImplementation 'org.apache.logging.log4j:log4j-core:2.13.1'
 }
 
 tasks.withType(dokka.getClass()) {

integration/kotlinx-coroutines-log4j/src/Log4jDiagnosticContext.kt

@@ -0,0 +1,156 @@
+/*
+ * Copyright 2016-2018 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license.
+ */
+
+package kotlinx.coroutines.log4j
+
+import kotlinx.coroutines.ThreadContextElement
+import org.apache.logging.log4j.ThreadContext
+import kotlin.coroutines.AbstractCoroutineContextElement
+import kotlin.coroutines.CoroutineContext
+
+/**
+ * Creates a new, immutable, [DiagnosticContext].
+ *
+ * See [DiagnosticContext] for usage instructions.
+ *
+ * If not modifying the [ThreadContext], this method is preferred over [MutableDiagnosticContext] as it performs fewer
+ * unnecessary allocations.
+ */
+public fun immutableDiagnosticContext(): DiagnosticContext = DiagnosticContext(ThreadContext.getImmutableContext())
+
+/**
+ * Enables the use of Log4J 2's [ThreadContext] with coroutines.
+ *
+ * See [DiagnosticContext] for usage instructions.
+ *
+ * @param mappedContext Mapped diagnostic context to apply for the duration of the corresponding [CoroutineContext].
+ */
+public class MutableDiagnosticContext private constructor(
+    // Local reference so we can mutate the state
+    private val mappedContext: MutableMap<String, String?>
+) : DiagnosticContext(mappedContext, ThreadContext.getImmutableContext()) {
+
+    /**
+     * Creates a new [MutableDiagnosticContext] populated with the current [ThreadContext].
+     *
+     * If not intending to modify the [ThreadContext], consider using [immutableDiagnosticContext] instead.
+     * [immutableDiagnosticContext] is preferred in this case as it performs fewer unnecessary allocations.
+     */
+    public constructor() : this(ThreadContext.getContext())
+
+    /**
+     * Adds an entry to the Log4J context map.
+     *
+     * The entry will remain as part of the diagnostic context for the duration of the current coroutine context.
+     *
+     * This is the coroutine-compatible equivalent of [ThreadContext.put].
+     *
+     * @param key Key of the entry to add to the diagnostic context.
+     * @param value Value of the entry to add to the diagnostic context.
+     * @return This instance.
+     */
+    public fun put(key: String, value: String?): MutableDiagnosticContext {
+        mappedContext[key] = value
+        return this
+    }
+
+    /**
+     * Adds all entries to the Log4J context map.
+     *
+     * The entries will remain as part of the diagnostic context for the duration of the current coroutine context.
+     *
+     * This is the coroutine-compatible equivalent of [ThreadContext.putAll].
+     *
+     * @param from Entries to add to the diagnostic context.
+     * @return This instance.
+     */
+    public fun putAll(from: Map<String, String?>): MutableDiagnosticContext {
+        mappedContext.putAll(from)
+        return this
+    }
+}
+
+/**
+ * Enables the use of Log4J 2's [ThreadContext] with coroutines.
+ *
+ * # Example
+ * The following example demonstrates usage of this class. All `assert`s pass. Note that only the mapped diagnostic
+ * context is supported.
+ *
+ * ```kotlin
+ * ThreadContext.put("kotlin", "rocks") // Put a value into the ThreadContext.
+ * launch(immutableDiagnosticContext()) { // The contents of the ThreadContext are captured into the newly created CoroutineContext.
+ *   assert(ThreadContext.get("kotlin") == "rocks")
+ *
+ *   withContext(MutableDiagnosticContext().put("kotlin", "is great") {
+ *     assert(ThreadContext.get("kotlin") == "is great")
+ *
+ *     launch(Dispatchers.IO) {
+ *       assert(ThreadContext.get("kotlin") == "is great") // The diagnostic context is inherited by child CoroutineContexts.
+ *     }
+ *   }
+ *   assert(ThreadContext.get("kotlin") == "rocks") // The ThreadContext is reset when the CoroutineContext exits.
+ * }
+ * ```
+ *
+ * ## Combine with others
+ * You may wish to combine this [ThreadContextElement] with other [CoroutineContext]s.
+ *
+ * ```kotlin
+ * launch(Dispatchers.IO + immutableDiagnosticContext()) { ... }
+ * ```
+ *
+ * # CloseableThreadContext
+ * [org.apache.logging.log4j.CloseableThreadContext] is useful for automatically cleaning up the [ThreadContext] after a
+ * block of code. The structured concurrency provided by coroutines offers the same functionality.
+ *
+ * In the following example, the modifications to the [ThreadContext] are cleaned up when the coroutine exits.
+ *
+ * ```kotlin
+ * ThreadContext.put("kotlin", "rocks")
+ *
+ * withContext(MutableDiagnosticContext().put("kotlin", "is awesome") {
+ *   assert(ThreadContext.get("kotlin") == "is awesome")
+ * }
+ * assert(ThreadContext.get("kotlin") == "rocks")
+ * ```
+ *
+ * @param mappedContextBefore Mapped diagnostic context to apply for the duration of the corresponding [CoroutineContext].
+ * @param mappedContextAfter Mapped diagnostic context to restore when the corresponding [CoroutineContext] exits.
+ */
+public open class DiagnosticContext internal constructor(
+    private val mappedContextBefore: Map<String, String?>,
+    private val mappedContextAfter: Map<String, String?> = mappedContextBefore
+) : ThreadContextElement<DiagnosticContext>, AbstractCoroutineContextElement(Key) {
+
+    /**
+     * Key of [DiagnosticContext] in [CoroutineContext].
+     */
+    public companion object Key : CoroutineContext.Key<DiagnosticContext>
+
+    /** @suppress */
+    final override fun updateThreadContext(context: CoroutineContext): DiagnosticContext {
+        setCurrent(mappedContextBefore)
+        return this
+    }
+
+    /** @suppress */
+    final override fun restoreThreadContext(context: CoroutineContext, oldState: DiagnosticContext) {
+        setCurrent(oldState.mappedContextAfter)
+    }
+
+    private fun setCurrent(map: Map<String, String?>) {
+        /*
+         * The logic here varies significantly from how CloseableThreadContext works. CloseableThreadContext has the
+         * luxury of assuming that it is appending new state to the existing state of the current thread. We cannot make
+         * this assumption. It is very realistic for us to be restoring a context to a thread that has loads of state
+         * that we are not at all interested in, due to the Log4J ThreadContext being implemented as a ThreadLocal.
+         *
+         * So, to make sure that the ThreadLocal belonging to the Thread servicing this CoroutineContext is has the
+         * correct state, we first clear everything existing, and then apply the desired state.
+         */
+        ThreadContext.clearMap()
+        ThreadContext.putAll(map)
+    }
+}