summaryrefslogtreecommitdiff
path: root/src
diff options
context:
space:
mode:
Diffstat (limited to 'src')
-rw-r--r--src/Propellor/CmdLine.hs2
-rw-r--r--src/Propellor/Message.hs1
-rw-r--r--src/Utility/ConcurrentOutput.hs13
-rw-r--r--src/wrapper.hs2
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"