[FLINK-40208] Add JobMdcRegistry for config-driven MDC enrichment

docs/layouts/shortcodes/generated/mdc_configuration.html

@@ -0,0 +1,18 @@
+<table class="configuration table table-bordered">
+    <thead>
+        <tr>
+            <th class="text-left" style="width: 20%">Key</th>
+            <th class="text-left" style="width: 15%">Default</th>
+            <th class="text-left" style="width: 10%">Type</th>
+            <th class="text-left" style="width: 55%">Description</th>
+        </tr>
+    </thead>
+    <tbody>
+        <tr>
+            <td><h5>mdc.job-configuration-to-mdc-keys</h5></td>
+            <td style="word-wrap: break-word;"></td>
+            <td>Map</td>
+            <td>Maps job configuration keys to MDC key names. At job start, each listed configuration key is looked up; if the value is present and non-blank it is emitted into MDC under the mapped name. Keys absent or blank in the job configuration are skipped.</td>
+        </tr>
+    </tbody>
+</table>

flink-core/src/main/java/org/apache/flink/configuration/MdcOptions.java

@@ -0,0 +1,48 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.flink.configuration;
+
+import org.apache.flink.annotation.PublicEvolving;
+
+import java.util.Collections;
+import java.util.Map;
+
+import static org.apache.flink.configuration.ConfigOptions.key;
+
+/** Configuration options for MDC (Mapped Diagnostic Context) enrichment. */
+@PublicEvolving
+public final class MdcOptions {
+
+    /**
+     * Maps job configuration keys to MDC key names. Keys absent or blank in the job configuration
+     * are skipped.
+     */
+    @PublicEvolving
+    public static final ConfigOption<Map<String, String>> JOB_CONFIGURATION_TO_MDC_KEYS =
+            key("mdc.job-configuration-to-mdc-keys")
+                    .mapType()
+                    .defaultValue(Collections.emptyMap())
+                    .withDescription(
+                            "Maps job configuration keys to MDC key names. "
+                                    + "At job start, each listed configuration key is looked up; "
+                                    + "if the value is present and non-blank it is emitted into MDC under the mapped name. "
+                                    + "Keys absent or blank in the job configuration are skipped.");
+
+    private MdcOptions() {}
+}

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

@@ -0,0 +1,72 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.flink.util;
+
+import org.apache.flink.annotation.Internal;
+import org.apache.flink.annotation.VisibleForTesting;
+import org.apache.flink.api.common.JobID;
+import org.apache.flink.configuration.Configuration;
+
+import javax.annotation.Nullable;
+import javax.annotation.concurrent.ThreadSafe;
+
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
+
+/**
+ * Process-wide registry mapping {@link JobID} to enriched MDC context, populated where the job
+ * {@link Configuration} is available and consulted by {@link MdcUtils#asContextData(JobID)}.
+ */
+@Internal
+@ThreadSafe
+public final class JobMdcRegistry {
+
+    private static final Map<JobID, Map<String, String>> REGISTRY = new ConcurrentHashMap<>();
+
+    private JobMdcRegistry() {}
+
+    /**
+     * Registers enriched MDC context if the configuration carries any MDC key mappings; clears any
+     * stale entry otherwise. Equivalent to {@link #unregister} when the config is unenriched.
+     */
+    public static void registerOrClear(final JobID jobID, final Configuration jobConfiguration) {
+        final Map<String, String> context = MdcUtils.asContextData(jobID, jobConfiguration);
+        if (context.size() > 1) {
+            REGISTRY.put(jobID, context);
+        } else {
+            unregister(jobID);
+        }
+    }
+
+    /** Remove the registered context for the job. */
+    public static void unregister(final JobID jobID) {
+        REGISTRY.remove(jobID);
+    }
+
+    /** Return the registered context for the job, or {@code null} if none. */
+    @Nullable
+    public static Map<String, String> lookup(final JobID jobID) {
+        return REGISTRY.get(jobID);
+    }
+
+    @VisibleForTesting
+    public static void clear() {
+        REGISTRY.clear();
+    }
+}

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

@@ -19,10 +19,13 @@
 package org.apache.flink.util;
 
 import org.apache.flink.api.common.JobID;
+import org.apache.flink.configuration.Configuration;
+import org.apache.flink.configuration.MdcOptions;
 
 import org.slf4j.MDC;
 
 import java.util.Collections;
+import java.util.HashMap;
 import java.util.Map;
 import java.util.concurrent.Callable;
 import java.util.concurrent.Executor;
@@ -31,7 +34,7 @@
 
 import static org.apache.flink.util.Preconditions.checkArgument;
 
-/** Utility class to manage common Flink attributes in {@link MDC} (only {@link JobID} ATM). */
+/** Utility class to manage common Flink attributes in {@link MDC}. */
 public class MdcUtils {
 
     public static final String JOB_ID = "flink-job-id";
@@ -103,6 +106,14 @@ public static ExecutorService scopeToJob(JobID jobID, ExecutorService delegate)
         return new MdcAwareExecutorService<>(delegate, asContextData(jobID));
     }
 
+    /** Wraps the given {@link ExecutorService} into one with a copy of the current context. */
+    public static ExecutorService propagate(ExecutorService delegate) {
+        checkArgument(!(delegate instanceof MdcAwareExecutorService));
+        final Map<String, String> context = MDC.getCopyOfContextMap();
+        return new MdcAwareExecutorService<>(
+                delegate, context != null ? context : Collections.emptyMap());
+    }
+
     /**
      * Wrap the given {@link ScheduledExecutorService} so that the given {@link JobID} is added
      * before it executes any submitted commands and removed afterward.
@@ -112,7 +123,38 @@ public static ScheduledExecutorService scopeToJob(JobID jobID, ScheduledExecutor
         return new MdcAwareScheduledExecutorService(ses, asContextData(jobID));
     }
 
+    /**
+     * Build MDC context for a job. Consults the {@link JobMdcRegistry} for enriched context
+     * registered where the job {@link Configuration} is available; falls back to the plain job ID
+     * entry.
+     */
     public static Map<String, String> asContextData(JobID jobID) {
+        final Map<String, String> registered = JobMdcRegistry.lookup(jobID);
+        if (registered != null) {
+            return registered;
+        }
         return Collections.singletonMap(JOB_ID, jobID.toHexString());
     }
+
+    /**
+     * Build MDC context from a job ID and job configuration, enriching with context entries
+     * configured via {@link MdcOptions#JOB_CONFIGURATION_TO_MDC_KEYS}.
+     */
+    public static Map<String, String> asContextData(
+            final JobID jobID, final Configuration jobConfiguration) {
+        final Map<String, String> mdcKeyMapping =
+                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 value = jobConfiguration.getString(entry.getKey(), null);
+            if (value != null && !value.isBlank()) {
+                context.put(entry.getValue(), value);
+            }
+        }
+        if (context.isEmpty()) {
+            return Collections.singletonMap(JOB_ID, jobID.toHexString());
+        }
+        context.put(JOB_ID, jobID.toHexString());
+        return Collections.unmodifiableMap(context);
+    }
 }

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

@@ -0,0 +1,109 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.flink.util;
+
+import org.apache.flink.api.common.JobID;
+import org.apache.flink.configuration.Configuration;
+
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.Arguments;
+import org.junit.jupiter.params.provider.MethodSource;
+
+import java.util.function.Consumer;
+import java.util.stream.Stream;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+/** Tests for {@link JobMdcRegistry}. */
+class JobMdcRegistryTest {
+
+    @AfterEach
+    void clearRegistry() {
+        JobMdcRegistry.clear();
+    }
+
+    @Test
+    void testEnrichedContextStoredOnRegister() {
+        final JobID jobID = new JobID();
+        JobMdcRegistry.registerOrClear(
+                jobID, MdcTestFixtures.enrichedConfiguration("val-1", "val-2"));
+        assertThat(JobMdcRegistry.lookup(jobID))
+                .containsEntry(MdcUtils.JOB_ID, jobID.toHexString())
+                .containsEntry("mdc-key-1", "val-1")
+                .containsEntry("mdc-key-2", "val-2")
+                .hasSize(3);
+    }
+
+    private static Stream<Arguments> clearingActions() {
+        return Stream.of(
+                // explicit unregister — also verifies idempotency (double unregister stays null)
+                Arguments.of(
+                        "after explicit unregister",
+                        (Consumer<JobID>)
+                                jobID -> {
+                                    JobMdcRegistry.unregister(jobID);
+                                    assertThat(JobMdcRegistry.lookup(jobID)).isNull();
+                                    JobMdcRegistry.unregister(jobID);
+                                }),
+                // unenriched config on a fresh job stores nothing
+                Arguments.of(
+                        "after registerOrClear with empty config (no prior entry)",
+                        (Consumer<JobID>)
+                                jobID ->
+                                        JobMdcRegistry.registerOrClear(jobID, new Configuration())),
+                // unenriched config overwrites an existing enriched entry
+                Arguments.of(
+                        "after registerOrClear with empty config (clears prior entry)",
+                        (Consumer<JobID>)
+                                jobID -> {
+                                    JobMdcRegistry.registerOrClear(
+                                            jobID, MdcTestFixtures.enrichedConfiguration("val-1"));
+                                    assertThat(JobMdcRegistry.lookup(jobID)).isNotNull();
+                                    JobMdcRegistry.registerOrClear(jobID, new Configuration());
+                                }));
+    }
+
+    @ParameterizedTest
+    @MethodSource("clearingActions")
+    void testLookupReturnsNullAfterRemoval(String scenario, Consumer<JobID> clearAction) {
+        final JobID jobID = new JobID();
+        clearAction.accept(jobID);
+        assertThat(JobMdcRegistry.lookup(jobID)).as(scenario).isNull();
+    }
+
+    @Test
+    void testLatestEnrichmentWinsOnReRegister() {
+        final JobID jobID = new JobID();
+        JobMdcRegistry.registerOrClear(jobID, MdcTestFixtures.enrichedConfiguration("val-first"));
+        JobMdcRegistry.registerOrClear(jobID, MdcTestFixtures.enrichedConfiguration("val-second"));
+        assertThat(JobMdcRegistry.lookup(jobID)).containsEntry("mdc-key-1", "val-second");
+    }
+
+    @Test
+    void testContextIsolatedPerJob() {
+        final JobID first = new JobID();
+        final JobID second = new JobID();
+        JobMdcRegistry.registerOrClear(first, MdcTestFixtures.enrichedConfiguration("val-first"));
+        JobMdcRegistry.registerOrClear(second, MdcTestFixtures.enrichedConfiguration("val-second"));
+        assertThat(JobMdcRegistry.lookup(first)).containsEntry("mdc-key-1", "val-first");
+        assertThat(JobMdcRegistry.lookup(second)).containsEntry("mdc-key-1", "val-second");
+    }
+}

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

@@ -0,0 +1,51 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.flink.util;
+
+import org.apache.flink.configuration.Configuration;
+import org.apache.flink.configuration.MdcOptions;
+
+import java.util.HashMap;
+import java.util.Map;
+
+/** Shared test fixtures for MDC-related tests. */
+final class MdcTestFixtures {
+
+    /** Returns a two-entry key mapping from generic job config keys to MDC key names. */
+    static Map<String, String> testKeyMapping() {
+        final Map<String, String> mapping = new HashMap<>();
+        mapping.put("job.key-1", "mdc-key-1");
+        mapping.put("job.key-2", "mdc-key-2");
+        return mapping;
+    }
+
+    static Configuration enrichedConfiguration(final String key1Value, final String key2Value) {
+        final Configuration conf = new Configuration();
+        conf.set(MdcOptions.JOB_CONFIGURATION_TO_MDC_KEYS, testKeyMapping());
+        conf.setString("job.key-1", key1Value);
+        conf.setString("job.key-2", key2Value);
+        return conf;
+    }
+
+    static Configuration enrichedConfiguration(final String key1Value) {
+        return enrichedConfiguration(key1Value, "val-2");
+    }
+
+    private MdcTestFixtures() {}
+}