implement entrypoint in Python

- make sure to set PYTHONUNBUFFERED
- implement tee with subprocess readlines
- forward all relevant signals from entrypoint to 'true' command
- add monitor process to ensure child shuts down if parent does (maybe not useful in docker?)

Author: minrk

repo2docker/buildpacks/base.py

@@ -181,6 +181,7 @@
 {% endif -%}
 
 # Add entrypoint
+ENV PYTHONUNBUFFERED=1
 COPY /repo2docker-entrypoint /usr/local/bin/repo2docker-entrypoint
 ENTRYPOINT ["/usr/local/bin/repo2docker-entrypoint"]
 

repo2docker/buildpacks/repo2docker-entrypoint

@@ -1,24 +1,150 @@
-#!/bin/bash -l
-# lightest possible entrypoint that ensures that
-# we use a login shell to get a fully configured shell environment
-# (e.g. sourcing /etc/profile.d, ~/.bashrc, and friends)
-
-# Setup a file descriptor (FD) that is connected to a tee process which
-# writes its input to $REPO_DIR/.jupyter-server-log.txt
-# We later use this FD as a place to redirect the output of the actual
-# command to. We can't add `tee` to the command directly as that will prevent
-# the container from exiting when `docker stop` is run.
-# See https://stackoverflow.com/a/55678435
-exec {log_fd}> >(exec tee $REPO_DIR/.jupyter-server-log.txt)
-
-if [[ ! -z "${R2D_ENTRYPOINT:-}" ]]; then
-    if [[ ! -x "$R2D_ENTRYPOINT" ]]; then
-        chmod u+x "$R2D_ENTRYPOINT"
-    fi
-    exec "$R2D_ENTRYPOINT" "$@" 2>&1 >&"$log_fd"
-else
-    exec "$@" 2>&1 >&"$log_fd"
-fi
-
-# Close the logging output again
-exec {log_fd}>&-
+#!/usr/bin/env python3
+
+# goals:
+# - load environment variables from a login shell (bash -l)
+# - preserve signal handling of subprocess (kill -TERM and friends)
+# - tee output to a log file
+
+import json
+import os
+import signal
+import subprocess
+import sys
+import time
+
+
+def get_login_env():
+    """Instantiate a login shell to retrieve environment variables
+
+    Serialize with Python to ensure proper escapes
+    """
+    p = subprocess.run(
+        [
+            "bash",
+            "-l",
+            "-c",
+            "python3 -c 'import os, json; print(json.dumps(dict(os.environ)))'",
+        ],
+        stdout=subprocess.PIPE,
+    )
+    if p.returncode:
+        print("Error getting login env")
+        return {}
+
+    last_line = p.stdout.splitlines()[-1]
+    try:
+        return json.loads(last_line)
+    except Exception as e:
+        print(f"Error getting login env: {e}", file=sys.stderr)
+        return {}
+
+
+def monitor_parent(parent_pid, child_pgid):
+    """Monitor parent_pid and shutdown child_pgid if parent goes away first"""
+    while True:
+        try:
+            os.kill(parent_pid, 0)
+        except ProcessLookupError:
+            # parent is gone, likely by SIGKILL
+            # send SIGKILL to child process group
+            try:
+                os.killpg(child_pgid, signal.SIGKILL)
+            except (ProcessLookupError, PermissionError):
+                # ignore if the child is already gone
+                pass
+            return
+        else:
+            time.sleep(1)
+
+
+# signals to be forwarded to the child
+SIGNALS = [
+    signal.SIGHUP,
+    signal.SIGINT,
+    # signal.SIGKILL,
+    signal.SIGQUIT,
+    signal.SIGTERM,
+    signal.SIGUSR1,
+    signal.SIGUSR2,
+    signal.SIGWINCH,
+]
+
+
+def main():
+
+    # load login shell environment
+    login_env = get_login_env()
+    env = os.environ.copy()
+    env.update(login_env)
+
+    # open log file to send output
+    log_file = open(
+        os.path.join(os.environ.get("REPO_DIR", "."), ".jupyter-server-log.txt"),
+        "a",
+    )
+
+    child = subprocess.Popen(
+        sys.argv[1:],
+        bufsize=1,
+        env=env,
+        start_new_session=True,
+        stdout=subprocess.PIPE,
+        stderr=subprocess.STDOUT,
+        universal_newlines=True,
+    )
+    child_pgid = os.getpgid(child.pid)
+
+    # if parent is forcefully shutdown,
+    # make sure child shuts down immediately as well
+    parent_pid = os.getpid()
+
+    monitor_pid = os.fork()
+    if monitor_pid == 0:
+        # child process, sibling of 'real' command
+        # avoid receiving signals sent to parent
+        os.setpgrp()
+        # terminate child if parent goes away,
+        # e.g. in ungraceful KILL not relayed to children
+        monitor_parent(parent_pid, child_pgid)
+        return
+
+    # hook up ~all signals so that every signal the parent gets,
+    # the children also get
+
+    def relay_signal(sig, frame):
+        """Relay a signal to children"""
+        print(f"Forwarding signal {sig} to {child_pgid}")
+        os.killpg(child_pgid, sig)
+
+    # question: maybe use all valid_signals() except a few, e.g. SIGCHLD?
+    # rather than opt-in list
+    for signum in SIGNALS:
+        signal.signal(signum, relay_signal)
+
+    # tee output from child to both our stdout and the log file
+    def tee(chunk):
+        for f in [sys.stdout, log_file]:
+            f.write(chunk)
+            f.flush()
+
+    while child.poll() is None:
+        tee(child.stdout.readline())
+
+    # flush the rest
+    chunk = child.stdout.read()
+    while chunk:
+        tee(chunk)
+        chunk = child.stdout.read()
+
+    # child exited, cleanup monitor
+    try:
+        os.kill(monitor_pid, signal.SIGKILL)
+    except ProcessLookupError:
+        pass
+
+    # preserve returncode
+    sys.exit(child.returncode)
+
+
+if __name__ == "__main__":
+    main()

tests/unit/test_env.py

@@ -3,12 +3,13 @@
 """
 import os
 import subprocess
+import sys
 import tempfile
 import time
 from getpass import getuser
 
 
-def test_env():
+def test_env(capfd):
     """
     Validate that you can define environment variables
 
@@ -42,32 +43,28 @@ def test_env():
                 # value
                 "--env",
                 "SPAM_2=",
-                # "--",
                 tmpdir,
+                "--",
                 "/bin/bash",
                 "-c",
                 # Docker exports all passed env variables, so we can
                 # just look at exported variables.
-                "export; sleep 1",
-                # "export; echo TIMDONE",
-                # "export",
+                "export",
             ],
-            universal_newlines=True,
-            stdout=subprocess.PIPE,
-            stderr=subprocess.PIPE,
         )
+    captured = capfd.readouterr()
+    print(captured.out, end="")
+    print(captured.err, file=sys.stderr, end="")
+
     assert result.returncode == 0
 
     # all docker output is returned by repo2docker on stderr
     # extract just the declare for better failure message formatting
     # stdout should be empty
     assert not result.stdout
 
-    print(result.stderr.split("\n"))
-    # assert False
-
     # stderr should contain lines of output
-    declares = [x for x in result.stderr.split("\n") if x.startswith("declare")]
+    declares = [x for x in captured.err.splitlines() if x.startswith("declare")]
     assert 'declare -x FOO="{}"'.format(ts) in declares
     assert 'declare -x BAR="baz"' in declares
     assert 'declare -x SPAM="eggs"' in declares