Fix pull async bug, add test

2 files changed, 84 insertions, 8 deletions

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 18c1ad1dc93..db10a85bb45 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
@@ -6,6 +6,7 @@ import com.github.dockerjava.api.command.ExecCreateCmdResponse;
 import com.github.dockerjava.api.command.ExecStartCmd;
 import com.github.dockerjava.api.command.InspectContainerResponse;
 import com.github.dockerjava.api.command.InspectExecResponse;
+import com.github.dockerjava.api.command.InspectImageResponse;
 import com.github.dockerjava.api.exception.DockerClientException;
 import com.github.dockerjava.api.exception.NotFoundException;
 import com.github.dockerjava.api.exception.NotModifiedException;
@@ -165,8 +166,10 @@ public class DockerImpl implements Docker {
             synchronized (monitor) {
                 if (scheduledPulls.contains(image)) return true;
                 if (imageIsDownloaded(image)) return false;
+
+                scheduledPulls.add(image);
                 dockerClient.pullImageCmd(image.asString()).exec(new ImagePullCallback(image));
-                return false;
+                return true;
             }
         } catch (RuntimeException e) {
             numberOfDockerDaemonFails.add();
@@ -183,14 +186,15 @@ public class DockerImpl implements Docker {
     /**
      * Check if a given image is already in the local registry
      */
-    private boolean imageIsDownloaded(final DockerImage dockerImage) {
+    boolean imageIsDownloaded(final DockerImage dockerImage) {
         return inspectImage(dockerImage).isPresent();
     }
 
-    private Optional<Image> inspectImage(DockerImage dockerImage) {
+    private Optional<InspectImageResponse> inspectImage(DockerImage dockerImage) {
         try {
-            return dockerClient.listImagesCmd().withShowAll(true)
-                    .withImageNameFilter(dockerImage.asString()).exec().stream().findFirst();
+            return Optional.of(dockerClient.inspectImageCmd(dockerImage.asString()).exec());
+        } catch (NotFoundException e) {
+            return Optional.empty();
         } catch (RuntimeException e) {
             numberOfDockerDaemonFails.add();
             throw new DockerException("Failed to inspect image '" + dockerImage.asString() + "'", e);
@@ -449,7 +453,7 @@ public class DockerImpl implements Docker {
 
         @Override
         public void onComplete() {
-            Optional<Image> image = inspectImage(dockerImage);
+            Optional<InspectImageResponse> 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);
@@ -515,7 +519,7 @@ public class DockerImpl implements Docker {
                 .withDockerCmdExecFactory(dockerFactory);
     }
 
-    private void setMetrics(MetricReceiverWrapper metricReceiver) {
+    void setMetrics(MetricReceiverWrapper metricReceiver) {
         Dimensions dimensions = new Dimensions.Builder().add("role", "docker").build();
         numberOfDockerDaemonFails = metricReceiver.declareCounter(MetricReceiverWrapper.APPLICATION_DOCKER, dimensions, "daemon.api_fails");
     }
diff --git a/docker-api/src/test/java/com/yahoo/vespa/hosted/dockerapi/DockerImplTest.java b/docker-api/src/test/java/com/yahoo/vespa/hosted/dockerapi/DockerImplTest.java
index 84d1ca82d17..12e52dde494 100644
--- a/docker-api/src/test/java/com/yahoo/vespa/hosted/dockerapi/DockerImplTest.java
+++ b/docker-api/src/test/java/com/yahoo/vespa/hosted/dockerapi/DockerImplTest.java
@@ -2,14 +2,22 @@
 package com.yahoo.vespa.hosted.dockerapi;
 
 import com.github.dockerjava.api.DockerClient;
+import com.github.dockerjava.api.async.ResultCallback;
 import com.github.dockerjava.api.command.ExecCreateCmd;
 import com.github.dockerjava.api.command.ExecCreateCmdResponse;
 import com.github.dockerjava.api.command.ExecStartCmd;
 import com.github.dockerjava.api.command.InspectExecCmd;
 import com.github.dockerjava.api.command.InspectExecResponse;
+import com.github.dockerjava.api.command.InspectImageCmd;
+import com.github.dockerjava.api.command.InspectImageResponse;
+import com.github.dockerjava.api.command.PullImageCmd;
+import com.github.dockerjava.api.exception.NotFoundException;
 import com.github.dockerjava.core.DefaultDockerClientConfig;
 import com.github.dockerjava.core.command.ExecStartResultCallback;
+import com.yahoo.metrics.simple.MetricReceiver;
+import com.yahoo.vespa.hosted.dockerapi.metrics.MetricReceiverWrapper;
 import org.junit.Test;
+import org.mockito.ArgumentCaptor;
 import org.mockito.Matchers;
 
 import java.io.IOException;
@@ -19,9 +27,11 @@ import java.security.NoSuchAlgorithmException;
 import java.security.UnrecoverableKeyException;
 
 import static org.hamcrest.CoreMatchers.is;
+import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertThat;
 import static org.junit.Assert.assertTrue;
 import static org.mockito.Matchers.any;
+import static org.mockito.Matchers.eq;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
 
@@ -72,7 +82,6 @@ public class DockerImplTest {
         clientConfig.getSSLConfig().getSSLContext();
     }
 
-
     private static DockerConfig createConfig(String uri, String caCertPath, String clientCertPath, String clientKeyPath) {
         DockerConfig.Builder configBuilder = new DockerConfig.Builder();
 
@@ -84,6 +93,7 @@ public class DockerImplTest {
         return new DockerConfig(configBuilder);
     }
 
+
     @Test
     public void testExecuteCompletes() throws Exception {
         final String containerId = "container-id";
@@ -119,4 +129,66 @@ public class DockerImplTest {
         final ProcessResult result = docker.executeInContainer(new ContainerName(containerId), command);
         assertThat(result.getExitStatus(), is(exitCode));
     }
+
+    @Test
+    @SuppressWarnings({"unchecked", "rawtypes"})
+    public void pullImageAsyncIfNeededSuccessfully() {
+        final DockerImage image = new DockerImage("test:1.2.3");
+
+        InspectImageResponse inspectImageResponse = mock(InspectImageResponse.class);
+        when(inspectImageResponse.getId()).thenReturn(image.asString());
+
+        InspectImageCmd imageInspectCmd = mock(InspectImageCmd.class);
+        when(imageInspectCmd.exec())
+                .thenThrow(new NotFoundException("Image not found"))
+                .thenReturn(inspectImageResponse);
+
+        // Array to make it final
+        ArgumentCaptor<ResultCallback> resultCallback = ArgumentCaptor.forClass(ResultCallback.class);
+        PullImageCmd pullImageCmd = mock(PullImageCmd.class);
+        when(pullImageCmd.exec(resultCallback.capture())).thenReturn(null);
+
+        final DockerClient dockerClient = mock(DockerClient.class);
+        when(dockerClient.inspectImageCmd(image.asString())).thenReturn(imageInspectCmd);
+        when(dockerClient.pullImageCmd(eq(image.asString()))).thenReturn(pullImageCmd);
+
+        final DockerImpl docker = new DockerImpl(dockerClient);
+        docker.setMetrics(new MetricReceiverWrapper(MetricReceiver.nullImplementation));
+        assertTrue("Should return true, we just scheduled the pull", docker.pullImageAsyncIfNeeded(image));
+        assertTrue("Should return true, the pull i still ongoing", docker.pullImageAsyncIfNeeded(image));
+
+        assertTrue(docker.imageIsDownloaded(image));
+        resultCallback.getValue().onComplete();
+        assertFalse(docker.pullImageAsyncIfNeeded(image));
+    }
+
+    @Test
+    @SuppressWarnings({"unchecked", "rawtypes"})
+    public void pullImageAsyncIfNeededWithError() {
+        final DockerImage image = new DockerImage("test:1.2.3");
+
+        InspectImageCmd imageInspectCmd = mock(InspectImageCmd.class);
+        when(imageInspectCmd.exec()).thenThrow(new NotFoundException("Image not found"));
+
+        // Array to make it final
+        ArgumentCaptor<ResultCallback> resultCallback = ArgumentCaptor.forClass(ResultCallback.class);
+        PullImageCmd pullImageCmd = mock(PullImageCmd.class);
+        when(pullImageCmd.exec(resultCallback.capture())).thenReturn(null);
+
+        final DockerClient dockerClient = mock(DockerClient.class);
+        when(dockerClient.inspectImageCmd(image.asString())).thenReturn(imageInspectCmd);
+        when(dockerClient.pullImageCmd(eq(image.asString()))).thenReturn(pullImageCmd);
+
+        final DockerImpl docker = new DockerImpl(dockerClient);
+        docker.setMetrics(new MetricReceiverWrapper(MetricReceiver.nullImplementation));
+        assertTrue("Should return true, we just scheduled the pull", docker.pullImageAsyncIfNeeded(image));
+        assertTrue("Should return true, the pull i still ongoing", docker.pullImageAsyncIfNeeded(image));
+
+        try {
+            resultCallback.getValue().onComplete();
+        } catch (Exception ignored) { }
+
+        assertFalse(docker.imageIsDownloaded(image));
+        assertTrue("Should return true, new pull scheduled", docker.pullImageAsyncIfNeeded(image));
+    }
 }