feat(jvm) add withCoroutineLoggingContext

Author: maniax89

src/jvmMain/kotlin/io/github/oshai/kotlinlogging/coroutines/KotlinLoggingAsyncMDC.kt

@@ -26,6 +26,19 @@ public suspend inline fun <T> withLoggingContextAsync(
     withContext(MDCContext()) { body() }
   }
 
+/**
+ * Use a pair inheriting coroutine context but not current MDC Context. Example:
+ * ```
+ * withLoggingContextAsync("userId" to userId) {
+ *   doSomething()
+ * }
+ * ```
+ */
+public suspend inline fun <T> withCoroutineLoggingContext(
+  pair: Pair<String, String?>,
+  crossinline body: suspend () -> T,
+): T = withCoroutineLoggingContext(mapOf(pair), body)
+
 /**
  * Use a varying number of pairs in an asynchronous MDC context. Example:
  * ```
@@ -43,6 +56,19 @@ public suspend inline fun <T> withLoggingContextAsync(
     withContext(MDCContext()) { body() }
   }
 
+/**
+ * Use a varying number of pairs inheriting coroutine context but not current MDC Context. Example:
+ * ```
+ * withCoroutineLoggingContext("userId" to userId) {
+ *   doSomething()
+ * }
+ * ```
+ */
+public suspend inline fun <T> withCoroutineLoggingContext(
+  vararg pair: Pair<String, String?>,
+  crossinline body: suspend () -> T,
+): T = withCoroutineLoggingContext(pair.toMap(), body)
+
 /**
  * Use a map in an asynchronous MDC context. Example:
  * ```
@@ -64,3 +90,21 @@ public suspend inline fun <T> withLoggingContextAsync(
   withLoggingContext(map, restorePrevious = restorePrevious) {
     withContext(MDCContext()) { body() }
   }
+
+/**
+ * Use a map inheriting coroutine context but not current MDC Context. Example:
+ * ```
+ * withCoroutineLoggingContext(mapOf("userId" to userId)) {
+ *   doSomething()
+ * }
+ * ```
+ */
+public suspend inline fun <T> withCoroutineLoggingContext(
+  map: Map<String, String?>,
+  crossinline body: suspend () -> T,
+): T {
+  val parent = currentCoroutineContext()[MDCContext]?.contextMap ?: emptyMap()
+  val merged = (parent + map.filterValues { it != null }).mapValues { it.value!! }
+
+  return withContext(MDCContext(merged)) { body() }
+}

src/jvmTest/kotlin/io/github/oshai/kotlinlogging/coroutines/KotlinLoggingAsyncMDCTest.kt

@@ -1,17 +1,16 @@
 package io.github.oshai.kotlinlogging.coroutines
 
-import kotlinx.coroutines.CompletableDeferred
+import java.util.concurrent.Executors
 import kotlin.test.*
+import kotlinx.coroutines.CompletableDeferred
 import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.asCoroutineDispatcher
 import kotlinx.coroutines.launch
-import kotlinx.coroutines.runBlocking
 import kotlinx.coroutines.test.runTest
 import org.apache.logging.log4j.*
 import org.apache.logging.log4j.core.config.*
 import org.junit.jupiter.api.BeforeEach
 import org.slf4j.*
-import java.util.concurrent.Executors
 
 @ExperimentalCoroutinesApi
 class KotlinLoggingAsyncMDCTest {
@@ -168,37 +167,126 @@ class KotlinLoggingAsyncMDCTest {
   fun `withLoggingContextAsync leaks context across coroutines`() = runTest {
     val dispatcher = Executors.newFixedThreadPool(1).asCoroutineDispatcher()
     // establish order of events with deferreds:
-    // 1. request A will start, set its context, and then wait for request B (before exiting withLoggingContextAsync)
+    // 1. request A will start, set its context, and then wait for request B (before exiting
+    // withLoggingContextAsync)
     // 2. request B will start, set its context, and then capture the MDC
     val requestAStarted = CompletableDeferred<Unit>()
     val requestBDone = CompletableDeferred<Unit>()
 
     val requestAAfterBlock = CompletableDeferred<Map<String, String>?>()
     val requestBAfterBlock = CompletableDeferred<Map<String, String>?>()
 
-    val jobA = launch(dispatcher) {
-      withLoggingContextAsync("foo" to "original") {
-        requestAStarted.complete(Unit)
-        requestBDone.await() // Suspend while B runs
+    val jobA =
+      launch(dispatcher) {
+        withLoggingContextAsync("foo" to "original") {
+          requestAStarted.complete(Unit)
+          requestBDone.await() // Suspend while B runs
+        }
+        // will be empty
+        requestAAfterBlock.complete(MDC.getCopyOfContextMap())
       }
-      // will be empty
-      requestAAfterBlock.complete(MDC.getCopyOfContextMap())
-    }
 
-    val jobB = launch(dispatcher) {
-      requestAStarted.await()
-      withLoggingContextAsync("foo" to "bar") { }
-      // Capture MDC state immediately after Request B's block, will NOT be empty due to the bug
-      requestBAfterBlock.complete(MDC.getCopyOfContextMap())
-      requestBDone.complete(Unit)
-    }
+    val jobB =
+      launch(dispatcher) {
+        requestAStarted.await()
+        withLoggingContextAsync("foo" to "bar") {}
+        // Capture MDC state immediately after Request B's block, will NOT be empty due to the bug
+        requestBAfterBlock.complete(MDC.getCopyOfContextMap())
+        requestBDone.complete(Unit)
+      }
 
     jobA.join()
     jobB.join()
 
     assertEquals(emptyMap(), requestAAfterBlock.await())
-    // THIS IS THE BUG: Request B sees "AAA" after its own block exits
+    // THIS IS THE BUG: Request B sees "original" after its own block exits
     // It should be null (no context) but it's leaking Request A's value
-    assertEquals(mapOf("foo" to "original"), requestBAfterBlock.await()) // This will FAIL
+    assertEquals(mapOf("foo" to "original"), requestBAfterBlock.await())
+  }
+
+  @Test
+  fun `withCoroutineLoggingContext does not leak context across coroutines`() = runTest {
+    val dispatcher = Executors.newFixedThreadPool(1).asCoroutineDispatcher()
+    // establish order of events with deferreds:
+    // 1. request A will start, set its context, and then wait for request B (before exiting
+    // withLoggingContextAsync)
+    // 2. request B will start, set its context, and then capture the MDC
+    val requestAStarted = CompletableDeferred<Unit>()
+    val requestBDone = CompletableDeferred<Unit>()
+
+    val requestAAfterBlock = CompletableDeferred<Map<String, String>?>()
+    val requestBAfterBlock = CompletableDeferred<Map<String, String>?>()
+
+    val jobA =
+      launch(dispatcher) {
+        withCoroutineLoggingContext("foo" to "original") {
+          requestAStarted.complete(Unit)
+          requestBDone.await() // Suspend while B runs
+        }
+        // will be empty
+        requestAAfterBlock.complete(MDC.getCopyOfContextMap())
+      }
+
+    val jobB =
+      launch(dispatcher) {
+        requestAStarted.await()
+        withCoroutineLoggingContext("foo" to "bar") {}
+        // Capture MDC state immediately after Request B's block, will NOT be empty due to the bug
+        requestBAfterBlock.complete(MDC.getCopyOfContextMap())
+        requestBDone.complete(Unit)
+      }
+
+    jobA.join()
+    jobB.join()
+
+    assertEquals(emptyMap(), requestAAfterBlock.await())
+    assertEquals(emptyMap(), requestBAfterBlock.await())
+  }
+
+  @Test
+  fun `map withCoroutineLoggingContext`() = runTest {
+    assertNull(MDC.get("a"))
+    assertNull(MDC.get("c"))
+    assertNull(MDC.get("e"))
+    assertNull(MDC.get("f"))
+    assertNull(MDC.get("k"))
+
+    // original MDC outside of coroutine context
+    MDC.put("e", "g")
+    MDC.put("k", "l")
+
+    withCoroutineLoggingContext(mapOf("a" to "b", "c" to "d", "e" to null, "f" to "h")) {
+      assertEquals("b", MDC.get("a"))
+      assertEquals("d", MDC.get("c"))
+      // does NOT inherit original MDC
+      assertNull(MDC.get("e"))
+      assertEquals("h", MDC.get("f"))
+      // does NOT inherit original MDC
+      assertNull(MDC.get("k"))
+
+      withCoroutineLoggingContext(mapOf("a" to "b", "e" to "i", "f" to "j")) {
+        assertEquals("b", MDC.get("a"))
+        // DOES inherit previous coroutine context
+        assertEquals("d", MDC.get("c"))
+        assertEquals("i", MDC.get("e"))
+        assertEquals("j", MDC.get("f"))
+        assertNull(MDC.get("k"))
+      }
+
+      assertEquals("b", MDC.get("a"))
+      assertEquals("d", MDC.get("c"))
+      assertNull(MDC.get("e"))
+      assertEquals("h", MDC.get("f"))
+      assertNull(MDC.get("k"))
+    }
+
+    assertNull(MDC.get("a"))
+    assertNull(MDC.get("c"))
+
+    // original MDC is restored
+    assertEquals("g", MDC.get("e"))
+    assertNull(MDC.get("f"))
+    // original MDC is restored
+    assertEquals("l", MDC.get("k"))
   }
 }