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 MAX_PROCS=$(expr $(ulimit -u) / 3 - $PROC_BUFFER)
readonly MAX_CMD_SIZE=350
readonly READER="$(dirname $(realpath $0))/fixed-read"
readonly TIMEOUT=600 # 10 minutes

readonly CONF_LIST="$1"
readonly SERVER_LIST="$2"
readonly OUT_FMT="${3:-$DEFAULT_OUT_FMT}"

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) 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 "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_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
    {
        printf "echo > %s\necho\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:]]*$//')
            #cmd_sanitized="$(echo $cmd | sed "$S_SPACE;$R_DASH;$S_QUOTE")"
            #out_file="$(printf "$OUT_FMT" $cmd_sanitized)"
            printf "$server: %s\n" "$cmd" >&2
            #printf "%s > %s\necho\n" "$cmd" "$out_file"
            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"
    clean_server "$loop"
    echo "Server '$server' finished!" >&2
}

clean_up() {
    for pid in $(pgrep -P $$); do
        pkill -P $pid
    done
    rm "$CMD_FEED"
    exit 2
}

main() {
    mkfifo "$CMD_FEED"
    trap clean_up 2 15
    cat "$CONF_LIST" | sed '/^[[:space:]]*$/d' | \
        xargs -d'\n' printf "%-$MAX_CMD_SIZE.${MAX_CMD_SIZE}s" > "$CMD_FEED" &
    for server in $(head -n$MAX_PROCS "$SERVER_LIST"); do
        run_server "$server" < "$CMD_FEED" > /dev/null &
    done
    for pid in $(pgrep -P $$); do
        wait $pid
    done
    clean_up
    echo "All jobs finished!"
}

if [ "$#" -eq 2 -o "$#" -eq 3 ]; then
    main
else
    help
fi