diff --git a/drove-executor/src/main/java/com/phonepe/drove/executor/utils/ImagePullProgressHandler.java b/drove-executor/src/main/java/com/phonepe/drove/executor/utils/ImagePullProgressHandler.java index ce04d161..de886281 100644 --- a/drove-executor/src/main/java/com/phonepe/drove/executor/utils/ImagePullProgressHandler.java +++ b/drove-executor/src/main/java/com/phonepe/drove/executor/utils/ImagePullProgressHandler.java @@ -21,6 +21,7 @@ import com.google.common.base.Strings; import lombok.extern.slf4j.Slf4j; import lombok.val; +import org.slf4j.MDC; import java.util.Map; import java.util.Objects; @@ -33,14 +34,19 @@ public class ImagePullProgressHandler extends ResultCallback.Adapter { private final Map downloadState = new ConcurrentHashMap<>(); + private final Map mdc; private final String image; - public ImagePullProgressHandler(String image) { + public ImagePullProgressHandler(Map mdc, String image) { + this.mdc = mdc; this.image = image; } @Override public void onNext(PullResponseItem responseItem) { + if (null != mdc) { + MDC.setContextMap(mdc); + } val layerId = responseItem.getId(); if (!Strings.isNullOrEmpty(layerId)) { val progressDetail = responseItem.getProgressDetail(); diff --git a/drove-executor/src/test/java/com/phonepe/drove/executor/ContainerHelperExtension.java b/drove-executor/src/test/java/com/phonepe/drove/executor/ContainerHelperExtension.java index 63ce2a2f..6d2529de 100644 --- a/drove-executor/src/test/java/com/phonepe/drove/executor/ContainerHelperExtension.java +++ b/drove-executor/src/test/java/com/phonepe/drove/executor/ContainerHelperExtension.java @@ -96,7 +96,7 @@ private void ensureImage(final String imageName) { } try { DOCKER_CLIENT.pullImageCmd(imageName) - .exec(new ImagePullProgressHandler(imageName)) + .exec(new ImagePullProgressHandler(null, imageName)) .awaitCompletion(); } catch (InterruptedException e) { diff --git a/drove-executor/src/test/java/com/phonepe/drove/executor/checker/CmdCheckerTest.java b/drove-executor/src/test/java/com/phonepe/drove/executor/checker/CmdCheckerTest.java index 99f014b2..fc1c0922 100644 --- a/drove-executor/src/test/java/com/phonepe/drove/executor/checker/CmdCheckerTest.java +++ b/drove-executor/src/test/java/com/phonepe/drove/executor/checker/CmdCheckerTest.java @@ -29,6 +29,7 @@ import lombok.SneakyThrows; import lombok.val; import org.junit.jupiter.api.Test; +import org.slf4j.MDC; import static com.phonepe.drove.common.CommonTestUtils.TASK_IMAGE_NAME; import static com.phonepe.drove.executor.ExecutorTestingUtils.DOCKER_CLIENT; @@ -99,7 +100,8 @@ private interface TestConsumer { @SneakyThrows void runTest(TestConsumer test) { DOCKER_CLIENT.pullImageCmd(TASK_IMAGE_NAME) - .exec(new ImagePullProgressHandler(TASK_IMAGE_NAME)) + .exec(new ImagePullProgressHandler(MDC.getCopyOfContextMap(), + TASK_IMAGE_NAME)) .awaitCompletion(); val containerId = DOCKER_CLIENT.createContainerCmd(TASK_IMAGE_NAME) .withHostConfig(new HostConfig().withAutoRemove(true))