repo2docker/repo2docker/buildpacks/repo2docker-entrypoint

#!/usr/bin/env python3
# note: must run on Python >= 3.5, which mainly means no f-strings

# 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("Error getting login env: {e}".format(e=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",
    )

    command = sys.argv[1:]
    r2d_entrypoint = os.environ.get("R2D_ENTRYPOINT")
    if r2d_entrypoint:
        command.insert(0, r2d_entrypoint)

    child = subprocess.Popen(
        command,
        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(
            "Forwarding signal {sig} to {child_pgid}".format(
                sig=sig, child_pgid=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()