SSH-Master-Thesis/src/distributer/distribute.sh

#!/bin/sh

readonly DEFAULT_OUT_FMT="%s.out"
readonly CMD_FEED="$(mktemp -u /tmp/distributer-XXX.fifo)"
readonly PROC_BUFFER=10
readonly PROCS_PER_SERVER=6
readonly MAX_PROCS=$(expr $(ulimit -u) / $PROCS_PER_SERVER - $PROC_BUFFER)
readonly MAX_CMD_SIZE=512
readonly READER="$(dirname $(realpath $0))/fixed-read"
readonly TIMEOUT=600 # 10 minutes

readonly MK_DIR_CMD="if [ ! -d \"$(dirname "$OUT_FMT")\" ]; then
                        mkdir -p \"$(dirname "$OUT_FMT")\" > /dev/null
                     fi"

readonly S_SPACE="s/[[:space:].\/]/_/g"
readonly R_DASH="s/-//g"
readonly S_QUOTE="s/'/\\\\'/g"

help() {
    printf "Usage: $(basename $0) [options] cmd_list server_list [out_file_fmt]\n" >&2
    printf "    cmd_list: A text file containing a list of commands to run.\n" >&2
    printf "    server_list: A text file containing a list of servers to\n" >&2
    printf "    connect to and run commands on.\n" >&2
    printf "    out_file_fmt: File name format to write the output of each\n" >&2
    printf "    command to (default: %s).\n\n" "$DEFAULT_OUT_FMT" >&2
    printf "Options:.\n" >&2
    printf "    -h -?:          Print this help mesage and exit. \n" >&2
    printf "    -d destination: Move files to destination as space becomes\n" >&2
    printf "      unavailable.  (Requires quota to be available)\n" >&2
    printf "    -c processor:   If present, instead of directly executing\n\n" >&2
    printf "      cmd_list, will launch a processor on each server and \n" >&2
    printf "      distribute cmds to the processors.\n\n" >&2
    printf "All commands will be allocated to the first available server.\n" >&2
    printf "Each command must be valid on every server.\n" >&2
    printf "The output will be saved to a text file on the remote systems.\n" >&2
    exit $1
}

clean_server() {
    cmd_feed="$1"
    rm "$cmd_feed"
    exit 2
}

run_processor() {
    server="$1"
    loop="$2"
    out_file="$3"
    {
        # Test that ssh is ready to receive.
        printf "echo\n"
        read -t $TIMEOUT line < "$loop"
        printf "Running processor: '$PROCESSOR'\n" "$out_file" >&2
        printf "$PROCESSOR\n" "$out_file"
        # Block until command completes
        while read -t $TIMEOUT line < "$loop" > /dev/null && \
            cmd=$("$READER" $MAX_CMD_SIZE)
        do
            cmd=$(printf "%s" "$cmd" |
                sed -e 's/^[[:space:]]*//' \
                    -e 's/[[:space:]]*$//')
            printf "%s\n" "$cmd" >&2
            printf "%s\n" "$cmd"
        done
    } | ssh -oBatchMode=yes -oStrictHostKeyChecking=no "$server" "sh" > "$loop"
}

run_cmd_list() {
    server="$1"
    loop="$2"
    out_file="$3"
    {
        # Test that ssh is ready to receive.
        printf "echo\n"
        # Block until command completes
        while read -t $TIMEOUT line < "$loop" > /dev/null && \
            cmd=$("$READER" $MAX_CMD_SIZE)
        do
            cmd=$(printf "%s" "$cmd" | 
                sed -e 's/^[[:space:]]*//' \
                    -e 's/[[:space:]]*$//')
            printf "%s\n" "$cmd" >&2
            printf "printf \"%s: \" >> %s\n" "$cmd"  "$out_file"
            printf "%s >> %s\n" "$cmd"  "$out_file"
            printf "echo\n"
        done
    } | ssh -oBatchMode=yes -oStrictHostKeyChecking=no "$server" "sh" > "$loop"
}

run_server() {
    server="$1"
    loop="$(mktemp -u /tmp/$(basename $CMD_FEED .fifo)-$server-XXX.fifo)"
    out_file=$(printf "$OUT_FMT" "$(basename "$loop" .fifo)")
    mkfifo "$loop"
    trap "clean_server $loop" 2 15
    if [ -n "$PROCESSOR" ]; then
        run_processor "$server" "$loop" "$out_file"
    else
        run_cmd_list "$server" "$loop" "$out_file"
    fi # > "$loop"
    #| ssh -oBatchMode=yes -oStrictHostKeyChecking=no "$server" "sh" > "$loop"
    clean_server "$loop"
    echo "Finished!" >&2
}

clean_up() {
    pkill -P $$
    rm "$CMD_FEED"
    exit 2
}

prepend() {
    while read -r line; do
        printf "%s: %s\n" "$1" "$line"
    done
}

promote_subordinate_master() {
    server="$(head -1 -)"
    until server="$(head -1 -)";\
        ssh -oBatchMode=yes -oStrictHostKeyChecking=no "$server" \
        "$(realpath $0) -t"; do
        sleep 1
    done
}

launch_children() {
    for server in $(head -n$MAX_PROCS -); do
        run_server "$server" < "$CMD_FEED" 2>&1 | prepend "$server" &
    done
    # promote_subordinate_master &
    while pgrep -P $$ > /dev/null; do
        wait
    done
}

main() {
    mkfifo "$CMD_FEED"
    trap clean_up 2 15
    trap "" 10
    cat "$CONF_LIST" | sed '/^[[:space:]]*$/d' | \
        xargs -d'\n' printf "%-$MAX_CMD_SIZE.${MAX_CMD_SIZE}s" > "$CMD_FEED" &
    cat "$SERVER_LIST" | launch_children
    clean_up
    echo "All jobs finished!"
}

OPTIND=1
while getopts "h?c:t" opt; do
    case "$opt" in
        h|\?)
            help 0
            ;;
        c)  readonly PROCESSOR="$OPTARG"
            ;;
        t)  launch_children
            ;;
        *)  echo "Unknown option '$opt'" >&2
            help 1
            ;;
    esac
done
shift $(expr $OPTIND - 1)
if [ "$#" -lt 2 -o "$#" -gt 3 ]; then
    help 1
fi
readonly CONF_LIST="$1"
readonly SERVER_LIST="$2"
readonly OUT_FMT="${3:-$DEFAULT_OUT_FMT}"

main