Skip to content

Commit 4b438b8

Browse files
maniax89maniax89
committed
fix(jvm): apply review feedback from copilot
1 parent 6ff2f9a commit 4b438b8

3 files changed

Lines changed: 93 additions & 80 deletions

File tree

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

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -104,7 +104,15 @@ public suspend inline fun <T> withCoroutineLoggingContext(
104104
crossinline body: suspend () -> T,
105105
): T {
106106
val parent = currentCoroutineContext()[MDCContext]?.contextMap ?: emptyMap()
107-
val merged = (parent + map.filterValues { it != null }).mapValues { it.value!! }
107+
val merged = parent.toMutableMap().apply {
108+
for((key, value) in map) {
109+
if (value == null) {
110+
remove(key)
111+
} else {
112+
put(key, value)
113+
}
114+
}
115+
}
108116

109117
return withContext(MDCContext(merged)) { body() }
110118
}

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

Lines changed: 81 additions & 76 deletions
Original file line numberDiff line numberDiff line change
@@ -163,98 +163,103 @@ class KotlinLoggingAsyncMDCTest {
163163
assertEquals("l", MDC.get("k"))
164164
}
165165

166+
166167
@Test
167-
fun `withLoggingContextAsync leaks context across coroutines`() = runTest {
168+
fun `withLoggingContextAsync leaks context across coroutines per GH issue 75`() = runTest {
168169
val dispatcher = Executors.newFixedThreadPool(1).asCoroutineDispatcher()
169-
// establish order of events with deferreds:
170-
// 1. request A will start, set its context, and then wait for request B (before exiting
171-
// withLoggingContextAsync)
172-
// 2. request B will start, set its context, and then capture the MDC
173-
val requestAStarted = CompletableDeferred<Unit>()
174-
val requestBDone = CompletableDeferred<Unit>()
175-
176-
val requestADuringBlock = CompletableDeferred<Map<String, String>?>()
177-
val requestAAfterBlock = CompletableDeferred<Map<String, String>?>()
178-
val requestBDuringBlock = CompletableDeferred<Map<String, String>?>()
179-
val requestBAfterBlock = CompletableDeferred<Map<String, String>?>()
180-
181-
val jobA =
182-
launch(dispatcher) {
183-
withLoggingContextAsync("foo" to "original") {
184-
requestAStarted.complete(Unit)
185-
requestADuringBlock.complete(MDC.getCopyOfContextMap())
186-
requestBDone.await() // Suspend while B runs
170+
dispatcher.use {
171+
// establish order of events with deferreds:
172+
// 1. request A will start, set its context, and then wait for request B (before exiting
173+
// withLoggingContextAsync)
174+
// 2. request B will start, set its context, and then capture the MDC
175+
val requestAStarted = CompletableDeferred<Unit>()
176+
val requestBDone = CompletableDeferred<Unit>()
177+
178+
val requestADuringBlock = CompletableDeferred<Map<String, String>>()
179+
val requestAAfterBlock = CompletableDeferred<Map<String, String>>()
180+
val requestBDuringBlock = CompletableDeferred<Map<String, String>>()
181+
val requestBAfterBlock = CompletableDeferred<Map<String, String>>()
182+
183+
val jobA =
184+
launch(dispatcher) {
185+
withLoggingContextAsync("foo" to "original") {
186+
requestAStarted.complete(Unit)
187+
requestADuringBlock.complete(MDC.getCopyOfContextMap() ?: emptyMap())
188+
requestBDone.await() // Suspend while B runs
189+
}
190+
// will be empty
191+
requestAAfterBlock.complete(MDC.getCopyOfContextMap() ?: emptyMap())
187192
}
188-
// will be empty
189-
requestAAfterBlock.complete(MDC.getCopyOfContextMap())
190-
}
191193

192-
val jobB =
193-
launch(dispatcher) {
194-
requestAStarted.await()
195-
withLoggingContextAsync("foo" to "bar") {
196-
requestBDuringBlock.complete(MDC.getCopyOfContextMap())
194+
val jobB =
195+
launch(dispatcher) {
196+
requestAStarted.await()
197+
withLoggingContextAsync("foo" to "bar") {
198+
requestBDuringBlock.complete(MDC.getCopyOfContextMap() ?: emptyMap())
199+
}
200+
// Capture MDC state immediately after Request B's block, will NOT be empty due to the bug
201+
requestBAfterBlock.complete(MDC.getCopyOfContextMap() ?: emptyMap())
202+
requestBDone.complete(Unit)
197203
}
198-
// Capture MDC state immediately after Request B's block, will NOT be empty due to the bug
199-
requestBAfterBlock.complete(MDC.getCopyOfContextMap())
200-
requestBDone.complete(Unit)
201-
}
202204

203-
jobA.join()
204-
jobB.join()
205+
jobA.join()
206+
jobB.join()
205207

206-
assertEquals(mapOf("foo" to "original"), requestADuringBlock.await())
207-
assertEquals(emptyMap(), requestAAfterBlock.await())
208-
assertEquals(mapOf("foo" to "bar"), requestBDuringBlock.await())
209-
// THIS IS THE BUG: Request B sees "original" after its own block exits
210-
// It should be null (no context) but it's leaking Request A's value
211-
assertEquals(mapOf("foo" to "original"), requestBAfterBlock.await())
208+
assertEquals(mapOf("foo" to "original"), requestADuringBlock.await())
209+
assertEquals(emptyMap(), requestAAfterBlock.await())
210+
assertEquals(mapOf("foo" to "bar"), requestBDuringBlock.await())
211+
// THIS IS THE BUG: Request B sees "original" after its own block exits
212+
// It should be null (no context) but it's leaking Request A's value
213+
assertEquals(mapOf("foo" to "original"), requestBAfterBlock.await())
214+
}
212215
}
213216

214217
@Test
215-
fun `withCoroutineLoggingContext does not leak context across coroutines`() = runTest {
218+
fun `withCoroutineLoggingContext binds logging context to current coroutine contextMap`() = runTest {
216219
val dispatcher = Executors.newFixedThreadPool(1).asCoroutineDispatcher()
217-
// establish order of events with deferreds:
218-
// 1. request A will start, set its context, and then wait for request B (before exiting
219-
// withLoggingContextAsync)
220-
// 2. request B will start, set its context, and then capture the MDC
221-
val requestAStarted = CompletableDeferred<Unit>()
222-
val requestBDone = CompletableDeferred<Unit>()
223-
224-
val requestADuringBlock = CompletableDeferred<Map<String, String>?>()
225-
val requestAAfterBlock = CompletableDeferred<Map<String, String>?>()
226-
val requestBDuringBlock = CompletableDeferred<Map<String, String>?>()
227-
val requestBAfterBlock = CompletableDeferred<Map<String, String>?>()
228-
229-
val jobA =
230-
launch(dispatcher) {
231-
withCoroutineLoggingContext("foo" to "original") {
232-
requestAStarted.complete(Unit)
233-
requestADuringBlock.complete(MDC.getCopyOfContextMap())
234-
requestBDone.await() // Suspend while B runs
220+
dispatcher.use {
221+
// establish order of events with deferreds:
222+
// 1. request A will start, set its context, and then wait for request B (before exiting
223+
// withLoggingContextAsync)
224+
// 2. request B will start, set its context, and then capture the MDC
225+
val requestAStarted = CompletableDeferred<Unit>()
226+
val requestBDone = CompletableDeferred<Unit>()
227+
228+
val requestADuringBlock = CompletableDeferred<Map<String, String>>()
229+
val requestAAfterBlock = CompletableDeferred<Map<String, String>>()
230+
val requestBDuringBlock = CompletableDeferred<Map<String, String>>()
231+
val requestBAfterBlock = CompletableDeferred<Map<String, String>>()
232+
233+
val jobA =
234+
launch(dispatcher) {
235+
withCoroutineLoggingContext("foo" to "original") {
236+
requestAStarted.complete(Unit)
237+
requestADuringBlock.complete(MDC.getCopyOfContextMap() ?: emptyMap())
238+
requestBDone.await() // Suspend while B runs
239+
}
240+
// will be empty
241+
requestAAfterBlock.complete(MDC.getCopyOfContextMap() ?: emptyMap())
235242
}
236-
// will be empty
237-
requestAAfterBlock.complete(MDC.getCopyOfContextMap())
238-
}
239243

240-
val jobB =
241-
launch(dispatcher) {
242-
requestAStarted.await()
243-
withCoroutineLoggingContext("foo" to "bar") {
244-
requestBDuringBlock.complete(MDC.getCopyOfContextMap())
244+
val jobB =
245+
launch(dispatcher) {
246+
requestAStarted.await()
247+
withCoroutineLoggingContext("foo" to "bar") {
248+
requestBDuringBlock.complete(MDC.getCopyOfContextMap() ?: emptyMap())
249+
}
250+
// will also be empty
251+
requestBAfterBlock.complete(MDC.getCopyOfContextMap() ?: emptyMap())
252+
requestBDone.complete(Unit)
245253
}
246-
// will also be empty
247-
requestBAfterBlock.complete(MDC.getCopyOfContextMap())
248-
requestBDone.complete(Unit)
249-
}
250254

251-
jobA.join()
252-
jobB.join()
255+
jobA.join()
256+
jobB.join()
253257

254-
assertEquals(mapOf("foo" to "original"), requestADuringBlock.await())
255-
assertEquals(emptyMap(), requestAAfterBlock.await())
256-
assertEquals(mapOf("foo" to "bar"), requestBDuringBlock.await())
257-
assertEquals(emptyMap(), requestBAfterBlock.await())
258+
assertEquals(mapOf("foo" to "original"), requestADuringBlock.await())
259+
assertEquals(emptyMap(), requestAAfterBlock.await())
260+
assertEquals(mapOf("foo" to "bar"), requestBDuringBlock.await())
261+
assertEquals(emptyMap(), requestBAfterBlock.await())
262+
}
258263
}
259264

260265
@Test

‎src/jvmTest/kotlin/io/github/oshai/kotlinlogging/coroutines/README.md‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -20,9 +20,9 @@ withLoggingContextAsync("userId" to "ANOTHER_USER_ID", restorePrevious = false)
2020

2121
## Coroutine-Safe wrappers
2222

23-
If you prefer to leverage coroutine context instead of thread local storage, you can use the `withCoroutineLoggingContext` wrapper.
24-
This will NOT inherit the current MDC set prior, but will carry existing coroutine context within nested declarations, restoring previous context
25-
at the logical coroutine level rather than thread level.
23+
If you prefer to leverage coroutine context to propogate MDC safely across suspensions, you can use the `withCoroutineLoggingContext` wrapper.
24+
This will NOT inherit the current MDC set on the calling thread, but will carry the logging context established by outer `withCoroutineLoggingContext` calls
25+
within nested declarations, restoring previous MDC for the current coroutine when the block completes.
2626

2727
```kotlin
2828
MDC.put("foo", "bar")

0 commit comments

Comments
 (0)