From 6564b14a41dd51a3a7dcd3b632802b14a014b484 Mon Sep 17 00:00:00 2001
From: Mitchell Rosen <mitchellwrosen@gmail.com>
Date: Tue, 25 Feb 2020 08:54:22 -0500
Subject: [PATCH] Wait for tee'd process stdout/stderr to hit EOF

---
 src/Database/Postgres/Temp/Internal/Core.hs | 34 ++++++++++++---------
 test/Main.hs                                |  4 +++
 2 files changed, 24 insertions(+), 14 deletions(-)

diff --git a/src/Database/Postgres/Temp/Internal/Core.hs b/src/Database/Postgres/Temp/Internal/Core.hs
index 1379c6b..73fce10 100644
--- a/src/Database/Postgres/Temp/Internal/Core.hs
+++ b/src/Database/Postgres/Temp/Internal/Core.hs
@@ -7,7 +7,7 @@ See 'startPlan' for more details.
 module Database.Postgres.Temp.Internal.Core where
 
 import           Control.Concurrent
-import           Control.Concurrent.Async (race_, withAsync)
+import           Control.Concurrent.Async (Async, race_, wait, withAsync)
 import           Control.Exception
 import           Control.Monad
 import qualified Data.ByteString.Char8 as BSC
@@ -123,18 +123,19 @@ waitForDB logger options = do
     Right () -> return ()
 
 -- Only useful if we believe the output is finite
-teeHandle :: Handle -> (Handle -> IO a) -> IO (a, String)
+teeHandle :: Handle -> (Handle -> Async String -> IO a) -> IO a
 teeHandle orig f =
   bracket createPipe (\(x, y) -> hClose x >> hClose y) $ \(readEnd, writeEnd) -> do
-    outputRef <- newIORef []
+    let readerLoop acc = do
+          eof <- hIsEOF readEnd
+          if eof then
+            pure acc
+          else do
+            theLine <- hGetLine readEnd
+            hPutStrLn orig theLine
+            readerLoop (acc <> theLine)
 
-    let readerLoop = forever $ do
-          theLine <- hGetLine readEnd
-          modifyIORef outputRef (<>theLine)
-          hPutStrLn orig theLine
-
-    res <- withAsync readerLoop $ \_ -> f writeEnd
-    (res,) <$> readIORef outputRef
+    withAsync (readerLoop "") $ f writeEnd
 
 -- | 'CompleteProcessConfig' contains the configuration necessary for starting a
 --   process. It is essentially a stripped down 'System.Process.CreateProcess'.
@@ -220,13 +221,18 @@ executeProcessAndTee
   -> CompleteProcessConfig
   -- ^ Process config
   -> IO (ExitCode, String, String)
-executeProcessAndTee name config = fmap (\((x, y), z) -> (x, z, y)) $
-  teeHandle (completeProcessConfigStdOut config) $ \newOut ->
-    teeHandle (completeProcessConfigStdErr config) $ \newErr ->
-      executeProcess name $ config
+executeProcessAndTee name config =
+  teeHandle (completeProcessConfigStdOut config) $ \newOut aOut ->
+    teeHandle (completeProcessConfigStdErr config) $ \newErr aErr -> do
+      result <- executeProcess name $ config
         { completeProcessConfigStdErr = newErr
         , completeProcessConfigStdOut = newOut
         }
+      for_ [newOut, newErr] hClose -- executeProcess doesn't close these
+      stdoutString <- wait aOut
+      stderrString <- wait aErr
+      pure (result, stdoutString, stderrString)
+
 -------------------------------------------------------------------------------
 -- PostgresProcess Life cycle management
 -------------------------------------------------------------------------------
diff --git a/test/Main.hs b/test/Main.hs
index 63676c0..1a7cb9a 100644
--- a/test/Main.hs
+++ b/test/Main.hs
@@ -811,6 +811,10 @@ spec = do
           PG.query_ conn "SELECT id FROM foo ORDER BY id ASC"
             `shouldReturn` [PG.Only (1 :: Int)]
 
+  it "executeProcessAndTee doesn't lose stderr" do
+    let config = CompleteProcessConfig [] ["--bogus-argument"] devNull devNull devNull
+    let expected = "could not find a \"initdb\" to executeinitdb: unrecognized option '--bogus-argument'Try \"initdb --help\" for more information."
+    executeProcessAndTee "initdb" config `shouldReturn` (ExitFailure 1, "", expected)
 
 main :: IO ()
 main = hspec spec