bs_parallel.sh

#!/bin/sh
############################################################ IDENT(1)
#
# $Title: Script for bootstrapping multiple hosts in parallel $
# $Copyright: 2017-2020 Devin Teske. All rights reserved. $
#
############################################################ INFORMATION
#
# Sometimes you just want to bootstrap a bunch of machines. This script helps
# you bootstrap them in parallel to minimize the amount of time it takes to get
# a large amount of machines bootstrapped.
#
############################################################ CONFIGURATION

#
# Bootstrap script
# NB: Exported for parallelized sh
#
export BSCMD="./bs.sh" # -b path

#
# Bootstrap script arguments (these go after the host name)
# NB: Exported for parallelized sh
#
export BSARGS=

#
# Bootstrap script options (these go before the host name)
#
export BSOPTS=

#
# Directory where logs are to be stored
# NB: Exported for parallelized sh
#
export LOGDIR="$HOME/logs" # -l path

#
# Maximum concurrent bootstrap scripts that can run in-parallel
#
MAXPARALLEL=9 # -n num

#
# Maximum seconds to randomly sleep before each bootstrap starts
# NB: Exported for parallelized sh
#
export MAXSLEEP=30 # -s num

#
# Test mode. If enabled (non-NULL value), don't actually do anything
# NB: Exported for parallelized sh
#
export TEST= # -t

############################################################ ENVIRONMENT

SUDO_PROMPT="[sudo] Password:"
SUDO_PROMPT_RE='\[sudo\] Password:'
export SUDO_PROMPT SUDO_PROMPT_RE

# OPTIONAL: Ansible become-user support
#?BECOME_PROMPT="SUDO password: "
#?BECOME_PROMPT_RE="SUDO password: "
#?export BECOME_PROMPT BECOME_PROMPT_RE

# OPTIONAL: Ansible vault support
#?VAULT_PROMPT="Vault password: "
#?VAULT_PROMPT_RE="$VAULT_PROMPT"
#?export VAULT_PROMPT VAULT_PROMPT_RE

# OPTIONAL: Kerberos support
#?KINIT_PROMPT="Password for admin@domain: "
#?KINIT_PROMPT_RE="$KINIT_PROMPT"
#?export KINIT_PROMPT KINIT_PROMPT_RE

############################################################ GLOBALS

pgm="${0##*/}" # program basename

#
# Global exit status
#
SUCCESS=0
FAILURE=1

#
# Command-line options
#
OPT_EXPECT_SUDO=        # -S

#
# Terminal Detection
#
stty size > /dev/null 2>&1 && TTY=1

#
# Inline routine for calculating time elapsed from $estart
#
export ELAPSED_DURATION='
    estop=$( date +%s )
    elapsed=$(( $estop - $estart ))
    hour=$(( $elapsed / 3600 ))
    min=$(( ($elapsed - $hour * 3600) / 60 ))
    sec=$(( $elapsed - $hour * 3600 - $min * 60 ))
    [ $hour -lt 10 ] && hour=0$hour
    [ $min -lt 10 ] && min=0$min
    [ $sec -lt 10 ] && sec=0$sec
    if [ $hour -gt 0 ]; then
        duration=$hour:$min:$sec
    else
        duration=$min:$sec
    fi
' # END-QUOTE

############################################################ FUNCTIONS

die()
{
    [ "$*" ] && fail "$*"
    exit $FAILURE
}

exec 3<&1
eval2(){ printf '\e[2m%s\e[m\n' "$*" >&3; eval "$@"; }
fail(){ printf "\e[1m+ \e[31mFAIL\e[m %s\n\n" "$*"; }
pass(){ printf "\e[1m+ \e[32mPASS\e[m %s\n\n" "$*"; }
step(){ printf "\e[32;1m==>\e[39m %s\e[m\n" "$*"; }
warn(){ printf '\e[33;1m!WARNING!\e[39m %s\e[m\n' "$*"; }

usage()
{
    local fmt="\t%-9s %s\n"
    exec >&2
    [ "$*" ] && echo "$pgm:" "$@"
    printf "Usage: %s [OPTIONS] host ...\n" "$pgm"
    printf "OPTIONS:\n"
    printf "$fmt" "-a args" \
        "Bootstrap command arguments to use after each host name."
    printf "$fmt" "-b path" \
        "Bootstrap command (default \`$BSCMD')."
    printf "$fmt" "-l path" \
        "Path to log directory (default \`$LOGDIR')."
    printf "$fmt" "-n num" \
        "Maximum number of parallel instances (default $MAXPARALLEL)."
    printf "$fmt" "-o opts" \
        "Bootstrap command options to use before each host name."
    printf "$fmt" "-S" \
        "Use tcl/expect to handle sudo prompts automatically."
    printf "$fmt" "-s num" \
        "Maximum seconds to randomly sleep (default $MAXSLEEP)."
    printf "$fmt" "-t" \
        "Test mode. Don't actually do anything but pretend."
    die
}

isset()
{
    eval [ \"\${$1+set}\" ]
}

sread()
{
    local OPTIND=1 OPTARG flag
    local prompt= retval

    while getopts p: flag; do
        case "$flag" in
        p) prompt="$OPTARG" ;;
        esac
    done
    shift $(( OPTIND - 1 ))

    ! isset "$1" || return $SUCCESS
    ! [ "$prompt" ] || printf "%s" "$prompt"

    trap 'stty echo' EXIT
    stty -echo
    retval=$?
    read -r $1
    stty echo
    trap - EXIT

    echo
    export $1
    return $retval
}

make_expect()
{
    local var=$1
    shift 1 # var
    exec 9<<-EOF
    set timeout 86400
    fconfigure stdout -buffering none
    #log_user 0
    spawn {*}\$argv
    #log_user 1
    while (1) { expect {$(
        for key_value in "$@"; do
            printf '\n\t\t'
            printf '%s $::env(%s) { send "$::env(%s)\\n" }' \
                -re ${key_value%%=*} ${key_value#*=}
        done )
        timeout { Time_Out; break }
        eof { break }
    } }
    lassign [wait] pid spawnid os_error_flag value
    exit \$value
    EOF
    eval $var='$( cat <&9 )'
}

############################################################ MAIN

#
# Command-line options
#
while getopts a:b:l:n:o:Ss:t flag; do
    case "$flag" in
    a) BSARGS="$OPTARG" ;;
    b) BSCMD="$OPTARG" ;;
    l) LOGDIR="$OPTARG" ;;
    n) MAXPARALLEL="$OPTARG" ;;
    o) BSOPTS="$OPTARG" ;;
    S) OPT_EXPECT_SUDO=1 ;;
    s) MAXSLEEP="$OPTARG" ;;
    t) TEST=1 ;;
    *) usage # NOTREACHED
    esac
done
shift $(( $OPTIND - 1 ))

#
# Validate `-b path' option
#
[ -e "$BSCMD" ] || die "$BSCMD: No such file or directory" # NOTREACHED
[ -d "$BSCMD" ] && die "$BSCMD: Is a directory" # NOTREACHED
[ -x "$BSCMD" ] || die "$BSCMD: Permission denied" # NOTREACHED

#
# Validate `-l path' option
#
[ -e "$LOGDIR" ] || mkdir -p "$LOGDIR" || die # NOTREACHED
[ -d "${LOGDIR%/}/" ] || die "$LOGDIR: Not a directory" # NOTREACHED

#
# Validate `-n num' option
#
case "$MAXPARALLEL" in
"") usage "-n flag requires an argument" ;; # NOTREACHED
*[!0-9]*) usage "$MAXPARALLEL: -n argument must be a number" ;; # NOTREACHED
esac

#
# Validate `-s num' option
#
case "$MAXSLEEP" in
"") usage "-s flag requires an argument" ;; # NOTREACHED
*[!0-9]*) usage "$MAXSLEEP: -s argument must be a number" ;; # NOTREACHED
esac

#
# Validate number of command-line arguments
#
[ $# -gt 0 ] || usage "missing host argument(s)" # NOTREACHED

#
# Provide some basic information
#
step "Run Details"
echo "Maximum parallel instances: $MAXPARALLEL"
echo "Maximum random sleep: $MAXSLEEP"
echo "Hosts: $*"
set -- $* || die
echo "Total Hosts: $#"
echo "Running (for each host): $BSCMD ${BSOPTS:+$BSOPTS }<host> $BSARGS"
echo "Logs go to: ${LOGDIR%/}/<host>.log"
if [ "$TEST" ]; then
    printf "${TTY:+\e[33m}INFO:${TTY:+\e[0m} %s\n" \
        "TEST MODE enabled (bootstrap actions will not be performed)"
else
    printf "${TTY:+\e[35m}%s${TTY:+\e[0m} " \
        "< Press ENTER to proceed or Ctrl-C to cancel >"
    read IGNORED
fi

#
# Tcl/Expect
#
if [ "$OPT_EXPECT_SUDO" ]; then
    errexit=
    case "$-" in
    *e*) errexit=1 ;;
    esac
    set -e # errexit
    step "Test sudo credentials"
    sread -p "$SUDO_PROMPT" _SP || die
    make_expect EXPECT_SUDO SUDO_PROMPT_RE=_SP
    echo "$EXPECT_SUDO" | eval2 expect -f- sudo /bin/true || die
    pass
    # OPTIONAL: Kerberos support
    #?step "Test kerberos credentials"
    #?sread -p "$KINIT_PROMPT" _KP || die
    #?make_expect EXPECT_KINIT KINIT_PROMPT_RE=_KP
    #?if ! eval2 klist \| grep -q admin; then
    #?  eval2 kdestroy
    #?  eval2 sleep 3
    #?fi
    #?echo "$EXPECT_KINIT" | eval2 expect -f- kinit admin || die
    #?pass
    # OPTIONAL: Ansible support
    #?if ! isset _VP; then
    #?  echo
    #?  warn "Using kerberos credentials for ansible vault"
    #?  echo
    #?  _VP="$_KP"
    #?  export _VP
    #?fi
    #?make_expect EXPECT_SUDO_OR_VAULT \
    #?  SUDO_PROMPT_RE=_SP BECOME_PROMPT_RE=_SP VAULT_PROMPT_RE=_VP
    [ "$errexit" ] || set +e
fi

#
# Configure EXIT trap to provide ending information
#
trap '
    echo "End parallel bootstrap: $( date )"
    eval "$ELAPSED_DURATION"
    echo "Elapsed parallel time: $duration"
' EXIT

#
# Parallel sh code
# NB: The first argument ($1) represents the host to be bootstrapped
#
export OPT_EXPECT_SUDO EXPECT_SUDO EXPECT_SUDO_OR_VAULT
exec 9<<'EOF'
    log="${LOGDIR%/}/$1.log"
    exec > "$log" 2>&1

    sleep=$(( $RANDOM * $MAXSLEEP / 32768 ))
    echo "Sleeping $sleep second(s) before starting bootstrap..."
    sleep $sleep

    echo "Bootstrap started: $( date )"
    estart=$( date +%s )

    echo "Running: $BSCMD ${BSOPTS:+$BSOPTS }$1 $BSARGS"
    trap '
        echo "Exit code: $?"
        echo "Bootstrap ended: $( date )"
        eval "$ELAPSED_DURATION"
        echo "Elapsed time: $duration"
    ' EXIT
    if [ "$TEST" ]; then
        echo "INFO: TEST MODE enabled (nothing done)"
    elif [ "$OPT_EXPECT_SUDO" ]; then
        # OPTIONAL: Ansible vault support
        #?echo "$EXPECT_SUDO_OR_VAULT" |
        echo "$EXPECT_SUDO" |
            expect -f- $BSCMD $BSOPTS $1 $BSARGS
    else
        $BSCMD $BSOPTS $1 $BSARGS
    fi
EOF
BS_PARALLEL_WORKER=$( cat <&9 )

#
# Run parallel instances of sh code, one instance per host to be bootstrapped
# NB: Worker code evaluated to prevent raw code appearing in ps(1)
#
step "Run"
echo "Start parallel bootstrap: $( date )"
export BS_PARALLEL_WORKER
estart=$( date +%s )
exec 3<&1 4<&2
( exec 5<&1 >&3 2>&4 3<&5 4>&- 5>&- # fd0-2 to TTY, fd3 to pipe
    ################################################## INFORMATION
    #
    # Shell co-routine for sending work to xargs and displaying status
    #
    ################################################## MAIN

    #
    # Send list of hosts to xargs
    #
    n=0
    for host in $*; do
        n=$(( $n + 1 ))

        log="${LOGDIR%/}/$host.log"
        rm -f "$log"

        unset done$n
        eval log$n=\"\$log\"

        echo $host >&3
    done

    #
    # Display status and elapsed time for completed host(s)
    #
    remainder=$n
    printf "%4s %6s %9s %s\n" "#" STATUS ELAPSED HOSTNAME
    while [ $remainder -gt 0 ]; do
        n=0
        alldone=1
        for host in $*; do
            n=$(( $n + 1 ))
            eval done=\$done$n
            [ "$done" ] && continue

            # This host is not done yet, check its status
            alldone=
            eval log=\"\$log$n\"
            status=$( awk '
                /^Exit code:/, ++fs { status = $NF }
                /^Elapsed time:/, ++fe { elapsed = $NF }
                fs && fe { exit found = 1 }
                END { print status, elapsed; exit !found }
            ' "$log" 2> /dev/null ) || continue

            # This host has recently completed, display status
            eval done$n=1
            if [ "$status" != "${status#"$SUCCESS "}" ]; then
                msg="[${TTY:+\e[32;1m} OK ${TTY:+\e[0m}]"
            else
                msg="[${TTY:+\e[31;1m}FAIL${TTY:+\e[0m}]"
            fi
            elapsed="${status#*[$IFS]}"
            printf "%4s $msg %9s %s\n" \
                "$remainder" "$elapsed" "$host"

            remainder=$(( $remainder - 1 ))
        done
        [ "$alldone" ] && break
        sleep 1
    done

    ######################################################################
) | xargs -rn1 -P$MAXPARALLEL sh -c 'eval "$BS_PARALLEL_WORKER"' /bin/sh

################################################################################
# END
################################################################################