Fix docker image pull retry

4 files changed, 29 insertions, 45 deletions

diff --git a/docker-api/src/main/java/com/yahoo/vespa/hosted/dockerapi/Docker.java b/docker-api/src/main/java/com/yahoo/vespa/hosted/dockerapi/Docker.java
index b512b7b5f90..865908bdc8e 100644
--- a/docker-api/src/main/java/com/yahoo/vespa/hosted/dockerapi/Docker.java
+++ b/docker-api/src/main/java/com/yahoo/vespa/hosted/dockerapi/Docker.java
@@ -6,7 +6,6 @@ import java.net.InetAddress;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
-import java.util.concurrent.CompletableFuture;
 
 /**
  * API to simplify the com.github.dockerjava API for clients,
@@ -60,9 +59,14 @@ public interface Docker {
 
     Optional<Container> getContainer(ContainerName containerName);
 
-    CompletableFuture<DockerImage> pullImageAsync(DockerImage image);
-
-    boolean imageIsDownloaded(DockerImage image);
+    /**
+     * Checks if the image is currently being pulled or is already pulled, if not, starts an async
+     * pull of the image
+     *
+     * @param image Docker image to pull
+     * @return true iff image being pulled, false otherwise
+     */
+    boolean pullImageAsyncIfNeeded(DockerImage image);
 
     void deleteImage(DockerImage dockerImage);
 
diff --git a/docker-api/src/main/java/com/yahoo/vespa/hosted/dockerapi/DockerImpl.java b/docker-api/src/main/java/com/yahoo/vespa/hosted/dockerapi/DockerImpl.java
index 938b277d90b..18c1ad1dc93 100644
--- a/docker-api/src/main/java/com/yahoo/vespa/hosted/dockerapi/DockerImpl.java
+++ b/docker-api/src/main/java/com/yahoo/vespa/hosted/dockerapi/DockerImpl.java
@@ -37,10 +37,11 @@ import java.time.Duration;
 import java.util.Arrays;
 import java.util.Collections;
 import java.util.HashMap;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
-import java.util.concurrent.CompletableFuture;
+import java.util.Set;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.TimeUnit;
 import java.util.logging.Logger;
@@ -63,7 +64,7 @@ public class DockerImpl implements Docker {
 
     private final Object monitor = new Object();
     @GuardedBy("monitor")
-    private final Map<DockerImage, CompletableFuture<DockerImage>> scheduledPulls = new HashMap<>();
+    private final Set<DockerImage> scheduledPulls = new HashSet<>();
 
     // Exposed for testing.
     final DockerClient dockerClient;
@@ -159,37 +160,30 @@ public class DockerImpl implements Docker {
     }
 
     @Override
-    public CompletableFuture<DockerImage> pullImageAsync(final DockerImage image) {
-        final CompletableFuture<DockerImage> completionListener;
-        synchronized (monitor) {
-            if (scheduledPulls.containsKey(image)) {
-                return scheduledPulls.get(image);
-            }
-            completionListener = new CompletableFuture<>();
-            scheduledPulls.put(image, completionListener);
-        }
-
+    public boolean pullImageAsyncIfNeeded(final DockerImage image) {
         try {
-            dockerClient.pullImageCmd(image.asString()).exec(new ImagePullCallback(image));
+            synchronized (monitor) {
+                if (scheduledPulls.contains(image)) return true;
+                if (imageIsDownloaded(image)) return false;
+                dockerClient.pullImageCmd(image.asString()).exec(new ImagePullCallback(image));
+                return false;
+            }
         } catch (RuntimeException e) {
             numberOfDockerDaemonFails.add();
             throw new DockerException("Failed to pull image '" + image.asString() + "'", e);
         }
-
-        return completionListener;
     }
 
-    private CompletableFuture<DockerImage> removeScheduledPoll(final DockerImage image) {
+    private void removeScheduledPoll(final DockerImage image) {
         synchronized (monitor) {
-            return scheduledPulls.remove(image);
+            scheduledPulls.remove(image);
         }
     }
 
     /**
      * Check if a given image is already in the local registry
      */
-    @Override
-    public boolean imageIsDownloaded(final DockerImage dockerImage) {
+    private boolean imageIsDownloaded(final DockerImage dockerImage) {
         return inspectImage(dockerImage).isPresent();
     }
 
@@ -448,7 +442,8 @@ public class DockerImpl implements Docker {
 
         @Override
         public void onError(Throwable throwable) {
-            removeScheduledPoll(dockerImage).completeExceptionally(throwable);
+            removeScheduledPoll(dockerImage);
+            throw new DockerClientException("Could not download image: " + dockerImage);
         }
 
 
@@ -457,10 +452,9 @@ public class DockerImpl implements Docker {
             Optional<Image> image = inspectImage(dockerImage);
             if (image.isPresent()) { // Download successful, update image GC with the newly downloaded image
                 dockerImageGC.ifPresent(imageGC -> imageGC.updateLastUsedTimeFor(image.get().getId()));
-                removeScheduledPoll(dockerImage).complete(dockerImage);
+                removeScheduledPoll(dockerImage);
             } else {
-                removeScheduledPoll(dockerImage).completeExceptionally(
-                        new DockerClientException("Could not download image: " + dockerImage));
+                throw new DockerClientException("Could not download image: " + dockerImage);
             }
         }
     }
diff --git a/docker-api/src/test/java/com/yahoo/vespa/hosted/dockerapi/DockerTest.java b/docker-api/src/test/java/com/yahoo/vespa/hosted/dockerapi/DockerTest.java
index 4e2567b9fa8..310fb4ffdd3 100644
--- a/docker-api/src/test/java/com/yahoo/vespa/hosted/dockerapi/DockerTest.java
+++ b/docker-api/src/test/java/com/yahoo/vespa/hosted/dockerapi/DockerTest.java
@@ -14,7 +14,6 @@ import java.util.concurrent.ExecutionException;
 
 import static org.hamcrest.CoreMatchers.is;
 import static org.junit.Assert.assertEquals;
-import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertThat;
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assume.assumeTrue;
@@ -30,21 +29,6 @@ public class DockerTest {
     private static final DockerImage dockerImage = new DockerImage("simple-ipv6-server:Dockerfile");
     private static final String MANAGER_NAME = "docker-test";
 
-    // It is ignored since it is a bit slow and unstable, at least on Mac.
-    @Ignore
-    @Test
-    public void testDockerImagePullDelete() throws ExecutionException, InterruptedException {
-        DockerImage dockerImage = new DockerImage("busybox:1.24.0");
-
-        // Pull the image and wait for the pull to complete
-        docker.pullImageAsync(dockerImage).get();
-        assertTrue("Failed to download " + dockerImage.asString() + " image", docker.imageIsDownloaded(dockerImage));
-
-        // Remove the image
-        docker.deleteImage(dockerImage);
-        assertFalse("Failed to delete " + dockerImage.asString() + " image", docker.imageIsDownloaded(dockerImage));
-    }
-
     // Ignored because the test is very slow (several minutes) when swap is enabled, to disable: (Linux)
     // $ sudo swapoff -a
     @Ignore
diff --git a/docker-api/src/test/java/com/yahoo/vespa/hosted/dockerapi/RunSystemTests.java b/docker-api/src/test/java/com/yahoo/vespa/hosted/dockerapi/RunSystemTests.java
index 5f090abec71..915b3b53867 100644
--- a/docker-api/src/test/java/com/yahoo/vespa/hosted/dockerapi/RunSystemTests.java
+++ b/docker-api/src/test/java/com/yahoo/vespa/hosted/dockerapi/RunSystemTests.java
@@ -151,9 +151,11 @@ public class RunSystemTests {
     }
 
     private void buildVespaSystestDockerImage(Docker docker, DockerImage vespaBaseImage) throws IOException, ExecutionException, InterruptedException {
-        if (!docker.imageIsDownloaded(vespaBaseImage)) {
+        if (docker.pullImageAsyncIfNeeded(vespaBaseImage)) {
             logger.info("Pulling " + vespaBaseImage.asString() + " (This may take a while)");
-            docker.pullImageAsync(vespaBaseImage).get();
+            while (docker.pullImageAsyncIfNeeded(vespaBaseImage)) {
+                Thread.sleep(5000);
+            };
         }
 
         Path systestBuildDirectory = pathToVespaRepoInHost.resolve("docker-api/src/test/resources/systest/");