-- File created: 2008-02-11 12:55:34 {-# LANGUAGE CPP #-} {-# OPTIONS_GHC -fglasgow-exts -frewrite-rules #-} -- for the rewrite rule ------------------------------------------------------------------------------- -- | -- Module : System.Process.Pipe -- Copyright : (c) Matti Niemenmaa 2008 -- License : BSD (see LICENSE.txt) -- -- Maintainer : Matti Niemenmaa <matti.niemenmaa+web@iki.fi> -- Stability : experimental -- Portability : portable -- -- Operations for piping data through multiple processes. -- -- 'pipe' is the most general function, with 'filePipe' and 'pipe\'' provided -- for convenience purposes. For the common case of piping between 'String's, -- the 'pipeString' wrapper and the 'word8ToString' and 'stringToWord8' helpers -- are included. -- -- Whenever specifying a path to a process, explicitly specifying the current -- directory is recommended for portability. That is: use \".\/foo\" instead of -- \"foo\", for instance. -- -- On Windows, appending \".exe\" to process paths is attempted if the -- invocation fails. ------------------------------------------------------------------------------- module System.Process.Pipe ( filePipe , Tap(..), Sink(..), bufferSize , pipe, pipe', pipeString , word8ToString, stringToWord8 ) where import Control.Monad (forM) import Data.Char (chr, ord) import Data.Maybe (fromJust) import Data.Word (Word8) import Foreign.Marshal.Alloc (allocaBytes) import Foreign.Marshal.Array (peekArray, pokeArray) import Foreign.Ptr (Ptr) import System.FilePath (dropFileName) import System.IO ( withBinaryFile, IOMode (ReadMode, WriteMode) , Handle, hClose , hGetContents, hPutStr) import System.Process ( CreateProcess(..), createProcess , CmdSpec (RawCommand) , StdStream (CreatePipe, Inherit, UseHandle) , ProcessHandle, waitForProcess) import System.Process.Pipe.Plumbing #if mingw32_HOST_OS import Data.ByteString (ByteString) import qualified Data.ByteString as BS import Foreign.Ptr (castPtr) #else import System.IO (hIsOpen, hGetBuf, hGetBufNonBlocking, hPutBuf) import System.IO.Error (isFullError) import System.Posix.Signals ( Signal, openEndedPipe , Handler (Ignore), installHandler) #endif type Proc = ( Handle, Handle, ProcessHandle) type MProc = (Maybe Handle, Maybe Handle, ProcessHandle) createProc :: FilePath -> StdStream -> StdStream -> (FilePath,[String]) -> IO MProc createProc wdir inp out (p,args) = do let proc = CreateProcess { cmdspec = undefined , cwd = Just wdir , env = Nothing , std_in = inp , std_out = out , std_err = Inherit , close_fds = True } (i,o,_,pid) <- createProcess proc { cmdspec = RawCommand p args } #if mingw32_HOST_OS `catch` const (createProcess proc { cmdspec = RawCommand (p ++ ".exe") args }) #endif return (i,o,pid) -- | Pipes the contents of the first file to the second file through all the -- programs named. -- -- The working directory used is the directory component of the path to the -- first file. filePipe :: [(FilePath,[String])] -> FilePath -> FilePath -> IO () filePipe progs infile outfile = do withBinaryFile outfile WriteMode $ \outhdl -> withBinaryFile infile ReadMode $ \inhdl -> pipe (dropFileName infile) progs inhdl outhdl return () -- | Pipes data from the 'Tap' to the 'Sink' through all the commands named, in -- the given working directory. -- -- Be careful! All IO is at the byte level: this means that piping even a -- String such as \"foo\" will result in the raw UTF-32 moving: the bytes (in -- my case; I believe this is implementation-dependent) in question are not the -- ASCII @[102, 111, 111]@ but rather @[102, 0, 0, 0, 111, 0, 0, 0, 111, 0, 0, -- 0]@. -- -- Note to Windows users: since 'hGetBufNonBlocking' doesn't work on Windows -- (it blocks despite its name, see -- <http://hackage.haskell.org/trac/ghc/ticket/806>), this pipeline uses a -- non-constant amount of space. The amount used is linear in the amount of -- data used at any point in the pipeline. So if you want to pipe 20 gibioctets -- of data to a program, you better make sure you have at least said amount of -- memory available. (In fact, ByteStrings are used, and their documentation -- suggests that you might want twice that, just in case.) -- -- In addition, the 'Tap' and 'Sink' classes are meant for the POSIX code: -- having to move data through the 'Ptr' 'Word8' types, 'bufferSize' bytes at a -- time, results in extra complexity. -- -- If you want to do something about the above, ideally fix the GHC ticket -- (probably nontrivial) and let me know so that I can activate the better code -- for Windows as well. Alternatively, feel free to code an implementation of -- this which works on Windows. pipe :: (Tap t, Sink s) => FilePath -> [(FilePath,[String])] -> t -> s -> IO (t,s) pipe wdir progs otap osink = do let cp = createProc wdir CreatePipe CreatePipe allocaBytes bufferSize $ \buf -> do ps <- forM progs $ \pr -> do (i,o,pid) <- cp pr return (fromJust i, fromJust o, pid) -- See 'pipeline' comment below for why this needs to be done differently. #if mingw32_HOST_OS -- Gather up all data from the tap until it's exhausted. let loop tap s = do exh <- exhausted tap if exh then return (tap,s) else do (tap',sz) <- flowOut tap buf bufferSize xs <- peekArray sz (castPtr buf) loop tap' (s `BS.append` BS.pack xs) (tap, s) <- loop otap BS.empty sink <- pipeline osink buf ps s return (tap,sink) #else let loop :: (Sink s, Tap t) => t -> s -> IO (t,s) loop tap sink = do exh <- exhausted tap if exh then return (tap,sink) else do (tap' , sz) <- flowOut tap buf bufferSize (sink', status) <- pipeline sink ps buf sz exh' <- exhausted tap' if isNeed status && exh' then do -- The first process wants more input, but there's no -- more to give. Hence what we do is enter the final -- pipeline: have the process close its stdin and deal -- with any leftover output. (sink'',_) <- finalPipeline sink' ps buf return (tap',sink'') else loop tap' sink' loop otap osink #endif #if mingw32_HOST_OS -- hGetBufNonBlocking doesn't work on Windows (see -- http://hackage.haskell.org/trac/ghc/ticket/806). I can't think of a way of -- doing a robust constant-space pipeline without it. Hence we use this silly -- implementation instead, which is crap but does the job. (This is how -- pipelines worked in Windows pre-NT anyway, AFAIK. (At least in DOS.)) -- -- Basically we just wait for the previous command to complete, gathering up -- all its output, and then give it all at once to the next one. And yes, this -- means that if some process outputs (or the input file contains) an infinite -- amount of data or more than fits in memory, you're screwed. pipeline :: Sink s => s -> Ptr Word8 -> [Proc] -> ByteString -> IO s pipeline sink buf ((i,o,p):ps) dat = do BS.hPut i dat hClose i dat' <- BS.hGetContents o waitForProcess p pipeline sink buf ps dat' pipeline osink buf [] dat = do let loop sink s = if BS.null s then return sink else do let (xs,ys) = BS.splitAt bufferSize s pokeArray (castPtr buf) (BS.unpack xs) sink' <- flowIn sink buf (BS.length xs) loop sink' ys loop osink dat #else data Need = Done | Need isNeed :: Need -> Bool isNeed Need = True isNeed _ = False -- This matches more than just EPIPE but it's the best we can do. onEPIPE :: IO a -> IO a -> IO a a `onEPIPE` b = a `catch` \e -> if isFullError e then b else ioError e pipeline, shoveDown :: Sink s => s -> [Proc] -> Ptr Word8 -> Int -> IO (s,Need) pipeline sink [] buf sz = toSink sink buf sz pipeline sink ps@((inp, out, _):_) buf sz = do -- Put the given data to the stdin of this process. -- -- We have to be careful with SIGPIPEs here. If hPutBuf fails, that's -- because the input pipe is either closed or broken. We try to close it, -- and if that fails then we know the pipe is broken. -- -- Simply piping infinite output to a process which takes only a bit of -- input ("yes | head", for instance) triggers both of the onEPIPE cases, as -- well as the one in finalize. brokenPipe <- withIgnoringSignal openEndedPipe $ (hPutBuf inp buf sz >> return False) `onEPIPE` (hClose inp >> return False) `onEPIPE` return True -- We can't block here, lest it be the case where the program has output -- e.g. bufferSize-1 bytes, but is waiting on input. If we block here then -- we're waiting for more output while the process is waiting for more -- input---deadlock! sz' <- hGetBufNonBlocking out buf bufferSize if sz' == 0 -- We got no output from this process. If the process's stdin is open, -- we request more input for it from the process above us. Otherwise -- we block, waiting for the process to complete its computation or -- whatever it's doing. We can do this safely now since the stdin is -- closed: there's no way to reach the deadlock outlined above. then do wantsMore <- hIsOpen inp if wantsMore && not brokenPipe then return (sink, Need) else blockingPipeline sink ps buf else shoveDown sink ps buf sz' shoveDown sink [] buf sz = toSink sink buf sz shoveDown sink ps@(p:rest) buf sz = do -- Pipe the output from the process above to the next process in the -- pipeline. (sink', below) <- pipeline sink rest buf sz case below of -- The process below wants more data: go back and try a nonblocking get -- from this process. Need -> pipeline sink' ps buf 0 -- The process below us says it wants nothing more ever again. Since it -- wants nothing from us, we have nothing to do either: finish up here -- and bubble the information up. Done -> finalize p >> return (sink', Done) blockingPipeline, finalPipeline :: Sink s => s -> [Proc] -> Ptr Word8 -> IO (s, Need) blockingPipeline sink [] buf = toSink sink buf 0 blockingPipeline sink ps@(p@(_,out,_):rest) buf = do sz <- hGetBuf out buf bufferSize if sz == 0 -- The blocking call returned 0: this means we've hit EOF, i.e. the -- process is done and will no longer output anything. Go down the -- pipeline, closing the stdin of the process below us and pushing any -- last output further down. After that, finish up here and tell the ones -- above to do the same. then do (sink', _) <- finalPipeline sink rest buf finalize p return (sink', Done) else shoveDown sink ps buf sz finalPipeline sink [] buf = toSink sink buf 0 finalPipeline sink ps@((inp,_,_):_) buf = do hClose inp -- Since the stdin is closed, we can jump straight to the blocking version -- of the pipeline: while we could go by way of the non-blocking one there's -- no need to do so. blockingPipeline sink ps buf -- Some data found its way all the way down the pipeline, so we put it in the -- output handle and let the processes know that we're ready for more. -- -- The special case for 0, while handled in hPutBuf, is my little hint to the -- optimizer that it should inline these calls where the 0 is given explicitly -- above. toSink :: Sink s => s -> Ptr Word8 -> Int -> IO (s, Need) toSink sink _ 0 = return (sink , Need) toSink sink buf sz = flowIn sink buf sz >>= \sink' -> return (sink', Need) finalize :: Proc -> IO () finalize (i,o,p) = do -- Once again the stdin pipe may be broken so we sidestep a possible -- SIGPIPE. withIgnoringSignal openEndedPipe (hClose i `onEPIPE` return ()) -- We close the stdout as soon as we no longer need to read anything from -- it, so that we don't have to wait for the process to finish outputting. -- With infinite-output processes (and possibly even just with buffered -- output) we'd be stuck forever in waitForProcess if we don't do this. -- -- This is also the reason why withIgnoringSignal openEndedPipe isn't -- wrapped around this whole thing already in the pipe function: we want the -- processes in the pipeline to handle (or die of) SIGPIPE (which this -- hClose causes on their end) by themselves. hClose o waitForProcess p return () withIgnoringSignal :: Signal -> IO a -> IO a withIgnoringSignal sig mx = do old <- installHandler sig Ignore Nothing x <- mx installHandler sig old Nothing return x #endif {-# RULES "pipe->handlePipe" pipe = handlePipe #-} -- Smarter way of piping Handle-to-Handle handlePipe :: FilePath -> [(FilePath,[String])] -> Handle -> Handle -> IO (Handle, Handle) handlePipe _ [] inhdl outhdl = do hGetContents inhdl >>= hPutStr outhdl return (inhdl, outhdl) handlePipe wdir (p:ps) inhdl outhdl = do let cp = createProc wdir f pids out [] = return (out, pids) f pids out (pr:prs) = do (i,_,pid) <- cp CreatePipe out pr f (pid:pids) (UseHandle . fromJust $ i) prs (inp, pids) <- f [] (UseHandle outhdl) (reverse ps) (_,_,pid) <- cp (UseHandle inhdl) inp p mapM_ waitForProcess (pid:pids) return (inhdl, outhdl) -- | A convenience function for when you don't care about the working -- directory. -- -- > pipe' = pipe "." pipe' :: (Tap t, Sink s) => [(FilePath,[String])] -> t -> s -> IO (t,s) pipe' = pipe "." -- | A convenience function for the common case of piping from a 'String' to a -- 'String'. This uses the 'word8ToString' and 'stringToWord8' functions and -- thus loses information if your 'Char's are non-ASCII. pipeString :: [(FilePath, [String])] -> String -> IO String pipeString progs s = do (_, out) <- pipe' progs (stringToWord8 s) ([] :: [Word8]) return (word8ToString out) -- | A helper function which converts a @['Word8']@ to a 'String' by mapping -- 'chr' over the octets. -- -- In most cases, when you wish to pipe data to a String, you do not want to -- interpret the results as the raw byte pattern of 'Char's, so you use -- @['Word8']@ as the 'Sink' type. This function handles the common case of -- ASCII data simply—if you're dealing with non-ASCII data you probably need -- to handle the results in a different way. word8ToString :: [Word8] -> String word8ToString = map (chr.fromIntegral) -- | The inverse of 'word8ToString'. Any 'Char's greater than 255 are -- truncated: once again, be careful with non-ASCII. stringToWord8 :: String -> [Word8] stringToWord8 = map (fromIntegral.ord)