{-# LANGUAGE DeriveDataTypeable #-}
{-# LANGUAGE ForeignFunctionInterface #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ScopedTypeVariables #-}
module Data.Conduit.Process.Unix
( -- * Starting processes
forkExecuteFile
, killProcess
, terminateProcess
, waitForProcess
, ProcessStatus (..)
, signalProcessHandle
, signalProcessHandleGroup
-- * Process tracking
-- $processTracker
-- ** Types
, ProcessTracker
, TrackedProcess
, ProcessTrackerException (..)
-- ** Functions
, initProcessTracker
, trackProcess
, untrackProcess
-- * Logging
, RotatingLog
, openRotatingLog
, forkExecuteLog
-- * Monitored process
, MonitoredProcess
, monitorProcess
, terminateMonitoredProcess
) where
import Control.Applicative ((<$>), (<*>))
import Control.Arrow ((***))
import Control.Concurrent (forkIO)
import Control.Concurrent (threadDelay)
import Control.Concurrent.MVar (MVar, modifyMVar, newMVar,
readMVar, swapMVar)
import Control.Exception (Exception, SomeException,
bracketOnError, finally,
handle, mask, mask_,
onException, throwIO, try)
import Control.Monad (unless, void, when, zipWithM_)
import Control.Monad.Trans.Class (lift)
import Data.ByteString (ByteString, append, concat,
null, singleton)
import qualified Data.ByteString.Char8 as S8
import Data.ByteString.Unsafe (unsafePackCStringFinalizer,
unsafeUseAsCString,
unsafeUseAsCStringLen)
import Data.Conduit (Sink, Source, yield, ($$))
import Data.Conduit.Binary (sinkHandle, sourceHandle)
import Data.Conduit.List (mapM_)
import qualified Data.Conduit.List as CL
import Data.Conduit.LogFile
import Data.IORef (IORef, newIORef, readIORef,
writeIORef)
import Data.Time (getCurrentTime)
import Data.Time (diffUTCTime)
import Data.Typeable (Typeable)
import Foreign.C.String
import Foreign.C.Types
import Foreign.Marshal.Alloc (allocaBytes, free,
mallocBytes)
import Foreign.Ptr (Ptr, castPtr, nullPtr)
import Foreign.Storable (pokeElemOff, sizeOf)
import Prelude (Bool (..), Either (..), IO,
Maybe (..), Monad (..), Show,
const, error, flip, fmap,
fromIntegral, fst, head, id,
length, map, maybe, show, snd,
($), ($!), (*), (.), (<),
(==))
import System.IO (hClose)
import System.Posix.IO.ByteString (FdOption (CloseOnExec),
closeFd, createPipe,
fdReadBuf, fdToHandle,
fdWriteBuf, setFdOption)
import System.Posix.Process.ByteString (ProcessStatus (..),
getProcessGroupIDOf,
getProcessStatus)
import System.Posix.Signals (Signal, sigKILL,
signalProcess,
signalProcessGroup)
import System.Posix.Types (CPid (..), Fd, ProcessID)
import System.Process
import System.Process.Internals (ProcessHandle (..),
ProcessHandle__ (..),
withProcessHandle_)
-- | Kill a process by sending it the KILL (9) signal.
--
-- Since 0.1.0
killProcess :: ProcessHandle -> IO ()
killProcess ph = withProcessHandle_ ph $ \p_ ->
case p_ of
ClosedHandle _ -> return p_
OpenHandle h -> do
signalProcess sigKILL h
return p_
signalProcessHandle :: Signal -> ProcessHandle -> IO ()
signalProcessHandle signal ph = withProcessHandle_ ph $ \p_ ->
case p_ of
ClosedHandle _ -> return p_
OpenHandle h -> do
signalProcess signal h
return p_
signalProcessHandleGroup :: Signal -> ProcessHandle -> IO ()
signalProcessHandleGroup signal ph = withProcessHandle_ ph $ \p_ ->
case p_ of
ClosedHandle _ -> return p_
OpenHandle h -> do
pgid <- getProcessGroupIDOf h
signalProcessGroup signal pgid
return p_
-- | Fork a new process and execute the given command.
--
-- This is a wrapper around with fork() and exec*() syscalls, set up to work
-- with @conduit@ datatypes for standard input, output, and error. If @Nothing@
-- is provided for any of those arguments, then the original file handles will
-- remain open to the child process.
--
-- If you would like to simply discard data provided by the child process,
-- provide @sinkNull@ for stdout and/or stderr. To provide an empty input
-- stream, use @return ()@.
--
-- Since 0.1.0
forkExecuteFile :: ByteString -- ^ command
-> [ByteString] -- ^ args
-> Maybe [(ByteString, ByteString)] -- ^ environment
-> Maybe ByteString -- ^ working directory
-> Maybe (Source IO ByteString) -- ^ stdin
-> Maybe (Sink ByteString IO ()) -- ^ stdout
-> Maybe (Sink ByteString IO ()) -- ^ stderr
-> IO ProcessHandle
forkExecuteFile cmd args menv mwdir mstdin mstdout mstderr = do
(min, mout, merr, ph) <- createProcess cp
case (,) <$> mstdin <*> min of
Just (source, h) -> void $ forkIO $ ignoreExceptions $
(source $$ sinkHandle h) `finally` hClose h
Nothing -> return ()
case (,) <$> mstdout <*> mout of
Just (sink, h) -> void $ forkIO $ ignoreExceptions $
(sourceHandle h $$ sink) `finally` hClose h
Nothing -> return ()
case (,) <$> mstderr <*> merr of
Just (sink, h) -> void $ forkIO $ ignoreExceptions $
(sourceHandle h $$ sink) `finally` hClose h
Nothing -> return ()
return ph
where
cp = CreateProcess
{ cmdspec = RawCommand (S8.unpack cmd) (map S8.unpack args)
, cwd = S8.unpack <$> mwdir
, env = map (S8.unpack *** S8.unpack) <$> menv
, std_in = maybe Inherit (const CreatePipe) mstdin
, std_out = maybe Inherit (const CreatePipe) mstdout
, std_err = maybe Inherit (const CreatePipe) mstderr
, close_fds = True
, create_group = True
}
ignoreExceptions :: IO () -> IO ()
ignoreExceptions = handle (\(_ :: SomeException) -> return ())
closeOnExec :: Fd -> IO ()
closeOnExec fd = setFdOption fd CloseOnExec True
withMString :: Maybe ByteString -> (CString -> IO a) -> IO a
withMString Nothing f = f nullPtr
withMString (Just bs) f = unsafeUseAsCString (bs `append` singleton 0) f
withEnv :: Maybe [(ByteString, ByteString)] -> (Ptr CString -> IO a) -> IO a
withEnv Nothing f = f nullPtr
withEnv (Just pairs) f =
withArgs (map toBS pairs) f
where
toBS (x, y) = concat [x, "=", y]
withArgs :: [ByteString] -> (Ptr CString -> IO a) -> IO a
withArgs bss0 f =
loop bss0 id
where
loop [] front = run (front [nullPtr])
loop (bs:bss) front =
unsafeUseAsCString (bs `append` singleton 0) $ \ptr ->
loop bss (front . (ptr:))
run ptrs = allocaBytes (length ptrs * sizeOf (head ptrs)) $ \res ->
zipWithM_ (pokeElemOff res) [0..] ptrs >> f res
-- $processTracker
--
-- Ensure that child processes are killed, regardless of how the parent process exits.
--
-- The technique used here is:
--
-- * Create a pipe.
--
-- * Fork a new child process that listens on the pipe.
--
-- * In the current process, send updates about processes that should be auto-killed.
--
-- * When the parent process dies, listening on the pipe in the child process will get an EOF.
--
-- * When the child process receives that EOF, it kills all processes it was told to auto-kill.
--
-- This code was originally written for Keter, but was moved to unix-process
-- conduit in the 0.2.1 release.
foreign import ccall unsafe "launch_process_tracker"
c_launch_process_tracker :: IO CInt
foreign import ccall unsafe "track_process"
c_track_process :: ProcessTracker -> CPid -> CInt -> IO ()
-- | Represents the child process which handles process cleanup.
--
-- Since 0.2.1
newtype ProcessTracker = ProcessTracker CInt
-- | Represents a child process which is currently being tracked by the cleanup
-- child process.
--
-- Since 0.2.1
data TrackedProcess = TrackedProcess !ProcessTracker !(IORef MaybePid)
data MaybePid = NoPid | Pid !CPid
-- | Fork off the child cleanup process.
--
-- This will ideally only be run once for your entire application.
--
-- Since 0.2.1
initProcessTracker :: IO ProcessTracker
initProcessTracker = do
i <- c_launch_process_tracker
if i == -1
then throwIO CannotLaunchProcessTracker
else return $! ProcessTracker i
-- | Since 0.2.1
data ProcessTrackerException = CannotLaunchProcessTracker
deriving (Show, Typeable)
instance Exception ProcessTrackerException
-- | Begin tracking the given process. If the 'ProcessHandle' refers to a
-- closed process, no tracking will occur. If the process is closed, then it
-- will be untracked automatically.
--
-- Note that you /must/ compile your program with @-threaded@; see
-- 'waitForProcess'.
--
-- Since 0.2.1
trackProcess :: ProcessTracker -> ProcessHandle -> IO TrackedProcess
trackProcess pt ph@(ProcessHandle mph) = mask_ $ do
mpid <- readMVar mph
mpid' <- case mpid of
ClosedHandle{} -> return NoPid
OpenHandle pid -> do
c_track_process pt pid 1
return $ Pid pid
ipid <- newIORef mpid'
let tp = TrackedProcess pt ipid
case mpid' of
NoPid -> return ()
Pid _ -> void $ forkIO $ do
void $ waitForProcess ph
untrackProcess tp
return $! tp
-- | Explicitly remove the given process from the tracked process list in the
-- cleanup process.
--
-- Since 0.2.1
untrackProcess :: TrackedProcess -> IO ()
untrackProcess (TrackedProcess pt ipid) = mask_ $ do
mpid <- readIORef ipid
case mpid of
NoPid -> return ()
Pid pid -> do
c_track_process pt pid 0
writeIORef ipid NoPid
-- | Fork and execute a subprocess, sending stdout and stderr to the specified
-- rotating log.
--
-- Since 0.2.1
forkExecuteLog :: ByteString -- ^ command
-> [ByteString] -- ^ args
-> Maybe [(ByteString, ByteString)] -- ^ environment
-> Maybe ByteString -- ^ working directory
-> Maybe (Source IO ByteString) -- ^ stdin
-> RotatingLog -- ^ both stdout and stderr will be sent to this location
-> IO ProcessHandle
forkExecuteLog cmd args menv mwdir mstdin rlog = bracketOnError
setupPipe
cleanupPipes
usePipes
where
setupPipe = bracketOnError
createPipe
(\(x, y) -> closeFd x `finally` closeFd y)
(\(x, y) -> (,) <$> fdToHandle x <*> fdToHandle y)
cleanupPipes (x, y) = hClose x `finally` hClose y
usePipes pipes@(readerH, writerH) = do
(min, _, _, ph) <- createProcess CreateProcess
{ cmdspec = RawCommand (S8.unpack cmd) (map S8.unpack args)
, cwd = S8.unpack <$> mwdir
, env = map (S8.unpack *** S8.unpack) <$> menv
, std_in = maybe Inherit (const CreatePipe) mstdin
, std_out = UseHandle writerH
, std_err = UseHandle writerH
, close_fds = True
, create_group = True
}
ignoreExceptions $ addAttachMessage pipes ph
void $ forkIO $ ignoreExceptions $
(sourceHandle readerH $$ CL.mapM_ (addChunk rlog)) `finally` hClose readerH
case (min, mstdin) of
(Just h, Just source) -> void $ forkIO $ ignoreExceptions $
(source $$ sinkHandle h) `finally` hClose h
(Nothing, Nothing) -> return ()
_ -> error $ "Invariant violated: Data.Conduit.Process.Unix.forkExecuteLog"
return ph
addAttachMessage pipes ph = withProcessHandle_ ph $ \p_ -> do
now <- getCurrentTime
case p_ of
ClosedHandle ec -> do
addChunk rlog $ S8.concat
[ "\n\n"
, S8.pack $ show now
, ": Process immediately died with exit code "
, S8.pack $ show ec
, "\n\n"
]
cleanupPipes pipes
OpenHandle h -> do
addChunk rlog $ S8.concat
[ "\n\n"
, S8.pack $ show now
, ": Attached new process "
, S8.pack $ show h
, "\n\n"
]
return p_
data Status = NeedsRestart | NoRestart | Running ProcessHandle
-- | Run the given command, restarting if the process dies.
monitorProcess
:: (ByteString -> IO ()) -- ^ log
-> ProcessTracker
-> Maybe S8.ByteString -- ^ setuid
-> S8.ByteString -- ^ executable
-> S8.ByteString -- ^ working directory
-> [S8.ByteString] -- ^ command line parameter
-> [(S8.ByteString, S8.ByteString)] -- ^ environment
-> RotatingLog
-> IO MonitoredProcess
monitorProcess log processTracker msetuid exec dir args env rlog = do
mstatus <- newMVar NeedsRestart
let loop mlast = do
next <- modifyMVar mstatus $ \status ->
case status of
NoRestart -> return (NoRestart, return ())
_ -> do
now <- getCurrentTime
case mlast of
Just last | diffUTCTime now last < 5 -> do
log $ "Process restarting too quickly, waiting before trying again: " `S8.append` exec
threadDelay $ 5 * 1000 * 1000
_ -> return ()
let (cmd, args') =
case msetuid of
Nothing -> (exec, args)
Just setuid -> ("sudo", "-E" : "-u" : setuid : "--" : exec : args)
res <- try $ forkExecuteLog
cmd
args'
(Just env)
(Just dir)
(Just $ return ())
rlog
case res of
Left e -> do
log $ "Data.Conduit.Process.Unix.monitorProcess: " `S8.append` S8.pack (show (e :: SomeException))
return (NeedsRestart, return ())
Right pid -> do
log $ "Process created: " `S8.append` exec
return (Running pid, do
void $ trackProcess processTracker pid
void $ waitForProcess pid
loop (Just now))
next
forkIO $ loop Nothing
return $ MonitoredProcess mstatus
-- | Abstract type containing information on a process which will be restarted.
newtype MonitoredProcess = MonitoredProcess (MVar Status)
-- | Terminate the process and prevent it from being restarted.
terminateMonitoredProcess :: MonitoredProcess -> IO ()
terminateMonitoredProcess (MonitoredProcess mstatus) = do
status <- swapMVar mstatus NoRestart
case status of
Running pid -> do
terminateProcess pid
threadDelay 1000000
killProcess pid
_ -> return ()