have to flush concurrent output before printing result when chaining

4 files changed, 21 insertions, 14 deletions

diff --git a/src/Propellor/Message.hs b/src/Propellor/Message.hs
index 7439c362..7df5104a 100644
--- a/src/Propellor/Message.hs
+++ b/src/Propellor/Message.hs
@@ -135,7 +135,7 @@ processChainOutput h = go Nothing
 				Just l -> case readish l of
 					Just r -> pure r
 					Nothing -> do
-						outputConcurrent l
+						outputConcurrent (l ++ "\n")
 						return FailedChange
 			Just s -> do
 				outputConcurrent $
diff --git a/src/Propellor/Property/Chroot.hs b/src/Propellor/Property/Chroot.hs
index 8b923aab..e72d1bd9 100644
--- a/src/Propellor/Property/Chroot.hs
+++ b/src/Propellor/Property/Chroot.hs
@@ -213,6 +213,7 @@ chain hostlist (ChrootChain hn loc systemdonly onconsole) =
 					then [Systemd.installed]
 					else map ignoreInfo $
 						hostProperties h
+			flushConcurrentOutput
 			putStrLn $ "\n" ++ show r
 chain _ _ = errorMessage "bad chain command"
 
diff --git a/src/Propellor/Property/Docker.hs b/src/Propellor/Property/Docker.hs
index 5f41209a..9082460f 100644
--- a/src/Propellor/Property/Docker.hs
+++ b/src/Propellor/Property/Docker.hs
@@ -540,6 +540,7 @@ init s = case toContainerId s of
 				warningMessage "Boot provision failed!"
 		void $ async $ job reapzombies
 		job $ do
+			flushConcurrentOutput
 			void $ tryIO $ ifM (inPath "bash")
 				( boolSystem "bash" [Param "-l"]
 				, boolSystem "/bin/sh" []
@@ -583,6 +584,7 @@ chain hostlist hn s = case toContainerId s of
 			r <- runPropellor h $ ensureProperties $
 				map ignoreInfo $
 					hostProperties h
+			flushConcurrentOutput
 			putStrLn $ "\n" ++ show r
 
 stopContainer :: ContainerId -> IO Bool
diff --git a/src/Utility/ConcurrentOutput.hs b/src/Utility/ConcurrentOutput.hs
index db0bae0a..3f28068a 100644
--- a/src/Utility/ConcurrentOutput.hs
+++ b/src/Utility/ConcurrentOutput.hs
@@ -5,6 +5,7 @@
 
 module Utility.ConcurrentOutput (
 	withConcurrentOutput,
+	flushConcurrentOutput,
 	outputConcurrent,
 	createProcessConcurrent,
 	waitForProcessConcurrent,
@@ -105,19 +106,22 @@ dropOutputLock = withLock $ void . takeTMVar
 -- This is necessary to ensure that buffered concurrent output actually
 -- gets displayed before the program exits.
 withConcurrentOutput :: IO a -> IO a
-withConcurrentOutput a = a `finally` drain
-  where
-	-- Wait for all outputThreads to finish. Then, take the output lock
-	-- to ensure that nothing is currently generating output, and flush
-	-- any buffered output.
-	drain = do
-		v <- outputThreads <$> getOutputHandle
-		atomically $ do
-			r <- takeTMVar v
-			if r == S.empty
-				then return ()
-				else retry
-		lockOutput $ return ()
+withConcurrentOutput a = a `finally` flushConcurrentOutput
+
+-- | Blocks until any processes started by `createProcessConcurrent` have
+-- finished, and any buffered output is displayed.
+flushConcurrentOutput :: IO ()
+flushConcurrentOutput = do
+	-- Wait for all outputThreads to finish.
+	v <- outputThreads <$> getOutputHandle
+	atomically $ do
+		r <- takeTMVar v
+		if r == S.empty
+			then return ()
+			else retry
+	-- Take output lock to ensure that nothing else is currently
+	-- generating output, and flush any buffered output.
+	lockOutput $ return ()
 
 -- | Displays a string to stdout, and flush output so it's displayed.
 --