Adds ability to only SIGTERM the immediate process

This is useful in certain situations (with multiprocessing apps such as
`gunicorn` or `celery`) where we want to signal to the "main" process
only and let *it* handle gracefully terminating the "worker" processes.

In case the "main" process does not gracefully shutdown the "worker"
processes in the given "timeout" we still make sure to send SIGKILL to
all "(-1)" remaining processes.

Author: denibertovic

README.md

@@ -22,7 +22,7 @@ repo](https://github.com/snoyberg/docker-testing#readme).
 
 ### Usage
 
-> pid1 [-e|--env ENV] [-u|--user USER] [-g|--group GROUP] [-w|--workdir DIR] [-t|--timeout TIMEOUT] COMMAND [ARG1 ARG2 ... ARGN]
+> pid1 [-e|--env ENV] [-u|--user USER] [-g|--group GROUP] [-w|--workdir DIR] [-t|--timeout TIMEOUT] [-s|--single] COMMAND [ARG1 ARG2 ... ARGN]
 
 Where:
 * `-e`, `--env` `ENV` - Override environment variable from given name=value
@@ -33,6 +33,7 @@ Where:
   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
+* `-s`, `--single` - flag if we should only send SIGTERM to the immediate child process
 
 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

@@ -17,7 +17,8 @@ options defaultEnv =
   , 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 ['t'] ["timeout"] (ReqArg (setRunExitTimeoutSec . read) "TIMEOUT") "timeout (in seconds) to wait for all child processes to exit" ]
+  , Option ['t'] ["timeout"] (ReqArg (setRunExitTimeoutSec . read) "TIMEOUT") "timeout (in seconds) to wait for all child processes to exit"
+  , Option ['s'] ["single"] (NoArg (setRunSignalImmediateChildOnly True)) "Flag if we should only send SIGTERM to the immediate child process" ]
   where optEnv env' kv =
           let kvp = fmap (drop 1) $ span (/= '=') kv in
             kvp:filter ((fst kvp /=) . fst) env'

pid1.cabal

@@ -1,5 +1,5 @@
 name:                pid1
-version:             0.1.2.0
+version:             0.1.3.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

@@ -7,10 +7,12 @@ module System.Process.PID1
     , getRunGroup
     , getRunUser
     , getRunWorkDir
+    , getRunSignalImmediateChildOnly
     , run
     , runWithOptions
     , setRunEnv
     , setRunExitTimeoutSec
+    , setRunSignalImmediateChildOnly
     , setRunGroup
     , setRunUser
     , setRunWorkDir
@@ -52,6 +54,7 @@ data RunOptions = RunOptions
     -- timeout (in seconds) to wait for all child processes to exit after
     -- receiving SIGTERM or SIGINT signal
   , runExitTimeoutSec :: Int
+  , runSignalImmediateChildOnly :: Bool
   } deriving Show
 
 -- | return default `RunOptions`
@@ -63,7 +66,8 @@ defaultRunOptions = RunOptions
   , runUser = Nothing
   , runGroup = Nothing
   , runWorkDir = Nothing
-  , runExitTimeoutSec = 5 }
+  , runExitTimeoutSec = 5
+  , runSignalImmediateChildOnly = False }
 
 -- | Get environment variable overrides for the given `RunOptions`
 --
@@ -127,6 +131,19 @@ getRunExitTimeoutSec = runExitTimeoutSec
 setRunExitTimeoutSec :: Int -> RunOptions -> RunOptions
 setRunExitTimeoutSec sec opts = opts { runExitTimeoutSec = sec }
 
+-- | Return boolean flag if we should only send SIGTERM to the immediate child process
+--
+--- @since 0.1.3.0
+getRunSignalImmediateChildOnly :: RunOptions -> Bool
+getRunSignalImmediateChildOnly = runSignalImmediateChildOnly
+
+-- | Set boolean flag if we should only send SIGTERM to the immediate child process
+--
+-- @since 0.1.3.0
+setRunSignalImmediateChildOnly :: Bool -> RunOptions -> RunOptions
+setRunSignalImmediateChildOnly x opts = opts { runSignalImmediateChildOnly = x }
+
+
 -- | Run the given command with specified arguments, with optional environment
 -- variable override (default is to use the current process's environment).
 --
@@ -168,22 +185,20 @@ runWithOptions opts cmd args = do
   for_ (runWorkDir opts) setCurrentDirectory
   let env' = runEnv opts
       timeout = runExitTimeoutSec opts
+      single = runSignalImmediateChildOnly opts
   -- check if we should act as pid1 or just exec the process
   myID <- getProcessID
   if myID == 1
-    then runAsPID1 cmd args env' timeout
+    then runAsPID1 single 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)] -> Int -> IO a
-runAsPID1 cmd args env' timeout = do
+runAsPID1 :: Bool -> FilePath -> [String] -> Maybe [(String, String)] -> Int -> IO a
+runAsPID1 single 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 timeout
 
     -- Helper function to start killing, used below
     let startKilling = void $ tryPutMVar killChildrenVar ()
@@ -206,6 +221,13 @@ runAsPID1 cmd args env' timeout = do
             ClosedHandle e -> assert False (exitWith e)
             OpenHandle pid -> return pid
 
+    _ <- forkIO $ do
+        takeMVar killChildrenVar
+        if single then
+          killImmediateChild child timeout
+        else
+          killAllChildren timeout
+
     -- Loop on reaping child processes
     reap startKilling child
 
@@ -267,6 +289,26 @@ killAllChildren timeout = do
             then return ()
             else throwIO e
 
+killImmediateChild :: CPid -> Int -> IO ()
+killImmediateChild cid timeout = do
+    -- Send immediate child processes the TERM signal
+    signalProcess sigTERM cid `catch` \e ->
+        if isDoesNotExistError e
+            then return ()
+            else throwIO e
+
+    -- 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 $ timeout * 1000 * 1000
+
+    -- OK, some children didn't exit. Now time to get serious!
+    signalProcess sigKILL (-1) `catch` \e ->
+        if isDoesNotExistError e
+            then return ()
+            else throwIO e
+
 -- | Convert a ProcessStatus to an ExitCode. In the case of a signal being the
 -- cause of termination, see 'signalToEC'.
 toExitCode :: ProcessStatus -> ExitCode