trc

#!/bin/bash

# TrivialRC - The minimalistic init system and process manager for containers
# Copyright (c) 2016 by Oleksii Tsvietnov, me@vorakl.name
# Version: 1.0.5

# --[ Environment ]--
# RC_DEBUG (true|false) [false]
#    Prints out all commands which are being executed
# RC_VERBOSE (true|false) [false]
#    Prints out service information
# RC_WAIT_POLICY (wait_all|wait_any|wait_forever) [wait_any]
#    - wait_all      quit after exiting the last command (back- or foreground)
#    - wait_any      quit after exiting any of command (including zero commands)
#    - wait_forever  will be waiting forever after exiting all commands.
#                    Usefull in case of daemons which are detaching and exiting

set +e # Do not exit on errors by default
if [ "$RC_DEBUG" = "true" ]; then
    set -x # Turns on Debug mode if true
fi
: ${RC_WAIT_POLICY:=wait_any} # Sets default value for a wait policy

main() {
    trap '_exitcode=$?; hook_main_exit; exit $_exitcode' EXIT

    export SELF_NAME=$(basename -s .sh $0) # Self name for logging purpose
    export DIR_NAME=$(dirname $0)
    export MAINPID="$BASHPID" # PID of the main process
    childpid="" # Collecting child processes
    ns="main" # Name Space
    local _wait_bg_cmd _file _exitcode=0 _exitcode2=0

    case "$RC_WAIT_POLICY" in
        wait_all)  _wait_bg_cmd="command -p wait";;
        wait_any)  _wait_bg_cmd="command -p wait";;
        *)         _wait_bg_cmd=":";;
    esac

    say "The default wait policy: $RC_WAIT_POLICY"

    # Reads commands from files to run on the background (in parallel)
    for _file in $(ls $DIR_NAME/trc.bg.* 2>/dev/null); 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 '_exitcode=$?; hook_sub_exit $_exitcode "$_file"; exit $_exitcode' EXIT

            # In case of exit on errors (set -e), catch exitcode and exit with it, which
            # will lead to triggering an EXIT trap
            set -e
            trap 'exit $?' ERR # Exit with the status code of a last command

            ns="bg"
            say "Running on the $(ns_long $ns): $_file"
            . $_file
        )&
        childpid="${childpid} $!"
    done

    # Checks for background tasks in the command line
    while [ "$1" = "-D" ]; do
        shift
        (
            # 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 '_exitcode=$?; hook_sub_exit $_exitcode "$1"; exit $_exitcode' EXIT

            # In case of exit on errors (set -e), catch exitcode and exit with it, which
            # will lead to triggering an EXIT trap
            set -e
            trap 'exit $?' ERR # Exit with the status code of a last command

            ns="bg"
            say "Running on the $(ns_long $ns): $1"
            eval "$1"
        )&
        childpid="${childpid} $!"
        shift
    done

    # Checks for foreground tasks in files (sequentially)
    for _file in $(ls $DIR_NAME/trc.fg.* 2>/dev/null); 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 '_exitcode=$?; hook_sub_exit $_exitcode "$_file"; exit $_exitcode' EXIT

            # In case of exit on errors (set -e), catch exitcode and exit with it, which
            # will lead to triggering an EXIT trap
            set -e
            trap 'exit $?' ERR # Exit with the status code of a last command

            ns="fg"
            say "Running on the $(ns_long $ns): $_file"
            . $_file
        )
        # Catch the exitcode of a foreground sub-program
        _exitcode=$?
    done

    # Checks for a foreground task in the command line (one parameter with all commands)
    if [ -n "$*" ]; then
        (
            # 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 '_exitcode=$?; hook_sub_exit $_exitcode "$@"; exit $_exitcode' EXIT

            # In case of exit on errors (set -e), catch exitcode and exit with it, which
            # will lead to trigger an EXIT trap
            set -e
            trap 'exit $?' ERR # Exit with the status code of a last command

            ns="fg"
            say "Running on the $(ns_long $ns): $@"
            "$@"
        )
        # Catch the exitcode of a foreground sub-program
        _exitcode=$?
    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
    $_wait_bg_cmd $childpid;
    _exitcode2=$?

    if [ $_exitcode2 -ne 0 ]; then
        _exitcode=$_exitcode2 # update exit code only if one of background processes has failed
    fi

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

    return $_exitcode
}

say() {
    if [ "$RC_VERBOSE" = "true" ]; then
        echo "$(date '+%Y-%m-%d %H:%M:%S') ($SECONDS sec) ${SELF_NAME} [$ns/$BASHPID]: $@"
    fi
}

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

hook_main_exit() {
    local _pid _file

    say "Going down. Running shutdown scripts..."
    # Checks for shutdown tasks in files (sequentially)
    for _file in $(ls $DIR_NAME/trc.sd.* 2>/dev/null); do
        (
            ns="sd"
            set -e # Exit on errors in the sub-shell
            say "Running the shutdown script: $_file"
            . $_file
        )
    done

    say "Handling of termination..."
    for _pid in $childpid; do
        if ps -p $_pid &> /dev/null; then
            warn " - terminating child <$_pid>"
            kill -TERM $_pid &> /dev/null
        fi
    done
    say "Exited."

    # It's not needed to exit with a proper exit code here.
    # The programm will exit with correct exit code from EXIT trap of main()
}

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

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

    say "Exiting on the $(ns_long $ns) ($_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 ps -p $MAINPID &> /dev/null; then
            warn " - terminating main <$MAINPID>"
            kill -TERM $MAINPID &> /dev/null
        fi
    fi
    # To prevent invoking error trap from error trap ;)
    # The correct exit code was taken before diving into this function
    return 0
}

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

ns_long() {
    local _ns

    case "$1" in
        fg) _ns="foreground";;
        bg) _ns="background";;
        sd) _ns="shutdown";;
         *) _ns="$1";;
    esac

    echo "$_ns"
}

main "$@"