[FLINK-40208][runtime] Add test coverage and blank-key validation for JobMdcRegistry

docs/content.zh/docs/ops/logging_context.md

@@ -1,6 +1,6 @@
 ---
 title: "Logging Context (MDC)"
-weight: 7
+weight: 9
 type: docs
 ---
 <!--

docs/content/docs/ops/logging_context.md

@@ -1,6 +1,6 @@
 ---
 title: "Logging Context (MDC)"
-weight: 7
+weight: 9
 type: docs
 ---
 <!--

flink-core/src/main/java/org/apache/flink/util/MdcUtils.java

@@ -26,6 +26,7 @@
 import org.slf4j.MDC;
 
 import javax.annotation.Nonnull;
+import javax.annotation.concurrent.ThreadSafe;
 
 import java.util.Collections;
 import java.util.HashMap;
@@ -38,6 +39,7 @@
 import static org.apache.flink.util.Preconditions.checkArgument;
 
 /** Utility class to manage common Flink attributes in {@link MDC}. */
+@ThreadSafe
 public class MdcUtils {
 
     public static final String JOB_ID = "flink-job-id";
@@ -162,9 +164,12 @@ public static Map<String, String> asContextData(
                 jobConfiguration.get(MdcOptions.JOB_CONFIGURATION_TO_MDC_KEYS);
         final Map<String, String> context = new HashMap<>();
         for (Map.Entry<String, String> entry : mdcKeyMapping.entrySet()) {
+            final String mdcKeyName = entry.getValue();
+            final String effectiveMdcKey =
+                    (mdcKeyName == null || mdcKeyName.isBlank()) ? entry.getKey() : mdcKeyName;
             final String value = jobConfiguration.getString(entry.getKey(), null);
             if (value != null && !value.isBlank()) {
-                context.put(entry.getValue(), value);
+                context.put(effectiveMdcKey, value);
             }
         }
         if (context.isEmpty()) {

flink-core/src/test/java/org/apache/flink/util/MdcUtilsTest.java

@@ -43,6 +43,7 @@
 import org.slf4j.MDC;
 
 import java.util.Collections;
+import java.util.HashMap;
 import java.util.Map;
 import java.util.concurrent.Executor;
 import java.util.concurrent.ExecutorService;
@@ -311,6 +312,30 @@ void testKeySkippedWhenValueAbsentOrBlank(final String scenario, final String co
                 .isEqualTo(Collections.singletonMap(MdcUtils.JOB_ID, jobID.toHexString()));
     }
 
+    private static Stream<Arguments> blankMdcKeyNameCases() {
+        return Stream.of(
+                Arguments.of("empty string", ""),
+                Arguments.of("whitespace only", "   "),
+                Arguments.of("null key name", (String) null));
+    }
+
+    @ParameterizedTest
+    @MethodSource("blankMdcKeyNameCases")
+    void testBlankMdcKeyNameFallsBackToConfigKey(final String scenario, final String mdcKeyName) {
+        final JobID jobID = new JobID();
+        final Configuration conf = new Configuration();
+        final Map<String, String> keyMapping = new HashMap<>();
+        keyMapping.put("job.key-1", mdcKeyName);
+        conf.set(MdcOptions.JOB_CONFIGURATION_TO_MDC_KEYS, keyMapping);
+        conf.setString("job.key-1", "val-1");
+
+        final Map<String, String> context = MdcUtils.asContextData(jobID, conf);
+        assertThat(context)
+                .as(scenario)
+                .containsEntry(MdcUtils.JOB_ID, jobID.toHexString())
+                .containsEntry("job.key-1", "val-1");
+    }
+
     // --- JobMdcRegistry integration: registry-first lookup ---
 
     @Test

flink-runtime/src/test/java/org/apache/flink/runtime/dispatcher/DispatcherTest.java

@@ -107,6 +107,7 @@
 import org.apache.flink.util.FlinkException;
 import org.apache.flink.util.InstantiationUtil;
 import org.apache.flink.util.JobMdcRegistry;
+import org.apache.flink.util.MdcUtils;
 import org.apache.flink.util.Preconditions;
 import org.apache.flink.util.concurrent.FutureUtils;
 
@@ -572,12 +573,7 @@ public void testCancellationOfNonCanceledTerminalJobFailsWithAppropriateExceptio
 
     @Test
     public void testJobMdcContextRegisteredOnSubmissionAndClearedOnTermination() throws Exception {
-        final Map<String, String> keyMapping = new HashMap<>();
-        keyMapping.put("job.key-1", "mdc-key-1");
-        keyMapping.put("job.key-2", "mdc-key-2");
-        jobGraph.getJobConfiguration().set(MdcOptions.JOB_CONFIGURATION_TO_MDC_KEYS, keyMapping);
-        jobGraph.getJobConfiguration().setString("job.key-1", "val-1");
-        jobGraph.getJobConfiguration().setString("job.key-2", "val-2");
+        configureJobWithMdcEnrichment();
 
         final CompletableFuture<JobManagerRunnerResult> resultFuture = new CompletableFuture<>();
         dispatcher =
@@ -611,6 +607,32 @@ public void testJobMdcContextRegisteredOnSubmissionAndClearedOnTermination() thr
         CommonTestUtils.waitUntilCondition(() -> JobMdcRegistry.lookup(jobId) == null);
     }
 
+    @Test
+    public void testJobMdcContextRegisteredOnRecovery() throws Exception {
+        configureJobWithMdcEnrichment();
+
+        jobMasterLeaderElection.isLeader(UUID.randomUUID());
+
+        final TestingJobMasterServiceLeadershipRunnerFactory runnerFactory =
+                new TestingJobMasterServiceLeadershipRunnerFactory();
+        dispatcher =
+                createTestingDispatcherBuilder()
+                        .setJobManagerRunnerFactory(runnerFactory)
+                        .setRecoveredJobs(Collections.singleton(jobGraph))
+                        .build(rpcService);
+        dispatcher.start();
+
+        // takeCreatedJobManagerRunner blocks until the runner is created,
+        // which happens AFTER registerOrClear in runRecoveredJob
+        runnerFactory.takeCreatedJobManagerRunner();
+
+        assertThat(JobMdcRegistry.lookup(jobId))
+                .containsEntry(MdcUtils.JOB_ID, jobId.toHexString())
+                .containsEntry("mdc-key-1", "val-1")
+                .containsEntry("mdc-key-2", "val-2")
+                .hasSize(3);
+    }
+
     @Test
     public void testNoHistoryServerArchiveCreatedForSuspendedJob() throws Exception {
         final CompletableFuture<Void> archiveAttemptFuture = new CompletableFuture<>();
@@ -675,6 +697,15 @@ private void mockApplicationFinished() throws Exception {
                 .get();
     }
 
+    private void configureJobWithMdcEnrichment() {
+        final Map<String, String> keyMapping = new HashMap<>();
+        keyMapping.put("job.key-1", "mdc-key-1");
+        keyMapping.put("job.key-2", "mdc-key-2");
+        jobGraph.getJobConfiguration().set(MdcOptions.JOB_CONFIGURATION_TO_MDC_KEYS, keyMapping);
+        jobGraph.getJobConfiguration().setString("job.key-1", "val-1");
+        jobGraph.getJobConfiguration().setString("job.key-2", "val-2");
+    }
+
     @Test
     public void testJobManagerRunnerInitializationFailureFailsJob() throws Exception {
         final TestingJobMasterServiceLeadershipRunnerFactory testingJobManagerRunnerFactory =