From 64c7a675be16d44fde291ee21a973e3685bb8827 Mon Sep 17 00:00:00 2001
From: Aleksandr Iushmanov <aiushmanov@confluent.io>
Date: Tue, 4 Aug 2026 10:13:48 +0100
Subject: [PATCH] [FLINK-40208] Add test coverage and blank-key validation for
 JobMdcRegistry

Generated-by: Claude Code
---
 docs/content.zh/docs/ops/logging_context.md   |   2 +-
 docs/content/docs/ops/logging_context.md      |   2 +-
 .../java/org/apache/flink/util/MdcUtils.java  |   7 +-
 .../org/apache/flink/util/MdcUtilsTest.java   |  25 +++
 .../runtime/dispatcher/DispatcherTest.java    |  43 +++-
 .../TaskExecutorSubmissionTest.java           | 190 ++++++++++++++++++
 .../TaskSubmissionTestEnvironment.java        |  14 +-
 7 files changed, 272 insertions(+), 11 deletions(-)

diff --git a/docs/content.zh/docs/ops/logging_context.md b/docs/content.zh/docs/ops/logging_context.md
index fb06fdeb4f1c44..fb89b9fd41e748 100644
--- a/docs/content.zh/docs/ops/logging_context.md
+++ b/docs/content.zh/docs/ops/logging_context.md
@@ -1,6 +1,6 @@
 ---
 title: "Logging Context (MDC)"
-weight: 7
+weight: 9
 type: docs
 ---
 <!--
diff --git a/docs/content/docs/ops/logging_context.md b/docs/content/docs/ops/logging_context.md
index fb06fdeb4f1c44..fb89b9fd41e748 100644
--- a/docs/content/docs/ops/logging_context.md
+++ b/docs/content/docs/ops/logging_context.md
@@ -1,6 +1,6 @@
 ---
 title: "Logging Context (MDC)"
-weight: 7
+weight: 9
 type: docs
 ---
 <!--
diff --git a/flink-core/src/main/java/org/apache/flink/util/MdcUtils.java b/flink-core/src/main/java/org/apache/flink/util/MdcUtils.java
index 979b8a2c94310e..47a7a5f1b33bc0 100644
--- a/flink-core/src/main/java/org/apache/flink/util/MdcUtils.java
+++ b/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()) {
diff --git a/flink-core/src/test/java/org/apache/flink/util/MdcUtilsTest.java b/flink-core/src/test/java/org/apache/flink/util/MdcUtilsTest.java
index 8075bffd9a1047..457bb982bfddf1 100644
--- a/flink-core/src/test/java/org/apache/flink/util/MdcUtilsTest.java
+++ b/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
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/dispatcher/DispatcherTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/dispatcher/DispatcherTest.java
index 5bac062a99dbfb..1c85c7c3b2c21a 100755
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/dispatcher/DispatcherTest.java
+++ b/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 =
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/taskexecutor/TaskExecutorSubmissionTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/taskexecutor/TaskExecutorSubmissionTest.java
index c13445e43b9f70..8e87b86fcadd67 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/taskexecutor/TaskExecutorSubmissionTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/taskexecutor/TaskExecutorSubmissionTest.java
@@ -18,12 +18,21 @@
 
 package org.apache.flink.runtime.taskexecutor;
 
+import org.apache.flink.api.common.ApplicationID;
 import org.apache.flink.api.common.ExecutionConfig;
 import org.apache.flink.api.common.JobID;
+import org.apache.flink.configuration.BlobServerOptions;
 import org.apache.flink.configuration.Configuration;
 import org.apache.flink.configuration.MdcOptions;
 import org.apache.flink.configuration.NettyShuffleEnvironmentOptions;
+import org.apache.flink.runtime.blob.BlobServer;
+import org.apache.flink.runtime.blob.BlobWriter;
+import org.apache.flink.runtime.blob.JobPermanentBlobService;
+import org.apache.flink.runtime.blob.NoOpTransientBlobService;
 import org.apache.flink.runtime.blob.PermanentBlobKey;
+import org.apache.flink.runtime.blob.TaskExecutorBlobService;
+import org.apache.flink.runtime.blob.TransientBlobService;
+import org.apache.flink.runtime.blob.VoidBlobStore;
 import org.apache.flink.runtime.clusterframework.types.AllocationID;
 import org.apache.flink.runtime.clusterframework.types.ResourceID;
 import org.apache.flink.runtime.deployment.InputGateDeploymentDescriptor;
@@ -56,9 +65,11 @@
 import org.apache.flink.runtime.taskexecutor.slot.TaskSlotTable;
 import org.apache.flink.runtime.taskmanager.Task;
 import org.apache.flink.runtime.testtasks.BlockingNoOpInvokable;
+import org.apache.flink.runtime.testutils.CommonTestUtils;
 import org.apache.flink.runtime.util.NettyShuffleDescriptorBuilder;
 import org.apache.flink.testutils.TestingUtils;
 import org.apache.flink.testutils.executor.TestExecutorExtension;
+import org.apache.flink.types.Either;
 import org.apache.flink.util.JobMdcRegistry;
 import org.apache.flink.util.MdcUtils;
 import org.apache.flink.util.NetUtils;
@@ -71,10 +82,14 @@
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.TestInfo;
 import org.junit.jupiter.api.extension.RegisterExtension;
+import org.junit.jupiter.api.io.TempDir;
 import org.mockito.Mockito;
 
+import java.io.File;
 import java.io.IOException;
+import java.net.InetSocketAddress;
 import java.net.URL;
+import java.nio.file.Path;
 import java.time.Duration;
 import java.util.Collection;
 import java.util.Collections;
@@ -84,11 +99,13 @@
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.TimeUnit;
 
 import static org.apache.flink.core.testutils.FlinkAssertions.assertThatFuture;
 import static org.apache.flink.runtime.executiongraph.ExecutionGraphTestUtils.createExecutionAttemptId;
 import static org.apache.flink.runtime.util.NettyShuffleDescriptorBuilder.createRemoteWithIdAndLocation;
 import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
 import static org.mockito.Mockito.doThrow;
 import static org.mockito.Mockito.mock;
 
@@ -179,6 +196,104 @@ void testJobMdcContextRegisteredOnSubmitTask() throws Exception {
         }
     }
 
+    /**
+     * Tests that submitting a task with offloaded (blob-stored) job information still populates the
+     * JobMdcRegistry after loadBigData resolves the offloaded data.
+     */
+    @Test
+    void testJobMdcContextRegisteredOnSubmitTaskWithOffloadedJobInfo(@TempDir Path tempDir)
+            throws Exception {
+        try (BlobServer blobServer = createBlobServer(tempDir)) {
+            final ExecutionAttemptID eid = createExecutionAttemptId();
+
+            final Configuration jobConfiguration = singleKeyMdcJobConfiguration();
+
+            final JobInformation jobInformation =
+                    new JobInformation(
+                            jobId,
+                            JobType.STREAMING,
+                            "offloaded-test",
+                            new SerializedValue<>(new ExecutionConfig()),
+                            jobConfiguration,
+                            Collections.emptyList(),
+                            Collections.emptyList());
+
+            final Either<SerializedValue<JobInformation>, PermanentBlobKey> offloadResult =
+                    BlobWriter.serializeAndTryOffload(jobInformation, jobId, blobServer);
+            assertThat(offloadResult.isRight())
+                    .as("job info must be offloaded (OFFLOAD_MINSIZE=0)")
+                    .isTrue();
+            final PermanentBlobKey jobInfoBlobKey = offloadResult.right();
+
+            final TaskDeploymentDescriptor tdd =
+                    createOffloadedTaskDeploymentDescriptor(eid, jobInfoBlobKey);
+
+            assertThatThrownBy(tdd::getJobInformation).isInstanceOf(IllegalStateException.class);
+
+            try (TaskSubmissionTestEnvironment env =
+                    new TaskSubmissionTestEnvironment.Builder(jobId)
+                            .setSlotSize(1)
+                            .setTaskExecutorBlobService(blobServiceFrom(blobServer))
+                            .build(EXECUTOR_EXTENSION.getExecutor())) {
+                TaskExecutorGateway tmGateway = env.getTaskExecutorGateway();
+                TaskSlotTable taskSlotTable = env.getTaskSlotTable();
+
+                taskSlotTable.allocateSlot(0, jobId, tdd.getAllocationId(), Duration.ofSeconds(60));
+                tmGateway.submitTask(tdd, env.getJobMasterId(), timeout).get();
+
+                assertThat(MdcUtils.asContextData(jobId))
+                        .containsEntry(MdcUtils.JOB_ID, jobId.toHexString())
+                        .containsEntry("mdc-key-1", "val-1")
+                        .hasSize(2);
+            }
+        }
+    }
+
+    /**
+     * Tests that releasing a job's resources via slot freeing unregisters its MDC context from the
+     * process-wide registry.
+     */
+    @Test
+    void testJobMdcContextUnregisteredOnJobRelease() throws Exception {
+        final ExecutionAttemptID eid = createExecutionAttemptId();
+
+        final TaskDeploymentDescriptor tdd =
+                createTestTaskDeploymentDescriptor(
+                        "test task",
+                        eid,
+                        FutureCompletingInvokable.class,
+                        singleKeyMdcJobConfiguration());
+        final AllocationID allocationId = tdd.getAllocationId();
+        final CompletableFuture<Void> taskFinishedFuture = new CompletableFuture<>();
+
+        try (TaskSubmissionTestEnvironment env =
+                new TaskSubmissionTestEnvironment.Builder(jobId)
+                        .setSlotSize(1)
+                        .addTaskManagerActionListener(
+                                eid, ExecutionState.FINISHED, taskFinishedFuture)
+                        .build(EXECUTOR_EXTENSION.getExecutor())) {
+            TaskExecutorGateway tmGateway = env.getTaskExecutorGateway();
+            TaskSlotTable taskSlotTable = env.getTaskSlotTable();
+
+            taskSlotTable.allocateSlot(0, jobId, allocationId, Duration.ofSeconds(60));
+            tmGateway.submitTask(tdd, env.getJobMasterId(), timeout).get();
+
+            assertThat(MdcUtils.asContextData(jobId))
+                    .containsEntry("mdc-key-1", "val-1")
+                    .hasSize(2);
+
+            taskFinishedFuture.get(timeout.toMillis(), TimeUnit.MILLISECONDS);
+            taskSlotTable.removeTask(eid);
+            tmGateway.freeSlot(allocationId, new Exception("test release"), timeout).get();
+
+            CommonTestUtils.waitUntilCondition(() -> JobMdcRegistry.lookup(jobId) == null);
+
+            assertThat(MdcUtils.asContextData(jobId))
+                    .containsEntry(MdcUtils.JOB_ID, jobId.toHexString())
+                    .hasSize(1);
+        }
+    }
+
     /**
      * Tests that the TaskManager sends a proper exception back to the sender if the submit task
      * message fails.
@@ -849,6 +964,81 @@ static TaskDeploymentDescriptor createTaskDeploymentDescriptor(
                 inputGates);
     }
 
+    private static Configuration singleKeyMdcJobConfiguration() {
+        final Configuration conf = new Configuration();
+        conf.set(MdcOptions.JOB_CONFIGURATION_TO_MDC_KEYS, Map.of("job.key-1", "mdc-key-1"));
+        conf.setString("job.key-1", "val-1");
+        return conf;
+    }
+
+    private TaskDeploymentDescriptor createOffloadedTaskDeploymentDescriptor(
+            ExecutionAttemptID eid, PermanentBlobKey jobInfoBlobKey) throws IOException {
+        final TaskInformation taskInformation =
+                new TaskInformation(
+                        eid.getJobVertexId(),
+                        "test task",
+                        1,
+                        1,
+                        FutureCompletingInvokable.class.getName(),
+                        new Configuration());
+        return new TaskDeploymentDescriptor(
+                jobId,
+                new TaskDeploymentDescriptor.Offloaded<>(jobInfoBlobKey),
+                new TaskDeploymentDescriptor.NonOffloaded<>(new SerializedValue<>(taskInformation)),
+                eid,
+                new AllocationID(),
+                null,
+                Collections.emptyList(),
+                Collections.emptyList());
+    }
+
+    private static BlobServer createBlobServer(Path tempDir) throws IOException {
+        final Configuration config = new Configuration();
+        config.set(BlobServerOptions.OFFLOAD_MINSIZE, 0);
+        final BlobServer blobServer = new BlobServer(config, tempDir.toFile(), new VoidBlobStore());
+        blobServer.start();
+        return blobServer;
+    }
+
+    private static TaskExecutorBlobService blobServiceFrom(final BlobServer blobServer) {
+        return new TaskExecutorBlobService() {
+            @Override
+            public JobPermanentBlobService getPermanentBlobService() {
+                return new JobPermanentBlobService() {
+                    @Override
+                    public void registerJob(JobID jobId, ApplicationID applicationId) {}
+
+                    @Override
+                    public void releaseJob(JobID jobId, ApplicationID applicationId) {}
+
+                    @Override
+                    public File getFile(JobID jobId, PermanentBlobKey key) throws IOException {
+                        return blobServer.getFile(jobId, key);
+                    }
+
+                    @Override
+                    public void close() {}
+                };
+            }
+
+            @Override
+            public TransientBlobService getTransientBlobService() {
+                return NoOpTransientBlobService.INSTANCE;
+            }
+
+            @Override
+            public void setBlobServerAddress(InetSocketAddress blobServerAddress) {}
+
+            @Override
+            public int getPort() {
+                return 0;
+            }
+
+            @Override
+            public void close() {}
+        };
+    }
+
     /** Test invokable which completes the given future when executed. */
     public static class FutureCompletingInvokable extends AbstractInvokable {
 
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/taskexecutor/TaskSubmissionTestEnvironment.java b/flink-runtime/src/test/java/org/apache/flink/runtime/taskexecutor/TaskSubmissionTestEnvironment.java
index b460cec83dd941..8b47556f49c133 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/taskexecutor/TaskSubmissionTestEnvironment.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/taskexecutor/TaskSubmissionTestEnvironment.java
@@ -89,8 +89,7 @@ class TaskSubmissionTestEnvironment implements AutoCloseable {
 
     private final HeartbeatServices heartbeatServices = new HeartbeatServicesImpl(1000L, 1000L);
     private final TestingRpcService testingRpcService;
-    private final TaskExecutorBlobService taskExecutorBlobService =
-            NoOpTaskExecutorBlobService.INSTANCE;
+    private final TaskExecutorBlobService taskExecutorBlobService;
     private final Duration timeout = Duration.ofMillis(10000L);
     private final TestingFatalErrorHandler testingFatalErrorHandler =
             new TestingFatalErrorHandler();
@@ -114,9 +113,11 @@ private TaskSubmissionTestEnvironment(
             @Nullable String metricQueryServiceAddress,
             TestingRpcService testingRpcService,
             ShuffleEnvironment<?, ?> shuffleEnvironment,
+            TaskExecutorBlobService taskExecutorBlobService,
             ScheduledExecutorService executor)
             throws Exception {
 
+        this.taskExecutorBlobService = taskExecutorBlobService;
         this.timerService =
                 new DefaultTimerService<>(
                         java.util.concurrent.Executors.newSingleThreadScheduledExecutor(),
@@ -360,6 +361,9 @@ public static final class Builder {
         private ResourceID resourceID = ResourceID.generate();
         @Nullable private String metricQueryServiceAddress;
 
+        private TaskExecutorBlobService taskExecutorBlobService =
+                NoOpTaskExecutorBlobService.INSTANCE;
+
         private List<Tuple3<ExecutionAttemptID, ExecutionState, CompletableFuture<Void>>>
                 taskManagerActionListeners = new ArrayList<>();
 
@@ -422,6 +426,11 @@ public Builder setResourceID(ResourceID resourceID) {
             return this;
         }
 
+        Builder setTaskExecutorBlobService(TaskExecutorBlobService taskExecutorBlobService) {
+            this.taskExecutorBlobService = taskExecutorBlobService;
+            return this;
+        }
+
         public TaskSubmissionTestEnvironment build(ScheduledExecutorService executorService)
                 throws Exception {
             final TestingRpcService testingRpcService = new TestingRpcService();
@@ -450,6 +459,7 @@ public TaskSubmissionTestEnvironment build(ScheduledExecutorService executorServ
                     metricQueryServiceAddress,
                     testingRpcService,
                     network,
+                    taskExecutorBlobService,
                     executorService);
         }
     }