feat(jvm) add withCoroutineLoggingContext

Author: Andrew Turgeon

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

@@ -26,6 +26,24 @@ 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()
+ * }
+ * ```
+ * ```
+ * withLoggingContextAsync("userId" to userId, restorePrevious = false) {
+ *   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 +61,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 +95,22 @@ 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

@@ -197,8 +197,91 @@ class KotlinLoggingAsyncMDCTest {
     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 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"))
   }
 }