Refactor

testcontainers · Mar 18, 2019 · 4c7f395 · 4c7f395
1 parent 516433e
commit 4c7f395
Show file tree

Hide file tree

Showing 2 changed files with 38 additions and 7 deletions.
diff --git a/core/src/main/java/org/testcontainers/images/LoggedPullImageResultCallback.java b/core/src/main/java/org/testcontainers/images/LoggedPullImageResultCallback.java
@@ -21,6 +21,7 @@ class LoggedPullImageResultCallback extends PullImageResultCallback {
     private Set<String> pulledLayers = new HashSet<>();
     private Map<String, Long> totalSizes = new HashMap<>();
     private Map<String, Long> currentSizes = new HashMap<>();
+    private boolean completed;
 
     LoggedPullImageResultCallback(final Logger logger) {
         this.logger = logger;
@@ -84,14 +85,21 @@ public void onNext(final PullResponseItem item) {
                 byteCountToDisplaySize(currentSize),
                 friendlyTotalSize);
         }
+
+        if (status != null && status.contains("complete")) {
+            completed = true;
+        }
     }
 
     @Override
     public void onComplete() {
         super.onComplete();
 
         long totalSize = totalLayerSize();
-        logger.info("Pull complete ({} layers, {})", allLayers.size(), byteCountToDisplaySize(totalSize));
+
+        if (completed) {
+            logger.info("Pull complete ({} layers, {})", allLayers.size(), byteCountToDisplaySize(totalSize));
+        }
     }
 
     private long downloadedLayerSize() {

diff --git a/core/src/main/java/org/testcontainers/images/TimeLimitedLoggedPullImageResultCallback.java b/core/src/main/java/org/testcontainers/images/TimeLimitedLoggedPullImageResultCallback.java
@@ -5,6 +5,7 @@
 import org.slf4j.Logger;
 
 import java.io.Closeable;
+import java.io.IOException;
 import java.time.Duration;
 import java.util.HashSet;
 import java.util.Set;
@@ -20,7 +21,7 @@
  * {@link PullImageResultCallback} with improved logging of pull progress and a 'watchdog' which will abort the pull
  * if progress is not being made, to prevent a hanging test
  */
-public class TimeLimitedLoggedPullImageResultCallback extends LoggedPullImageResultCallback {
+class TimeLimitedLoggedPullImageResultCallback extends LoggedPullImageResultCallback {
 
     private static final AtomicInteger THREAD_ID = new AtomicInteger(0);
     private static final ScheduledExecutorService PROGRESS_WATCHDOG_EXECUTOR =
@@ -33,7 +34,10 @@ public class TimeLimitedLoggedPullImageResultCallback extends LoggedPullImageRes
     private static final Duration PULL_PAUSE_TOLERANCE = Duration.ofSeconds(30);
     private final Logger logger;
 
+    // A future which, if it ever fires, will kill the pull
     private ScheduledFuture<?> nextCheckForProgress;
+
+    // All threads that are 'awaiting' this pull
     private Set<Thread> waitingThreads = new HashSet<>();
 
     TimeLimitedLoggedPullImageResultCallback(Logger logger) {
@@ -47,6 +51,12 @@ public PullImageResultCallback awaitCompletion() throws InterruptedException {
         return super.awaitCompletion();
     }
 
+    @Override
+    public boolean awaitCompletion(long timeout, TimeUnit timeUnit) throws InterruptedException {
+        waitingThreads.add(Thread.currentThread());
+        return super.awaitCompletion(timeout, timeUnit);
+    }
+
     @Override
     public void onNext(PullResponseItem item) {
         if (item.getProgressDetail() != null) {
@@ -81,13 +91,26 @@ public void onComplete() {
      */
     private void resetProgressWatchdog(boolean isFinished) {
         if (nextCheckForProgress != null && ! nextCheckForProgress.isCancelled()) {
-            nextCheckForProgress.cancel(true);
+            nextCheckForProgress.cancel(false);
         }
         if (!isFinished) {
-            nextCheckForProgress = PROGRESS_WATCHDOG_EXECUTOR.schedule(() -> {
-                logger.error("Docker image pull has not made progress in {}s - aborting pull", PULL_PAUSE_TOLERANCE.getSeconds());
-                waitingThreads.forEach(Thread::interrupt);
-            }, PULL_PAUSE_TOLERANCE.getSeconds(), TimeUnit.SECONDS);
+            nextCheckForProgress = PROGRESS_WATCHDOG_EXECUTOR.schedule(
+                this::abortPull,
+                PULL_PAUSE_TOLERANCE.getSeconds(),
+                TimeUnit.SECONDS
+            );
+        }
+    }
+
+    private void abortPull() {
+        logger.error("Docker image pull has not made progress in {}s - aborting pull", PULL_PAUSE_TOLERANCE.getSeconds());
+        // Interrupt any threads that are waiting, before closing streams, because the stream can take
+        //  an indeterminate amount of time to close
+        waitingThreads.forEach(Thread::interrupt);
+        try {
+            close();
+        } catch (IOException ignored) {
+            // no action
         }
     }
 }