join back threads before ending concurrent output so display works

I didn't really want to put allowConcurrentOutput in CmdLine.Action,
but there were dep loops and that was the best place available.

1 files changed, 26 insertions, 1 deletions

diff --git a/CmdLine/Action.hs b/CmdLine/Action.hs
index 73cffec76..2579196c9 100644
--- a/CmdLine/Action.hs
+++ b/CmdLine/Action.hs
@@ -5,6 +5,8 @@
  - Licensed under the GNU GPL version 3 or higher.
  -}
 
+{-# LANGUAGE CPP #-}
+
 module CmdLine.Action where
 
 import Common.Annex
@@ -12,6 +14,7 @@ import qualified Annex
 import Annex.Concurrent
 import Types.Command
 import qualified Annex.Queue
+import Messages.Concurrent
 import Messages.Internal
 import Types.Messages
 
@@ -19,6 +22,10 @@ import Control.Concurrent.Async
 import Control.Exception (throwIO)
 import Data.Either
 
+#ifdef WITH_CONCURRENTOUTPUT
+import qualified System.Console.Regions as Regions
+#endif
+
 {- Runs a command, starting with the check stage, and then
  - the seek stage. Finishes by running the continutation, and 
  - then showing a count of any failures. -}
@@ -71,7 +78,9 @@ commandAction a = withOutputType go
  -}
 finishCommandActions :: Annex ()
 finishCommandActions = do
-	l <- liftIO . drainTo 0 =<< Annex.getState Annex.workers
+	ws <- Annex.getState Annex.workers
+	Annex.changeState $ \s -> s { Annex.workers = [] }
+	l <- liftIO $ drainTo 0 ws
 	forM_ (lefts l) mergeState
 
 {- Wait for Asyncs from the list to finish, replacing them with their
@@ -138,3 +147,19 @@ callCommandAction = start
 	skip = return True
 	failure = showEndFail >> return False
 	status r = showEndResult r >> return r
+
+{- Do concurrent output when that has been requested. -}
+allowConcurrentOutput :: Annex a -> Annex a
+#ifdef WITH_CONCURRENTOUTPUT
+allowConcurrentOutput a = go =<< Annex.getState Annex.concurrentjobs
+  where
+	go Nothing = a
+	go (Just n) = Regions.displayConsoleRegions $
+		bracket_ (setup n) cleanup a
+	setup = Annex.setOutput . ConcurrentOutput
+	cleanup = do
+		finishCommandActions
+		Annex.setOutput NormalOutput
+#else
+allowConcurrentOutput = id
+#endif