This is an automated email from the ASF dual-hosted git repository. rkhachatryan pushed a commit to branch master in repository https://gitbox.apache.org/repos/asf/flink.git
commit b198cddc893246626a1a577ea7d3b208a849a483 Author: Aleksandr Iushmanov <[email protected]> AuthorDate: Wed Jul 22 10:54:55 2026 +0100 [FLINK-40208] Add JobMdcRegistry for config-driven MDC enrichment Introduce JobMdcRegistry, a process-wide registry mapping JobID to an enriched MDC context built from job configuration. Add MdcOptions with the mdc.job-configuration-to-mdc-keys config option (@PublicEvolving). Wire in mdc enrichment on the job/task submission paths. Generated-by: Claude Code --- .../shortcodes/generated/mdc_configuration.html | 18 ++ .../org/apache/flink/configuration/MdcOptions.java | 51 ++++ .../java/org/apache/flink/util/JobMdcRegistry.java | 75 +++++ .../main/java/org/apache/flink/util/MdcUtils.java | 36 ++- .../org/apache/flink/util/JobMdcRegistryTest.java | 109 +++++++ .../org/apache/flink/util/MdcTestFixtures.java | 51 ++++ .../java/org/apache/flink/util/MdcUtilsTest.java | 327 ++++++++++++++++++--- .../flink/runtime/dispatcher/Dispatcher.java | 11 +- .../flink/runtime/taskexecutor/TaskExecutor.java | 35 ++- .../flink/runtime/dispatcher/DispatcherTest.java | 46 +++ .../taskexecutor/TaskExecutorSubmissionTest.java | 83 +++++- .../apache/flink/test/misc/JobIDLoggingITCase.java | 49 ++- 12 files changed, 828 insertions(+), 63 deletions(-) diff --git a/docs/layouts/shortcodes/generated/mdc_configuration.html b/docs/layouts/shortcodes/generated/mdc_configuration.html new file mode 100644 index 00000000000..ff70817f280 --- /dev/null +++ b/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. The job ID is always added to MDC under the key 'flink-job-id' regardless of this setting. Example: 'pipeline.name:pipeline-name' maps the job configuration key 'pipeline.name' to the MDC key 'pipeline-name'.</td> + </tr> + </tbody> +</table> diff --git a/flink-core/src/main/java/org/apache/flink/configuration/MdcOptions.java b/flink-core/src/main/java/org/apache/flink/configuration/MdcOptions.java new file mode 100644 index 00000000000..c575f689728 --- /dev/null +++ b/flink-core/src/main/java/org/apache/flink/configuration/MdcOptions.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.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. " + + "The job ID is always added to MDC under the key 'flink-job-id' regardless of this setting. " + + "Example: 'pipeline.name:pipeline-name' maps the job configuration key " + + "'pipeline.name' to the MDC key 'pipeline-name'."); + + private MdcOptions() {} +} diff --git a/flink-core/src/main/java/org/apache/flink/util/JobMdcRegistry.java b/flink-core/src/main/java/org/apache/flink/util/JobMdcRegistry.java new file mode 100644 index 00000000000..a297eb522b3 --- /dev/null +++ b/flink-core/src/main/java/org/apache/flink/util/JobMdcRegistry.java @@ -0,0 +1,75 @@ +/* + * 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. The returned map is + * unmodifiable. + */ + @Nullable + public static Map<String, String> lookup(final JobID jobID) { + return REGISTRY.get(jobID); + } + + @VisibleForTesting + public static void clear() { + REGISTRY.clear(); + } +} 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 b076c64c783..935dfa79505 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 @@ -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 java.util.concurrent.ScheduledExecutorService; 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"; @@ -112,7 +115,38 @@ public class MdcUtils { 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); + } } diff --git a/flink-core/src/test/java/org/apache/flink/util/JobMdcRegistryTest.java b/flink-core/src/test/java/org/apache/flink/util/JobMdcRegistryTest.java new file mode 100644 index 00000000000..2bbb906644a --- /dev/null +++ b/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"); + } +} diff --git a/flink-core/src/test/java/org/apache/flink/util/MdcTestFixtures.java b/flink-core/src/test/java/org/apache/flink/util/MdcTestFixtures.java new file mode 100644 index 00000000000..fa85241e0c2 --- /dev/null +++ b/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() {} +} 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 92d99a15c86..117a51b74bc 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 @@ -19,6 +19,9 @@ 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.apache.flink.core.testutils.ManuallyTriggeredScheduledExecutorService; import org.apache.flink.testutils.logging.LoggerAuditingExtension; import org.apache.flink.util.MdcUtils.MdcCloseable; import org.apache.flink.util.concurrent.Executors; @@ -26,19 +29,32 @@ import org.apache.flink.util.function.ThrowingConsumer; import org.apache.logging.log4j.core.LogEvent; import org.assertj.core.api.AbstractObjectAssert; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.RegisterExtension; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.Arguments; +import org.junit.jupiter.params.provider.MethodSource; import org.slf4j.Logger; import org.slf4j.LoggerFactory; +import org.slf4j.MDC; import java.util.Collections; +import java.util.Map; +import java.util.concurrent.Executor; +import java.util.concurrent.ExecutorService; import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicReference; +import java.util.function.Function; +import java.util.stream.Stream; import static org.apache.flink.util.MdcUtils.asContextData; import static org.apache.flink.util.MdcUtils.wrapCallable; import static org.apache.flink.util.MdcUtils.wrapRunnable; import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; import static org.slf4j.event.Level.DEBUG; /** Tests for the {@link MdcUtils}. */ @@ -50,21 +66,69 @@ class MdcUtilsTest { public final LoggerAuditingExtension loggerExtension = new LoggerAuditingExtension(MdcUtilsTest.class, DEBUG); + @BeforeEach + @AfterEach + void clearMdcAndRegistry() { + MDC.clear(); + JobMdcRegistry.clear(); + } + @Test void testJobIDAsContext() { JobID jobID = new JobID(); assertThat(MdcUtils.asContextData(jobID)) - .isEqualTo(Collections.singletonMap("flink-job-id", jobID.toHexString())); + .isEqualTo(Collections.singletonMap(MdcUtils.JOB_ID, jobID.toHexString())); } - @Test - void testMdcCloseableAddsJobId() throws Exception { - assertJobIDLogged( - jobID -> { - try (MdcCloseable ignored = MdcUtils.withContext(asContextData(jobID))) { - LOGGER.warn("ignore"); - } - }); + private static Stream<Arguments> wrappingMechanisms() { + return Stream.of( + Arguments.of( + "MdcCloseable", + (ThrowingConsumer<JobID, Exception>) + jobID -> { + try (MdcCloseable ignored = + MdcUtils.withContext(asContextData(jobID))) { + LOGGER.warn("ignore"); + } + }), + Arguments.of( + "wrapRunnable", + (ThrowingConsumer<JobID, Exception>) + jobID -> + wrapRunnable(asContextData(jobID), LOGGING_RUNNABLE).run()), + Arguments.of( + "wrapCallable", + (ThrowingConsumer<JobID, Exception>) + jobID -> + wrapCallable( + asContextData(jobID), + () -> { + LOGGER.info("ignore"); + return null; + }) + .call()), + Arguments.of( + "scopeToJob(Executor)", + (ThrowingConsumer<JobID, Exception>) + jobID -> + MdcUtils.scopeToJob(jobID, Executors.directExecutor()) + .execute(LOGGING_RUNNABLE)), + Arguments.of( + "scopeToJob(ExecutorService)", + (ThrowingConsumer<JobID, Exception>) + jobID -> + MdcUtils.scopeToJob( + jobID, Executors.newDirectExecutorService()) + .submit(LOGGING_RUNNABLE) + .get())); + } + + @ParameterizedTest + @MethodSource("wrappingMechanisms") + void testJobIdLoggedByWrappingMechanism( + final String scenario, final ThrowingConsumer<JobID, Exception> action) + throws Exception { + assertJobIDLogged(scenario, jobID -> action.accept(jobID)); } @Test @@ -78,67 +142,240 @@ class MdcUtilsTest { } @Test - void testWrapRunnable() throws Exception { - assertJobIDLogged(jobID -> wrapRunnable(asContextData(jobID), LOGGING_RUNNABLE).run()); + void testScopeScheduledExecutorService() throws Exception { + ScheduledExecutorService ses = + java.util.concurrent.Executors.newSingleThreadScheduledExecutor(); + try { + assertJobIDLogged( + jobID -> + MdcUtils.scopeToJob(jobID, ses) + .schedule(LOGGING_RUNNABLE, 1L, TimeUnit.MILLISECONDS) + .get()); + } finally { + ses.shutdownNow(); + } + } + + // --- asContextData(JobID, Configuration): map-based extraction --- + + private static Stream<Arguments> configurationBranches() { + return Stream.of( + // both keys present + Arguments.of( + Map.of("job.key-1", "mdc-key-1", "job.key-2", "mdc-key-2"), + "val-1", + "val-2", + 3), + // only key-1 + Arguments.of(Map.of("job.key-1", "mdc-key-1"), "val-1", null, 2), + // only key-2 + Arguments.of(Map.of("job.key-2", "mdc-key-2"), null, "val-2", 2), + // empty map → only job-id + Arguments.of(Collections.emptyMap(), null, null, 1)); + } + + @ParameterizedTest + @MethodSource("configurationBranches") + void testContextEntriesExtractedFromConfiguration( + final Map<String, String> keyMapping, + final String key1Value, + final String key2Value, + final int expectedSize) { + final JobID jobID = new JobID(); + final Configuration conf = new Configuration(); + conf.set(MdcOptions.JOB_CONFIGURATION_TO_MDC_KEYS, keyMapping); + if (key1Value != null) { + conf.setString("job.key-1", key1Value); + } + if (key2Value != null) { + conf.setString("job.key-2", key2Value); + } + + final Map<String, String> context = MdcUtils.asContextData(jobID, conf); + + assertContextEntries(context, jobID, key1Value, key2Value, expectedSize); } @Test - void testWrapCallable() throws Exception { - assertJobIDLogged( - jobID -> - wrapCallable( - asContextData(jobID), - () -> { - LOGGER.info("ignore"); - return null; - }) - .call()); + void testMappingTargetingJobIdKeyIsIgnored() { + final JobID jobID = new JobID(); + final Configuration conf = new Configuration(); + conf.set(MdcOptions.JOB_CONFIGURATION_TO_MDC_KEYS, Map.of("job.key-1", MdcUtils.JOB_ID)); + conf.setString("job.key-1", "user-supplied-value"); + + final Map<String, String> context = MdcUtils.asContextData(jobID, conf); + + assertThat(context) + .containsEntry(MdcUtils.JOB_ID, jobID.toHexString()) + .doesNotContainEntry(MdcUtils.JOB_ID, "user-supplied-value"); } @Test - void testScopeExecutor() throws Exception { - assertJobIDLogged( - jobID -> - MdcUtils.scopeToJob(jobID, Executors.directExecutor()) - .execute(LOGGING_RUNNABLE)); + void testConfigContextIsUnmodifiable() { + final JobID jobID = new JobID(); + 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"); + + final Map<String, String> context = MdcUtils.asContextData(jobID, conf); + + assertThatThrownBy(() -> context.put("extra", "value")) + .isInstanceOf(UnsupportedOperationException.class); } + private static Stream<Arguments> skippedValueCases() { + return Stream.of(Arguments.of("blank value", " "), Arguments.of("missing key", null)); + } + + @ParameterizedTest + @MethodSource("skippedValueCases") + void testKeySkippedWhenValueAbsentOrBlank(final String scenario, final String configValue) { + final JobID jobID = new JobID(); + final Configuration conf = new Configuration(); + conf.set(MdcOptions.JOB_CONFIGURATION_TO_MDC_KEYS, Map.of("job.key-1", "mdc-key-1")); + if (configValue != null) { + conf.setString("job.key-1", configValue); + } + + final Map<String, String> context = MdcUtils.asContextData(jobID, conf); + + assertThat(context) + .as(scenario) + .isEqualTo(Collections.singletonMap(MdcUtils.JOB_ID, jobID.toHexString())); + } + + // --- JobMdcRegistry integration: registry-first lookup --- + @Test - void testScopeExecutorService() throws Exception { - assertJobIDLogged( - jobID -> - MdcUtils.scopeToJob(jobID, Executors.newDirectExecutorService()) - .submit(LOGGING_RUNNABLE) - .get()); + void testAsContextDataUsesRegistry() { + final JobID jobID = new JobID(); + JobMdcRegistry.registerOrClear( + jobID, MdcTestFixtures.enrichedConfiguration("val-1", "val-2")); + + assertThat(MdcUtils.asContextData(jobID)) + .containsEntry(MdcUtils.JOB_ID, jobID.toHexString()) + .containsEntry("mdc-key-1", "val-1") + .containsEntry("mdc-key-2", "val-2") + .hasSize(3); } @Test - void testScopeScheduledExecutorService() throws Exception { - ScheduledExecutorService ses = - java.util.concurrent.Executors.newSingleThreadScheduledExecutor(); - try { - assertJobIDLogged( - jobID -> - MdcUtils.scopeToJob(jobID, ses) - .schedule(LOGGING_RUNNABLE, 1L, TimeUnit.MILLISECONDS) - .get()); - } finally { - ses.shutdownNow(); + void testMdcRestoredAfterScopeCloses() { + final JobID jobID = new JobID(); + final Configuration conf = new Configuration(); + conf.set( + MdcOptions.JOB_CONFIGURATION_TO_MDC_KEYS, + Map.of("job.key-1", "mdc-key-1", "job.key-2", "mdc-key-2")); + conf.setString("job.key-1", "scoped-val-1"); + conf.setString("job.key-2", "scoped-val-2"); + + try (MdcCloseable ignored = MdcUtils.withContext(MdcUtils.asContextData(jobID, conf))) { + assertThat(MDC.get("mdc-key-1")).isEqualTo("scoped-val-1"); + } + assertThat(MDC.get("mdc-key-1")).isNull(); + assertThat(MDC.get("mdc-key-2")).isNull(); + } + + private static Stream<Arguments> jobScopedRunners() { + return Stream.of( + Arguments.of( + (Function<JobID, ThrowingConsumer<Runnable, Exception>>) + jobID -> { + final Executor wrapped = + MdcUtils.scopeToJob(jobID, Executors.directExecutor()); + return wrapped::execute; + }), + Arguments.of( + (Function<JobID, ThrowingConsumer<Runnable, Exception>>) + jobID -> { + final ExecutorService wrapped = + MdcUtils.scopeToJob( + jobID, Executors.newDirectExecutorService()); + return action -> + wrapped.submit( + () -> { + action.run(); + return null; + }) + .get(); + }), + Arguments.of( + (Function<JobID, ThrowingConsumer<Runnable, Exception>>) + jobID -> { + final ManuallyTriggeredScheduledExecutorService ses = + new ManuallyTriggeredScheduledExecutorService(); + final ScheduledExecutorService wrapped = + MdcUtils.scopeToJob(jobID, ses); + return action -> { + wrapped.schedule(action, 0L, TimeUnit.MILLISECONDS); + ses.triggerScheduledTasks(); + }; + })); + } + + @ParameterizedTest + @MethodSource("jobScopedRunners") + void testScopeToJobCapturesEnrichedContextAtConstructionTime( + final Function<JobID, ThrowingConsumer<Runnable, Exception>> runnerFactory) + throws Exception { + final JobID jobID = new JobID(); + + // Register before scopeToJob — context is baked at construction time + JobMdcRegistry.registerOrClear( + jobID, MdcTestFixtures.enrichedConfiguration("val-1", "val-2")); + final ThrowingConsumer<Runnable, Exception> runner = runnerFactory.apply(jobID); + final AtomicReference<Map<String, String>> captured = new AtomicReference<>(); + final Runnable capture = () -> captured.set(MDC.getCopyOfContextMap()); + + runner.accept(capture); + assertThat(captured.get()) + .containsEntry(MdcUtils.JOB_ID, jobID.toHexString()) + .containsEntry("mdc-key-1", "val-1") + .containsEntry("mdc-key-2", "val-2") + .hasSize(3); + } + + // --- helpers --- + + private static void assertContextEntries( + final Map<String, String> context, + final JobID jobID, + final String expectedKey1, + final String expectedKey2, + final int expectedSize) { + assertThat(context) + .containsEntry(MdcUtils.JOB_ID, jobID.toHexString()) + .hasSize(expectedSize); + if (expectedKey1 != null) { + assertThat(context).containsEntry("mdc-key-1", expectedKey1); + } + if (expectedKey2 != null) { + assertThat(context).containsEntry("mdc-key-2", expectedKey2); } } private void assertJobIDLogged(ThrowingConsumer<JobID, Exception> action) throws Exception { + assertJobIDLogged(null, action); + } + + private void assertJobIDLogged(String scenario, ThrowingConsumer<JobID, Exception> action) + throws Exception { JobID jobID = new JobID(); action.accept(jobID); - assertJobIdLogged(jobID); + assertJobIdLogged(scenario, jobID); } private void assertJobIdLogged(JobID jobId) { + assertJobIdLogged(null, jobId); + } + + private void assertJobIdLogged(String scenario, JobID jobId) { AbstractObjectAssert<?, Object> extracting = assertThat(loggerExtension.getEvents()) .singleElement() .extracting(LogEvent::getContextData) - .extracting(m -> m.getValue("flink-job-id")); + .extracting(m -> m.getValue(MdcUtils.JOB_ID)) + .as(scenario); if (jobId == null) { extracting.isNull(); } else { diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/dispatcher/Dispatcher.java b/flink-runtime/src/main/java/org/apache/flink/runtime/dispatcher/Dispatcher.java index 6d53fa0e221..90f1f732b2a 100644 --- a/flink-runtime/src/main/java/org/apache/flink/runtime/dispatcher/Dispatcher.java +++ b/flink-runtime/src/main/java/org/apache/flink/runtime/dispatcher/Dispatcher.java @@ -119,6 +119,7 @@ import org.apache.flink.streaming.api.graph.StreamingJobGraphGenerator; import org.apache.flink.util.CollectionUtil; import org.apache.flink.util.ExceptionUtils; import org.apache.flink.util.FlinkException; +import org.apache.flink.util.JobMdcRegistry; import org.apache.flink.util.MdcUtils; import org.apache.flink.util.MdcUtils.MdcCloseable; import org.apache.flink.util.Preconditions; @@ -572,6 +573,7 @@ public abstract class Dispatcher extends FencedRpcEndpoint<DispatcherId> initJobClientExpiredTime(recoveredJob); final JobID jobId = recoveredJob.getJobID(); + JobMdcRegistry.registerOrClear(jobId, recoveredJob.getJobConfiguration()); try (MdcCloseable ignored = MdcUtils.withContext(MdcUtils.asContextData(jobId))) { if (wrapIntoApplication) { internalSubmitApplication(new SingleJobApplication(recoveredJob, true)).get(); @@ -581,6 +583,7 @@ public abstract class Dispatcher extends FencedRpcEndpoint<DispatcherId> ExecutionType.RECOVERY, recoveredJob.getApplicationId().orElse(null)); } catch (Throwable throwable) { + JobMdcRegistry.unregister(recoveredJob.getJobID()); onFatalError( new DispatcherException( String.format("Could not start recovered job %s.", jobId), throwable)); @@ -834,7 +837,9 @@ public abstract class Dispatcher extends FencedRpcEndpoint<DispatcherId> @Override public CompletableFuture<Acknowledge> submitJob(ExecutionPlan executionPlan, Duration timeout) { final JobID jobID = executionPlan.getJobID(); - try (MdcCloseable ignored = MdcUtils.withContext(MdcUtils.asContextData(jobID))) { + try (MdcCloseable ignored = + MdcUtils.withContext( + MdcUtils.asContextData(jobID, executionPlan.getJobConfiguration()))) { log.info("Received job submission '{}' ({}).", executionPlan.getName(), jobID); } return isInGloballyTerminalState(jobID) @@ -1276,6 +1281,7 @@ public abstract class Dispatcher extends FencedRpcEndpoint<DispatcherId> final String jobName = executionPlan.getName(); final ApplicationID applicationId = executionPlan.getApplicationId().orElse(null); + JobMdcRegistry.registerOrClear(jobId, executionPlan.getJobConfiguration()); log.info( "Submitting job '{}' ({}) with associated application ({}).", jobName, @@ -1321,6 +1327,7 @@ public abstract class Dispatcher extends FencedRpcEndpoint<DispatcherId> ExceptionUtils.stripCompletionException( terminationThrowable); log.error("Failed to submit job {}.", jobId, strippedThrowable); + JobMdcRegistry.unregister(jobId); throw new CompletionException( new JobSubmissionException( jobId, "Failed to submit job.", strippedThrowable)); @@ -1437,6 +1444,8 @@ public abstract class Dispatcher extends FencedRpcEndpoint<DispatcherId> jobTerminationFuture, (thread, throwable) -> fatalErrorHandler.onFatalError(throwable)); registerJobManagerRunnerTerminationFuture(jobId, jobTerminationFuture); + jobTerminationFuture.whenComplete( + (ignored, ignoredThrowable) -> JobMdcRegistry.unregister(jobId)); } @Nullable diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/taskexecutor/TaskExecutor.java b/flink-runtime/src/main/java/org/apache/flink/runtime/taskexecutor/TaskExecutor.java index e5b9fc9b79d..166935cbd8e 100644 --- a/flink-runtime/src/main/java/org/apache/flink/runtime/taskexecutor/TaskExecutor.java +++ b/flink-runtime/src/main/java/org/apache/flink/runtime/taskexecutor/TaskExecutor.java @@ -146,6 +146,7 @@ import org.apache.flink.util.CollectionUtil; import org.apache.flink.util.ExceptionUtils; import org.apache.flink.util.FlinkException; import org.apache.flink.util.FlinkExpectedException; +import org.apache.flink.util.JobMdcRegistry; import org.apache.flink.util.MdcUtils; import org.apache.flink.util.MdcUtils.MdcCloseable; import org.apache.flink.util.OptionalConsumer; @@ -661,9 +662,20 @@ public class TaskExecutor extends RpcEndpoint implements TaskExecutorGateway { TaskDeploymentDescriptor tdd, JobMasterId jobMasterId, Duration timeout) { final JobID jobId = tdd.getJobId(); - // todo: consider adding task info - try (MdcCloseable ignored = MdcUtils.withContext(MdcUtils.asContextData(jobId))) { - + JobInformation jobInformation = null; + try { + jobInformation = tdd.getJobInformation(); + } catch (IllegalStateException ignored) { + // Expected when job information is offloaded to blob storage and not yet loaded. + } catch (IOException | ClassNotFoundException e) { + log.debug("Could not deserialize job information for early MDC enrichment", e); + } + try (MdcCloseable ignored = + MdcUtils.withContext( + jobInformation == null + ? MdcUtils.asContextData(jobId) + : MdcUtils.asContextData( + jobId, jobInformation.getJobConfiguration()))) { final ExecutionAttemptID executionAttemptID = tdd.getExecutionAttemptId(); final JobTable.Connection jobManagerConnection = @@ -716,18 +728,18 @@ public class TaskExecutor extends RpcEndpoint implements TaskExecutorGateway { } // deserialize the pre-serialized information - final JobInformation jobInformation; final TaskInformation taskInformation; final JobManagerTaskRestore taskRestore; try { - jobInformation = tdd.getJobInformation(); + if (jobInformation == null) { + jobInformation = tdd.getJobInformation(); + } taskInformation = tdd.getTaskInformation(); taskRestore = tdd.getTaskRestore(); } catch (IOException | ClassNotFoundException e) { throw new TaskSubmissionException( "Could not deserialize the job or task information.", e); } - if (!jobId.equals(jobInformation.getJobId())) { throw new TaskSubmissionException( "Inconsistent job ID information inside TaskDeploymentDescriptor (" @@ -736,7 +748,7 @@ public class TaskExecutor extends RpcEndpoint implements TaskExecutorGateway { + jobInformation.getJobId() + ")"); } - + JobMdcRegistry.registerOrClear(jobId, jobInformation.getJobConfiguration()); TaskManagerJobMetricGroup jobGroup = taskManagerMetricGroup.addJob( jobInformation.getJobId(), jobInformation.getJobName()); @@ -757,13 +769,7 @@ public class TaskExecutor extends RpcEndpoint implements TaskExecutorGateway { new RpcTaskOperatorEventGateway( jobManagerConnection.getJobManagerGateway(), executionAttemptID, - (t) -> - runAsync( - () -> - failTask( - jobInformation.getJobId(), - executionAttemptID, - t))); + (t) -> runAsync(() -> failTask(jobId, executionAttemptID, t))); TaskManagerActions taskManagerActions = jobManagerConnection.getTaskManagerActions(); CheckpointResponder checkpointResponder = jobManagerConnection.getCheckpointResponder(); @@ -2078,6 +2084,7 @@ public class TaskExecutor extends RpcEndpoint implements TaskExecutorGateway { taskInformationCache.clearCacheForGroup(jobId); shuffleDescriptorsCache.clearCacheForGroup(jobId); fileMergingManager.releaseMergingSnapshotManagerForJob(jobId); + JobMdcRegistry.unregister(jobId); } private void scheduleResultPartitionCleanup(JobID jobId) { 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 ad4f10be633..5bac062a99d 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 @@ -25,6 +25,7 @@ import org.apache.flink.api.common.JobStatus; import org.apache.flink.api.common.operators.ResourceSpec; import org.apache.flink.configuration.Configuration; import org.apache.flink.configuration.JobManagerOptions; +import org.apache.flink.configuration.MdcOptions; import org.apache.flink.configuration.PipelineOptions; import org.apache.flink.core.execution.SavepointFormatType; import org.apache.flink.core.failure.FailureEnricher; @@ -105,6 +106,7 @@ import org.apache.flink.streaming.api.graph.ExecutionPlan; import org.apache.flink.util.ExceptionUtils; import org.apache.flink.util.FlinkException; import org.apache.flink.util.InstantiationUtil; +import org.apache.flink.util.JobMdcRegistry; import org.apache.flink.util.Preconditions; import org.apache.flink.util.concurrent.FutureUtils; @@ -133,8 +135,10 @@ import java.util.ArrayList; import java.util.Arrays; import java.util.Collection; import java.util.Collections; +import java.util.HashMap; import java.util.Iterator; import java.util.List; +import java.util.Map; import java.util.Optional; import java.util.Queue; import java.util.UUID; @@ -202,6 +206,7 @@ public class DispatcherTest extends AbstractDispatcherTest { @After public void tearDown() throws Exception { + JobMdcRegistry.clear(); if (dispatcher != null) { RpcUtils.terminateRpcEndpoint(dispatcher); } @@ -565,6 +570,47 @@ public class DispatcherTest extends AbstractDispatcherTest { .withCauseOfType(FlinkJobTerminatedWithoutCancellationException.class); } + @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"); + + final CompletableFuture<JobManagerRunnerResult> resultFuture = new CompletableFuture<>(); + dispatcher = + createAndStartDispatcher( + heartbeatServices, + haServices, + new FinishingJobManagerRunnerFactory(resultFuture, () -> {})); + jobMasterLeaderElection.isLeader(UUID.randomUUID()); + final DispatcherGateway dispatcherGateway = + dispatcher.getSelfGateway(DispatcherGateway.class); + + submitApplication(); + dispatcherGateway.submitJob(jobGraph, TIMEOUT).get(); + + assertThat(JobMdcRegistry.lookup(jobId)) + .containsEntry("mdc-key-1", "val-1") + .containsEntry("mdc-key-2", "val-2"); + + resultFuture.complete( + JobManagerRunnerResult.forSuccess( + new ExecutionGraphInfo( + new ArchivedExecutionGraphBuilder() + .setJobID(jobId) + .setState(JobStatus.FINISHED) + .build()))); + mockApplicationFinished(); + dispatcher.getJobTerminationFuture(jobId, TIMEOUT).get(); + + // the unregistration callback is an independent dependent of the termination future, + // so poll instead of asserting immediately + CommonTestUtils.waitUntilCondition(() -> JobMdcRegistry.lookup(jobId) == null); + } + @Test public void testNoHistoryServerArchiveCreatedForSuspendedJob() throws Exception { final CompletableFuture<Void> archiveAttemptFuture = new CompletableFuture<>(); 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 225629876b0..c13445e43b9 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 @@ -21,6 +21,7 @@ package org.apache.flink.runtime.taskexecutor; import org.apache.flink.api.common.ExecutionConfig; import org.apache.flink.api.common.JobID; import org.apache.flink.configuration.Configuration; +import org.apache.flink.configuration.MdcOptions; import org.apache.flink.configuration.NettyShuffleEnvironmentOptions; import org.apache.flink.runtime.blob.PermanentBlobKey; import org.apache.flink.runtime.clusterframework.types.AllocationID; @@ -58,11 +59,14 @@ import org.apache.flink.runtime.testtasks.BlockingNoOpInvokable; import org.apache.flink.runtime.util.NettyShuffleDescriptorBuilder; import org.apache.flink.testutils.TestingUtils; import org.apache.flink.testutils.executor.TestExecutorExtension; +import org.apache.flink.util.JobMdcRegistry; +import org.apache.flink.util.MdcUtils; import org.apache.flink.util.NetUtils; import org.apache.flink.util.Preconditions; import org.apache.flink.util.SerializedValue; import org.apache.flink.util.concurrent.FutureUtils; +import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.TestInfo; @@ -74,7 +78,9 @@ import java.net.URL; import java.time.Duration; import java.util.Collection; import java.util.Collections; +import java.util.HashMap; import java.util.List; +import java.util.Map; import java.util.concurrent.CompletableFuture; import java.util.concurrent.ExecutionException; import java.util.concurrent.ScheduledExecutorService; @@ -104,6 +110,11 @@ class TaskExecutorSubmissionTest { this.testInfo = testInfo; } + @AfterEach + void clearJobMdcRegistry() { + JobMdcRegistry.clear(); + } + /** * Tests that we can submit a task to the TaskManager given that we've allocated a slot there. */ @@ -133,6 +144,41 @@ class TaskExecutorSubmissionTest { } } + /** Tests that a successful task submission registers the job's MDC context. */ + @Test + void testJobMdcContextRegisteredOnSubmitTask() throws Exception { + final ExecutionAttemptID eid = createExecutionAttemptId(); + + final Map<String, String> keyMapping = new HashMap<>(); + keyMapping.put("job.key-1", "mdc-key-1"); + keyMapping.put("job.key-2", "mdc-key-2"); + final Configuration jobConfiguration = new Configuration(); + jobConfiguration.set(MdcOptions.JOB_CONFIGURATION_TO_MDC_KEYS, keyMapping); + jobConfiguration.setString("job.key-1", "val-1"); + jobConfiguration.setString("job.key-2", "val-2"); + + final TaskDeploymentDescriptor tdd = + createTestTaskDeploymentDescriptor( + "test task", eid, FutureCompletingInvokable.class, jobConfiguration); + + try (TaskSubmissionTestEnvironment env = + new TaskSubmissionTestEnvironment.Builder(jobId) + .setSlotSize(1) + .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(JobMdcRegistry.lookup(jobId)) + .containsEntry(MdcUtils.JOB_ID, jobId.toHexString()) + .containsEntry("mdc-key-1", "val-1") + .containsEntry("mdc-key-2", "val-2") + .hasSize(3); + } + } + /** * Tests that the TaskManager sends a proper exception back to the sender if the submit task * message fails. @@ -697,6 +743,41 @@ class TaskExecutorSubmissionTest { List<ResultPartitionDeploymentDescriptor> producedPartitions, List<InputGateDeploymentDescriptor> inputGates) throws IOException { + return createTestTaskDeploymentDescriptor( + taskName, + eid, + abstractInvokable, + maxNumberOfSubtasks, + producedPartitions, + inputGates, + new Configuration()); + } + + private TaskDeploymentDescriptor createTestTaskDeploymentDescriptor( + String taskName, + ExecutionAttemptID eid, + Class<? extends AbstractInvokable> abstractInvokable, + Configuration jobConfiguration) + throws IOException { + return createTestTaskDeploymentDescriptor( + taskName, + eid, + abstractInvokable, + 1, + Collections.emptyList(), + Collections.emptyList(), + jobConfiguration); + } + + private TaskDeploymentDescriptor createTestTaskDeploymentDescriptor( + String taskName, + ExecutionAttemptID eid, + Class<? extends AbstractInvokable> abstractInvokable, + int maxNumberOfSubtasks, + List<ResultPartitionDeploymentDescriptor> producedPartitions, + List<InputGateDeploymentDescriptor> inputGates, + Configuration jobConfiguration) + throws IOException { Preconditions.checkNotNull(producedPartitions); Preconditions.checkNotNull(inputGates); return createTaskDeploymentDescriptor( @@ -707,7 +788,7 @@ class TaskExecutorSubmissionTest { taskName, maxNumberOfSubtasks, 1, - new Configuration(), + jobConfiguration, new Configuration(), abstractInvokable.getName(), producedPartitions, diff --git a/flink-tests/src/test/java/org/apache/flink/test/misc/JobIDLoggingITCase.java b/flink-tests/src/test/java/org/apache/flink/test/misc/JobIDLoggingITCase.java index 103dd74f58e..e493000f47e 100644 --- a/flink-tests/src/test/java/org/apache/flink/test/misc/JobIDLoggingITCase.java +++ b/flink-tests/src/test/java/org/apache/flink/test/misc/JobIDLoggingITCase.java @@ -27,6 +27,7 @@ import org.apache.flink.api.connector.source.lib.NumberSequenceSource; import org.apache.flink.client.program.ClusterClient; import org.apache.flink.configuration.Configuration; import org.apache.flink.configuration.JobManagerOptions; +import org.apache.flink.configuration.MdcOptions; import org.apache.flink.core.execution.CheckpointType; import org.apache.flink.runtime.checkpoint.CheckpointCoordinator; import org.apache.flink.runtime.checkpoint.CheckpointException; @@ -50,6 +51,7 @@ import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.RegisterExtension; import java.time.Duration; +import java.util.Map; import java.util.concurrent.ExecutionException; import static java.util.Arrays.asList; @@ -228,7 +230,50 @@ class JobIDLoggingITCase { ".* finished asynchronous part of checkpoint .*")); } + @Test + void testEnrichedMdcLogging(@InjectClusterClient ClusterClient<?> clusterClient) + throws Exception { + final Configuration enrichmentConfig = new Configuration(); + enrichmentConfig.set( + MdcOptions.JOB_CONFIGURATION_TO_MDC_KEYS, Map.of("job.key-1", "mdc-key-1")); + enrichmentConfig.setString("job.key-1", "val-1"); + + final JobID jobID = runJob(clusterClient, enrichmentConfig); + clusterClient.cancel(jobID).get(); + + assertKeyPresent( + "mdc-key-1", + "val-1", + jobMasterLogging, + asList("Initializing job .*", "Starting execution of job .*"), + "Registration at ResourceManager.*", + "Registration with ResourceManager.*", + "Resolved ResourceManager address.*"); + + assertKeyPresent( + "mdc-key-1", + "val-1", + taskExecutorLogging, + asList("Received task .*"), + "TaskManager received a checkpoint confirmation for unknown task.*", + "TaskManager received an aborted checkpoint for unknown task.*", + "Un-registering task.*", + "Successful registration.*", + "Establish JobManager connection.*", + "Offer reserved slots.*", + ".*ResourceManager.*", + "Operator event.*", + "Recovered slot allocation snapshots.*", + ".*heartbeat.*", + ".*leadership.*"); + } + private static JobID runJob(ClusterClient<?> clusterClient) throws Exception { + return runJob(clusterClient, new Configuration()); + } + + private static JobID runJob(ClusterClient<?> clusterClient, Configuration jobConfig) + throws Exception { StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment(); env.fromSource( @@ -239,7 +284,9 @@ class JobIDLoggingITCase { .withTimestampAssigner((r, t) -> (long) r), "Source-42441337") .addSink(new DiscardingSink<>()); - JobID jobId = clusterClient.submitJob(env.getStreamGraph().getJobGraph()).get(); + var jobGraph = env.getStreamGraph().getJobGraph(); + jobGraph.getJobConfiguration().addAll(jobConfig); + JobID jobId = clusterClient.submitJob(jobGraph).get(); Deadline deadline = Deadline.fromNow(Duration.ofMinutes(5)); while (deadline.hasTimeLeft() && clusterClient.listJobs().get().stream()
