Source

udon / src / Process.hs

Full commit
module Process
    ( runCommand
    , runCommandWithFd
    , system
    , systemWithInput
    , forkDaemon
    , waitChild
    , waitChild'
    , killProcess
    , parseSignal

    -- re-export
    , ProcessStatus
    , ProcessID
    ) where

import Control.Concurrent (threadDelay, forkIO)
import Data.Char (isDigit)
import Data.Set (member, fromList)
import System.Exit(ExitCode(ExitSuccess))
import System.IO (Handle, hSetBuffering, BufferMode(NoBuffering, LineBuffering), hPutStr, hClose, hFlush)
import System.IO.Error (try)
import System.Posix.Directory (changeWorkingDirectory, getWorkingDirectory)
import System.Posix.Env (getEnvironment)
import System.Posix.IO (stdInput, stdOutput, stdError, closeFd, createPipe, dupTo, openFd, OpenMode(ReadWrite), defaultFileFlags, fdToHandle)
import System.Posix.Process (forkProcess, executeFile, createSession, ProcessStatus, getProcessStatus, exitImmediately)
import System.Posix.Signals (signalProcess, Signal, unblockSignals, fullSignalSet)
import System.Posix.Types (Fd, ProcessID)
import System.Posix.Unistd (getSysVar, SysVar(OpenFileLimit))

import qualified System.Posix.Signals as S

runCommand :: String                 -- ^ Command name
           -> [String]               -- ^ Command arguments
           -> String                 -- ^ Current working directory
           -> [(String, String)]     -- ^ Environmental variables
           -> IO (Fd, Fd, ProcessID) -- ^ (stdout, stderr, Child's process ID)
runCommand cmd args cwd envs = do
    (stdout_r, stdout_w) <- createPipe
    (stderr_r, stderr_w) <- createPipe
    pid <- runCommandWithFd cmd args cwd envs Nothing (Just stdout_w) (Just stderr_w)
    closeFd stdout_w
    closeFd stderr_w
    return (stdout_r, stderr_r, pid)

runCommandWithFd :: String -> [String] -> String -> [(String, String)] -> Maybe Fd -> Maybe Fd -> Maybe Fd -> IO ProcessID
runCommandWithFd cmd args cwd envs mstdin mstdout mstderr = do
    devNull <- openFd "/dev/null" ReadWrite Nothing defaultFileFlags
    let stdin  = maybe devNull id mstdin
        stdout = maybe devNull id mstdout
        stderr = maybe devNull id mstderr
    pid <- forkProcess $ execCommand cmd args cwd envs stdin stdout stderr
    closeFd devNull
    return pid

system :: String           -- ^ Command name
       -> [String]         -- ^ Command arguments
       -> IO ProcessStatus
system cmd args = do
    cwd <- getWorkingDirectory
    envs <- getEnvironment
    pid <- runCommandWithFd cmd args cwd envs (Just stdInput) (Just stdOutput) (Just stdError)
    Just ps <- waitChild pid
    return ps

systemWithInput :: String           -- ^ Command name
                -> [String]         -- ^ Command arguments
                -> String           -- ^ Input data
                -> IO ProcessStatus -- ^ (stdout, stderr, Child's process ID)
systemWithInput cmd args inp = do
    cwd <- getWorkingDirectory
    envs <- getEnvironment
    (stdin_r, stdin_w)  <- createPipe
    pid <- runCommandWithFd cmd args cwd envs (Just stdin_r) (Just stdOutput) (Just stdError)
    closeFd stdin_r
    forkIO $ do
        h <- fdToHandle stdin_w
        hPutStr h inp
        hFlush h
        hClose h
    Just ps <- waitChild pid
    return ps

execCommand :: String -> [String] -> String -> [(String, String)] -> Fd -> Fd -> Fd -> IO ()
execCommand cmd args cwd envs stdin stdout stderr = do
    dupTo stdin stdInput

    dupTo stdout stdOutput
    dupTo stderr stdError

    closeNotStdFds
    createSession
    changeWorkingDirectory cwd

    unblockSignals fullSignalSet

    executeFile cmd True args (Just envs)
    return ()

-- | Run IO action as daemon
forkDaemon :: IO () -> IO ()
forkDaemon action = do
    pid <- forkProcess $ do
        createSession
        forkProcess $ do
            changeWorkingDirectory "/"
            devNull <- openFd "/dev/null" ReadWrite Nothing defaultFileFlags
            dupTo devNull stdInput
            dupTo devNull stdOutput
            dupTo devNull stdError
            closeNotStdFds

            action

        exitImmediately ExitSuccess
    getProcessStatus True False pid
    return ()

-- | close All Fds greater than 3
closeNotStdFds :: IO ()
closeNotStdFds = do
    n <- getSysVar OpenFileLimit
    mapM_ forceCloseFd [3..n]
  where
    forceCloseFd i = try (closeFd $ fromInteger i) >> return ()

waitChild :: ProcessID -> IO (Maybe ProcessStatus)
waitChild pid = do
    m <- getProcessStatus False False pid
    case m of
        Just ps -> return m
        Nothing -> waitChild_ 0.05 pid

waitChild_ :: Double -> ProcessID -> IO (Maybe ProcessStatus)
waitChild_ wait pid = do
    threadDelay (truncate $ wait * 1000)
    m <- getProcessStatus False False pid
    case m of
        Just ps -> return m
        Nothing -> waitChild_ (min (wait * 2) 1.0) pid

-- | wait child process
-- caution! this function blocks entire thread
waitChild' :: ProcessID -> IO (Maybe ProcessStatus)
waitChild' pid = getProcessStatus True False pid

killProcess :: Signal -> ProcessID -> IO ()
killProcess = signalProcess

parseSignal :: String -> Maybe Signal
parseSignal cs | all isDigit cs = let sig = read cs
                                  in if sig `member` knownSignals
                                       then Just sig
                                       else Nothing
parseSignal ('S':'I':'G':cs) = parseSignal cs
parseSignal "ABRT"  = Just S.sigABRT
parseSignal "FPE"   = Just S.sigFPE
parseSignal "PIPE"  = Just S.sigPIPE
parseSignal "STOP"  = Just S.sigSTOP
parseSignal "TTIN"  = Just S.sigTTIN
parseSignal "VTALR" = Just S.sigVTALRM
parseSignal "ALRM"  = Just S.sigALRM
parseSignal "HUP"   = Just S.sigHUP
parseSignal "POLL"  = Just S.sigPOLL
parseSignal "SYS"   = Just S.sigSYS
parseSignal "TTOU"  = Just S.sigTTOU
parseSignal "XCPU"  = Just S.sigXCPU
parseSignal "BUS"   = Just S.sigBUS
parseSignal "ILL"   = Just S.sigILL
parseSignal "PROF"  = Just S.sigPROF
parseSignal "TERM"  = Just S.sigTERM
parseSignal "URG"   = Just S.sigURG
parseSignal "XFSZ"  = Just S.sigXFSZ
parseSignal "CHLD"  = Just S.sigCHLD
parseSignal "INT"   = Just S.sigINT
parseSignal "QUIT"  = Just S.sigQUIT
parseSignal "TRAP"  = Just S.sigTRAP
parseSignal "USR1"  = Just S.sigUSR1
parseSignal "CONT"  = Just S.sigCONT
parseSignal "KILL"  = Just S.sigKILL
parseSignal "SEGV"  = Just S.sigSEGV
parseSignal "TSTP"  = Just S.sigTSTP
parseSignal "USR2"  = Just S.sigUSR2
parseSignal _       = Nothing


knownSignals = fromList $
    [ S.sigABRT
    , S.sigFPE
    , S.sigPIPE
    , S.sigSTOP
    , S.sigTTIN
    , S.sigVTALRM
    , S.sigALRM
    , S.sigHUP
    , S.sigPOLL
    , S.sigSYS
    , S.sigTTOU
    , S.sigXCPU
    , S.sigBUS
    , S.sigILL
    , S.sigPROF
    , S.sigTERM
    , S.sigURG
    , S.sigXFSZ
    , S.sigCHLD
    , S.sigINT
    , S.sigQUIT
    , S.sigTRAP
    , S.sigUSR1
    , S.sigCONT
    , S.sigKILL
    , S.sigSEGV
    , S.sigTSTP
    , S.sigUSR2
    ]