trc

#!/bin/bash

# TrivialRC
# The minimalistic RC system and process manager for containers and applications
# Copyright (c) 2016, 2017 by vorakl

trc_version="v1.1.12"

set +e # Do not exit on errors by default
if [[ "${RC_DEBUG}" = "true" ]]; then
    set -x # Turns on Debug mode
fi

main() {
    check_ver_usage "$1"
    check_req_bins ls ps pkill sleep date basename dirname

    # It's possible to exit from the last `halt` command with ${_exit_status} which comes from main() 
    # and rewrite by this an exit status of the `halt` stage which always has a priority against main's exit status
    trap '_exit_status=$?; trap '_ignore_mt_ex=1' EXIT; debug "exit-trap (exitcode=${_exit_status})"; hook_main_exit' EXIT
    trap '_es_mt_s=$?; trap '_ignore_mt_s=1' HUP INT ILL TRAP ABRT BUS FPE USR1 SEGV USR2 PIPE ALRM TERM XCPU XFSZ VTALRM SYS; debug "sig-trap (exitcode=${_es_mt_s})"; exit ${_es_mt_s}' HUP INT ILL TRAP ABRT BUS FPE USR1 SEGV USR2 PIPE ALRM TERM XCPU XFSZ VTALRM SYS
    trap '_es_mt_er=$?; trap '_ignore_mt_er=1' ERR; if [[ "${RC_WAIT_POLICY}" = "wait_err" ]]; then debug "err-trap (exitcode=${_es_mt_er})"; exit ${_es_mt_er}; fi' ERR

    local _wait_bg_cmd _file _es_m_f=0 _es_m_d=0

    export SELF_NAME=$(run -p basename $0 .sh) # Self name for logging purpose
    export MAINPID="${BASHPID}" # PID of the main process
    : ${RC_WAIT_POLICY:=wait_any}; export RC_WAIT_POLICY 

    work_dir="${PWD}"
    boot_stage_left=0
    halt_cmds_exist=0
    halt_cmds=()
    child_pids="" # Collects of child processes
    ns="main" # Name Space

    case "${RC_WAIT_POLICY}" in
        wait_all)  _wait_bg_cmd="run wait";;
        wait_any)  _wait_bg_cmd="run wait";;
        wait_err)  _wait_bg_cmd="run wait";;
        *)         _wait_bg_cmd=":";;
    esac

    log "The wait policy: ${RC_WAIT_POLICY}"

    if [[ "$1" = "-w" ]] || [[ "$1" = "--workdir" ]]; then
        shift
        work_dir="$1"
        shift
    fi
    if [[ -n "${RC_WORK_DIR}" ]]; then
        work_dir="${RC_WORK_DIR}"
    else
        RC_WORK_DIR="${work_dir}"
        export RC_WORK_DIR
    fi
    log "Work dir: ${RC_WORK_DIR}"

    debug "Looking for \`boot\` scripts and commands..."
    {
        # Reads commands from files to run before everything
        for _file in $(run -p ls ${work_dir}/trc.boot.* ${work_dir}/trc.d/boot.* 2>/dev/null || true); do
            log "Launching in the $(ns_long ${ns}): ${_file}"
            set -e 
            . ${_file}
            set +e
        done

        # Checks for boot tasks in the command line (all commands and their parameters should be supplied as one parameter)
        while [[ "$1" = "-B" ]]; do
            shift
            log "Launching in the $(ns_long ${ns}): $1"
            set -e 
            eval "$1"
            set +e
            shift
        done
    }
    boot_stage_left=1

    debug "Looking for \`halt\` commands..."
    # Collects commands for the `halt` stage from a command line to be run on exit from the main exit trap
    while [[ "$1" = "-H" ]]; do
        shift
            halt_cmds_exist=1
            halt_cmds[${#halt_cmds[*]}]="$1"
        shift
    done

    debug "Looking for \`async\` scripts and commands..."
    # Reads commands from files to run in the background (in parallel)
    for _file in $(run -p ls ${work_dir}/trc.async.* ${work_dir}/trc.d/async.* 2>/dev/null || true); do
        (
            # Run this on any exit, catching the exitcode of the "main" command,
            # printing additional info and finishing up this sub-program with the right exitcode
            trap '_es_dt_ex=$?; trap '_ignore_mt_ex=1' EXIT; debug "exit-trap (exitcode=${_es_dt_ex})"; hook_sub_exit ${_es_dt_ex} "${_file}"' EXIT

            # In case of exit on errors (set -e) or by a signal, catch exitcode and exit with it, which
            # will lead to triggering an EXIT trap
            trap '_es_dt_s=$?; trap '_ignore_dt_s=1' HUP INT ILL TRAP ABRT BUS FPE USR1 SEGV USR2 PIPE ALRM TERM XCPU XFSZ VTALRM SYS; debug "sig-trap (exitcode=$_es_dt_s)"; exit ${_es_dt_s}' HUP INT ILL TRAP ABRT BUS FPE USR1 SEGV USR2 PIPE ALRM TERM XCPU XFSZ VTALRM SYS
            trap '_es_dt_er=$?; trap '_ignore_dt_er=1' ERR; debug "err-trap (exitcode=$_es_dt_er)"; exit ${_es_dt_er}' ERR

            ns="async"
            log "Launching in the $(ns_long ${ns}): ${_file}"
            . ${_file}
        )& 
        child_pids="${child_pids} $!"
    done

    # Checks for background tasks in the command line (all commands and their parameters should be supplied as one parameter)
    while [[ "$1" = "-D" ]]; do
        shift
        (
            trap '_es_dt_ex=$?; trap '_ignore_dt_ex=1' EXIT; debug "exit-trap (exitcode=${_es_dt_ex})"; hook_sub_exit ${_es_dt_ex} "$1"' EXIT
            trap '_es_dt_s=$?; trap '_ignore_dt_s=1' HUP INT ILL TRAP ABRT BUS FPE USR1 SEGV USR2 PIPE ALRM TERM XCPU XFSZ VTALRM SYS; debug "sig-trap (exitcode=$_es_dt_s)"; exit ${_es_dt_s}' HUP INT ILL TRAP ABRT BUS FPE USR1 SEGV USR2 PIPE ALRM TERM XCPU XFSZ VTALRM SYS
            trap '_es_dt_er=$?; trap '_ignore_dt_er=1' ERR; debug "err-trap (exitcode=$_es_dt_er)"; exit ${_es_dt_er}' ERR

            ns="async"
            log "Launching in the $(ns_long ${ns}): $1"
            eval "$1"
        )& 
        child_pids="${child_pids} $!"
        shift
    done

    debug "Looking for \`sync\` scripts and commands..." 
    # Checks for foreground tasks in files (sequentially)
    for _file in $(run -p ls ${work_dir}/trc.sync.* ${work_dir}/trc.d/sync.* 2>/dev/null || true); do
        (
            trap '_es_ft_ex=$?; trap '_ignore_ft_ex=1' EXIT; debug "exit-trap (exitcode=${_es_ft_ex})"; hook_sub_exit ${_es_ft_ex} "${_file}"' EXIT
            trap '_es_ft_s=$?; trap '_ignore_ft_s=1' HUP INT ILL TRAP ABRT BUS FPE USR1 SEGV USR2 PIPE ALRM TERM XCPU XFSZ VTALRM SYS; debug "sig-trap (exitcode=$_es_ft_s)"; exit ${_es_ft_s}' HUP INT ILL TRAP ABRT BUS FPE USR1 SEGV USR2 PIPE ALRM TERM XCPU XFSZ VTALRM SYS
            trap '_es_ft_er=$?; trap '_ignore_ft_er=1' ERR; debug "err-trap (exitcode=$_es_ft_er)"; exit ${_es_ft_er}' ERR

            ns="sync"
            log "Launching in the $(ns_long ${ns}): ${_file}"
            . ${_file}
        )
        # Catch the exitcode of a foreground sub-program
        _es_m_f=$?
    done

    # Checks for foreground tasks in the command line (all commands and their parameters should be supplied as one parameter)
    while [[ "$1" = "-F" ]]; do
        shift
        (
            trap '_es_ft_ex=$?; trap '_ignore_ft_ex=1' EXIT; debug "exit-trap (exitcode=${_es_ft_ex})"; hook_sub_exit ${_es_ft_ex} "$1"' EXIT
            trap '_es_ft_s=$?; trap '_ignore_ft_s=1' HUP INT ILL TRAP ABRT BUS FPE USR1 SEGV USR2 PIPE ALRM TERM XCPU XFSZ VTALRM SYS; debug "sig-trap (exitcode=$_es_ft_s)"; exit ${_es_ft_s}' HUP INT ILL TRAP ABRT BUS FPE USR1 SEGV USR2 PIPE ALRM TERM XCPU XFSZ VTALRM SYS
            trap '_es_ft_er=$?; trap '_ignore_ft_er=1' ERR; debug "err-trap (exitcode=$_es_ft_er)"; exit ${_es_ft_er}' ERR

            ns="sync"
            log "Launching in the $(ns_long ${ns}): $1"
            eval "$1"
        )
        _es_m_f=$?
        shift
    done

    # Checks for a foreground task in the command line (without any quotation!)
    if [[ -n "$*" ]]; then
        (
            set -e
            trap '_es_ft_ex=$?; trap '_ignore_ft_ex=1' EXIT; debug "exit-trap (exitcode=${_es_ft_ex})"; hook_sub_exit ${_es_ft_ex} "$@"' EXIT
            trap '_es_ft_s=$?; trap '_ignore_ft_s=1' HUP INT ILL TRAP ABRT BUS FPE USR1 SEGV USR2 PIPE ALRM TERM XCPU XFSZ VTALRM SYS; debug "sig-trap (exitcode=$_es_ft_s)"; exit ${_es_ft_s}' HUP INT ILL TRAP ABRT BUS FPE USR1 SEGV USR2 PIPE ALRM TERM XCPU XFSZ VTALRM SYS
            trap '_es_ft_er=$?; trap '_ignore_ft_er=1' ERR; debug "err-trap (exitcode=$_es_ft_er)"; exit ${_es_ft_er}' ERR

            ns="sync"
            log "Launching in the $(ns_long ${ns}): $@"
            "$@"
        )
        # Catch the exitcode of a foreground sub-program
        _es_m_f=$?
    fi

    # Wait for all background processes and exit with a status of the last one
    # or with 128+SIGNAL in case of getting a signal
    # this is a work around for not trigerring ERR trap in main() two times in case of non 0 exist status
    ${_wait_bg_cmd} ${child_pids} || _es_m_d=$?

    if [ ${_es_m_d} -ne 0 ]; then
        _es_m_f=${_es_m_d} # update exit code only if one of background processes has failed
    fi

    if [[ "${RC_WAIT_POLICY}" = "wait_forever" ]]; then
        infinite_loop
    fi

    return ${_es_m_f}
}

say() {
    if [[ "${RC_VERBOSE}" = "true" ]]; then
        run echo "$@"
    fi
}

log() {
    say "$(run -p date '+%Y-%m-%d %H:%M:%S') ${SELF_NAME} [${ns}/${BASHPID}]: $@"
}

warn() {
    log "$@" >&2
}

err() {
    warn "$@"
    exit 1
}

debug() {
    if [[ "${RC_VERBOSE_EXTRA}" = "true" ]]; then
        log " - $@"
    fi
}

hook_main_exit() {
    local _pid _cpids _file _cmd _es_mtt_h=0 _term_timeout=5

    # If it fails in the `boot` stage, it appears here with 'set -e'. So, reset to +e again
    set +e
    trap 'true' ERR

    log "Trying to terminate sub-processes..."
    for _pid in ${child_pids}; do
        if run -p ps -p ${_pid} &> /dev/null; then
            log "terminating the child process <pid=${_pid}>"

            # Removes all unexpected sub-processes
            _cpids="$(run -p pgrep -P ${_pid} -d,)"
            run kill -TERM ${_pid} $(echo "${_cpids}" | tr ',' ' ') &> /dev/null

            while run -p ps -p ${_pid},${_cpids} &> /dev/null && [[ ${_term_timeout} -gt 0 ]]
            do
                run -p sleep 1
                : $((_term_timeout--))
            done

            if run -p ps -p ${_pid},${_cpids} &>/dev/null; then
                run kill -KILL ${_pid} $(run echo "${_cpids}" | run -p tr ',' ' ') &> /dev/null
            fi
        fi
    done
    debug "removing all unexpected sub-processes"
    if ! run -p pkill -KILL -P ${MAINPID} &> /dev/null; then
        true
    fi

    # --- `halt`
    if [[ boot_stage_left -eq 1 ]]; then
        debug "Looking for \`halt\` scripts..."
        # Checks for shutdown tasks in files (sequentially)
        for _file in $(run -p ls ${work_dir}/trc.halt.* ${work_dir}/trc.d/halt.* 2>/dev/null || true); do
            halt_cmds_exist=1
            (
                set -e # Exit on errors in the sub-shell
                ns="halt"
                log "Running the shutdown script: ${_file}"
                . ${_file}
            )
            _es_mtt_h=$?
            log "Exiting from the shutdown script (exitcode=${_es_mtt_h}): ${_file}"
        done

        # Runs shutdown tasks from the command line if any
        for _cmd in "${halt_cmds[@]}"; do
            (
                set -e # Exit on errors in the sub-shell

                ns="halt"
                log "Running the shutdown command: ${_cmd}"
                eval "${_cmd}"
            )
            _es_mtt_h=$?
            log "Exiting from the shutdown command (exitcode=${_es_mtt_h}): ${_cmd}"
        done    
    fi 

    if [[ ${halt_cmds_exist} -eq 1 ]]; then
        _exit_status=${_es_mtt_h}
    fi

    log "Exited (exitcode=${_exit_status})"
    exit ${_exit_status}
}

hook_sub_exit() {
    set +e # do not stop or errors anyway

    local _rc=$1 # Getting the exit code for a logging purpose only
    shift

    log "Exiting in the $(ns_long ${ns}) (exitcode=${_rc}): $@"

    if [[ "${RC_WAIT_POLICY}" = "wait_any" ]]; then 
        # If exiting from a bg process and don't need to wait other processes, let's stop the main
        if run -p ps -p ${MAINPID} &> /dev/null; then
            debug "terminating the main process <pid=${MAINPID}>"
            run kill -TERM ${MAINPID} &> /dev/null
        fi
    fi

    exit ${_rc}
}

infinite_loop() {
    log "Activated infinite loop! To stop, press <Ctrl+C> or send SIGTERM..."
    while true; do
        run -p sleep 1
    done
}

ns_long() {
    local _ns

    case "$1" in
        fg|sync)  _ns="foreground";;
        bg|async) _ns="background";;
        sd|halt)  _ns="shutdown";;
        main)     _ns="boot";;
         *)       _ns="$1";;
    esac

    echo "${_ns}"
}

check_ver_usage() {
    case "$1" in
        -v|--version) print_version ;;
        -h|--help)    print_usage ;;
         *)           : ;;
    esac
}

print_usage() {
    local _self=$0
    run echo "Usage: ${_self} [-w|--workdir 'dir'] [-B 'cmds' [...]] [-H 'cmds' [...]] [-D 'cmds' [...]] [-F 'cmds' [...]] [cmd [args]]"
    run echo "Examples:"
    run echo -e " $ ${_self} -B 'name=\$(id -un); echo booting...' -H 'echo halting...' -F 'echo Hello, \${name}!'"
    run echo -e " $ RC_WAIT_POLICY=wait_all ${_self} -D 'echo Hello' -D 'sleep 2; echo World' echo waiting..."
    run echo -e " $ RC_VERBOSE=true ${_self} -F 'echo -n \"Hello \"; echo World'"
    run echo -e " $ ${_self} --workdir /opt/app"
    exit 0
}

print_version() {
    run echo "trc (TrivialRC) $trc_version"
    exit 0
}

run() {
    builtin command "$@"
}

check_req_bins() {
    local _cmd

    for _cmd in "$@"; do
        run -v ${_cmd} &> /dev/null || { echo "There is no the required command: ${_cmd}" >&2; exit 1; }
    done
}

main "$@"