Use Awaitility to Poll the Image Pull

core/src/main/java/org/testcontainers/images/RemoteDockerImage.java

@@ -11,6 +11,9 @@
 import lombok.SneakyThrows;
 import lombok.ToString;
 import lombok.With;
+import org.awaitility.Awaitility;
+import org.awaitility.pollinterval.IterativePollInterval;
+import org.awaitility.pollinterval.PollInterval;
 import org.slf4j.Logger;
 import org.testcontainers.DockerClientFactory;
 import org.testcontainers.containers.ContainerFetchException;
@@ -21,9 +24,11 @@
 
 import java.time.Duration;
 import java.time.Instant;
+import java.util.concurrent.Callable;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.Future;
+import java.util.concurrent.atomic.AtomicReference;
 
 @ToString
 @AllArgsConstructor(access = AccessLevel.PACKAGE)
@@ -65,7 +70,7 @@ public RemoteDockerImage(@NonNull Future<String> imageFuture) {
     @SneakyThrows({ InterruptedException.class, ExecutionException.class })
     protected final String resolve() {
         final DockerImageName imageName = getImageName();
-        Logger logger = DockerLoggerFactory.getLogger(imageName.toString());
+        final Logger logger = DockerLoggerFactory.getLogger(imageName.toString());
         try {
             if (!imagePullPolicy.shouldPull(imageName)) {
                 return imageName.asCanonicalNameString();
@@ -77,55 +82,88 @@ protected final String resolve() {
                 imageName
             );
 
-            Exception lastFailure = null;
+            final Instant startedAt = Instant.now();
             final Instant lastRetryAllowed = Instant.now().plus(PULL_RETRY_TIME_LIMIT);
-
-            Instant startedAt = Instant.now();
-            while (Instant.now().isBefore(lastRetryAllowed)) {
-                try {
-                    PullImageCmd pullImageCmd = dockerClient
-                        .pullImageCmd(imageName.getUnversionedPart())
-                        .withTag(imageName.getVersionPart());
-
-                    try {
-                        pullImageCmd.exec(new TimeLimitedLoggedPullImageResultCallback(logger)).awaitCompletion();
-                    } catch (DockerClientException e) {
-                        // Try to fallback to x86
-                        pullImageCmd
-                            .withPlatform("linux/amd64")
-                            .exec(new TimeLimitedLoggedPullImageResultCallback(logger))
-                            .awaitCompletion();
-                    }
-                    String dockerImageName = imageName.asCanonicalNameString();
-                    logger.info("Image {} pull took {}", dockerImageName, Duration.between(startedAt, Instant.now()));
-
-                    LocalImagesCache.INSTANCE.refreshCache(imageName);
-
-                    return dockerImageName;
-                } catch (InterruptedException | InternalServerErrorException e) {
-                    // these classes of exception often relate to timeout/connection errors so should be retried
-                    lastFailure = e;
-                    logger.warn(
-                        "Retrying pull for image: {} ({}s remaining)",
-                        imageName,
-                        Duration.between(Instant.now(), lastRetryAllowed).getSeconds()
-                    );
-                }
+            final AtomicReference<Exception> lastFailure = new AtomicReference<>();
+            final PullImageCmd pullImageCmd = dockerClient
+                .pullImageCmd(imageName.getUnversionedPart())
+                .withTag(imageName.getVersionPart());
+            final AtomicReference<String> dockerImageName = new AtomicReference<>();
+
+            // The following poll interval in ms: 50, 100, 200, 400, 800....
+            // Results in ~70 requests in over 2 minutes
+            final PollInterval interval = IterativePollInterval
+                .iterative(duration -> duration.multipliedBy(2))
+                .startDuration(Duration.ofMillis(50));
+
+            Awaitility
+                .await()
+                .pollInSameThread()
+                .pollDelay(Duration.ZERO) // start checking immediately
+                .atMost(PULL_RETRY_TIME_LIMIT)
+                .pollInterval(interval)
+                .until(
+                    tryImagePullCommand(pullImageCmd, logger, dockerImageName, imageName, lastFailure, lastRetryAllowed)
+                );
+
+            if (dockerImageName.get() == null) {
+                final Exception lastException = lastFailure.get();
+                logger.error(
+                    "Failed to pull image: {}. Please check output of `docker pull {}`",
+                    imageName,
+                    imageName,
+                    lastException
+                );
+                throw new ContainerFetchException("Failed to pull image: " + imageName, lastException);
             }
 
-            logger.error(
-                "Failed to pull image: {}. Please check output of `docker pull {}`",
-                imageName,
-                imageName,
-                lastFailure
-            );
-
-            throw new ContainerFetchException("Failed to pull image: " + imageName, lastFailure);
+            logger.info("Image {} pull took {}", dockerImageName.get(), Duration.between(startedAt, Instant.now()));
+            LocalImagesCache.INSTANCE.refreshCache(imageName);
+            return dockerImageName.get();
         } catch (DockerClientException e) {
             throw new ContainerFetchException("Failed to get Docker client for " + imageName, e);
         }
     }
 
+    private Callable<Boolean> tryImagePullCommand(
+        PullImageCmd pullImageCmd,
+        Logger logger,
+        AtomicReference<String> dockerImageName,
+        DockerImageName imageName,
+        AtomicReference<Exception> lastFailure,
+        Instant lastRetryAllowed
+    ) {
+        return () -> {
+            try {
+                pullImage(pullImageCmd, logger);
+                dockerImageName.set(imageName.asCanonicalNameString());
+                return true;
+            } catch (InterruptedException | InternalServerErrorException e) {
+                // these classes of exception often relate to timeout/connection errors so should be retried
+                lastFailure.set(e);
+                logger.warn(
+                    "Retrying pull for image: {} ({}s remaining)",
+                    imageName,
+                    Duration.between(Instant.now(), lastRetryAllowed).getSeconds()
+                );
+                return false;
+            }
+        };
+    }
+
+    private TimeLimitedLoggedPullImageResultCallback pullImage(PullImageCmd pullImageCmd, Logger logger)
+        throws InterruptedException {
+        try {
+            return pullImageCmd.exec(new TimeLimitedLoggedPullImageResultCallback(logger)).awaitCompletion();
+        } catch (DockerClientException e) {
+            // Try to fallback to x86
+            return pullImageCmd
+                .withPlatform("linux/amd64")
+                .exec(new TimeLimitedLoggedPullImageResultCallback(logger))
+                .awaitCompletion();
+        }
+    }
+
     private DockerImageName getImageName() throws InterruptedException, ExecutionException {
         final DockerImageName specifiedImageName = imageNameFuture.get();