Merge pull request #7 from mikemckibben/feature/timeout

Adds support for `-t` child process kill timeout option

ChangeLog.md

@@ -1,3 +1,9 @@
+## 0.1.2.0
+
+* Removes support for ',' separated list of environment variables
+  for `-e` command line option
+* Adds support for setting child processes wait timeout on SIGTERM or SIGINT
+
 ## 0.1.1.0
 
 * Adds support for setuid and setguid when running command

README.md

@@ -22,17 +22,17 @@ repo](https://github.com/snoyberg/docker-testing#readme).
 
 ### Usage
 
-> pid1 [-e|--env LIST][-u|--user USER] [-g|--group GROUP] [-w|--workdir DIR] COMMAND [ARG1 ARG2 ... ARGN] 
+> pid1 [-e|--env ENV] [-u|--user USER] [-g|--group GROUP] [-w|--workdir DIR] [-t|--timeout TIMEOUT] COMMAND [ARG1 ARG2 ... ARGN]
 
 Where:
-* `-e`, `--env` `LIST` - Override environment variables. Comma separated
-  key=value pairs of environment variables to override in the existing
-  environment.
+* `-e`, `--env` `ENV` - Override environment variable from given name=value
+  pair. Can be specified multiple times to set multiple environment variables.
 * `-u`, `--user` `USER` - The username the process will setuid before executing
   COMMAND
 * `-g`, `--group` `GROUP` - The group name the process will setgid before
   executing COMMAND
 * `-w`, `--workdir` `DIR` - chdir to `DIR` before executing COMMAND
+* `-t`, `--timeout` `TIMEOUT` - timeout (in seconds) to wait for all child processes to exit
 
 The recommended use case for this executable is to embed it in a Docker image.
 Assuming you've placed it at `/sbin/pid1`, the two commonly recommended usages

app/Main.hs

@@ -3,7 +3,6 @@ module Main (main) where
 -- entrypoint. It will handle reaping orphans and handling TERM and
 -- INT signals.
 
-import Data.List (foldl')
 import Data.Maybe (fromMaybe)
 import System.Process.PID1
 import System.Environment
@@ -14,18 +13,15 @@ import System.Exit (exitFailure)
 -- | `GetOpt` command line options
 options :: [(String, String)] -> [OptDescr (RunOptions -> RunOptions)]
 options defaultEnv =
-  [ Option ['e'] ["env"] (ReqArg (\opt opts -> setRunEnv (optEnvList (getRunEnv opts) opt) opts) "LIST") "set environment variables from list of comma separated name=value pairs. Can be specified multiple times"
+  [ Option ['e'] ["env"] (ReqArg (\opt opts -> setRunEnv (optEnvList (getRunEnv opts) opt) opts) "ENV") "override environment variable from given name=value pair. Can be specified multiple times to set multiple environment variables"
   , Option ['u'] ["user"] (ReqArg setRunUser "USER") "run command as user"
   , Option ['g'] ["group"] (ReqArg setRunGroup "GROUP") "run command as group"
-  , Option ['w'] ["workdir"] (ReqArg setRunWorkDir "DIR") "command working directory"]
+  , Option ['w'] ["workdir"] (ReqArg setRunWorkDir "DIR") "command working directory"
+  , Option ['t'] ["timeout"] (ReqArg (setRunExitTimeoutSec . read) "TIMEOUT") "timeout (in seconds) to wait for all child processes to exit" ]
   where optEnv env' kv =
           let kvp = fmap (drop 1) $ span (/= '=') kv in
             kvp:filter ((fst kvp /=) . fst) env'
-        split [] = []
-        split s = case fmap (drop 1) $ span (/= ',') s of
-          ("", xs') -> split xs'
-          (x, xs') -> x:split xs'
-        optEnvList env' s = foldl' optEnv (fromMaybe defaultEnv env') $ split s
+        optEnvList = optEnv . fromMaybe defaultEnv
 
 main :: IO ()
 main = do

pid1.cabal

@@ -1,5 +1,5 @@
 name:                pid1
-version:             0.1.1.0
+version:             0.1.2.0
 synopsis:            Do signal handling and orphan reaping for Unix PID1 init processes
 description:         Please see README.md or view Haddocks at <https://www.stackage.org/package/pid1>
 homepage:            https://github.com/fpco/pid1#readme

src/System/Process/PID1.hs

@@ -3,12 +3,14 @@ module System.Process.PID1
     ( RunOptions
     , defaultRunOptions
     , getRunEnv
+    , getRunExitTimeoutSec
     , getRunGroup
     , getRunUser
     , getRunWorkDir
     , run
     , runWithOptions
     , setRunEnv
+    , setRunExitTimeoutSec
     , setRunGroup
     , setRunUser
     , setRunWorkDir
@@ -47,13 +49,21 @@ data RunOptions = RunOptions
   , runGroup :: Maybe String
     -- optional working directory
   , runWorkDir :: Maybe FilePath
+    -- timeout (in seconds) to wait for all child processes to exit after
+    -- receiving SIGTERM or SIGINT signal
+  , runExitTimeoutSec :: Int
   } deriving Show
 
 -- | return default `RunOptions`
 --
 -- @since 0.1.1.0
 defaultRunOptions :: RunOptions
-defaultRunOptions = RunOptions { runEnv = Nothing, runUser = Nothing, runGroup = Nothing, runWorkDir = Nothing }
+defaultRunOptions = RunOptions
+  { runEnv = Nothing
+  , runUser = Nothing
+  , runGroup = Nothing
+  , runWorkDir = Nothing
+  , runExitTimeoutSec = 5 }
 
 -- | Get environment variable overrides for the given `RunOptions`
 --
@@ -103,6 +113,20 @@ getRunWorkDir = runWorkDir
 setRunWorkDir :: FilePath -> RunOptions -> RunOptions
 setRunWorkDir dir opts = opts { runWorkDir = Just dir }
 
+-- | Return the timeout (in seconds) timeout (in seconds) to wait for all child
+-- processes to exit after receiving SIGTERM or SIGINT signal
+--
+-- @since 0.1.2.0
+getRunExitTimeoutSec :: RunOptions -> Int
+getRunExitTimeoutSec = runExitTimeoutSec
+
+-- | Set the timeout in seconds for the process reaper to wait for all child
+-- processes to exit after receiving SIGTERM or SIGINT signal
+--
+-- @since 0.1.2.0
+setRunExitTimeoutSec :: Int -> RunOptions -> RunOptions
+setRunExitTimeoutSec sec opts = opts { runExitTimeoutSec = sec }
+
 -- | Run the given command with specified arguments, with optional environment
 -- variable override (default is to use the current process's environment).
 --
@@ -143,22 +167,23 @@ runWithOptions opts cmd args = do
     setUserID $ userID entry
   for_ (runWorkDir opts) setCurrentDirectory
   let env' = runEnv opts
+      timeout = runExitTimeoutSec opts
   -- check if we should act as pid1 or just exec the process
   myID <- getProcessID
   if myID == 1
-    then runAsPID1 cmd args env'
+    then runAsPID1 cmd args env' timeout
     else executeFile cmd True args env'
 
 -- | Run as a child with signal handling and orphan reaping.
-runAsPID1 :: FilePath -> [String] -> Maybe [(String, String)] -> IO a
-runAsPID1 cmd args env' = do
+runAsPID1 :: FilePath -> [String] -> Maybe [(String, String)] -> Int -> IO a
+runAsPID1 cmd args env' timeout = do
     -- Set up an MVar to indicate we're ready to start killing all
     -- children processes. Then start a thread waiting for that
     -- variable to be filled and do the actual killing.
     killChildrenVar <- newEmptyMVar
     _ <- forkIO $ do
         takeMVar killChildrenVar
-        killAllChildren
+        killAllChildren timeout
 
     -- Helper function to start killing, used below
     let startKilling = void $ tryPutMVar killChildrenVar ()
@@ -222,19 +247,19 @@ reap startKilling child = do
                     startKilling
                 | otherwise -> return ()
 
-killAllChildren :: IO ()
-killAllChildren = do
+killAllChildren :: Int -> IO ()
+killAllChildren timeout = do
     -- Send all children processes the TERM signal
     signalProcess sigTERM (-1) `catch` \e ->
         if isDoesNotExistError e
             then return ()
             else throwIO e
 
-    -- Wait for five seconds. We don't need to put in any logic about
+    -- Wait for `timeout` seconds. We don't need to put in any logic about
     -- whether there are still child processes; if all children have
     -- exited, then the reap loop will exit and our process will shut
     -- down.
-    threadDelay $ 5 * 1000 * 1000
+    threadDelay $ timeout * 1000 * 1000
 
     -- OK, some children didn't exit. Now time to get serious!
     signalProcess sigKILL (-1) `catch` \e ->