Pull retry policy for pulling remote images

core/src/main/java/org/testcontainers/containers/Container.java

@@ -11,6 +11,7 @@
 import org.testcontainers.containers.wait.strategy.WaitStrategy;
 import org.testcontainers.images.ImagePullPolicy;
 import org.testcontainers.images.builder.Transferable;
+import org.testcontainers.images.retry.ImagePullRetryPolicy;
 import org.testcontainers.utility.LogUtils;
 import org.testcontainers.utility.MountableFile;
 
@@ -297,6 +298,13 @@ default SELF withEnv(String key, Function<Optional<String>, String> mapper) {
      */
     SELF withImagePullPolicy(ImagePullPolicy policy);
 
+    /**
+     * Set the image retry on pull error policy of the container
+     * @param policy the image pull retry policy
+     * @return this
+     */
+    SELF withImagePullRetryPolicy(ImagePullRetryPolicy policy);
+
     /**
      * Map a resource (file or directory) on the classpath to a path inside the container.
      * This will only work if you are running your tests outside a Docker container.

core/src/main/java/org/testcontainers/containers/GenericContainer.java

@@ -48,6 +48,7 @@
 import org.testcontainers.images.ImagePullPolicy;
 import org.testcontainers.images.RemoteDockerImage;
 import org.testcontainers.images.builder.Transferable;
+import org.testcontainers.images.retry.ImagePullRetryPolicy;
 import org.testcontainers.lifecycle.Startable;
 import org.testcontainers.lifecycle.Startables;
 import org.testcontainers.utility.Base58;
@@ -1181,6 +1182,12 @@ public SELF withImagePullPolicy(ImagePullPolicy imagePullPolicy) {
         return self();
     }
 
+    @Override
+    public SELF withImagePullRetryPolicy(ImagePullRetryPolicy policy) {
+        this.image = this.image.withImagePullRetryPolicy(policy);
+        return self();
+    }
+
     /**
      * {@inheritDoc}
      */

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

@@ -3,7 +3,6 @@
 import com.github.dockerjava.api.DockerClient;
 import com.github.dockerjava.api.command.PullImageCmd;
 import com.github.dockerjava.api.exception.DockerClientException;
-import com.github.dockerjava.api.exception.InternalServerErrorException;
 import com.github.dockerjava.api.exception.NotFoundException;
 import com.google.common.util.concurrent.Futures;
 import lombok.AccessLevel;
@@ -18,6 +17,8 @@
 import org.slf4j.Logger;
 import org.testcontainers.DockerClientFactory;
 import org.testcontainers.containers.ContainerFetchException;
+import org.testcontainers.images.retry.ImagePullRetryPolicy;
+import org.testcontainers.images.retry.PullRetryPolicy;
 import org.testcontainers.utility.DockerImageName;
 import org.testcontainers.utility.DockerLoggerFactory;
 import org.testcontainers.utility.ImageNameSubstitutor;
@@ -46,6 +47,9 @@ public class RemoteDockerImage extends LazyFuture<String> {
     @With
     ImagePullPolicy imagePullPolicy = PullPolicy.defaultPolicy();
 
+    @With
+    private ImagePullRetryPolicy imagePullRetryPolicy = PullRetryPolicy.defaultRetryPolicy();
+
     @With
     private ImageNameSubstitutor imageNameSubstitutor = ImageNameSubstitutor.instance();
 
@@ -87,7 +91,6 @@ protected final String resolve() {
             );
 
             final Instant startedAt = Instant.now();
-            final Instant lastRetryAllowed = Instant.now().plus(PULL_RETRY_TIME_LIMIT);
             final AtomicReference<Exception> lastFailure = new AtomicReference<>();
             final PullImageCmd pullImageCmd = dockerClient
                 .pullImageCmd(imageName.getUnversionedPart())
@@ -100,14 +103,16 @@ protected final String resolve() {
                 .iterative(duration -> duration.multipliedBy(2))
                 .startDuration(Duration.ofMillis(50));
 
+            imagePullRetryPolicy.pullStarted();
+
             Awaitility
                 .await()
                 .pollInSameThread()
                 .pollDelay(Duration.ZERO) // start checking immediately
                 .atMost(PULL_RETRY_TIME_LIMIT)
                 .pollInterval(interval)
                 .until(
-                    tryImagePullCommand(pullImageCmd, logger, dockerImageName, imageName, lastFailure, lastRetryAllowed)
+                    tryImagePullCommand(pullImageCmd, logger, dockerImageName, imageName, lastFailure, imagePullRetryPolicy)
                 );
 
             if (dockerImageName.get() == null) {
@@ -135,23 +140,23 @@ private Callable<Boolean> tryImagePullCommand(
         AtomicReference<String> dockerImageName,
         DockerImageName imageName,
         AtomicReference<Exception> lastFailure,
-        Instant lastRetryAllowed
+        ImagePullRetryPolicy imagePullRetryPolicy
     ) {
         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;
-            }
+            boolean pull;
+
+            do {
+                try {
+                    pullImage(pullImageCmd, logger);
+                    dockerImageName.set(imageName.asCanonicalNameString());
+                    return true;
+                } catch (Exception e) {
+                    lastFailure.set(e);
+                    pull = imagePullRetryPolicy.shouldRetry(imageName, e);
+                }
+            } while (pull);
+
+            return false;
         };
     }
 

core/src/main/java/org/testcontainers/images/retry/DefaultPullRetryPolicy.java

@@ -0,0 +1,41 @@
+package org.testcontainers.images.retry;
+
+import com.github.dockerjava.api.exception.InternalServerErrorException;
+import lombok.ToString;
+import lombok.extern.slf4j.Slf4j;
+import org.testcontainers.utility.DockerImageName;
+import org.testcontainers.utility.TestcontainersConfiguration;
+
+import java.time.Duration;
+
+/**
+ * Default pull retry policy.
+ *
+ * Will retry on <code>InterruptedException</code> and <code>InternalServerErrorException</code>
+ * exceptions for a time limit of two minutes.
+ */
+@Slf4j
+@ToString
+public class DefaultPullRetryPolicy extends LimitedDurationPullRetryPolicy {
+
+    private static final Duration PULL_RETRY_TIME_LIMIT = Duration.ofSeconds(
+        TestcontainersConfiguration.getInstance().getImagePullTimeout()
+    );
+
+    public DefaultPullRetryPolicy() {
+        super(PULL_RETRY_TIME_LIMIT);
+    }
+
+    @Override
+    public boolean shouldRetry(DockerImageName imageName, Throwable error) {
+        if (!mayRetry(error)) {
+            return false;
+        }
+
+        return super.shouldRetry(imageName, error);
+    }
+
+    private boolean mayRetry(Throwable error) {
+        return error instanceof InterruptedException || error instanceof InternalServerErrorException;
+    }
+}

core/src/main/java/org/testcontainers/images/retry/FailFastPullRetryPolicy.java

@@ -0,0 +1,18 @@
+package org.testcontainers.images.retry;
+
+import lombok.ToString;
+import lombok.extern.slf4j.Slf4j;
+import org.testcontainers.utility.DockerImageName;
+
+/**
+ * Fail-fast, i.e. not retry, pull policy
+ */
+@Slf4j
+@ToString
+public class FailFastPullRetryPolicy implements ImagePullRetryPolicy {
+
+    @Override
+    public boolean shouldRetry(DockerImageName imageName, Throwable error) {
+        return false;
+    }
+}

core/src/main/java/org/testcontainers/images/retry/ImagePullRetryPolicy.java

@@ -0,0 +1,9 @@
+package org.testcontainers.images.retry;
+
+import org.testcontainers.utility.DockerImageName;
+
+public interface ImagePullRetryPolicy {
+    default void pullStarted() {}
+
+    boolean shouldRetry(DockerImageName imageName, Throwable error);
+}

core/src/main/java/org/testcontainers/images/retry/LimitedDurationPullRetryPolicy.java

@@ -0,0 +1,59 @@
+package org.testcontainers.images.retry;
+
+import lombok.Getter;
+import lombok.ToString;
+import lombok.extern.slf4j.Slf4j;
+import org.testcontainers.utility.DockerImageName;
+
+import java.time.Duration;
+import java.time.Instant;
+import java.util.Objects;
+
+/**
+ * An ImagePullRetryPolicy which will retry a failed image pull if the time elapsed since the
+ * pull started is less than or equal to the configured {@link #maxAllowedDuration}.
+ *
+ */
+@Slf4j
+@ToString
+public class LimitedDurationPullRetryPolicy implements ImagePullRetryPolicy {
+
+    @Getter
+    private final Duration maxAllowedDuration;
+
+    private Instant lastRetryAllowed;
+
+    public LimitedDurationPullRetryPolicy(Duration maxAllowedDuration) {
+        Objects.requireNonNull(maxAllowedDuration, "maxAllowedDuration should not be null");
+
+        if (maxAllowedDuration.isNegative()) {
+            throw new IllegalArgumentException("maxAllowedDuration should not be negative");
+        }
+
+        this.maxAllowedDuration = maxAllowedDuration;
+    }
+
+    @Override
+    public void pullStarted() {
+        this.lastRetryAllowed = Instant.now().plus(maxAllowedDuration);
+    }
+
+    @Override
+    public boolean shouldRetry(DockerImageName imageName, Throwable error) {
+        if (lastRetryAllowed == null) {
+            throw new IllegalStateException("lastRetryAllowed is null. Please, check that pullStarted has been called.");
+        }
+
+        if (Instant.now().isBefore(lastRetryAllowed)) {
+            log.warn(
+                "Retrying pull for image: {} ({}s remaining)",
+                imageName,
+                Duration.between(Instant.now(), lastRetryAllowed).getSeconds()
+            );
+
+            return true;
+        }
+
+        return false;
+    }
+}

core/src/main/java/org/testcontainers/images/retry/NoOfAttemptsPullRetryPolicy.java

@@ -0,0 +1,41 @@
+package org.testcontainers.images.retry;
+
+import lombok.Getter;
+import lombok.ToString;
+import lombok.extern.slf4j.Slf4j;
+import org.testcontainers.utility.DockerImageName;
+
+import java.util.concurrent.atomic.AtomicInteger;
+
+/**
+ * An ImagePullRetryPolicy which will retry a failed image pull if the number of attempts
+ * is less than or equal to the configured {@link #maxAllowedNoOfAttempts}.
+ *
+ */
+@Slf4j
+@ToString
+public class NoOfAttemptsPullRetryPolicy implements ImagePullRetryPolicy {
+
+    @Getter
+    private final int maxAllowedNoOfAttempts;
+
+    private final AtomicInteger currentNoOfAttempts = new AtomicInteger(0);
+
+    public NoOfAttemptsPullRetryPolicy(int maxAllowedNoOfAttempts) {
+        if (maxAllowedNoOfAttempts < 0) {
+            throw new IllegalArgumentException("maxAllowedNoOfAttempts should not be negative");
+        }
+
+        this.maxAllowedNoOfAttempts = maxAllowedNoOfAttempts;
+    }
+
+    @Override
+    public void pullStarted() {
+        currentNoOfAttempts.set(0);
+    }
+
+    @Override
+    public boolean shouldRetry(DockerImageName imageName, Throwable error) {
+        return currentNoOfAttempts.incrementAndGet() <= maxAllowedNoOfAttempts;
+    }
+}

core/src/main/java/org/testcontainers/images/retry/PullRetryPolicy.java

@@ -0,0 +1,46 @@
+package org.testcontainers.images.retry;
+
+import lombok.experimental.UtilityClass;
+
+import java.time.Duration;
+
+/**
+ * Convenience class with logic for building common {@link ImagePullRetryPolicy} instances.
+ *
+ */
+@UtilityClass
+public class PullRetryPolicy {
+
+    /**
+     * Convenience method for returning the {@link FailFastPullRetryPolicy} failFast image pull retry policy
+     * @return {@link ImagePullRetryPolicy}
+     */
+    public static ImagePullRetryPolicy failFast() {
+        return new FailFastPullRetryPolicy();
+    }
+
+    /**
+     * Convenience method for returning the {@link NoOfAttemptsPullRetryPolicy} number of attempts based image pull
+     * retry policy.
+     * @return {@link ImagePullRetryPolicy}
+     */
+    public static ImagePullRetryPolicy noOfAttempts(int allowedNoOfAttempts) {
+        return new NoOfAttemptsPullRetryPolicy(allowedNoOfAttempts);
+    }
+
+    /**
+     * Convenience method for returning the {@link LimitedDurationPullRetryPolicy} duration image pull retry policy
+     * @return {@link ImagePullRetryPolicy}
+     */
+    public static ImagePullRetryPolicy limitedDuration(Duration maxAllowedDuration) {
+        return new LimitedDurationPullRetryPolicy(maxAllowedDuration);
+    }
+
+    /**
+     * Convenience method for returning the {@link DefaultPullRetryPolicy} default image pull retry policy.
+     * @return {@link ImagePullRetryPolicy}
+     */
+    public static ImagePullRetryPolicy defaultRetryPolicy() {
+        return new DefaultPullRetryPolicy();
+    }
+}