Coordinate image removals after stopping containers (#1745)

testcontainers · Aug 25, 2019 · aab9aa0 · aab9aa0
1 parent da97377
commit aab9aa0
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 22 deletions.
diff --git a/core/src/main/java/org/testcontainers/images/builder/ImageFromDockerfile.java b/core/src/main/java/org/testcontainers/images/builder/ImageFromDockerfile.java
@@ -22,6 +22,7 @@
 import org.testcontainers.utility.Base58;
 import org.testcontainers.utility.DockerLoggerFactory;
 import org.testcontainers.utility.LazyFuture;
+import org.testcontainers.utility.ResourceReaper;
 
 import java.io.IOException;
 import java.io.PipedInputStream;
@@ -42,26 +43,6 @@ public class ImageFromDockerfile extends LazyFuture<String> implements
         StringsTrait<ImageFromDockerfile>,
         DockerfileTrait<ImageFromDockerfile> {
 
-    private static final Set<String> imagesToDelete = Sets.newConcurrentHashSet();
-
-    static {
-        Runtime.getRuntime().addShutdownHook(new Thread(DockerClientFactory.TESTCONTAINERS_THREAD_GROUP, () -> {
-            DockerClient dockerClientForCleaning = DockerClientFactory.instance().client();
-            try {
-                for (String dockerImageName : imagesToDelete) {
-                    log.info("Removing image tagged {}", dockerImageName);
-                    try {
-                        dockerClientForCleaning.removeImageCmd(dockerImageName).withForce(true).exec();
-                    } catch (Throwable e) {
-                        log.warn("Unable to delete image " + dockerImageName, e);
-                    }
-                }
-            } catch (DockerClientException e) {
-                throw new RuntimeException(e);
-            }
-        }));
-    }
-
     private final String dockerImageName;
 
     private boolean deleteOnExit = true;
@@ -102,7 +83,7 @@ protected final String resolve() {
         DockerClient dockerClient = DockerClientFactory.instance().client();
         try {
             if (deleteOnExit) {
-                imagesToDelete.add(dockerImageName);
+                ResourceReaper.instance().registerImageForCleanup(dockerImageName);
             }
 
             BuildImageResultCallback resultCallback = new BuildImageResultCallback() {

diff --git a/core/src/main/java/org/testcontainers/utility/ResourceReaper.java b/core/src/main/java/org/testcontainers/utility/ResourceReaper.java
@@ -11,6 +11,8 @@
 import com.github.dockerjava.api.model.Volume;
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Throwables;
+import com.google.common.collect.Sets;
+
 import lombok.SneakyThrows;
 import lombok.extern.slf4j.Slf4j;
 import org.slf4j.Logger;
@@ -51,7 +53,8 @@ public final class ResourceReaper {
     private static ResourceReaper instance;
     private final DockerClient dockerClient;
     private Map<String, String> registeredContainers = new ConcurrentHashMap<>();
-    private Set<String> registeredNetworks = Collections.newSetFromMap(new ConcurrentHashMap<>());
+    private Set<String> registeredNetworks = Sets.newConcurrentHashSet();
+    private Set<String> registeredImages = Sets.newConcurrentHashSet();
     private AtomicBoolean hookIsSet = new AtomicBoolean(false);
 
     private ResourceReaper() {
@@ -164,6 +167,7 @@ public synchronized static ResourceReaper instance() {
     public synchronized void performCleanup() {
         registeredContainers.forEach(this::stopContainer);
         registeredNetworks.forEach(this::removeNetwork);
+        registeredImages.forEach(this::removeImage);
     }
 
     /**
@@ -337,6 +341,20 @@ public void unregisterNetwork(String identifier) {
     public void unregisterContainer(String identifier) {
         registeredContainers.remove(identifier);
     }
+
+    public void registerImageForCleanup(String dockerImageName) {
+        setHook();
+        registeredImages.add(dockerImageName);
+    }
+
+    private void removeImage(String dockerImageName) {
+        LOGGER.trace("Removing image tagged {}", dockerImageName);
+        try {
+            dockerClient.removeImageCmd(dockerImageName).withForce(true).exec();
+        } catch (Throwable e) {
+            LOGGER.warn("Unable to delete image " + dockerImageName, e);
+        }
+    }
 
     private void setHook() {
         if (hookIsSet.compareAndSet(false, true)) {