need withConcurrentOutput to flush any buffered concurrent output

4 files changed, 16 insertions, 2 deletions

diff --git a/src/Propellor/CmdLine.hs b/src/Propellor/CmdLine.hs
index 9f798166..4bca3986 100644
--- a/src/Propellor/CmdLine.hs
+++ b/src/Propellor/CmdLine.hs
@@ -89,7 +89,7 @@ processCmdLine = go =<< getArgs
 
 -- | Runs propellor on hosts, as controlled by command-line options.
 defaultMain :: [Host] -> IO ()
-defaultMain hostlist = do
+defaultMain hostlist = withConcurrentOutput $ do
 	Shim.cleanEnv
 	checkDebugMode
 	cmdline <- processCmdLine
diff --git a/src/Propellor/Message.hs b/src/Propellor/Message.hs
index 6d541b9a..7439c362 100644
--- a/src/Propellor/Message.hs
+++ b/src/Propellor/Message.hs
@@ -16,6 +16,7 @@ module Propellor.Message (
 	processChainOutput,
 	messagesDone,
 	createProcessConcurrent,
+	withConcurrentOutput,
 ) where
 
 import System.Console.ANSI
diff --git a/src/Utility/ConcurrentOutput.hs b/src/Utility/ConcurrentOutput.hs
index 1ca92d90..c6550b84 100644
--- a/src/Utility/ConcurrentOutput.hs
+++ b/src/Utility/ConcurrentOutput.hs
@@ -1,6 +1,7 @@
 -- | Concurrent output handling.
 
 module Utility.ConcurrentOutput (
+	withConcurrentOutput,
 	outputConcurrent,
 	createProcessConcurrent,
 ) where
@@ -113,6 +114,18 @@ updateOutputLocker l = do
 	putMVar lcker l
 	modifyMVar_ lcker (const $ return l)
 
+-- | Use this around any IO actions that use `outputConcurrent`
+-- or `createProcessConcurrent`
+--
+-- 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
+	-- Just taking the output lock is enough to ensure that anything
+	-- that was buffering output has had a chance to flush its buffer.
+	drain = lockOutput (return ())
+
 -- | Displays a string to stdout, and flush output so it's displayed.
 --
 -- Uses locking to ensure that the whole string is output atomically
diff --git a/src/wrapper.hs b/src/wrapper.hs
index e367fe69..0cfe319d 100644
--- a/src/wrapper.hs
+++ b/src/wrapper.hs
@@ -50,7 +50,7 @@ netrepo :: String
 netrepo = "https://github.com/joeyh/propellor.git"
 
 main :: IO ()
-main = do
+main = withConcurrentOutput $ do
 	args <- getArgs
 	home <- myHomeDir
 	let propellordir = home </> ".propellor"