1{-# LANGUAGE DataKinds #-}
2
3-- | The "System.Process.Typed" module from @typed-process@, but with
4-- added conduit helpers.
5module Data.Conduit.Process.Typed
6  ( -- * Conduit specific stuff
7    createSink
8  , createSinkClose
9  , createSource
10    -- * Running a process with logging
11  , withLoggedProcess_
12    -- * Reexports
13  , module System.Process.Typed
14  ) where
15
16import System.Process.Typed
17import qualified System.Process.Typed as P
18import Data.Conduit (ConduitM, (.|), runConduit)
19import qualified Data.Conduit.Binary as CB
20import Control.Monad.IO.Unlift
21import qualified Data.ByteString as S
22import qualified Data.Conduit.List as CL
23import qualified Data.ByteString.Lazy as BL
24import Data.IORef (IORef, newIORef, readIORef, modifyIORef)
25import Control.Exception (throwIO, catch)
26import Control.Concurrent.Async (concurrently)
27import System.IO (hSetBuffering, BufferMode (NoBuffering), hClose)
28
29-- | Provide input to a process by writing to a conduit. The sink provided here
30-- will leave the pipe to the child open after the stream ends. This allows the
31-- sink to be used multiple times, but may result in surprising behavior. You
32-- may prefer 'createSinkClose', see
33-- <https://github.com/snoyberg/conduit/issues/434>.
34--
35-- @since 1.2.1
36createSink :: MonadIO m => StreamSpec 'STInput (ConduitM S.ByteString o m ())
37createSink =
38  (\h -> liftIO (hSetBuffering h NoBuffering) >> CB.sinkHandle h)
39  `fmap` createPipe
40
41-- | Like 'createSink', but closes the pipe to the child process as soon as it
42-- runs out of data.
43--
44-- @since 1.3.5
45createSinkClose :: MonadIO m => StreamSpec 'STInput (ConduitM S.ByteString o m ())
46createSinkClose =
47  (\h -> liftIO (hSetBuffering h NoBuffering) >> CB.sinkHandle h >> liftIO (hClose h))
48  `fmap` createPipe
49
50-- | Read output from a process by read from a conduit.
51--
52-- @since 1.2.1
53createSource :: MonadIO m => StreamSpec 'STOutput (ConduitM i S.ByteString m ())
54createSource =
55  (\h -> liftIO (hSetBuffering h NoBuffering) >> CB.sourceHandle h)
56  `fmap` createPipe
57
58-- | Internal function: like 'createSource', but stick all chunks into
59-- the 'IORef'.
60createSourceLogged
61  :: MonadIO m
62  => IORef ([S.ByteString] -> [S.ByteString])
63  -> StreamSpec 'STOutput (ConduitM i S.ByteString m ())
64createSourceLogged ref =
65    -- We do not add a cleanup action to close the handle, since in
66    -- withLoggedProcess_ we attempt to read from the handle twice
67    (\h ->
68       (  CB.sourceHandle h
69       .| CL.iterM (\bs -> liftIO $ modifyIORef ref (. (bs:))))
70    )
71    `fmap` createPipe
72
73-- | Run a process, throwing an exception on a failure exit code. This
74-- will store all output from stdout and stderr in memory for better
75-- error messages. Note that this will require unbounded memory usage,
76-- so caveat emptor.
77--
78-- This will ignore any previous settings for the stdout and stderr
79-- streams, and instead force them to use 'createSource'.
80--
81-- @since 1.2.3
82withLoggedProcess_
83  :: MonadUnliftIO m
84  => ProcessConfig stdin stdoutIgnored stderrIgnored
85  -> (Process stdin (ConduitM () S.ByteString m ()) (ConduitM () S.ByteString m ()) -> m a)
86  -> m a
87withLoggedProcess_ pc inner = withUnliftIO $ \u -> do
88  stdoutBuffer <- newIORef id
89  stderrBuffer <- newIORef id
90  let pc' = setStdout (createSourceLogged stdoutBuffer)
91          $ setStderr (createSourceLogged stderrBuffer) pc
92  -- withProcessWait vs Term doesn't actually matter here, since we
93  -- call checkExitCode inside regardless. But still, Wait is the
94  -- safer function to use in general.
95  P.withProcessWait pc' $ \p -> do
96    a <- unliftIO u $ inner p
97    let drain src = unliftIO u (runConduit (src .| CL.sinkNull))
98    ((), ()) <- drain (getStdout p) `concurrently`
99                drain (getStderr p)
100    checkExitCode p `catch` \ece -> do
101      stdout <- readIORef stdoutBuffer
102      stderr <- readIORef stderrBuffer
103      throwIO ece
104        { eceStdout = BL.fromChunks $ stdout []
105        , eceStderr = BL.fromChunks $ stderr []
106        }
107    return a
108