{-# 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 = forall a. IO a -> a
unsafePerformIO forall a b. (a -> b) -> a -> b
$ TMVar Lock
-> TMVar OutputBuffer
-> TMVar OutputBuffer
-> TMVar Integer
-> OutputHandle
OutputHandle
	forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a. IO (TMVar a)
newEmptyTMVarIO
	forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall a. a -> IO (TMVar a)
newTMVarIO ([OutputBufferedActivity] -> OutputBuffer
OutputBuffer [])
	forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall a. a -> IO (TMVar a)
newTMVarIO ([OutputBufferedActivity] -> OutputBuffer
OutputBuffer [])
	forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> 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 = forall (m :: * -> *) a c b. MonadMask m => m a -> m c -> m b -> m b
bracket_ (forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO IO ()
takeOutputLock) (forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO IO ()
dropOutputLock)

-- | Blocks until we have the output lock.
takeOutputLock :: IO ()
takeOutputLock :: IO ()
takeOutputLock = forall (f :: * -> *) a. Functor f => f a -> f ()
void 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 = forall a. STM a -> IO a
atomically 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 <- forall a. (TMVar Lock -> STM a) -> IO a
withLock forall a b. (a -> b) -> a -> b
$ \TMVar Lock
l -> do
		Maybe Lock
v <- forall a. TMVar a -> STM (Maybe a)
tryTakeTMVar TMVar Lock
l
		case Maybe Lock
v of
			Just Lock
Locked
				| Bool
block -> forall a. STM a
retry
				| Bool
otherwise -> do
					-- Restore value we took.
					forall a. TMVar a -> a -> STM ()
putTMVar TMVar Lock
l Lock
Locked
					forall (m :: * -> *) a. Monad m => a -> m a
return Bool
False
			Maybe Lock
Nothing -> do
				forall a. TMVar a -> a -> STM ()
putTMVar TMVar Lock
l Lock
Locked
				forall (m :: * -> *) a. Monad m => a -> m a
return Bool
True
	forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when Bool
locked forall a b. (a -> b) -> a -> b
$ do
		(OutputBuffer
outbuf, OutputBuffer
errbuf) <- forall a. STM a -> IO a
atomically forall a b. (a -> b) -> a -> b
$ (,)
			forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a. TMVar a -> a -> STM a
swapTMVar (OutputHandle -> TMVar OutputBuffer
outputBuffer OutputHandle
globalOutputHandle) ([OutputBufferedActivity] -> OutputBuffer
OutputBuffer [])
			forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> 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
	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 = forall a. (TMVar Lock -> STM a) -> IO a
withLock forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) a. Functor f => f a -> f ()
void forall b c a. (b -> c) -> (a -> b) -> a -> c
. 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 forall (m :: * -> *) a b. MonadMask m => m a -> m b -> m a
`finally` 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
	forall a. STM a -> IO a
atomically forall a b. (a -> b) -> a -> b
$ do
		Integer
r <- forall a. TMVar a -> STM a
takeTMVar (OutputHandle -> TMVar Integer
outputThreads OutputHandle
globalOutputHandle)
		if Integer
r forall a. Ord a => a -> a -> Bool
<= Integer
0
			then forall a. TMVar a -> a -> STM ()
putTMVar (OutputHandle -> TMVar Integer
outputThreads OutputHandle
globalOutputHandle) Integer
r
			else forall a. STM a
retry
	-- Take output lock to wait for anything else that might be
	-- currently generating output.
	forall (m :: * -> *) a. (MonadIO m, MonadMask m) => m a -> m a
lockOutput forall a b. (a -> b) -> a -> b
$ 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 = 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 = forall v. Outputable v => v -> Text
toOutput forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> Text
L.toStrict

instance Outputable String where
	toOutput :: String -> Text
toOutput = forall v. Outputable v => v -> Text
toOutput 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 = 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 = 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 <- forall a. IO a -> IO (Async a)
async forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a c b.
MonadMask m =>
m a -> (a -> m c) -> (a -> m b) -> m b
bracket IO Bool
setup Bool -> IO ()
cleanup Bool -> IO ()
go
	forall a. Async a -> IO a
wait Async ()
worker
  where
	setup :: IO Bool
setup = IO Bool
tryTakeOutputLock
	cleanup :: Bool -> IO ()
cleanup Bool
False = 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 (forall v. Outputable v => v -> Text
toOutput v
v)
		Handle -> IO ()
hFlush Handle
h
	go Bool
False = do
		OutputBuffer
oldbuf <- forall a. STM a -> IO a
atomically forall a b. (a -> b) -> a -> b
$ forall a. TMVar a -> STM a
takeTMVar TMVar OutputBuffer
bv
		OutputBuffer
newbuf <- OutputBufferedActivity -> OutputBuffer -> IO OutputBuffer
addOutputBuffer (Text -> OutputBufferedActivity
Output (forall v. Outputable v => v -> Text
toOutput v
v)) OutputBuffer
oldbuf
		forall a. STM a -> IO a
atomically forall a b. (a -> b) -> a -> b
$ 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) =
		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 ()
_ <- forall a. IO a -> IO (Async a)
async forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) a. Functor f => f a -> f ()
void forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a.
MonadCatch m =>
m a -> m (Either IOException a)
tryIO forall a b. (a -> b) -> a -> b
$ ConcurrentProcessHandle -> IO ExitCode
P.waitForProcess ConcurrentProcessHandle
h
		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
		forall (m :: * -> *) a b. MonadCatch m => m a -> m b -> m a
`onException` IO ()
dropOutputLock
	IO ()
registerOutputThread
	-- Wait for the process to exit and drop the lock.
	Async ()
_ <- forall a. IO a -> IO (Async a)
async forall a b. (a -> b) -> a -> b
$ do
		forall (f :: * -> *) a. Functor f => f a -> f ()
void forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a.
MonadCatch m =>
m a -> m (Either IOException a)
tryIO forall a b. (a -> b) -> a -> b
$ ConcurrentProcessHandle -> IO ExitCode
P.waitForProcess ConcurrentProcessHandle
h
		IO ()
unregisterOutputThread
		IO ()
dropOutputLock
	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
		{ std_out :: StdStream
P.std_out = StdStream -> StdStream
rediroutput (CreateProcess -> StdStream
P.std_out CreateProcess
p)
		, std_err :: StdStream
P.std_err = StdStream -> StdStream
rediroutput (CreateProcess -> StdStream
P.std_err CreateProcess
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'
		forall (m :: * -> *) a b. 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
		, forall {a}. StdStream -> Maybe a -> Maybe a
mungeret (CreateProcess -> StdStream
P.std_out CreateProcess
p) Maybe Handle
stdout_h
		, 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 ()
_ <- forall a. IO a -> IO (Async a)
async forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) a. Functor f => f a -> f ()
void forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a.
MonadCatch m =>
m a -> m (Either IOException a)
tryIO 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 (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 (forall {a}. StdStream -> Maybe a -> Maybe a
mungebuf (CreateProcess -> StdStream
P.std_err CreateProcess
p) Maybe Handle
stderr_h)
	forall (f :: * -> *) a. Functor f => f a -> f ()
void forall a b. (a -> b) -> a -> b
$ forall a. IO a -> IO (Async a)
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]
	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 = forall a. Maybe a
Nothing
	mungeret :: StdStream -> Maybe a -> Maybe a
mungeret StdStream
ss Maybe a
mh
		| StdStream -> Bool
willOutput StdStream
ss = 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
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: OutputBuffer -> OutputBuffer -> Bool
$c/= :: OutputBuffer -> OutputBuffer -> Bool
== :: OutputBuffer -> OutputBuffer -> Bool
$c== :: 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
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: OutputBufferedActivity -> OutputBufferedActivity -> Bool
$c/= :: OutputBufferedActivity -> OutputBufferedActivity -> Bool
== :: OutputBufferedActivity -> OutputBufferedActivity -> Bool
$c== :: OutputBufferedActivity -> OutputBufferedActivity -> Bool
Eq)

data AtEnd = AtEnd
	deriving AtEnd -> AtEnd -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: AtEnd -> AtEnd -> Bool
$c/= :: AtEnd -> AtEnd -> Bool
== :: AtEnd -> AtEnd -> Bool
$c== :: 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 <- forall a. a -> IO (MVar a)
newMVar ([OutputBufferedActivity] -> OutputBuffer
OutputBuffer [])
	TMVar BufSig
bufsig <- forall a. STM a -> IO a
atomically forall a. STM (TMVar a)
newEmptyTMVar
	TMVar AtEnd
bufend <- forall a. STM a -> IO a
atomically forall a. STM (TMVar a)
newEmptyTMVar
	forall (f :: * -> *) a. Functor f => f a -> f ()
void forall a b. (a -> b) -> a -> b
$ forall a. IO a -> IO (Async a)
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
	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
				forall a. MVar a -> (a -> IO a) -> IO ()
modifyMVar_ MVar OutputBuffer
buf 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 = forall a. STM a -> IO a
atomically forall a b. (a -> b) -> a -> b
$ forall a. TMVar a -> a -> STM ()
putTMVar TMVar AtEnd
bufend AtEnd
AtEnd
	changed :: IO ()
changed = forall a. STM a -> IO a
atomically forall a b. (a -> b) -> a -> b
$ do
		forall (f :: * -> *) a. Functor f => f a -> f ()
void forall a b. (a -> b) -> a -> b
$ forall a. TMVar a -> STM (Maybe a)
tryTakeTMVar TMVar BufSig
bufsig
		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
	forall a. STM a -> IO a
atomically forall a b. (a -> b) -> a -> b
$ forall a. TMVar a -> a -> STM ()
putTMVar TMVar Integer
v forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Enum a => a -> a
succ forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< 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
	forall a. STM a -> IO a
atomically forall a b. (a -> b) -> a -> b
$ forall a. TMVar a -> a -> STM ()
putTMVar TMVar Integer
v forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Enum a => a -> a
pred forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< 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 <- forall a. STM a -> IO a
atomically forall a. STM (TMVar a)
newEmptyTMVar
	Async ()
worker1 <- forall a. IO a -> IO (Async a)
async forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a. (MonadIO m, MonadMask m) => m a -> m a
lockOutput forall a b. (a -> b) -> a -> b
$
		forall (m :: * -> *) a. Monad m => m Bool -> (m a, m a) -> m a
ifM (forall a. STM a -> IO a
atomically forall a b. (a -> b) -> a -> b
$ forall a. TMVar a -> a -> STM Bool
tryPutTMVar TMVar ()
activitysig ())
			( forall (f :: * -> *) a. Functor f => f a -> f ()
void forall a b. (a -> b) -> a -> b
$ 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
			, forall (m :: * -> *). Monad m => m ()
noop -- buffers already moved to global
			)
	Async ()
worker2 <- forall a. IO a -> IO (Async a)
async forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) a. Functor f => f a -> f ()
void forall a b. (a -> b) -> a -> b
$ forall {a}. TMVar () -> Async a -> IO ()
globalbuf TMVar ()
activitysig Async ()
worker1
	forall (f :: * -> *) a. Functor f => f a -> f ()
void forall a b. (a -> b) -> a -> b
$ forall a. IO a -> IO (Async a)
async forall a b. (a -> b) -> a -> b
$ do
		forall (f :: * -> *) a. Functor f => f a -> f ()
void forall a b. (a -> b) -> a -> b
$ forall a. Async a -> IO (Either SomeException a)
waitCatch Async ()
worker1
		forall (f :: * -> *) a. Functor f => f a -> f ()
void forall a b. (a -> b) -> a -> b
$ 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 <- forall a. STM a -> IO a
atomically forall a b. (a -> b) -> a -> b
$
			(forall a b. b -> Either a b
Right forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a. TMVar a -> STM a
takeTMVar TMVar BufSig
bufsig)
				forall a. STM a -> STM a -> STM a
`orElse`
			(forall a b. a -> Either a b
Left forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a. TMVar a -> STM a
takeTMVar TMVar AtEnd
bufend)
		OutputBuffer
l <- forall a. MVar a -> IO a
takeMVar MVar OutputBuffer
buf
		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 -> forall (m :: * -> *) a. Monad m => a -> m a
return ()
	globalbuf :: TMVar () -> Async a -> IO ()
globalbuf TMVar ()
activitysig Async a
worker1 = do
		Bool
ok <- forall a. STM a -> IO a
atomically forall a b. (a -> b) -> a -> b
$ do
			-- signal we're going to handle it
			-- (returns false if the displaybuf already did)
			Bool
ok <- forall a. TMVar a -> a -> STM Bool
tryPutTMVar TMVar ()
activitysig ()
			-- wait for end of all buffers
			forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when Bool
ok forall a b. (a -> b) -> a -> b
$
				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) -> forall a. TMVar a -> STM a
takeTMVar TMVar AtEnd
bufend) [(StdHandle, MVar OutputBuffer, TMVar BufSig, TMVar AtEnd)]
ts
			forall (m :: * -> *) a. Monad m => a -> m a
return Bool
ok
		forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when Bool
ok 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 <- 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 forall a b. (a -> b) -> a -> b
$ \(StdHandle
outh, MVar OutputBuffer
buf, TMVar BufSig
_bufsig, TMVar AtEnd
_bufend) ->
				(StdHandle
outh,) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a. MVar a -> IO a
takeMVar MVar OutputBuffer
buf
			forall a. STM a -> IO a
atomically forall a b. (a -> b) -> a -> b
$
				forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [(StdHandle, OutputBuffer)]
bs 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
			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' forall a. Ord a => a -> a -> Bool
<= Int
1048576 = forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ [OutputBufferedActivity] -> OutputBuffer
OutputBuffer (Text -> OutputBufferedActivity
Output Text
t' 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
		forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ [OutputBufferedActivity] -> OutputBuffer
OutputBuffer (OutputBufferedActivity
i forall a. a -> [a] -> [a]
: [OutputBufferedActivity]
other)
  where
	!t' :: Text
t' = [Text] -> Text
T.concat (forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe OutputBufferedActivity -> Maybe Text
getOutput [OutputBufferedActivity]
this) forall a. Semigroup a => a -> a -> a
<> Text
t
	!([OutputBufferedActivity]
this, [OutputBufferedActivity]
other) = 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'' -> forall a. a -> Maybe a
Just Text
t''
		OutputBufferedActivity
_ -> forall a. Maybe a
Nothing
addOutputBuffer OutputBufferedActivity
v (OutputBuffer [OutputBufferedActivity]
buf) = forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ [OutputBufferedActivity] -> OutputBuffer
OutputBuffer (OutputBufferedActivity
vforall 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 (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) <- forall a. TMVar a -> STM a
takeTMVar TMVar OutputBuffer
bv
	forall a. TMVar a -> a -> STM ()
putTMVar TMVar OutputBuffer
bv ([OutputBufferedActivity] -> OutputBuffer
OutputBuffer ([OutputBufferedActivity]
newbuf 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 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 forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a. TMVar a -> STM a
takeTMVar TMVar OutputBuffer
bv
		forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (OutputBuffer
selected forall a. Eq a => a -> a -> Bool
== [OutputBufferedActivity] -> OutputBuffer
OutputBuffer [])
			forall a. STM a
retry
		forall a. TMVar a -> a -> STM ()
putTMVar TMVar OutputBuffer
bv OutputBuffer
rest
		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) = 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
&& Text -> Char
T.last Text
t 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) = 
	forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ (forall a. [a] -> [a]
reverse [OutputBufferedActivity]
l) 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 forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< String -> IO Text
T.readFile String
tmp
			forall (f :: * -> *) a. Functor f => f a -> f ()
void forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a. MonadCatch m => m a -> m (Maybe a)
tryWhenExists 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 = forall (f :: * -> *) a. Functor f => f a -> f ()
void forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a.
MonadCatch m =>
m a -> m (Either IOException a)
tryIO forall a b. (a -> b) -> a -> b
$ do
		Handle -> Text -> IO ()
T.hPutStr Handle
outh Text
t
		Handle -> IO ()
hFlush Handle
outh