{-# LANGUAGE RecordWildCards #-} -- | A wrapping of createProcess to provide a more flexible interface. module General.Process( Buffer, newBuffer, readBuffer, process, ProcessOpts(..), Source(..), Destination(..) ) where import Control.Applicative import Control.Concurrent import Control.DeepSeq import Control.Exception.Extra as C import Control.Monad.Extra import Data.List.Extra import Data.Maybe import Foreign.C.Error import System.Exit import System.IO.Extra import System.Info.Extra import System.Process import System.Time.Extra import Data.Unique import Data.IORef import qualified Data.ByteString.Internal as BS(createAndTrim) import qualified Data.ByteString.Char8 as BS import qualified Data.ByteString.Lazy as LBS import General.Extra import Prelude import GHC.IO.Exception (IOErrorType(..), IOException(..)) --------------------------------------------------------------------- -- BUFFER ABSTRACTION data Buffer a = Buffer Unique (IORef [a]) instance Eq (Buffer a) where Buffer x _ == Buffer y _ = x == y instance Ord (Buffer a) where compare (Buffer x _) (Buffer y _) = compare x y newBuffer :: IO (Buffer a) newBuffer = liftM2 Buffer newUnique (newIORef []) addBuffer :: Buffer a -> a -> IO () addBuffer (Buffer _ ref) x = atomicModifyIORef ref $ \xs -> (x:xs, ()) readBuffer :: Buffer a -> IO [a] readBuffer (Buffer _ ref) = reverse <$> readIORef ref --------------------------------------------------------------------- -- OPTIONS data Source = SrcFile FilePath | SrcString String | SrcBytes LBS.ByteString data Destination = DestEcho | DestFile FilePath | DestString (Buffer String) | DestBytes (Buffer BS.ByteString) deriving (Eq,Ord) isDestString DestString{} = True; isDestString _ = False isDestBytes DestBytes{} = True; isDestBytes _ = False data ProcessOpts = ProcessOpts {poCommand :: CmdSpec ,poCwd :: Maybe FilePath ,poEnv :: Maybe [(String, String)] ,poTimeout :: Maybe Double ,poStdin :: [Source] ,poStdout :: [Destination] ,poStderr :: [Destination] ,poAsync :: Bool } --------------------------------------------------------------------- -- IMPLEMENTATION -- | If two buffers can be replaced by one and a copy, do that (only if they start empty) optimiseBuffers :: ProcessOpts -> IO (ProcessOpts, IO ()) optimiseBuffers po@ProcessOpts{..} = return (po{poStdout = nubOrd poStdout, poStderr = nubOrd poStderr}, return ()) stdStream :: (FilePath -> Handle) -> [Destination] -> [Destination] -> StdStream stdStream file [DestEcho] other = Inherit stdStream file [DestFile x] other | other == [DestFile x] || DestFile x `notElem` other = UseHandle $ file x stdStream file _ _ = CreatePipe stdIn :: (FilePath -> Handle) -> [Source] -> (StdStream, Handle -> IO ()) stdIn file [] = (Inherit, const $ return ()) stdIn file [SrcFile x] = (UseHandle $ file x, const $ return ()) stdIn file src = (,) CreatePipe $ \h -> ignoreSigPipe $ do forM_ src $ \x -> case x of SrcString x -> hPutStr h x SrcBytes x -> LBS.hPutStr h x SrcFile x -> LBS.hPutStr h =<< LBS.hGetContents (file x) hClose h ignoreSigPipe :: IO () -> IO () ignoreSigPipe = handleIO $ \e -> case e of IOError {ioe_type=ResourceVanished, ioe_errno=Just ioe} | Errno ioe == ePIPE -> return () _ -> throwIO e withTimeout :: Maybe Double -> IO () -> IO a -> IO a withTimeout Nothing stop go = go withTimeout (Just s) stop go = bracket (forkIO $ sleep s >> stop) killThread $ const go cmdSpec :: CmdSpec -> CreateProcess cmdSpec (ShellCommand x) = shell x cmdSpec (RawCommand x xs) = proc x xs forkWait :: IO a -> IO (IO a) forkWait a = do res <- newEmptyMVar _ <- mask $ \restore -> forkIO $ try_ (restore a) >>= putMVar res return $ takeMVar res >>= either throwIO return abort :: ProcessHandle -> IO () abort pid = do interruptProcessGroupOf pid sleep 5 -- give the process a few seconds grace period to die nicely -- seems to happen with some GHC 7.2 compiled binaries with FFI etc terminateProcess pid withFiles :: IOMode -> [FilePath] -> ((FilePath -> Handle) -> IO a) -> IO a withFiles mode files act = withs (map (`withFile` mode) files) $ \handles -> act $ \x -> fromJust $ lookup x $ zip files handles -- General approach taken from readProcessWithExitCode process :: ProcessOpts -> IO (ProcessHandle, ExitCode) process po = do (ProcessOpts{..}, flushBuffers) <- optimiseBuffers po let outFiles = nubOrd [x | DestFile x <- poStdout ++ poStderr] let inFiles = nubOrd [x | SrcFile x <- poStdin] withFiles WriteMode outFiles $ \outHandle -> withFiles ReadMode inFiles $ \inHandle -> do let cp = (cmdSpec poCommand){cwd = poCwd, env = poEnv, create_group = isJust poTimeout, close_fds = True ,std_in = fst $ stdIn inHandle poStdin ,std_out = stdStream outHandle poStdout poStderr, std_err = stdStream outHandle poStderr poStdout} withCreateProcessCompat cp $ \inh outh errh pid -> withTimeout poTimeout (abort pid) $ do let streams = [(outh, stdout, poStdout) | Just outh <- [outh], CreatePipe <- [std_out cp]] ++ [(errh, stderr, poStderr) | Just errh <- [errh], CreatePipe <- [std_err cp]] wait <- forM streams $ \(h, hh, dest) -> do -- no point tying the streams together if one is being streamed directly let isTied = not (poStdout `disjoint` poStderr) && length streams == 2 let isBinary = not $ any isDestString dest && not (any isDestBytes dest) when isTied $ hSetBuffering h LineBuffering when (DestEcho `elem` dest) $ do buf <- hGetBuffering hh case buf of BlockBuffering{} -> return () _ -> hSetBuffering h buf if isBinary then do hSetBinaryMode h True dest <- return $ for dest $ \d -> case d of DestEcho -> BS.hPut hh DestFile x -> BS.hPut (outHandle x) DestString x -> addBuffer x . (if isWindows then replace "\r\n" "\n" else id) . BS.unpack DestBytes x -> addBuffer x forkWait $ whileM $ do src <- bsHGetSome h 4096 mapM_ ($ src) dest notM $ hIsEOF h else if isTied then do dest <- return $ for dest $ \d -> case d of DestEcho -> hPutStrLn hh DestFile x -> hPutStrLn (outHandle x) DestString x -> addBuffer x . (++ "\n") forkWait $ whileM $ ifM (hIsEOF h) (return False) $ do src <- hGetLine h mapM_ ($ src) dest return True else do src <- hGetContents h wait1 <- forkWait $ C.evaluate $ rnf src waits <- forM dest $ \d -> case d of DestEcho -> forkWait $ hPutStr hh src DestFile x -> forkWait $ hPutStr (outHandle x) src DestString x -> do addBuffer x src; return $ return () return $ sequence_ $ wait1 : waits whenJust inh $ snd $ stdIn inHandle poStdin if poAsync then return (pid, ExitSuccess) else do sequence_ wait flushBuffers res <- waitForProcess pid whenJust outh hClose whenJust errh hClose return (pid, res) --------------------------------------------------------------------- -- COMPATIBILITY -- available in bytestring-0.9.1.10, GHC 7.8 and above -- implementation copied below bsHGetSome :: Handle -> Int -> IO BS.ByteString bsHGetSome h i = BS.createAndTrim i $ \p -> hGetBufSome h p i -- available in process-1.4.3.0, GHC ??? (Nov 2015) -- logic copied directly (apart from Ctrl-C handling magic using internal pieces) withCreateProcessCompat :: CreateProcess -> (Maybe Handle -> Maybe Handle -> Maybe Handle -> ProcessHandle -> IO a) -> IO a withCreateProcessCompat cp act = bracketOnError (createProcess cp) cleanup (\(m_in, m_out, m_err, ph) -> act m_in m_out m_err ph) where cleanup (inh, outh, errh, pid) = do terminateProcess pid whenJust inh $ ignoreSigPipe . hClose whenJust outh hClose whenJust errh hClose forkIO $ void $ waitForProcess pid