Merge pull request #6 from chriskuehl/process-groups

Add process group support

Dockerfile

@@ -3,7 +3,7 @@ FROM debian:jessie
 MAINTAINER Chris Kuehl <ckuehl@yelp.com>
 RUN DEBIAN_FRONTEND=noninteractive apt-get update && \
     apt-get install -y --no-install-recommends \
-        build-essential devscripts equivs && \
+        build-essential devscripts equivs procps psmisc && \
     apt-get clean
 WORKDIR /mnt
 

Makefile

@@ -1,7 +1,9 @@
+CFLAGS=-std=gnu99 -static -Wall -Werror
+
 DOCKER_RUN_TEST := docker run -v $(PWD):/mnt:ro
 DOCKER_DEB_TEST := sh -euxc ' \
 	apt-get update \
-	&& apt-get install -y --no-install-recommends procps \
+	&& apt-get install -y --no-install-recommends procps psmisc \
 	&& (which timeout || apt-get install -y --no-install-recommends timeout) \
 	&& dpkg -i /mnt/dist/*.deb \
 	&& cd /mnt \
@@ -9,7 +11,7 @@ DOCKER_DEB_TEST := sh -euxc ' \
 '
 DOCKER_PYTHON_TEST := sh -uexc ' \
 	apt-get update \
-	&& apt-get install -y --no-install-recommends python-pip build-essential procps \
+	&& apt-get install -y --no-install-recommends python-pip build-essential procps psmisc \
 	&& (which timeout || apt-get install -y --no-install-recommends timeout) \
 	&& tmp=$$(mktemp -d) \
 	&& cp -r /mnt/* "$$tmp" \
@@ -22,7 +24,7 @@ DOCKER_PYTHON_TEST := sh -uexc ' \
 
 .PHONY: build
 build:
-	$(CC) -static -Wall -Werror -o dumb-init dumb-init.c
+	$(CC) $(CFLAGS) -o dumb-init dumb-init.c
 
 .PHONY: clean
 clean: clean-tox

dumb-init.c

@@ -5,9 +5,10 @@
  * Usage:
  *   ./dumb-init python -c 'while True: pass'
  *
- * To get debug output on stderr, run with DUMB_INIT_DEBUG=1.
+ * To get debug output on stderr, run with DUMB_INIT_DEBUG=1
  */
 
+#include <errno.h>
 #include <signal.h>
 #include <stdio.h>
 #include <stdlib.h>
@@ -16,18 +17,24 @@
 #include <sys/wait.h>
 #include <unistd.h>
 
+#define DEBUG(...) do { \
+    if (debug) { \
+        fprintf(stderr, __VA_ARGS__); \
+    } \
+} while (0)
+
 pid_t child = -1;
 char debug = 0;
+char use_process_group = 1;
 
 void signal_handler(int signum) {
-    if (debug)
+    DEBUG("Received signal %d.\n", signum);
-        fprintf(stderr, "Received signal %d.\n", signum);
 
     if (child > 0) {
-        kill(child, signum);
+        kill(use_process_group ? -child : child, signum);
-
+        DEBUG("Forwarded signal to child.\n");
-        if (debug)
+    } else {
-            fprintf(stderr, "Forwarded signal to child.\n");
+        DEBUG("Didn't forward signal, no child exists yet.");
     }
 }
 
@@ -49,14 +56,18 @@ void print_help(char *argv[]) {
         "\n"
         "The proxy dies when your process dies, so it must not double-fork or do other\n"
         "weird things (this is basically a requirement for doing things sanely in\n"
-        "Docker anyway).\n",
+        "Docker anyway).\n"
+        "\n"
+        "By default, dumb-init starts a process group and kills all processes in it.\n"
+        "This is usually useful behavior, but if for some reason you wish to disable\n"
+        "it, run with DUMB_INIT_PROCESS_GROUP=0.\n",
         argv[0]
     );
 }
 
 int main(int argc, char *argv[]) {
     int signum, exit_status, status = 0;
-    char *debug_env;
+    char *debug_env, *pgroup_env;
 
     if (argc < 2) {
         print_help(argv);
@@ -64,9 +75,16 @@ int main(int argc, char *argv[]) {
     }
 
     debug_env = getenv("DUMB_INIT_DEBUG");
-    if (debug_env && strcmp(debug_env, "1") == 0)
+    if (debug_env && strcmp(debug_env, "1") == 0) {
         debug = 1;
+        DEBUG("Running in debug mode.\n");
+    }
 
+    pgroup_env = getenv("DUMB_INIT_PROCESS_GROUP");
+    if (pgroup_env && strcmp(pgroup_env, "0") == 0) {
+        use_process_group = 0;
+        DEBUG("Not running in process group mode.\n");
+    }
 
     /* register signal handlers */
     for (signum = 1; signum < 32; signum++) {
@@ -88,16 +106,29 @@ int main(int argc, char *argv[]) {
     }
 
     if (child == 0) {
+        if (use_process_group) {
+            pid_t result = setpgid(0, 0);
+            if (result != 0) {
+                fprintf(
+                    stderr,
+                    "Unable to create process group (errno=%d %s). Exiting.\n",
+                    errno,
+                    strerror(errno)
+                );
+                exit(1);
+            }
+            DEBUG("Set process group ID of child to its own PID.\n");
+        }
+
         execvp(argv[1], &argv[1]);
     } else {
-        if (debug)
+        DEBUG("Child spawned with PID %d.\n", child);
-            fprintf(stderr, "Child spawned with PID %d.\n", child);
 
+        /* wait for child to exit */
         waitpid(child, &status, 0);
         exit_status = WEXITSTATUS(status);
 
-        if (debug)
+        DEBUG("Child exited with status %d. Goodbye.\n", exit_status);
-            fprintf(stderr, "Child exited with status %d, goodbye.\n", exit_status);
 
         return exit_status;
     }

test

@@ -11,9 +11,15 @@ fi
 echo "Running with dumb-init at '$dumb_init_bin'"
 
 run_tests() {
+    export DUMB_INIT_PROCESS_GROUP
+    for DUMB_INIT_PROCESS_GROUP in 0 1; do
         ./test-proxies-signals "$dumb_init_bin"
         ./test-exit-status "$dumb_init_bin"
         ./test-help-message "$dumb_init_bin"
+    done
+
+    DUMB_INIT_PROCESS_GROUP=0 ./test-pgroup "$dumb_init_bin" 4
+    DUMB_INIT_PROCESS_GROUP=1 ./test-pgroup "$dumb_init_bin" 0
 }
 
 cd tests

tests/lib/print-signals

@@ -1,4 +1,8 @@
-#!/bin/bash -eux
+#!/bin/sh -eux
+# XXX: We use /bin/sh instead of /bin/bash since some old versions of bash
+# exhibit an issue where they seem to receive the same signal twice.
+# With /bin/sh, this does not seem to happen.
+
 # Print received signals into a file, one per line
 file="$1"
 

tests/test-help-message

@@ -6,15 +6,15 @@ dumb_init="$1"
 
 status=$($dumb_init > /dev/null 2>&1; echo $?)
 
-if [ "$status" -ne 0 ]; then
+if [ "$status" -eq 0 ]; then
-    msg=$($dumb_init 2>&1 || true)
-    msg_len=${#msg}
-
-    if [ "$msg_len" -le 50 ]; then
-        echo "Error: Expected dumb-init with no arguments to print a useful message, but it was only ${msg_len} chars long."
-        exit 1
-    fi
-else
     echo "Error: Expected dumb-init with no arguments to return nonzero, but it returned ${status}."
     exit 1
 fi
+
+msg=$($dumb_init 2>&1 || true)
+msg_len=${#msg}
+
+if [ "$msg_len" -le 50 ]; then
+    echo "Error: Expected dumb-init with no arguments to print a useful message, but it was only ${msg_len} chars long."
+    exit 1
+fi

tests/test-pgroup

@@ -0,0 +1,33 @@
+#!/bin/bash -eux
+# dumb-init should proxy signals to a process group rooted at its child when
+# requested.
+dumb_init="$1"
+after_count="$2"
+
+$dumb_init sh -c "yes 'oh, hi' | tail & yes error | tail >&2" &
+pid="$!"
+
+sleep 1
+pstree -p "$pid"
+pids=$(pstree -p "$pid" | grep -Po '(\d+)' | grep -Po '\d+')
+
+# ensure processes are running
+child_count=$(ps -o pid= $pids | wc -l) || true
+
+if [ "$child_count" -ne 6 ]; then
+    echo "Error: Expected 6 children, instead we had ${child_count}."
+    exit 1
+fi
+
+# ensure processes are dead after signal
+kill -TERM "$pid"
+sleep 1
+child_count=$(ps -o pid= $pids | wc -l) || true
+
+if [ "$child_count" -ne "$after_count" ]; then
+    echo "Error: Expected $after_count children, instead we had ${child_count}."
+    exit 1
+fi
+
+echo 'Killing any leftover processes.'
+xargs kill -9 <<< "$pids" || true