not log context when explicitly constructing the MessageWithContext w…

…hen conf is disabled:
anchovYu · Aug 13, 2024 · d5e37d8 · d5e37d8
1 parent 70b814b
commit d5e37d8
Show file tree

Hide file tree

Showing 2 changed files with 47 additions and 5 deletions.
diff --git a/common/utils/src/main/scala/org/apache/spark/internal/Logging.scala b/common/utils/src/main/scala/org/apache/spark/internal/Logging.scala
@@ -151,9 +151,7 @@ trait Logging {
       args.foreach { mdc =>
         val value = if (mdc.value != null) mdc.value.toString else null
         sb.append(value)
-        if (Logging.isStructuredLoggingEnabled) {
-          context.put(mdc.key.name, value)
-        }
+        context.put(mdc.key.name, value)
 
         if (processedParts.hasNext) {
           sb.append(StringContext.processEscapes(processedParts.next()))
@@ -165,7 +163,13 @@ trait Logging {
   }
 
   protected def withLogContext(context: java.util.HashMap[String, String])(body: => Unit): Unit = {
-    val threadContext = CloseableThreadContext.putAll(context)
+    // put into context only when structured logging is enabled
+    val updatedContext = if (Logging.isStructuredLoggingEnabled) {
+      context
+    } else {
+      new java.util.HashMap[String, String]()
+    }
+    val threadContext = CloseableThreadContext.putAll(updatedContext)
     try {
       body
     } finally {

diff --git a/common/utils/src/test/scala/org/apache/spark/util/StructuredLoggingSuite.scala b/common/utils/src/test/scala/org/apache/spark/util/StructuredLoggingSuite.scala
@@ -40,7 +40,7 @@ trait LoggingSuiteBase
   }
 
   // Return the newly added log contents in the log file after executing the function `f`
-  private def captureLogOutput(f: () => Unit): String = {
+  protected def captureLogOutput(f: () => Unit): String = {
     val content = if (logFile.exists()) {
       Files.readString(logFile.toPath)
     } else {
@@ -438,6 +438,44 @@ class StructuredLoggingSuite extends LoggingSuiteBase {
     assert((log"\r" + log"\n" + log"\t" + log"\b").message == "\r\n\t\b")
     assert((log"\r${MDC(LogKeys.EXECUTOR_ID, 1)}\n".message == "\r1\n"))
   }
+
+  test("disabled structured logging won't log context") {
+    Logging.disableStructuredLogging()
+    val expectedPatternWithoutContext = compactAndToRegexPattern(
+      s"""
+        {
+          "ts": "<timestamp>",
+          "level": "INFO",
+          "msg": "Lost executor 1.",
+          "logger": "$className"
+        }""")
+
+    Seq(
+      () => logInfo(log"Lost executor ${MDC(LogKeys.EXECUTOR_ID, "1")}."),
+      () => logInfo( // blocked when explicitly constructing the MessageWithContext
+        MessageWithContext(
+          "Lost executor 1.",
+          new java.util.HashMap[String, String] { put(LogKeys.EXECUTOR_ID.name, "1") }
+        )
+      )
+    ).foreach { f =>
+      val logOutput = captureLogOutput(f)
+      assert(expectedPatternWithoutContext.r.matches(logOutput))
+    }
+    Logging.enableStructuredLogging()
+  }
+
+  test("setting to MDC gets logged") {
+    val mdcPattern = s""""${LogKeys.DATA.name}":"some-data""""
+
+    org.slf4j.MDC.put(LogKeys.DATA.name, "some-data")
+    val logOutputWithMDCSet = captureLogOutput(() => logInfo(msgWithMDC))
+    assert(mdcPattern.r.findFirstIn(logOutputWithMDCSet).isDefined)
+
+    org.slf4j.MDC.remove(LogKeys.DATA.name)
+    val logOutputWithoutMDCSet = captureLogOutput(() => logInfo(msgWithMDC))
+    assert(mdcPattern.r.findFirstIn(logOutputWithoutMDCSet).isEmpty)
+  }
 }
 
 object CustomLogKeys {