Adds additional operations for working with MDC in suspending

jclyne · svc-squareup-copybara · commit b9426c89ab4c · 2025-02-06T21:18:57.000Z
and non-suspending contexts

GitOrigin-RevId: 2f1a75b27aa07cf9508a61ae9e10f60370aa8e3d
diff --git a/gradle/libs.versions.toml b/gradle/libs.versions.toml
@@ -151,6 +151,7 @@ kotlinStdLibJdk8 = { module = "org.jetbrains.kotlin:kotlin-stdlib-jdk8", version
 kotlinTest = { module = "org.jetbrains.kotlin:kotlin-test", version.ref = "kotlin" }
 kotlinxHtml = { module = "org.jetbrains.kotlinx:kotlinx-html-jvm", version = "0.12.0" }
 kotlinxCoroutinesCore = { module = "org.jetbrains.kotlinx:kotlinx-coroutines-core", version = "1.10.1" }
+kotlinxCoroutinesSlf4j = { module = "org.jetbrains.kotlinx:kotlinx-coroutines-slf4j", version = "1.10.1" }
 kotlinxCoroutinesTest = { module = "org.jetbrains.kotlinx:kotlinx-coroutines-test", version = "1.10.1" }
 kubernetesClient = { module = "io.kubernetes:client-java", version = "18.0.1" }
 kubernetesClientApi = { module = "io.kubernetes:client-java-api", version = "18.0.1" }
diff --git a/misk-api/api/misk-api.api b/misk-api/api/misk-api.api
@@ -50,7 +50,9 @@ public final class misk/client/NetworkInterceptorWrapper : okhttp3/Interceptor {
 public abstract interface class misk/logging/Mdc {
 	public abstract fun clear ()V
 	public abstract fun get (Ljava/lang/String;)Ljava/lang/String;
+	public abstract fun getCopyOfContextMap ()Ljava/util/Map;
 	public abstract fun put (Ljava/lang/String;Ljava/lang/String;)V
+	public abstract fun setContextMap (Ljava/util/Map;)V
 }
 
 public abstract interface class misk/scope/ActionScoped {
diff --git a/misk-api/src/main/kotlin/misk/logging/Mdc.kt b/misk-api/src/main/kotlin/misk/logging/Mdc.kt
@@ -1,10 +1,17 @@
 package misk.logging
 
+typealias MdcContextMap = Map<String, String>
+
 interface Mdc {
   fun put(key: String, value: String?)
 
   fun get(key: String): String?
 
   fun clear()
+
+  fun setContextMap(context: MdcContextMap)
+
+  fun getCopyOfContextMap(): MdcContextMap?
+
 }
 
diff --git a/misk/api/misk.api b/misk/api/misk.api
@@ -835,7 +835,19 @@ public final class misk/logging/MiskMdc : misk/logging/Mdc {
 	public static final field INSTANCE Lmisk/logging/MiskMdc;
 	public fun clear ()V
 	public fun get (Ljava/lang/String;)Ljava/lang/String;
+	public fun getCopyOfContextMap ()Ljava/util/Map;
 	public fun put (Ljava/lang/String;Ljava/lang/String;)V
+	public fun setContextMap (Ljava/util/Map;)V
+}
+
+public final class misk/logging/ScopedMdcKt {
+	public static final fun withMdc (Lmisk/logging/Mdc;Ljava/lang/String;Ljava/lang/String;Lkotlin/jvm/functions/Function0;)V
+	public static final fun withMdc (Lmisk/logging/Mdc;[Lkotlin/Pair;Lkotlin/jvm/functions/Function0;)V
+}
+
+public final class misk/logging/coroutines/ScopedMdcKt {
+	public static final fun withMdc (Lmisk/logging/Mdc;Ljava/lang/String;Ljava/lang/String;Lkotlin/jvm/functions/Function1;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
+	public static final fun withMdc (Lmisk/logging/Mdc;[Lkotlin/Pair;Lkotlin/jvm/functions/Function1;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
 }
 
 public final class misk/monitoring/JvmMetrics {
diff --git a/misk/build.gradle.kts b/misk/build.gradle.kts
@@ -54,6 +54,7 @@ dependencies {
   implementation(libs.kotlinReflect)
   implementation(libs.kotlinStdLibJdk8)
   implementation(libs.kotlinxCoroutinesCore)
+  implementation(libs.kotlinxCoroutinesSlf4j)
   implementation(libs.moshiAdapters)
   implementation(libs.okio)
   implementation(libs.openTracingConcurrent)
@@ -81,6 +82,7 @@ dependencies {
   testImplementation(libs.junitParams)
   testImplementation(libs.kotestAssertions)
   testImplementation(libs.kotlinTest)
+  testImplementation(libs.kotlinxCoroutinesCore)
   testImplementation(libs.kotlinxCoroutinesTest)
   testImplementation(libs.logbackClassic)
   testImplementation(libs.okHttpMockWebServer)
diff --git a/misk/src/main/kotlin/misk/logging/DynamicMdcContext.kt b/misk/src/main/kotlin/misk/logging/DynamicMdcContext.kt
diff --git a/misk/src/main/kotlin/misk/logging/MiskMdc.kt b/misk/src/main/kotlin/misk/logging/MiskMdc.kt
@@ -13,4 +13,14 @@ object MiskMdc : Mdc {
   override fun clear() {
     MDC.clear()
   }
+
+  override fun setContextMap(context: MdcContextMap) {
+    if (context.isNotEmpty()) {
+      MDC.setContextMap(context)
+    } else {
+      MDC.clear()
+    }
+  }
+
+  override fun getCopyOfContextMap(): MdcContextMap? = MDC.getCopyOfContextMap()
 }
diff --git a/misk/src/main/kotlin/misk/logging/ScopedMdc.kt b/misk/src/main/kotlin/misk/logging/ScopedMdc.kt
@@ -0,0 +1,23 @@
+package misk.logging
+
+import org.slf4j.MDC
+
+/**
+ * Adds the given key, value pair to the MDC for the duration of the block.
+ */
+inline fun Mdc.withMdc(key: String, value: String, block: () -> Unit) =
+  withMdc(key to value, block = block)
+
+/**
+ * Adds the given tags to the MDC for the duration of the block.
+ */
+inline fun Mdc.withMdc(vararg tags: Pair<String, String>, block: () -> Unit) {
+  val oldState = getCopyOfContextMap()
+  return try {
+    tags.forEach { (key, value) -> put(key, value) }
+    block()
+  } finally {
+    oldState?.let { setContextMap(it) } ?: clear()
+  }
+}
+
diff --git a/misk/src/main/kotlin/misk/logging/coroutines/ScopedMdc.kt b/misk/src/main/kotlin/misk/logging/coroutines/ScopedMdc.kt
@@ -0,0 +1,35 @@
+package misk.logging.coroutines
+
+import kotlinx.coroutines.slf4j.MDCContext
+import kotlinx.coroutines.withContext
+import misk.logging.Mdc
+import mu.KotlinLogging
+import wisp.logging.getLogger
+import kotlin.coroutines.coroutineContext
+
+/**
+ * Adds the given key, value pair to the MDC for the duration of the block.
+ * This is coroutine safe, so the additions will be added to the coroutine context
+ */
+suspend inline fun Mdc.withMdc(key: String, value: String, crossinline block: suspend () -> Unit) =
+  withMdc(key to value, block = block)
+
+/**
+ * Adds the given tags to the MDC for the duration of the block.
+ * This is coroutine safe, so the additions will be added to the coroutine context
+ */
+suspend inline fun Mdc.withMdc(
+  vararg tags: Pair<String, String>,
+  crossinline block: suspend () -> Unit
+) {
+  if(coroutineContext[MDCContext] == null) {
+    KotlinLogging.logger("misk.logging.coroutines.ScopedMdc").warn {
+      "MDCContext is not present in the coroutine context, this is required to restore the previous MDC state"
+    }
+  }
+  tags.forEach { (key, value) -> put(key, value) }
+  return withContext(MDCContext()) {
+    block()
+  }
+}
+
diff --git a/misk/src/main/kotlin/misk/web/actions/WebActions.kt b/misk/src/main/kotlin/misk/web/actions/WebActions.kt
@@ -2,11 +2,11 @@ package misk.web.actions
 
 import kotlinx.coroutines.launch
 import kotlinx.coroutines.runBlocking
+import kotlinx.coroutines.slf4j.MDCContext
 import misk.ApplicationInterceptor
 import misk.Chain
 import misk.grpc.GrpcMessageSinkChannel
 import misk.grpc.GrpcMessageSourceChannel
-import misk.logging.DynamicMdcContext
 import misk.scope.ActionScope
 import misk.web.HttpCall
 import misk.web.RealChain
@@ -40,7 +40,7 @@ internal fun WebAction.asChain(
       } else {
         // Handle suspending invocation, this includes building out the context to propagate MDC
         // and action scope.
-        val context = DynamicMdcContext() +
+        val context = MDCContext() +
           if (scope.inScope()) {
             scope.asContextElement()
           } else {
diff --git a/misk/src/test/kotlin/misk/logging/DynamicMdcContextTest.kt b/misk/src/test/kotlin/misk/logging/DynamicMdcContextTest.kt
diff --git a/misk/src/test/kotlin/misk/logging/ScopedMdcKtTest.kt b/misk/src/test/kotlin/misk/logging/ScopedMdcKtTest.kt
@@ -0,0 +1,79 @@
+package misk.logging
+
+import jakarta.inject.Inject
+import kotlinx.coroutines.delay
+import kotlinx.coroutines.slf4j.MDCContext
+import kotlinx.coroutines.test.runTest
+import misk.MiskTestingServiceModule
+import misk.testing.MiskTest
+import misk.testing.MiskTestModule
+import org.junit.jupiter.api.Test
+import kotlin.test.assertEquals
+import kotlin.test.assertNull
+import misk.logging.coroutines.withMdc as withMdcCoroutines
+
+@MiskTest(startService = false)
+internal class ScopedMdcKtTest {
+  @MiskTestModule
+  val module = MiskTestingServiceModule()
+
+  @Inject
+  lateinit var mdc: Mdc
+
+  @Test
+  fun `test withMdc in a coroutine for key value pairs`() = runTest(MDCContext()) {
+    val tags = (1..3).map { "key$it" to "value$it" }.toTypedArray()
+    mdc.withMdcCoroutines(*tags) {
+      tags.assertTags()
+      delay(100)
+      tags.assertTags()
+    }
+    tags.asserMissingTags()
+  }
+
+  @Test
+  fun `test withMdc for key value pairs`()  {
+    val tags = (1..3).map { "key$it" to "value$it" }.toTypedArray()
+    mdc.withMdc(*tags) {
+      tags.assertTags()
+    }
+    tags.asserMissingTags()
+  }
+
+  @Test
+  fun `test withMdc in a coroutine for key value pair overrides`() = runTest(MDCContext()) {
+    val tags = (1..3).map { "key$it" to "value$it" }.toTypedArray()
+    mdc.withMdcCoroutines(*tags) {
+      tags.assertTags()
+      delay(100)
+      tags.assertTags()
+      val updatedTags = tags.map { if (it.first == "key1"){it.first to it.second+"00"} else {it} }.toTypedArray()
+      mdc.withMdcCoroutines(*updatedTags) {
+        updatedTags.assertTags()
+        delay(100)
+        updatedTags.assertTags()
+      }
+      tags.forEach { it.asserTag() }
+    }
+    tags.asserMissingTags()
+  }
+
+  @Test
+  fun `test withMdc for key value pair overrides`()  {
+    val tags = (1..3).map { "key$it" to "value$it" }.toTypedArray()
+    mdc.withMdc(*tags) {
+      tags.assertTags()
+      val updatedTags = tags.map { if (it.first == "key1"){it.first to it.second+"00"} else {it} }.toTypedArray()
+      mdc.withMdc(*updatedTags) {
+        updatedTags.assertTags()
+      }
+      tags.forEach { it.asserTag() }
+    }
+    tags.asserMissingTags()
+  }
+  
+  fun Pair<String, String>.asserTag() = assertEquals(second, mdc.get(first))
+  fun Array<Pair<String, String>>.assertTags() = forEach { it.asserTag() }
+  fun Pair<String, String>.asserMissingTag() = assertNull( mdc.get(first))
+  fun Array<Pair<String, String>>.asserMissingTags() = forEach { it.asserMissingTag() }
+}

Original file line number	Diff line number	Diff line change
`@@ -50,7 +50,9 @@ public final class misk/client/NetworkInterceptorWrapper : okhttp3/Interceptor {`
`50`	`50`	`public abstract interface class misk/logging/Mdc {`
`51`	`51`	`public abstract fun clear ()V`
`52`	`52`	`public abstract fun get (Ljava/lang/String;)Ljava/lang/String;`
	`53`	`+ public abstract fun getCopyOfContextMap ()Ljava/util/Map;`
`53`	`54`	`public abstract fun put (Ljava/lang/String;Ljava/lang/String;)V`
	`55`	`+ public abstract fun setContextMap (Ljava/util/Map;)V`
`54`	`56`	`}`
`55`	`57`
`56`	`58`	`public abstract interface class misk/scope/ActionScoped {`