Routed image pull progress to instance log

PhonePe · Sep 25, 2024 · 3e0581f · 3e0581f
1 parent 64454aa
commit 3e0581f
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 3 deletions.
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
@@ -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<PullResponseItem> {
 
     private final Map<String, Long> downloadState = new ConcurrentHashMap<>();
+    private final Map<String, String> mdc;
     private final String image;
 
-    public ImagePullProgressHandler(String image) {
+    public ImagePullProgressHandler(Map<String, String> 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
@@ -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
@@ -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<T> {
     @SneakyThrows
     void runTest(TestConsumer<String> 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))