Avoid connection leaks in LogUtils (#2834)

testcontainers · Jun 6, 2020 · 0035c31 · 0035c31
1 parent ddae79c
commit 0035c31
Show file tree

Hide file tree

Showing 2 changed files with 44 additions and 19 deletions.
diff --git a/core/src/main/java/org/testcontainers/containers/wait/strategy/LogMessageWaitStrategy.java b/core/src/main/java/org/testcontainers/containers/wait/strategy/LogMessageWaitStrategy.java
@@ -1,34 +1,53 @@
 package org.testcontainers.containers.wait.strategy;
 
+import com.github.dockerjava.api.command.LogContainerCmd;
+import lombok.SneakyThrows;
 import org.testcontainers.DockerClientFactory;
 import org.testcontainers.containers.ContainerLaunchException;
+import org.testcontainers.containers.output.FrameConsumerResultCallback;
 import org.testcontainers.containers.output.OutputFrame;
 import org.testcontainers.containers.output.WaitingConsumer;
-import org.testcontainers.utility.LogUtils;
 
+import java.io.IOException;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
 import java.util.function.Predicate;
 
+import static org.testcontainers.containers.output.OutputFrame.OutputType.STDERR;
+import static org.testcontainers.containers.output.OutputFrame.OutputType.STDOUT;
+
 public class LogMessageWaitStrategy extends AbstractWaitStrategy {
 
     private String regEx;
 
     private int times = 1;
 
     @Override
+    @SneakyThrows(IOException.class)
     protected void waitUntilReady() {
         WaitingConsumer waitingConsumer = new WaitingConsumer();
-        LogUtils.followOutput(DockerClientFactory.instance().client(), waitStrategyTarget.getContainerId(), waitingConsumer);
 
-        Predicate<OutputFrame> waitPredicate = outputFrame ->
-            // (?s) enables line terminator matching (equivalent to Pattern.DOTALL)
-            outputFrame.getUtf8String().matches("(?s)" + regEx);
+        LogContainerCmd cmd = DockerClientFactory.instance().client().logContainerCmd(waitStrategyTarget.getContainerId())
+            .withFollowStream(true)
+            .withSince(0)
+            .withStdOut(true)
+            .withStdErr(true);
+
+        try (FrameConsumerResultCallback callback = new FrameConsumerResultCallback()) {
+            callback.addConsumer(STDOUT, waitingConsumer);
+            callback.addConsumer(STDERR, waitingConsumer);
+
+            cmd.exec(callback);
+
+            Predicate<OutputFrame> waitPredicate = outputFrame ->
+                // (?s) enables line terminator matching (equivalent to Pattern.DOTALL)
+                outputFrame.getUtf8String().matches("(?s)" + regEx);
 
-        try {
-            waitingConsumer.waitUntil(waitPredicate, startupTimeout.getSeconds(), TimeUnit.SECONDS, times);
-        } catch (TimeoutException e) {
-            throw new ContainerLaunchException("Timed out waiting for log output matching '" + regEx + "'");
+            try {
+                waitingConsumer.waitUntil(waitPredicate, startupTimeout.getSeconds(), TimeUnit.SECONDS, times);
+            } catch (TimeoutException e) {
+                throw new ContainerLaunchException("Timed out waiting for log output matching '" + regEx + "'");
+            }
         }
     }
 

diff --git a/core/src/main/java/org/testcontainers/utility/LogUtils.java b/core/src/main/java/org/testcontainers/utility/LogUtils.java
@@ -2,12 +2,15 @@
 
 import com.github.dockerjava.api.DockerClient;
 import com.github.dockerjava.api.command.LogContainerCmd;
+import lombok.SneakyThrows;
 import lombok.experimental.UtilityClass;
 import org.testcontainers.containers.output.FrameConsumerResultCallback;
 import org.testcontainers.containers.output.OutputFrame;
 import org.testcontainers.containers.output.ToStringConsumer;
 import org.testcontainers.containers.output.WaitingConsumer;
 
+import java.io.Closeable;
+import java.io.IOException;
 import java.util.function.Consumer;
 
 import static org.testcontainers.containers.output.OutputFrame.OutputType.STDERR;
@@ -59,6 +62,7 @@ public void followOutput(DockerClient dockerClient,
      * @param types        types of {@link OutputFrame} to receive
      * @return all previous output frames (stdout/stderr being separated by newline characters)
      */
+    @SneakyThrows(IOException.class)
     public String getOutput(DockerClient dockerClient,
                             String containerId,
                             OutputFrame.OutputType... types) {
@@ -73,17 +77,19 @@ public String getOutput(DockerClient dockerClient,
 
         final ToStringConsumer consumer = new ToStringConsumer();
         final WaitingConsumer wait = new WaitingConsumer();
-        attachConsumer(dockerClient, containerId, consumer.andThen(wait), false, types);
-
-        wait.waitUntilEnd();
-        return consumer.toUtf8String();
+        try (Closeable closeable = attachConsumer(dockerClient, containerId, consumer.andThen(wait), false, types)) {
+            wait.waitUntilEnd();
+            return consumer.toUtf8String();
+        }
     }
 
-    private static void attachConsumer(DockerClient dockerClient,
-                                       String containerId,
-                                       Consumer<OutputFrame> consumer,
-                                       boolean followStream,
-                                       OutputFrame.OutputType... types) {
+    private static Closeable attachConsumer(
+        DockerClient dockerClient,
+        String containerId,
+        Consumer<OutputFrame> consumer,
+        boolean followStream,
+        OutputFrame.OutputType... types
+    ) {
 
         final LogContainerCmd cmd = dockerClient.logContainerCmd(containerId)
             .withFollowStream(followStream)
@@ -96,6 +102,6 @@ private static void attachConsumer(DockerClient dockerClient,
             if (type == STDERR) cmd.withStdErr(true);
         }
 
-        cmd.exec(callback);
+        return cmd.exec(callback);
     }
 }