HandleReader.hs

{-# LANGUAGE ScopedTypeVariables #-}
module What4.Utils.HandleReader where

import           Control.Monad (unless)
import           Data.IORef
import           Data.Text (Text)
import qualified Data.Text as Text
import qualified Data.Text.Lazy as LazyText
import qualified Data.Text.IO as Text
import           Control.Exception(bracket,catch,IOException)
import           Control.Concurrent(ThreadId,forkIO,killThread)
import           Control.Concurrent.Chan(Chan,newChan,readChan,writeChan)
import           System.IO(Handle,hClose)
import           System.IO.Streams( OutputStream, InputStream )
import qualified System.IO.Streams as Streams


teeInputStream :: InputStream a -> OutputStream a -> IO (InputStream a)
teeInputStream i o = Streams.makeInputStream go
  where
  go = do x <- Streams.read i
          Streams.write x o
          return x

teeOutputStream :: OutputStream a -> OutputStream a -> IO (OutputStream a)
teeOutputStream o aux = Streams.makeOutputStream go
  where
  go x =
    do Streams.write x aux
       Streams.write x o

lineBufferedOutputStream :: Text -> OutputStream Text -> IO (OutputStream Text)
lineBufferedOutputStream prefix out =
    do ref <- newIORef mempty
       Streams.makeOutputStream (con ref)
 where
 newl = Text.pack "\n"

 con ref mx =
   do start <- readIORef ref
      case mx of
        Nothing ->
          do unless (Text.null start) (Streams.write (Just (prefix <> start)) out)
             Streams.write Nothing out
        Just x -> go ref (start <> x)

 go ref x =
   let (ln, x') = Text.break (== '\n') x in
   if Text.null x' then
     -- Flush
     do Streams.write (Just mempty) out
        writeIORef ref x
   else
     do Streams.write (Just (prefix <> ln <> newl)) out
        go ref (Text.drop 1 x')

demuxProcessHandles ::
  Handle {- ^ stdin for process -} ->
  Handle {- ^ stdout for process -} ->
  Handle {- ^ stderr for process -} ->
  Maybe (Text, Handle) {- optional handle to echo ouput; text argument is a line-comment prefix  -} ->
  IO ( OutputStream Text, InputStream Text, HandleReader )
demuxProcessHandles in_h out_h err_h Nothing =
  do in_str  <- Streams.encodeUtf8 =<< Streams.handleToOutputStream in_h
     out_str <- Streams.decodeUtf8 =<< Streams.handleToInputStream out_h
     err_reader <- startHandleReader err_h Nothing
     return (in_str, out_str, err_reader)
demuxProcessHandles in_h out_h err_h (Just (comment_prefix, aux_h)) =
  do aux_str <- Streams.lockingOutputStream =<< Streams.encodeUtf8 =<< Streams.handleToOutputStream aux_h
     in_str  <- Streams.encodeUtf8 =<< Streams.handleToOutputStream in_h
     out_str <- Streams.decodeUtf8 =<< Streams.handleToInputStream out_h

     in_aux <- lineBufferedOutputStream mempty aux_str
     in_str' <- teeOutputStream in_str in_aux

     out_aux <- lineBufferedOutputStream comment_prefix aux_str
     out_str' <- teeInputStream out_str out_aux

     err_reader <- startHandleReader err_h . Just
                    =<< lineBufferedOutputStream comment_prefix aux_str

     return (in_str', out_str', err_reader)


{- | Wrapper to help with reading from another process's
     standard out and stderr.

We want to be able to read from another process's stderr and stdout without
causing the process to stall because 'stdout' or 'stderr' becomes full.  This
data type will read from either of the handles, and buffer as much data
as needed in the queue.  It then provides a line-based method for reading
that data as strict bytestrings. -}
data HandleReader = HandleReader { hrChan :: !(Chan (Maybe Text))
                                 , hrHandle :: !Handle
                                 , hrThreadId :: !ThreadId
                                 }

streamLines :: Chan (Maybe Text) -> Handle -> Maybe (OutputStream Text) -> IO ()
streamLines c h Nothing = go
 where
 go = do ln <- Text.hGetLine h
         writeChan c (Just ln)
         go
streamLines c h (Just auxstr) = go
 where
 go = do ln <- Text.hGetLine h
         Streams.write (Just ln) auxstr
         writeChan c (Just ln)
         go

-- | Create a new handle reader for reading the given handle.
startHandleReader :: Handle -> Maybe (OutputStream Text) -> IO HandleReader
startHandleReader h auxOutput = do
  c <- newChan
  let handle_err (_e :: IOException) = writeChan c Nothing
  tid <- forkIO $ streamLines c h auxOutput `catch` handle_err

  return $! HandleReader { hrChan     = c
                         , hrHandle   = h
                         , hrThreadId = tid
                         }


-- | Stop the handle reader; cannot be used afterwards.
stopHandleReader :: HandleReader -> IO ()
stopHandleReader hr = do
  killThread (hrThreadId hr)
  hClose (hrHandle hr)

-- | Run an execution with a handle reader and stop it wheen down
withHandleReader :: Handle -> Maybe (OutputStream Text) -> (HandleReader -> IO a) -> IO a
withHandleReader h auxOut = bracket (startHandleReader h auxOut) stopHandleReader

readNextLine :: HandleReader -> IO (Maybe Text)
readNextLine hr = do
  mr <- readChan (hrChan hr)
  case mr of
    -- Write back 'Nothing' because thread should have terminated.
    Nothing -> writeChan (hrChan hr) Nothing
    Just{} -> return()
  return mr

readAllLines :: HandleReader -> IO LazyText.Text
readAllLines hr = go LazyText.empty
  where go :: LazyText.Text -> IO LazyText.Text
        go prev = do
          mr <- readNextLine hr
          case mr of
            Nothing -> return prev
            Just e -> go $! prev `LazyText.append` (LazyText.fromStrict e)
                                 `LazyText.snoc` '\n'