{-# LANGUAGE BangPatterns, TypeSynonymInstances, FlexibleInstances, TupleSections #-}
{-# OPTIONS_GHC -O2 #-}
{- Building this module with -O0 causes streams not to fuse and too much
 - memory to be used. -}

-- | 
-- Copyright: 2015 Joey Hess <id@joeyh.name>
-- License: BSD-2-clause
-- 
-- Concurrent output handling, internals.
--
-- May change at any time.

module System.Console.Concurrent.Internal where

import System.IO
import System.Directory
import System.Exit
import Control.Monad
import Control.Monad.IO.Class (liftIO, MonadIO)
import System.IO.Unsafe (unsafePerformIO)
import Control.Concurrent
import Control.Concurrent.STM
import Control.Concurrent.Async
import Data.Maybe
import Data.List
import Data.Monoid
import qualified System.Process as P
import qualified Data.Text as T
import qualified Data.Text.IO as T
import qualified Data.Text.Lazy as L
import Control.Applicative
import Prelude

import Utility.Monad
import Utility.Exception

data OutputHandle = OutputHandle
	{ OutputHandle -> TMVar Lock
outputLock :: TMVar Lock
	, OutputHandle -> TMVar OutputBuffer
outputBuffer :: TMVar OutputBuffer
	, OutputHandle -> TMVar OutputBuffer
errorBuffer :: TMVar OutputBuffer
	, OutputHandle -> TMVar Integer
outputThreads :: TMVar Integer
	}

data Lock = Locked

-- | A shared global variable for the OutputHandle.
{-# NOINLINE globalOutputHandle #-}
globalOutputHandle :: OutputHandle
globalOutputHandle :: OutputHandle
globalOutputHandle = IO OutputHandle -> OutputHandle
forall a. IO a -> a
unsafePerformIO (IO OutputHandle -> OutputHandle)
-> IO OutputHandle -> OutputHandle
forall a b. (a -> b) -> a -> b
$ TMVar Lock
-> TMVar OutputBuffer
-> TMVar OutputBuffer
-> TMVar Integer
-> OutputHandle
OutputHandle
	(TMVar Lock
 -> TMVar OutputBuffer
 -> TMVar OutputBuffer
 -> TMVar Integer
 -> OutputHandle)
-> IO (TMVar Lock)
-> IO
     (TMVar OutputBuffer
      -> TMVar OutputBuffer -> TMVar Integer -> OutputHandle)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IO (TMVar Lock)
forall a. IO (TMVar a)
newEmptyTMVarIO
	IO
  (TMVar OutputBuffer
   -> TMVar OutputBuffer -> TMVar Integer -> OutputHandle)
-> IO (TMVar OutputBuffer)
-> IO (TMVar OutputBuffer -> TMVar Integer -> OutputHandle)
forall a b. IO (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> OutputBuffer -> IO (TMVar OutputBuffer)
forall a. a -> IO (TMVar a)
newTMVarIO ([OutputBufferedActivity] -> OutputBuffer
OutputBuffer [])
	IO (TMVar OutputBuffer -> TMVar Integer -> OutputHandle)
-> IO (TMVar OutputBuffer) -> IO (TMVar Integer -> OutputHandle)
forall a b. IO (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> OutputBuffer -> IO (TMVar OutputBuffer)
forall a. a -> IO (TMVar a)
newTMVarIO ([OutputBufferedActivity] -> OutputBuffer
OutputBuffer [])
	IO (TMVar Integer -> OutputHandle)
-> IO (TMVar Integer) -> IO OutputHandle
forall a b. IO (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Integer -> IO (TMVar Integer)
forall a. a -> IO (TMVar a)
newTMVarIO Integer
0

-- | Holds a lock while performing an action. This allows the action to
-- perform its own output to the console, without using functions from this
-- module.
--
-- While this is running, other threads that try to lockOutput will block.
-- Any calls to `outputConcurrent` and `createProcessConcurrent` will not
-- block, but the output will be buffered and displayed only once the
-- action is done.
lockOutput :: (MonadIO m, MonadMask m) => m a -> m a
lockOutput :: forall (m :: * -> *) a. (MonadIO m, MonadMask m) => m a -> m a
lockOutput = m () -> m () -> m a -> m a
forall (m :: * -> *) a c b.
(HasCallStack, MonadMask m) =>
m a -> m c -> m b -> m b
bracket_ (IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO IO ()
takeOutputLock) (IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO IO ()
dropOutputLock)

-- | Blocks until we have the output lock.
takeOutputLock :: IO ()
takeOutputLock :: IO ()
takeOutputLock = IO Bool -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO Bool -> IO ()) -> IO Bool -> IO ()
forall a b. (a -> b) -> a -> b
$ Bool -> IO Bool
takeOutputLock' Bool
True

-- | Tries to take the output lock, without blocking.
tryTakeOutputLock :: IO Bool
tryTakeOutputLock :: IO Bool
tryTakeOutputLock = Bool -> IO Bool
takeOutputLock' Bool
False

withLock :: (TMVar Lock -> STM a) -> IO a
withLock :: forall a. (TMVar Lock -> STM a) -> IO a
withLock TMVar Lock -> STM a
a = STM a -> IO a
forall a. STM a -> IO a
atomically (STM a -> IO a) -> STM a -> IO a
forall a b. (a -> b) -> a -> b
$ TMVar Lock -> STM a
a (OutputHandle -> TMVar Lock
outputLock OutputHandle
globalOutputHandle)

takeOutputLock' :: Bool -> IO Bool
takeOutputLock' :: Bool -> IO Bool
takeOutputLock' Bool
block = do
	Bool
locked <- (TMVar Lock -> STM Bool) -> IO Bool
forall a. (TMVar Lock -> STM a) -> IO a
withLock ((TMVar Lock -> STM Bool) -> IO Bool)
-> (TMVar Lock -> STM Bool) -> IO Bool
forall a b. (a -> b) -> a -> b
$ \TMVar Lock
l -> do
		Maybe Lock
v <- TMVar Lock -> STM (Maybe Lock)
forall a. TMVar a -> STM (Maybe a)
tryTakeTMVar TMVar Lock
l
		case Maybe Lock
v of
			Just Lock
Locked
				| Bool
block -> STM Bool
forall a. STM a
retry
				| Bool
otherwise -> do
					-- Restore value we took.
					TMVar Lock -> Lock -> STM ()
forall a. TMVar a -> a -> STM ()
putTMVar TMVar Lock
l Lock
Locked
					Bool -> STM Bool
forall a. a -> STM a
forall (m :: * -> *) a. Monad m => a -> m a
return Bool
False
			Maybe Lock
Nothing -> do
				TMVar Lock -> Lock -> STM ()
forall a. TMVar a -> a -> STM ()
putTMVar TMVar Lock
l Lock
Locked
				Bool -> STM Bool
forall a. a -> STM a
forall (m :: * -> *) a. Monad m => a -> m a
return Bool
True
	Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when Bool
locked (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
		(OutputBuffer
outbuf, OutputBuffer
errbuf) <- STM (OutputBuffer, OutputBuffer) -> IO (OutputBuffer, OutputBuffer)
forall a. STM a -> IO a
atomically (STM (OutputBuffer, OutputBuffer)
 -> IO (OutputBuffer, OutputBuffer))
-> STM (OutputBuffer, OutputBuffer)
-> IO (OutputBuffer, OutputBuffer)
forall a b. (a -> b) -> a -> b
$ (,)
			(OutputBuffer -> OutputBuffer -> (OutputBuffer, OutputBuffer))
-> STM OutputBuffer
-> STM (OutputBuffer -> (OutputBuffer, OutputBuffer))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TMVar OutputBuffer -> OutputBuffer -> STM OutputBuffer
forall a. TMVar a -> a -> STM a
swapTMVar (OutputHandle -> TMVar OutputBuffer
outputBuffer OutputHandle
globalOutputHandle) ([OutputBufferedActivity] -> OutputBuffer
OutputBuffer [])
			STM (OutputBuffer -> (OutputBuffer, OutputBuffer))
-> STM OutputBuffer -> STM (OutputBuffer, OutputBuffer)
forall a b. STM (a -> b) -> STM a -> STM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> TMVar OutputBuffer -> OutputBuffer -> STM OutputBuffer
forall a. TMVar a -> a -> STM a
swapTMVar (OutputHandle -> TMVar OutputBuffer
errorBuffer OutputHandle
globalOutputHandle) ([OutputBufferedActivity] -> OutputBuffer
OutputBuffer [])
		StdHandle -> OutputBuffer -> IO ()
emitOutputBuffer StdHandle
StdOut OutputBuffer
outbuf
		StdHandle -> OutputBuffer -> IO ()
emitOutputBuffer StdHandle
StdErr OutputBuffer
errbuf
	Bool -> IO Bool
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return Bool
locked

-- | Only safe to call after taking the output lock.
dropOutputLock :: IO ()
dropOutputLock :: IO ()
dropOutputLock = (TMVar Lock -> STM ()) -> IO ()
forall a. (TMVar Lock -> STM a) -> IO a
withLock ((TMVar Lock -> STM ()) -> IO ())
-> (TMVar Lock -> STM ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ STM Lock -> STM ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (STM Lock -> STM ())
-> (TMVar Lock -> STM Lock) -> TMVar Lock -> STM ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TMVar Lock -> STM Lock
forall a. TMVar a -> STM a
takeTMVar

-- | Use this around any actions that use `outputConcurrent`
-- or `createProcessConcurrent`, unless 
-- `System.Console.Regions.displayConsoleRegions` is being used.
--
-- This is necessary to ensure that buffered concurrent output actually
-- gets displayed before the program exits.
withConcurrentOutput :: (MonadIO m, MonadMask m) => m a -> m a
withConcurrentOutput :: forall (m :: * -> *) a. (MonadIO m, MonadMask m) => m a -> m a
withConcurrentOutput m a
a = m a
a m a -> m () -> m a
forall (m :: * -> *) a b.
(HasCallStack, MonadMask m) =>
m a -> m b -> m a
`finally` IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO IO ()
flushConcurrentOutput

-- | Blocks until any processes started by `createProcessConcurrent` have
-- finished, and any buffered output is displayed. Also blocks while
-- `lockOutput` is is use.
--
-- `withConcurrentOutput` calls this at the end, so you do not normally
-- need to use this.
flushConcurrentOutput :: IO ()
flushConcurrentOutput :: IO ()
flushConcurrentOutput = do
	STM () -> IO ()
forall a. STM a -> IO a
atomically (STM () -> IO ()) -> STM () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
		Integer
r <- TMVar Integer -> STM Integer
forall a. TMVar a -> STM a
takeTMVar (OutputHandle -> TMVar Integer
outputThreads OutputHandle
globalOutputHandle)
		if Integer
r Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
<= Integer
0
			then TMVar Integer -> Integer -> STM ()
forall a. TMVar a -> a -> STM ()
putTMVar (OutputHandle -> TMVar Integer
outputThreads OutputHandle
globalOutputHandle) Integer
r
			else STM ()
forall a. STM a
retry
	-- Take output lock to wait for anything else that might be
	-- currently generating output.
	IO () -> IO ()
forall (m :: * -> *) a. (MonadIO m, MonadMask m) => m a -> m a
lockOutput (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ () -> IO ()
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ()

-- | Values that can be output.
class Outputable v where
	toOutput :: v -> T.Text

instance Outputable T.Text where
	toOutput :: Text -> Text
toOutput = Text -> Text
forall a. a -> a
id

-- | Note that using a lazy Text as an Outputable value 
-- will buffer it all in memory.
instance Outputable L.Text where
	toOutput :: Text -> Text
toOutput = Text -> Text
forall v. Outputable v => v -> Text
toOutput (Text -> Text) -> (Text -> Text) -> Text -> Text
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> Text
L.toStrict

instance Outputable String where
	toOutput :: String -> Text
toOutput = Text -> Text
forall v. Outputable v => v -> Text
toOutput (Text -> Text) -> (String -> Text) -> String -> Text
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> Text
T.pack

-- | Displays a value to stdout.
--
-- Uses locking to ensure that the whole output occurs atomically
-- even when other threads are concurrently generating output.
--
-- No newline is appended to the value, so if you want a newline, be sure
-- to include it yourself.
--
-- When something else is writing to the console at the same time, this does
-- not block. It buffers the value, so it will be displayed once the other
-- writer is done.
--
-- When outputConcurrent is used within a call to
-- `System.Console.Regions.displayConsoleRegions`, the output is displayed
-- above the currently open console regions. Only lines ending in a newline
-- are displayed in this case (it uses `waitCompleteLines`).
outputConcurrent :: Outputable v => v -> IO ()
outputConcurrent :: forall v. Outputable v => v -> IO ()
outputConcurrent = StdHandle -> v -> IO ()
forall v. Outputable v => StdHandle -> v -> IO ()
outputConcurrent' StdHandle
StdOut

-- | Like `outputConcurrent`, but displays to stderr.
--
-- (Does not throw an exception.)
errorConcurrent :: Outputable v => v -> IO ()
errorConcurrent :: forall v. Outputable v => v -> IO ()
errorConcurrent = StdHandle -> v -> IO ()
forall v. Outputable v => StdHandle -> v -> IO ()
outputConcurrent' StdHandle
StdErr

outputConcurrent' :: Outputable v => StdHandle -> v -> IO ()
outputConcurrent' :: forall v. Outputable v => StdHandle -> v -> IO ()
outputConcurrent' StdHandle
stdh v
v = do
	-- Use a worker thread. This is so any async exception that
	-- is thrown to the current thread does not affect
	-- tryTakeOutputLock, which is not async exception safe.
	Async ()
worker <- IO () -> IO (Async ())
forall a. IO a -> IO (Async a)
async (IO () -> IO (Async ())) -> IO () -> IO (Async ())
forall a b. (a -> b) -> a -> b
$ IO Bool -> (Bool -> IO ()) -> (Bool -> IO ()) -> IO ()
forall (m :: * -> *) a c b.
(HasCallStack, MonadMask m) =>
m a -> (a -> m c) -> (a -> m b) -> m b
bracket IO Bool
setup Bool -> IO ()
cleanup Bool -> IO ()
go
	Async () -> IO ()
forall a. Async a -> IO a
wait Async ()
worker
  where
	setup :: IO Bool
setup = IO Bool
tryTakeOutputLock
	cleanup :: Bool -> IO ()
cleanup Bool
False = () -> IO ()
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ()
	cleanup Bool
True = IO ()
dropOutputLock
	go :: Bool -> IO ()
go Bool
True = do
		Handle -> Text -> IO ()
T.hPutStr Handle
h (v -> Text
forall v. Outputable v => v -> Text
toOutput v
v)
		Handle -> IO ()
hFlush Handle
h
	go Bool
False = do
		OutputBuffer
oldbuf <- STM OutputBuffer -> IO OutputBuffer
forall a. STM a -> IO a
atomically (STM OutputBuffer -> IO OutputBuffer)
-> STM OutputBuffer -> IO OutputBuffer
forall a b. (a -> b) -> a -> b
$ TMVar OutputBuffer -> STM OutputBuffer
forall a. TMVar a -> STM a
takeTMVar TMVar OutputBuffer
bv
		OutputBuffer
newbuf <- OutputBufferedActivity -> OutputBuffer -> IO OutputBuffer
addOutputBuffer (Text -> OutputBufferedActivity
Output (v -> Text
forall v. Outputable v => v -> Text
toOutput v
v)) OutputBuffer
oldbuf
		STM () -> IO ()
forall a. STM a -> IO a
atomically (STM () -> IO ()) -> STM () -> IO ()
forall a b. (a -> b) -> a -> b
$ TMVar OutputBuffer -> OutputBuffer -> STM ()
forall a. TMVar a -> a -> STM ()
putTMVar TMVar OutputBuffer
bv OutputBuffer
newbuf
	h :: Handle
h = StdHandle -> Handle
toHandle StdHandle
stdh
	bv :: TMVar OutputBuffer
bv = StdHandle -> TMVar OutputBuffer
bufferFor StdHandle
stdh

-- | This alias is provided to avoid breaking backwards compatibility.
type ConcurrentProcessHandle = P.ProcessHandle

-- | Same as `P.waitForProcess`; provided to avoid breaking backwards
-- compatibility.
waitForProcessConcurrent :: ConcurrentProcessHandle -> IO ExitCode
waitForProcessConcurrent :: ConcurrentProcessHandle -> IO ExitCode
waitForProcessConcurrent = ConcurrentProcessHandle -> IO ExitCode
P.waitForProcess

-- | Wrapper around `System.Process.createProcess` that prevents 
-- multiple processes that are running concurrently from writing
-- to stdout/stderr at the same time.
--
-- If the process does not output to stdout or stderr, it's run
-- by createProcess entirely as usual. Only processes that can generate
-- output are handled specially:
--
-- A process is allowed to write to stdout and stderr in the usual
-- way, assuming it can successfully take the output lock.
--
-- When the output lock is held (ie, by another concurrent process,
-- or because `outputConcurrent` is being called at the same time),
-- the process is instead run with its stdout and stderr
-- redirected to a buffer. The buffered output will be displayed as soon
-- as the output lock becomes free.
--
-- Note that the the process is waited for by a background thread,
-- so unlike createProcess, neglecting to call waitForProcess will not
-- result in zombie processess.
createProcessConcurrent :: P.CreateProcess -> IO (Maybe Handle, Maybe Handle, Maybe Handle, P.ProcessHandle) 
createProcessConcurrent :: CreateProcess
-> IO
     (Maybe Handle, Maybe Handle, Maybe Handle, ConcurrentProcessHandle)
createProcessConcurrent CreateProcess
p
	| StdStream -> Bool
willOutput (CreateProcess -> StdStream
P.std_out CreateProcess
p) Bool -> Bool -> Bool
|| StdStream -> Bool
willOutput (CreateProcess -> StdStream
P.std_err CreateProcess
p) =
		IO Bool
-> (IO
      (Maybe Handle, Maybe Handle, Maybe Handle,
       ConcurrentProcessHandle),
    IO
      (Maybe Handle, Maybe Handle, Maybe Handle,
       ConcurrentProcessHandle))
-> IO
     (Maybe Handle, Maybe Handle, Maybe Handle, ConcurrentProcessHandle)
forall (m :: * -> *) a. Monad m => m Bool -> (m a, m a) -> m a
ifM IO Bool
tryTakeOutputLock
			( CreateProcess
-> IO
     (Maybe Handle, Maybe Handle, Maybe Handle, ConcurrentProcessHandle)
fgProcess CreateProcess
p
			, CreateProcess
-> IO
     (Maybe Handle, Maybe Handle, Maybe Handle, ConcurrentProcessHandle)
bgProcess CreateProcess
p
			)
	| Bool
otherwise = do
		r :: (Maybe Handle, Maybe Handle, Maybe Handle, ConcurrentProcessHandle)
r@(Maybe Handle
_, Maybe Handle
_, Maybe Handle
_, ConcurrentProcessHandle
h) <- CreateProcess
-> IO
     (Maybe Handle, Maybe Handle, Maybe Handle, ConcurrentProcessHandle)
P.createProcess CreateProcess
p
		Async ()
_ <- IO () -> IO (Async ())
forall a. IO a -> IO (Async a)
async (IO () -> IO (Async ())) -> IO () -> IO (Async ())
forall a b. (a -> b) -> a -> b
$ IO (Either IOException ExitCode) -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO (Either IOException ExitCode) -> IO ())
-> IO (Either IOException ExitCode) -> IO ()
forall a b. (a -> b) -> a -> b
$ IO ExitCode -> IO (Either IOException ExitCode)
forall (m :: * -> *) a.
MonadCatch m =>
m a -> m (Either IOException a)
tryIO (IO ExitCode -> IO (Either IOException ExitCode))
-> IO ExitCode -> IO (Either IOException ExitCode)
forall a b. (a -> b) -> a -> b
$ ConcurrentProcessHandle -> IO ExitCode
P.waitForProcess ConcurrentProcessHandle
h
		(Maybe Handle, Maybe Handle, Maybe Handle, ConcurrentProcessHandle)
-> IO
     (Maybe Handle, Maybe Handle, Maybe Handle, ConcurrentProcessHandle)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe Handle, Maybe Handle, Maybe Handle, ConcurrentProcessHandle)
r

-- | Wrapper around `System.Process.createProcess` that makes sure a process
-- is run in the foreground, with direct access to stdout and stderr.
-- Useful when eg, running an interactive process.
--
-- Note that the the process is waited for by a background thread,
-- so unlike createProcess, neglecting to call waitForProcess will not
-- result in zombie processess.
createProcessForeground :: P.CreateProcess -> IO (Maybe Handle, Maybe Handle, Maybe Handle, P.ProcessHandle)
createProcessForeground :: CreateProcess
-> IO
     (Maybe Handle, Maybe Handle, Maybe Handle, ConcurrentProcessHandle)
createProcessForeground CreateProcess
p = do
	IO ()
takeOutputLock
	CreateProcess
-> IO
     (Maybe Handle, Maybe Handle, Maybe Handle, ConcurrentProcessHandle)
fgProcess CreateProcess
p

fgProcess :: P.CreateProcess -> IO (Maybe Handle, Maybe Handle, Maybe Handle, P.ProcessHandle)
fgProcess :: CreateProcess
-> IO
     (Maybe Handle, Maybe Handle, Maybe Handle, ConcurrentProcessHandle)
fgProcess CreateProcess
p = do
	r :: (Maybe Handle, Maybe Handle, Maybe Handle, ConcurrentProcessHandle)
r@(Maybe Handle
_, Maybe Handle
_, Maybe Handle
_, ConcurrentProcessHandle
h) <- CreateProcess
-> IO
     (Maybe Handle, Maybe Handle, Maybe Handle, ConcurrentProcessHandle)
P.createProcess CreateProcess
p
		IO
  (Maybe Handle, Maybe Handle, Maybe Handle, ConcurrentProcessHandle)
-> IO ()
-> IO
     (Maybe Handle, Maybe Handle, Maybe Handle, ConcurrentProcessHandle)
forall (m :: * -> *) a b.
(HasCallStack, MonadCatch m) =>
m a -> m b -> m a
`onException` IO ()
dropOutputLock
	IO ()
registerOutputThread
	-- Wait for the process to exit and drop the lock.
	Async ()
_ <- IO () -> IO (Async ())
forall a. IO a -> IO (Async a)
async (IO () -> IO (Async ())) -> IO () -> IO (Async ())
forall a b. (a -> b) -> a -> b
$ do
		IO (Either IOException ExitCode) -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO (Either IOException ExitCode) -> IO ())
-> IO (Either IOException ExitCode) -> IO ()
forall a b. (a -> b) -> a -> b
$ IO ExitCode -> IO (Either IOException ExitCode)
forall (m :: * -> *) a.
MonadCatch m =>
m a -> m (Either IOException a)
tryIO (IO ExitCode -> IO (Either IOException ExitCode))
-> IO ExitCode -> IO (Either IOException ExitCode)
forall a b. (a -> b) -> a -> b
$ ConcurrentProcessHandle -> IO ExitCode
P.waitForProcess ConcurrentProcessHandle
h
		IO ()
unregisterOutputThread
		IO ()
dropOutputLock
	(Maybe Handle, Maybe Handle, Maybe Handle, ConcurrentProcessHandle)
-> IO
     (Maybe Handle, Maybe Handle, Maybe Handle, ConcurrentProcessHandle)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe Handle, Maybe Handle, Maybe Handle, ConcurrentProcessHandle)
r

bgProcess :: P.CreateProcess -> IO (Maybe Handle, Maybe Handle, Maybe Handle, P.ProcessHandle)
bgProcess :: CreateProcess
-> IO
     (Maybe Handle, Maybe Handle, Maybe Handle, ConcurrentProcessHandle)
bgProcess CreateProcess
p = do
	let p' :: CreateProcess
p' = CreateProcess
p
		{ P.std_out = rediroutput (P.std_out p)
		, P.std_err = rediroutput (P.std_err p)
		}
	IO ()
registerOutputThread
	(Maybe Handle
stdin_h, Maybe Handle
stdout_h, Maybe Handle
stderr_h, ConcurrentProcessHandle
h) <- CreateProcess
-> IO
     (Maybe Handle, Maybe Handle, Maybe Handle, ConcurrentProcessHandle)
P.createProcess CreateProcess
p'
		IO
  (Maybe Handle, Maybe Handle, Maybe Handle, ConcurrentProcessHandle)
-> IO ()
-> IO
     (Maybe Handle, Maybe Handle, Maybe Handle, ConcurrentProcessHandle)
forall (m :: * -> *) a b.
(HasCallStack, MonadCatch m) =>
m a -> m b -> m a
`onException` IO ()
unregisterOutputThread
	let r :: (Maybe Handle, Maybe Handle, Maybe Handle, ConcurrentProcessHandle)
r =
		( Maybe Handle
stdin_h
		, StdStream -> Maybe Handle -> Maybe Handle
forall {a}. StdStream -> Maybe a -> Maybe a
mungeret (CreateProcess -> StdStream
P.std_out CreateProcess
p) Maybe Handle
stdout_h
		, StdStream -> Maybe Handle -> Maybe Handle
forall {a}. StdStream -> Maybe a -> Maybe a
mungeret (CreateProcess -> StdStream
P.std_err CreateProcess
p) Maybe Handle
stderr_h
		, ConcurrentProcessHandle
h
		)
	-- Wait for the process for symmetry with fgProcess,
	-- which does the same.
	Async ()
_ <- IO () -> IO (Async ())
forall a. IO a -> IO (Async a)
async (IO () -> IO (Async ())) -> IO () -> IO (Async ())
forall a b. (a -> b) -> a -> b
$ IO (Either IOException ExitCode) -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO (Either IOException ExitCode) -> IO ())
-> IO (Either IOException ExitCode) -> IO ()
forall a b. (a -> b) -> a -> b
$ IO ExitCode -> IO (Either IOException ExitCode)
forall (m :: * -> *) a.
MonadCatch m =>
m a -> m (Either IOException a)
tryIO (IO ExitCode -> IO (Either IOException ExitCode))
-> IO ExitCode -> IO (Either IOException ExitCode)
forall a b. (a -> b) -> a -> b
$ ConcurrentProcessHandle -> IO ExitCode
P.waitForProcess ConcurrentProcessHandle
h
	(StdHandle, MVar OutputBuffer, TMVar BufSig, TMVar AtEnd)
outbuf <- StdHandle
-> Maybe Handle
-> IO (StdHandle, MVar OutputBuffer, TMVar BufSig, TMVar AtEnd)
setupOutputBuffer StdHandle
StdOut (StdStream -> Maybe Handle -> Maybe Handle
forall {a}. StdStream -> Maybe a -> Maybe a
mungebuf (CreateProcess -> StdStream
P.std_out CreateProcess
p) Maybe Handle
stdout_h)
	(StdHandle, MVar OutputBuffer, TMVar BufSig, TMVar AtEnd)
errbuf <- StdHandle
-> Maybe Handle
-> IO (StdHandle, MVar OutputBuffer, TMVar BufSig, TMVar AtEnd)
setupOutputBuffer StdHandle
StdErr (StdStream -> Maybe Handle -> Maybe Handle
forall {a}. StdStream -> Maybe a -> Maybe a
mungebuf (CreateProcess -> StdStream
P.std_err CreateProcess
p) Maybe Handle
stderr_h)
	IO (Async ()) -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO (Async ()) -> IO ()) -> IO (Async ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ IO () -> IO (Async ())
forall a. IO a -> IO (Async a)
async (IO () -> IO (Async ())) -> IO () -> IO (Async ())
forall a b. (a -> b) -> a -> b
$ [(StdHandle, MVar OutputBuffer, TMVar BufSig, TMVar AtEnd)]
-> IO ()
bufferWriter [(StdHandle, MVar OutputBuffer, TMVar BufSig, TMVar AtEnd)
outbuf, (StdHandle, MVar OutputBuffer, TMVar BufSig, TMVar AtEnd)
errbuf]
	(Maybe Handle, Maybe Handle, Maybe Handle, ConcurrentProcessHandle)
-> IO
     (Maybe Handle, Maybe Handle, Maybe Handle, ConcurrentProcessHandle)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe Handle, Maybe Handle, Maybe Handle, ConcurrentProcessHandle)
r
  where
	rediroutput :: StdStream -> StdStream
rediroutput StdStream
ss
		| StdStream -> Bool
willOutput StdStream
ss = StdStream
P.CreatePipe
		| Bool
otherwise = StdStream
ss
	mungebuf :: StdStream -> Maybe a -> Maybe a
mungebuf StdStream
ss Maybe a
mh
		| StdStream -> Bool
willOutput StdStream
ss = Maybe a
mh
		| Bool
otherwise = Maybe a
forall a. Maybe a
Nothing
	mungeret :: StdStream -> Maybe a -> Maybe a
mungeret StdStream
ss Maybe a
mh
		| StdStream -> Bool
willOutput StdStream
ss = Maybe a
forall a. Maybe a
Nothing
		| Bool
otherwise = Maybe a
mh

willOutput :: P.StdStream -> Bool
willOutput :: StdStream -> Bool
willOutput StdStream
P.Inherit = Bool
True
willOutput StdStream
_ = Bool
False

-- | Buffered output.
data OutputBuffer = OutputBuffer [OutputBufferedActivity]
	deriving (OutputBuffer -> OutputBuffer -> Bool
(OutputBuffer -> OutputBuffer -> Bool)
-> (OutputBuffer -> OutputBuffer -> Bool) -> Eq OutputBuffer
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: OutputBuffer -> OutputBuffer -> Bool
== :: OutputBuffer -> OutputBuffer -> Bool
$c/= :: OutputBuffer -> OutputBuffer -> Bool
/= :: OutputBuffer -> OutputBuffer -> Bool
Eq)

data StdHandle = StdOut | StdErr

toHandle :: StdHandle -> Handle
toHandle :: StdHandle -> Handle
toHandle StdHandle
StdOut = Handle
stdout
toHandle StdHandle
StdErr = Handle
stderr

bufferFor :: StdHandle -> TMVar OutputBuffer
bufferFor :: StdHandle -> TMVar OutputBuffer
bufferFor StdHandle
StdOut = OutputHandle -> TMVar OutputBuffer
outputBuffer OutputHandle
globalOutputHandle
bufferFor StdHandle
StdErr = OutputHandle -> TMVar OutputBuffer
errorBuffer OutputHandle
globalOutputHandle

data OutputBufferedActivity
	= Output T.Text
	| InTempFile
		{ OutputBufferedActivity -> String
tempFile :: FilePath
		, OutputBufferedActivity -> Bool
endsInNewLine :: Bool
		}
	deriving (OutputBufferedActivity -> OutputBufferedActivity -> Bool
(OutputBufferedActivity -> OutputBufferedActivity -> Bool)
-> (OutputBufferedActivity -> OutputBufferedActivity -> Bool)
-> Eq OutputBufferedActivity
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: OutputBufferedActivity -> OutputBufferedActivity -> Bool
== :: OutputBufferedActivity -> OutputBufferedActivity -> Bool
$c/= :: OutputBufferedActivity -> OutputBufferedActivity -> Bool
/= :: OutputBufferedActivity -> OutputBufferedActivity -> Bool
Eq)

data AtEnd = AtEnd
	deriving AtEnd -> AtEnd -> Bool
(AtEnd -> AtEnd -> Bool) -> (AtEnd -> AtEnd -> Bool) -> Eq AtEnd
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: AtEnd -> AtEnd -> Bool
== :: AtEnd -> AtEnd -> Bool
$c/= :: AtEnd -> AtEnd -> Bool
/= :: AtEnd -> AtEnd -> Bool
Eq

data BufSig = BufSig

setupOutputBuffer :: StdHandle -> Maybe Handle -> IO (StdHandle, MVar OutputBuffer, TMVar BufSig, TMVar AtEnd)
setupOutputBuffer :: StdHandle
-> Maybe Handle
-> IO (StdHandle, MVar OutputBuffer, TMVar BufSig, TMVar AtEnd)
setupOutputBuffer StdHandle
h Maybe Handle
fromh = do
	MVar OutputBuffer
buf <- OutputBuffer -> IO (MVar OutputBuffer)
forall a. a -> IO (MVar a)
newMVar ([OutputBufferedActivity] -> OutputBuffer
OutputBuffer [])
	TMVar BufSig
bufsig <- STM (TMVar BufSig) -> IO (TMVar BufSig)
forall a. STM a -> IO a
atomically STM (TMVar BufSig)
forall a. STM (TMVar a)
newEmptyTMVar
	TMVar AtEnd
bufend <- STM (TMVar AtEnd) -> IO (TMVar AtEnd)
forall a. STM a -> IO a
atomically STM (TMVar AtEnd)
forall a. STM (TMVar a)
newEmptyTMVar
	IO (Async ()) -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO (Async ()) -> IO ()) -> IO (Async ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ IO () -> IO (Async ())
forall a. IO a -> IO (Async a)
async (IO () -> IO (Async ())) -> IO () -> IO (Async ())
forall a b. (a -> b) -> a -> b
$ Maybe Handle
-> MVar OutputBuffer -> TMVar BufSig -> TMVar AtEnd -> IO ()
outputDrainer Maybe Handle
fromh MVar OutputBuffer
buf TMVar BufSig
bufsig TMVar AtEnd
bufend
	(StdHandle, MVar OutputBuffer, TMVar BufSig, TMVar AtEnd)
-> IO (StdHandle, MVar OutputBuffer, TMVar BufSig, TMVar AtEnd)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (StdHandle
h, MVar OutputBuffer
buf, TMVar BufSig
bufsig, TMVar AtEnd
bufend)

-- Drain output from the handle, and buffer it.
outputDrainer :: Maybe Handle -> MVar OutputBuffer -> TMVar BufSig -> TMVar AtEnd -> IO ()
outputDrainer :: Maybe Handle
-> MVar OutputBuffer -> TMVar BufSig -> TMVar AtEnd -> IO ()
outputDrainer Maybe Handle
mfromh MVar OutputBuffer
buf TMVar BufSig
bufsig TMVar AtEnd
bufend = case Maybe Handle
mfromh of
	Maybe Handle
Nothing -> IO ()
atend
	Just Handle
fromh -> Handle -> IO ()
go Handle
fromh
  where
	go :: Handle -> IO ()
go Handle
fromh = do
		Text
t <- Handle -> IO Text
T.hGetChunk Handle
fromh
		if Text -> Bool
T.null Text
t
			then do
				IO ()
atend
				Handle -> IO ()
hClose Handle
fromh
			else do
				MVar OutputBuffer -> (OutputBuffer -> IO OutputBuffer) -> IO ()
forall a. MVar a -> (a -> IO a) -> IO ()
modifyMVar_ MVar OutputBuffer
buf ((OutputBuffer -> IO OutputBuffer) -> IO ())
-> (OutputBuffer -> IO OutputBuffer) -> IO ()
forall a b. (a -> b) -> a -> b
$ OutputBufferedActivity -> OutputBuffer -> IO OutputBuffer
addOutputBuffer (Text -> OutputBufferedActivity
Output Text
t)
				IO ()
changed
				Handle -> IO ()
go Handle
fromh
	atend :: IO ()
atend = STM () -> IO ()
forall a. STM a -> IO a
atomically (STM () -> IO ()) -> STM () -> IO ()
forall a b. (a -> b) -> a -> b
$ TMVar AtEnd -> AtEnd -> STM ()
forall a. TMVar a -> a -> STM ()
putTMVar TMVar AtEnd
bufend AtEnd
AtEnd
	changed :: IO ()
changed = STM () -> IO ()
forall a. STM a -> IO a
atomically (STM () -> IO ()) -> STM () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
		STM (Maybe BufSig) -> STM ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (STM (Maybe BufSig) -> STM ()) -> STM (Maybe BufSig) -> STM ()
forall a b. (a -> b) -> a -> b
$ TMVar BufSig -> STM (Maybe BufSig)
forall a. TMVar a -> STM (Maybe a)
tryTakeTMVar TMVar BufSig
bufsig
		TMVar BufSig -> BufSig -> STM ()
forall a. TMVar a -> a -> STM ()
putTMVar TMVar BufSig
bufsig BufSig
BufSig

registerOutputThread :: IO ()
registerOutputThread :: IO ()
registerOutputThread = do
	let v :: TMVar Integer
v = OutputHandle -> TMVar Integer
outputThreads OutputHandle
globalOutputHandle
	STM () -> IO ()
forall a. STM a -> IO a
atomically (STM () -> IO ()) -> STM () -> IO ()
forall a b. (a -> b) -> a -> b
$ TMVar Integer -> Integer -> STM ()
forall a. TMVar a -> a -> STM ()
putTMVar TMVar Integer
v (Integer -> STM ()) -> (Integer -> Integer) -> Integer -> STM ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Integer -> Integer
forall a. Enum a => a -> a
succ (Integer -> STM ()) -> STM Integer -> STM ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< TMVar Integer -> STM Integer
forall a. TMVar a -> STM a
takeTMVar TMVar Integer
v
	
unregisterOutputThread :: IO ()
unregisterOutputThread :: IO ()
unregisterOutputThread = do
	let v :: TMVar Integer
v = OutputHandle -> TMVar Integer
outputThreads OutputHandle
globalOutputHandle
	STM () -> IO ()
forall a. STM a -> IO a
atomically (STM () -> IO ()) -> STM () -> IO ()
forall a b. (a -> b) -> a -> b
$ TMVar Integer -> Integer -> STM ()
forall a. TMVar a -> a -> STM ()
putTMVar TMVar Integer
v (Integer -> STM ()) -> (Integer -> Integer) -> Integer -> STM ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Integer -> Integer
forall a. Enum a => a -> a
pred (Integer -> STM ()) -> STM Integer -> STM ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< TMVar Integer -> STM Integer
forall a. TMVar a -> STM a
takeTMVar TMVar Integer
v

-- Wait to lock output, and once we can, display everything 
-- that's put into the buffers, until the end.
--
-- If end is reached before lock is taken, instead add the command's
-- buffers to the global outputBuffer and errorBuffer.
bufferWriter :: [(StdHandle, MVar OutputBuffer, TMVar BufSig, TMVar AtEnd)] -> IO ()
bufferWriter :: [(StdHandle, MVar OutputBuffer, TMVar BufSig, TMVar AtEnd)]
-> IO ()
bufferWriter [(StdHandle, MVar OutputBuffer, TMVar BufSig, TMVar AtEnd)]
ts = do
	TMVar ()
activitysig <- STM (TMVar ()) -> IO (TMVar ())
forall a. STM a -> IO a
atomically STM (TMVar ())
forall a. STM (TMVar a)
newEmptyTMVar
	Async ()
worker1 <- IO () -> IO (Async ())
forall a. IO a -> IO (Async a)
async (IO () -> IO (Async ())) -> IO () -> IO (Async ())
forall a b. (a -> b) -> a -> b
$ IO () -> IO ()
forall (m :: * -> *) a. (MonadIO m, MonadMask m) => m a -> m a
lockOutput (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$
		IO Bool -> (IO (), IO ()) -> IO ()
forall (m :: * -> *) a. Monad m => m Bool -> (m a, m a) -> m a
ifM (STM Bool -> IO Bool
forall a. STM a -> IO a
atomically (STM Bool -> IO Bool) -> STM Bool -> IO Bool
forall a b. (a -> b) -> a -> b
$ TMVar () -> () -> STM Bool
forall a. TMVar a -> a -> STM Bool
tryPutTMVar TMVar ()
activitysig ())
			( IO [()] -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO [()] -> IO ()) -> IO [()] -> IO ()
forall a b. (a -> b) -> a -> b
$ ((StdHandle, MVar OutputBuffer, TMVar BufSig, TMVar AtEnd)
 -> IO ())
-> [(StdHandle, MVar OutputBuffer, TMVar BufSig, TMVar AtEnd)]
-> IO [()]
forall (t :: * -> *) a b.
Traversable t =>
(a -> IO b) -> t a -> IO (t b)
mapConcurrently (StdHandle, MVar OutputBuffer, TMVar BufSig, TMVar AtEnd) -> IO ()
displaybuf [(StdHandle, MVar OutputBuffer, TMVar BufSig, TMVar AtEnd)]
ts
			, IO ()
forall (m :: * -> *). Monad m => m ()
noop -- buffers already moved to global
			)
	Async ()
worker2 <- IO () -> IO (Async ())
forall a. IO a -> IO (Async a)
async (IO () -> IO (Async ())) -> IO () -> IO (Async ())
forall a b. (a -> b) -> a -> b
$ IO () -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ TMVar () -> Async () -> IO ()
forall {a}. TMVar () -> Async a -> IO ()
globalbuf TMVar ()
activitysig Async ()
worker1
	IO (Async ()) -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO (Async ()) -> IO ()) -> IO (Async ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ IO () -> IO (Async ())
forall a. IO a -> IO (Async a)
async (IO () -> IO (Async ())) -> IO () -> IO (Async ())
forall a b. (a -> b) -> a -> b
$ do
		IO (Either SomeException ()) -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO (Either SomeException ()) -> IO ())
-> IO (Either SomeException ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ Async () -> IO (Either SomeException ())
forall a. Async a -> IO (Either SomeException a)
waitCatch Async ()
worker1
		IO (Either SomeException ()) -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO (Either SomeException ()) -> IO ())
-> IO (Either SomeException ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ Async () -> IO (Either SomeException ())
forall a. Async a -> IO (Either SomeException a)
waitCatch Async ()
worker2
		IO ()
unregisterOutputThread
  where
	displaybuf :: (StdHandle, MVar OutputBuffer, TMVar BufSig, TMVar AtEnd) -> IO ()
displaybuf v :: (StdHandle, MVar OutputBuffer, TMVar BufSig, TMVar AtEnd)
v@(StdHandle
outh, MVar OutputBuffer
buf, TMVar BufSig
bufsig, TMVar AtEnd
bufend) = do
		Either AtEnd BufSig
change <- STM (Either AtEnd BufSig) -> IO (Either AtEnd BufSig)
forall a. STM a -> IO a
atomically (STM (Either AtEnd BufSig) -> IO (Either AtEnd BufSig))
-> STM (Either AtEnd BufSig) -> IO (Either AtEnd BufSig)
forall a b. (a -> b) -> a -> b
$
			(BufSig -> Either AtEnd BufSig
forall a b. b -> Either a b
Right (BufSig -> Either AtEnd BufSig)
-> STM BufSig -> STM (Either AtEnd BufSig)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TMVar BufSig -> STM BufSig
forall a. TMVar a -> STM a
takeTMVar TMVar BufSig
bufsig)
				STM (Either AtEnd BufSig)
-> STM (Either AtEnd BufSig) -> STM (Either AtEnd BufSig)
forall a. STM a -> STM a -> STM a
`orElse`
			(AtEnd -> Either AtEnd BufSig
forall a b. a -> Either a b
Left (AtEnd -> Either AtEnd BufSig)
-> STM AtEnd -> STM (Either AtEnd BufSig)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TMVar AtEnd -> STM AtEnd
forall a. TMVar a -> STM a
takeTMVar TMVar AtEnd
bufend)
		OutputBuffer
l <- MVar OutputBuffer -> IO OutputBuffer
forall a. MVar a -> IO a
takeMVar MVar OutputBuffer
buf
		MVar OutputBuffer -> OutputBuffer -> IO ()
forall a. MVar a -> a -> IO ()
putMVar MVar OutputBuffer
buf ([OutputBufferedActivity] -> OutputBuffer
OutputBuffer [])
		StdHandle -> OutputBuffer -> IO ()
emitOutputBuffer StdHandle
outh OutputBuffer
l
		case Either AtEnd BufSig
change of
			Right BufSig
BufSig -> (StdHandle, MVar OutputBuffer, TMVar BufSig, TMVar AtEnd) -> IO ()
displaybuf (StdHandle, MVar OutputBuffer, TMVar BufSig, TMVar AtEnd)
v
			Left AtEnd
AtEnd -> () -> IO ()
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ()
	globalbuf :: TMVar () -> Async a -> IO ()
globalbuf TMVar ()
activitysig Async a
worker1 = do
		Bool
ok <- STM Bool -> IO Bool
forall a. STM a -> IO a
atomically (STM Bool -> IO Bool) -> STM Bool -> IO Bool
forall a b. (a -> b) -> a -> b
$ do
			-- signal we're going to handle it
			-- (returns false if the displaybuf already did)
			Bool
ok <- TMVar () -> () -> STM Bool
forall a. TMVar a -> a -> STM Bool
tryPutTMVar TMVar ()
activitysig ()
			-- wait for end of all buffers
			Bool -> STM () -> STM ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when Bool
ok (STM () -> STM ()) -> STM () -> STM ()
forall a b. (a -> b) -> a -> b
$
				((StdHandle, MVar OutputBuffer, TMVar BufSig, TMVar AtEnd)
 -> STM AtEnd)
-> [(StdHandle, MVar OutputBuffer, TMVar BufSig, TMVar AtEnd)]
-> STM ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (\(StdHandle
_outh, MVar OutputBuffer
_buf, TMVar BufSig
_bufsig, TMVar AtEnd
bufend) -> TMVar AtEnd -> STM AtEnd
forall a. TMVar a -> STM a
takeTMVar TMVar AtEnd
bufend) [(StdHandle, MVar OutputBuffer, TMVar BufSig, TMVar AtEnd)]
ts
			Bool -> STM Bool
forall a. a -> STM a
forall (m :: * -> *) a. Monad m => a -> m a
return Bool
ok
		Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when Bool
ok (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
			-- add all of the command's buffered output to the
			-- global output buffer, atomically
			[(StdHandle, OutputBuffer)]
bs <- [(StdHandle, MVar OutputBuffer, TMVar BufSig, TMVar AtEnd)]
-> ((StdHandle, MVar OutputBuffer, TMVar BufSig, TMVar AtEnd)
    -> IO (StdHandle, OutputBuffer))
-> IO [(StdHandle, OutputBuffer)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [(StdHandle, MVar OutputBuffer, TMVar BufSig, TMVar AtEnd)]
ts (((StdHandle, MVar OutputBuffer, TMVar BufSig, TMVar AtEnd)
  -> IO (StdHandle, OutputBuffer))
 -> IO [(StdHandle, OutputBuffer)])
-> ((StdHandle, MVar OutputBuffer, TMVar BufSig, TMVar AtEnd)
    -> IO (StdHandle, OutputBuffer))
-> IO [(StdHandle, OutputBuffer)]
forall a b. (a -> b) -> a -> b
$ \(StdHandle
outh, MVar OutputBuffer
buf, TMVar BufSig
_bufsig, TMVar AtEnd
_bufend) ->
				(StdHandle
outh,) (OutputBuffer -> (StdHandle, OutputBuffer))
-> IO OutputBuffer -> IO (StdHandle, OutputBuffer)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MVar OutputBuffer -> IO OutputBuffer
forall a. MVar a -> IO a
takeMVar MVar OutputBuffer
buf
			STM () -> IO ()
forall a. STM a -> IO a
atomically (STM () -> IO ()) -> STM () -> IO ()
forall a b. (a -> b) -> a -> b
$
				[(StdHandle, OutputBuffer)]
-> ((StdHandle, OutputBuffer) -> STM ()) -> STM ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [(StdHandle, OutputBuffer)]
bs (((StdHandle, OutputBuffer) -> STM ()) -> STM ())
-> ((StdHandle, OutputBuffer) -> STM ()) -> STM ()
forall a b. (a -> b) -> a -> b
$ \(StdHandle
outh, OutputBuffer
b) -> 
					StdHandle -> OutputBuffer -> STM ()
bufferOutputSTM' StdHandle
outh OutputBuffer
b
			-- worker1 might be blocked waiting for the output
			-- lock, and we've already done its job, so cancel it
			Async a -> IO ()
forall a. Async a -> IO ()
cancel Async a
worker1

-- Adds a value to the OutputBuffer. When adding Output to a Handle,
-- it's cheaper to combine it with any already buffered Output to that
-- same Handle.
--
-- When the total buffered Output exceeds 1 mb in size, it's moved out of
-- memory, to a temp file. This should only happen rarely, but is done to
-- avoid some verbose process unexpectedly causing excessive memory use.
addOutputBuffer :: OutputBufferedActivity -> OutputBuffer -> IO OutputBuffer
addOutputBuffer :: OutputBufferedActivity -> OutputBuffer -> IO OutputBuffer
addOutputBuffer (Output Text
t) (OutputBuffer [OutputBufferedActivity]
buf)
	| Text -> Int
T.length Text
t' Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
1048576 = OutputBuffer -> IO OutputBuffer
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (OutputBuffer -> IO OutputBuffer)
-> OutputBuffer -> IO OutputBuffer
forall a b. (a -> b) -> a -> b
$ [OutputBufferedActivity] -> OutputBuffer
OutputBuffer (Text -> OutputBufferedActivity
Output Text
t' OutputBufferedActivity
-> [OutputBufferedActivity] -> [OutputBufferedActivity]
forall a. a -> [a] -> [a]
: [OutputBufferedActivity]
other)
	| Bool
otherwise = do
		String
tmpdir <- IO String
getTemporaryDirectory
		(String
tmp, Handle
h) <- String -> String -> IO (String, Handle)
openTempFile String
tmpdir String
"output.tmp"
		let !endnl :: Bool
endnl = Text -> Bool
endsNewLine Text
t'
		let i :: OutputBufferedActivity
i = InTempFile
			{ tempFile :: String
tempFile = String
tmp
			, endsInNewLine :: Bool
endsInNewLine = Bool
endnl
			}
		Handle -> Text -> IO ()
T.hPutStr Handle
h Text
t'
		Handle -> IO ()
hClose Handle
h
		OutputBuffer -> IO OutputBuffer
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (OutputBuffer -> IO OutputBuffer)
-> OutputBuffer -> IO OutputBuffer
forall a b. (a -> b) -> a -> b
$ [OutputBufferedActivity] -> OutputBuffer
OutputBuffer (OutputBufferedActivity
i OutputBufferedActivity
-> [OutputBufferedActivity] -> [OutputBufferedActivity]
forall a. a -> [a] -> [a]
: [OutputBufferedActivity]
other)
  where
	!t' :: Text
t' = [Text] -> Text
T.concat ((OutputBufferedActivity -> Maybe Text)
-> [OutputBufferedActivity] -> [Text]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe OutputBufferedActivity -> Maybe Text
getOutput [OutputBufferedActivity]
this) Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
t
	!([OutputBufferedActivity]
this, [OutputBufferedActivity]
other) = (OutputBufferedActivity -> Bool)
-> [OutputBufferedActivity]
-> ([OutputBufferedActivity], [OutputBufferedActivity])
forall a. (a -> Bool) -> [a] -> ([a], [a])
partition OutputBufferedActivity -> Bool
isOutput [OutputBufferedActivity]
buf
	isOutput :: OutputBufferedActivity -> Bool
isOutput OutputBufferedActivity
v = case OutputBufferedActivity
v of
		Output Text
_ -> Bool
True
		OutputBufferedActivity
_ -> Bool
False
	getOutput :: OutputBufferedActivity -> Maybe Text
getOutput OutputBufferedActivity
v = case OutputBufferedActivity
v of
		Output Text
t'' -> Text -> Maybe Text
forall a. a -> Maybe a
Just Text
t''
		OutputBufferedActivity
_ -> Maybe Text
forall a. Maybe a
Nothing
addOutputBuffer OutputBufferedActivity
v (OutputBuffer [OutputBufferedActivity]
buf) = OutputBuffer -> IO OutputBuffer
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (OutputBuffer -> IO OutputBuffer)
-> OutputBuffer -> IO OutputBuffer
forall a b. (a -> b) -> a -> b
$ [OutputBufferedActivity] -> OutputBuffer
OutputBuffer (OutputBufferedActivity
vOutputBufferedActivity
-> [OutputBufferedActivity] -> [OutputBufferedActivity]
forall a. a -> [a] -> [a]
:[OutputBufferedActivity]
buf)

-- | Adds a value to the output buffer for later display.
--
-- Note that buffering large quantities of data this way will keep it
-- resident in memory until it can be displayed. While `outputConcurrent`
-- uses temp files if the buffer gets too big, this STM function cannot do
-- so.
bufferOutputSTM :: Outputable v => StdHandle -> v -> STM ()
bufferOutputSTM :: forall v. Outputable v => StdHandle -> v -> STM ()
bufferOutputSTM StdHandle
h v
v = StdHandle -> OutputBuffer -> STM ()
bufferOutputSTM' StdHandle
h ([OutputBufferedActivity] -> OutputBuffer
OutputBuffer [Text -> OutputBufferedActivity
Output (v -> Text
forall v. Outputable v => v -> Text
toOutput v
v)])

bufferOutputSTM' :: StdHandle -> OutputBuffer -> STM ()
bufferOutputSTM' :: StdHandle -> OutputBuffer -> STM ()
bufferOutputSTM' StdHandle
h (OutputBuffer [OutputBufferedActivity]
newbuf) = do
	(OutputBuffer [OutputBufferedActivity]
buf) <- TMVar OutputBuffer -> STM OutputBuffer
forall a. TMVar a -> STM a
takeTMVar TMVar OutputBuffer
bv
	TMVar OutputBuffer -> OutputBuffer -> STM ()
forall a. TMVar a -> a -> STM ()
putTMVar TMVar OutputBuffer
bv ([OutputBufferedActivity] -> OutputBuffer
OutputBuffer ([OutputBufferedActivity]
newbuf [OutputBufferedActivity]
-> [OutputBufferedActivity] -> [OutputBufferedActivity]
forall a. [a] -> [a] -> [a]
++ [OutputBufferedActivity]
buf))
  where
	bv :: TMVar OutputBuffer
bv = StdHandle -> TMVar OutputBuffer
bufferFor StdHandle
h

-- | A STM action that waits for some buffered output to become
-- available, and returns it.
--
-- The function can select a subset of output when only some is desired;
-- the fst part is returned and the snd is left in the buffer.
--
-- This will prevent it from being displayed in the usual way, so you'll
-- need to use `emitOutputBuffer` to display it yourself.
outputBufferWaiterSTM :: (OutputBuffer -> (OutputBuffer, OutputBuffer)) -> STM (StdHandle, OutputBuffer)
outputBufferWaiterSTM :: (OutputBuffer -> (OutputBuffer, OutputBuffer))
-> STM (StdHandle, OutputBuffer)
outputBufferWaiterSTM OutputBuffer -> (OutputBuffer, OutputBuffer)
selector = StdHandle -> STM (StdHandle, OutputBuffer)
waitgetbuf StdHandle
StdOut STM (StdHandle, OutputBuffer)
-> STM (StdHandle, OutputBuffer) -> STM (StdHandle, OutputBuffer)
forall a. STM a -> STM a -> STM a
`orElse` StdHandle -> STM (StdHandle, OutputBuffer)
waitgetbuf StdHandle
StdErr
  where
	waitgetbuf :: StdHandle -> STM (StdHandle, OutputBuffer)
waitgetbuf StdHandle
h = do
		let bv :: TMVar OutputBuffer
bv = StdHandle -> TMVar OutputBuffer
bufferFor StdHandle
h
		(OutputBuffer
selected, OutputBuffer
rest) <- OutputBuffer -> (OutputBuffer, OutputBuffer)
selector (OutputBuffer -> (OutputBuffer, OutputBuffer))
-> STM OutputBuffer -> STM (OutputBuffer, OutputBuffer)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TMVar OutputBuffer -> STM OutputBuffer
forall a. TMVar a -> STM a
takeTMVar TMVar OutputBuffer
bv
		Bool -> STM () -> STM ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (OutputBuffer
selected OutputBuffer -> OutputBuffer -> Bool
forall a. Eq a => a -> a -> Bool
== [OutputBufferedActivity] -> OutputBuffer
OutputBuffer [])
			STM ()
forall a. STM a
retry
		TMVar OutputBuffer -> OutputBuffer -> STM ()
forall a. TMVar a -> a -> STM ()
putTMVar TMVar OutputBuffer
bv OutputBuffer
rest
		(StdHandle, OutputBuffer) -> STM (StdHandle, OutputBuffer)
forall a. a -> STM a
forall (m :: * -> *) a. Monad m => a -> m a
return (StdHandle
h, OutputBuffer
selected)

waitAnyBuffer :: OutputBuffer -> (OutputBuffer, OutputBuffer)
waitAnyBuffer :: OutputBuffer -> (OutputBuffer, OutputBuffer)
waitAnyBuffer OutputBuffer
b = (OutputBuffer
b, [OutputBufferedActivity] -> OutputBuffer
OutputBuffer [])

-- | Use with `outputBufferWaiterSTM` to make it only return buffered
-- output that ends with a newline. Anything buffered without a newline
-- is left in the buffer.
waitCompleteLines :: OutputBuffer -> (OutputBuffer, OutputBuffer)
waitCompleteLines :: OutputBuffer -> (OutputBuffer, OutputBuffer)
waitCompleteLines (OutputBuffer [OutputBufferedActivity]
l) = 
	let ([OutputBufferedActivity]
selected, [OutputBufferedActivity]
rest) = (OutputBufferedActivity -> Bool)
-> [OutputBufferedActivity]
-> ([OutputBufferedActivity], [OutputBufferedActivity])
forall a. (a -> Bool) -> [a] -> ([a], [a])
span OutputBufferedActivity -> Bool
completeline [OutputBufferedActivity]
l
	in ([OutputBufferedActivity] -> OutputBuffer
OutputBuffer [OutputBufferedActivity]
selected, [OutputBufferedActivity] -> OutputBuffer
OutputBuffer [OutputBufferedActivity]
rest)
  where
	completeline :: OutputBufferedActivity -> Bool
completeline (v :: OutputBufferedActivity
v@(InTempFile {})) = OutputBufferedActivity -> Bool
endsInNewLine OutputBufferedActivity
v
	completeline (Output Text
b) = Text -> Bool
endsNewLine Text
b

endsNewLine :: T.Text -> Bool
endsNewLine :: Text -> Bool
endsNewLine Text
t = Bool -> Bool
not (Text -> Bool
T.null Text
t) Bool -> Bool -> Bool
&& HasCallStack => Text -> Char
Text -> Char
T.last Text
t Char -> Char -> Bool
forall a. Eq a => a -> a -> Bool
== Char
'\n'

-- | Emits the content of the OutputBuffer to the Handle
--
-- If you use this, you should use `lockOutput` to ensure you're the only
-- thread writing to the console.
emitOutputBuffer :: StdHandle -> OutputBuffer -> IO ()
emitOutputBuffer :: StdHandle -> OutputBuffer -> IO ()
emitOutputBuffer StdHandle
stdh (OutputBuffer [OutputBufferedActivity]
l) = 
	[OutputBufferedActivity]
-> (OutputBufferedActivity -> IO ()) -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ ([OutputBufferedActivity] -> [OutputBufferedActivity]
forall a. [a] -> [a]
reverse [OutputBufferedActivity]
l) ((OutputBufferedActivity -> IO ()) -> IO ())
-> (OutputBufferedActivity -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \OutputBufferedActivity
ba -> case OutputBufferedActivity
ba of
		Output Text
t -> Text -> IO ()
emit Text
t
		InTempFile String
tmp Bool
_ -> do
			Text -> IO ()
emit (Text -> IO ()) -> IO Text -> IO ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< String -> IO Text
T.readFile String
tmp
			IO (Maybe ()) -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO (Maybe ()) -> IO ()) -> IO (Maybe ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ IO () -> IO (Maybe ())
forall (m :: * -> *) a. MonadCatch m => m a -> m (Maybe a)
tryWhenExists (IO () -> IO (Maybe ())) -> IO () -> IO (Maybe ())
forall a b. (a -> b) -> a -> b
$ String -> IO ()
removeFile String
tmp
  where
	outh :: Handle
outh = StdHandle -> Handle
toHandle StdHandle
stdh
	emit :: Text -> IO ()
emit Text
t = IO (Either IOException ()) -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO (Either IOException ()) -> IO ())
-> IO (Either IOException ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ IO () -> IO (Either IOException ())
forall (m :: * -> *) a.
MonadCatch m =>
m a -> m (Either IOException a)
tryIO (IO () -> IO (Either IOException ()))
-> IO () -> IO (Either IOException ())
forall a b. (a -> b) -> a -> b
$ do
		Handle -> Text -> IO ()
T.hPutStr Handle
outh Text
t
		Handle -> IO ()
hFlush Handle
outh