parallel.bash

parallel.bash

#!/usr/bin/env bash
## parallel.bash
## -------------
## A bash script that executes provided tasks in parallel, with added handling
## for errors, logs, and signals (within reason).
##
## Features:
## - Executes multiple tasks in parallel with configurable concurrency
## - Provides interactive progress display with spinners
## - Captures stdout/stderr for each task in separate log files
## - Handles signals gracefully (SIGINT, SIGTERM)
## - Provides detailed error reporting for failed tasks
##
## Example:
## ```sh
## ./parallel.bash -m 2 \
##    'echo "sleeping1"; sleep 5' \
##    'echo "sleeping3"; sleep 2' \
##    'echo "sleeping3"; sleep 3' \
##    'bash -c "echo \"this is a forced error\" >&2; sleep 6; kill -ABRT $$"' \
##    'echo "sleeping3"; sleep 9' \
##    'echo "sleeping4"; sleep 4'
## ```

set -o nounset
export LC_CTYPE=en_US.UTF-8

# -----------------------------------------------------------------------------
# Help and usage information
# -----------------------------------------------------------------------------
usage() {
    cat <<EOF
Usage: $0 [-d directory] [-m integer] [-i boolean] [task_arg1] [task_arg2] ...

Options:
  -d directory     Specify directory for logs (default uses mktemp)
  -m integer       Maximum tasks that can run in parallel (default is hw.ncpu)
  -i boolean       Force disable interactive mode (no spinners or live updates)

Arguments:
  task_argN        Commands to execute in parallel (required, multiple allowed)

Examples:
  $0 -m 2 'echo "Task 1"; sleep 2' 'echo "Task 2"; sleep 3'
  $0 -d /tmp/logs 'make test' 'make lint' 'make build'

EOF
    exit 1
}

# -----------------------------------------------------------------------------
# Parse command line arguments
# -----------------------------------------------------------------------------
while getopts ":d:m:i" opt; do
    case "${opt}" in
    d) LOGS_DIR="${OPTARG%/}/" ;;      # Remove trailing slash if present, then add one
    m) PARALLEL=${OPTARG} ;;           # Number of parallel tasks
    i) ALLOW_INTERACTIVE=false ;;      # Disable interactive mode
    \?)
        printf "Invalid option: -%s\\n" "${OPTARG}" >&2
        usage
        ;;
    :)
        printf "Option -%s requires an argument.\\n" "${OPTARG}" >&2
        usage
        ;;
    esac
done
shift $((OPTIND - 1))                  # Remove processed options
(($# > 0)) || usage                    # Ensure at least one task is provided

# -----------------------------------------------------------------------------
# Constants and global state
# -----------------------------------------------------------------------------
# Constants
readonly LOGS_DIR=${LOGS_DIR:-${TMPDIR-:/tmp/}parallel.bash/$(uuidgen)/}  # Log directory
readonly FAILURE_FILE="${LOGS_DIR}failure.index"                          # File to track failures
declare -ir PARALLEL=${PARALLEL:-$(sysctl -n hw.ncpu)}                    # Max parallel tasks
declare -r ALLOW_INTERACTIVE=${ALLOW_INTERACTIVE:-true}                   # Interactive mode flag
declare -ir FAILURE_SUMMARY_MAX_LOG_LINES=5                               # Max lines to show in error summary

# Mutable global state
declare interactive=false              # Whether we're in interactive mode
declare -i columns=128                 # Terminal width
declare -i trapped_code=0              # Signal code if trapped

# Set up interactive mode if allowed and output is a terminal
$ALLOW_INTERACTIVE && [[ -t 1 ]] && interactive=true

if $interactive; then
    columns=$(tput cols)
    # ANSI color codes (only if terminal and not disabled)
    readonly RESET='\033[0m'
    readonly BOLD='\033[1m'
    readonly CYAN='\033[0;36m'
    readonly RED='\033[0;31m'
    readonly GREEN='\033[0;32m'
    readonly YELLOW='\033[0;33m'
    readonly BLUE='\033[0;34m'
    readonly GRAY='\033[0;90m'
else
    readonly RESET=''
    readonly BOLD=''
    readonly CYAN=''
    readonly RED=''
    readonly GREEN=''
    readonly YELLOW=''
    readonly BLUE=''
    readonly GRAY=''
fi

# -----------------------------------------------------------------------------
# Signal handlers
# -----------------------------------------------------------------------------
# Clean up on exit
handle_exit() {
    $interactive && printf '\e[?25h'   # Show the cursor if it was hidden
    rm -f "${FAILURE_FILE}"            # Remove failure tracking file
}

# Handle signals
handle_trap() { trapped_code=$1; }     # Store the signal code
handle_main_err() { handle_trap 1; }   # Handle ERR signal
handle_sigint() { handle_trap 130; }   # Handle Ctrl+C (SIGINT)
handle_sigterm() { handle_trap 143; }  # Handle SIGTERM

# Set up signal traps
trap handle_exit EXIT
trap handle_main_err ERR
trap handle_sigint INT
trap handle_sigterm TERM

# -----------------------------------------------------------------------------
# Helper functions
# -----------------------------------------------------------------------------
# Truncate a string to fit within a given width
print_truncated() {
    declare str=$1                     # String to truncate
    declare -i maxlen=$(($2-1))        # Max length (minus 1 for ellipsis)
    ((${#str} > maxlen)) && str="${str:0:maxlen}…"  # Truncate if needed
    printf '%s' "$str"                 # Print the result
}

# -----------------------------------------------------------------------------
# Main execution function
# -----------------------------------------------------------------------------
main() {
    # ---------------------------------------------------------------------
    # Initialize state
    # ---------------------------------------------------------------------
    # Immutable state
    declare -ar tasks=("$@")           # Array of tasks to execute
    declare -ar spinner=(⠋ ⠙ ⠹ ⠸ ⠼ ⠴ ⠦ ⠧ ⠇ ⠏)  # Spinner animation frames
    declare -ir timeout_seconds=3600   # Max execution time (1 hour)
    declare -ir term_timeout_seconds=3 # Grace period for termination
    declare -ir task_count=${#tasks[@]} # Number of tasks
    declare -ir spinner_count=${#spinner[@]} # Number of spinner frames
    declare -ir start_time=$SECONDS    # Script start time

    # Mutable state
    declare term_started=false         # Whether termination has started
    declare term_expired=false         # Whether grace period has expired
    declare -i executing_count=0       # Number of currently executing tasks
    declare -i spinner_index=0         # Current spinner frame
    declare -i term_start_time=0       # When termination started
    declare -a remaining_indexes=()    # Indexes of tasks not yet completed
    declare -a task_pids=()            # PIDs of running tasks
    declare -a task_statuses=()        # Status of each task

    # Initialize task arrays
    for ((i = 0; i < task_count; i++)); do
        remaining_indexes[i]=$i        # All tasks start as remaining
        task_pids[i]=0                 # No PID yet
        task_statuses[i]="pending"     # All tasks start as pending
    done

    # ---------------------------------------------------------------------
    # Task management functions
    # ---------------------------------------------------------------------
    # Start a new task
    start_task() {
        declare -ir index=$1 # Task index
        declare -r task=${tasks[index]} # Task command
        declare -r logs_dir="${LOGS_DIR}task-${index}/" # Log directory for this task
        declare -r out_log="${logs_dir}/out.log" # stdout log
        declare -r err_log="${logs_dir}/err.log" # stderr log

        mkdir -p "$logs_dir"   # Create log directory

        # Log the task command to both stdout and stderr logs
        printf 'parallel.bash task:%d: %s\n\n' $index "$task" | tee -a "$out_log" >>"$err_log"

        # Execute the task in background and capture its PID
        (${SHELL:-/usr/bin/env bash} -c "$task" || printf '%d\n' $index >>"$FAILURE_FILE") >>"$out_log" 2>>"$err_log" &
        task_pids[index]=$!   # Store the PID
        ((executing_count++)) # Increment executing count
    }

    # Process a single task (check status, start/stop as needed)
    process_task() {
        declare -ir index=$1 # Task index
        declare -ir pid=${task_pids[index]} # Task PID
        declare -r status=${task_statuses[index]} # Current task status

        # Task not started yet
        if ((pid == 0)); then
            if $term_started; then # If termination has started
                task_statuses[index]='cancelled'
                return 1 # Task is done (cancelled before start)
            fi
            # Start the task if we have capacity
            ((executing_count < PARALLEL)) && start_task $index
            return 0 # Task is still pending or just started
        fi

        # Task is running (PID exists)
        if kill -0 $pid 2>/dev/null; then  # Check if process is running
            if $term_started; then         # If termination has started
                if $term_expired; then     # If grace period expired
                    task_statuses[index]='cancelled'
                    return 1 # Task is done (forcibly cancelled)
                elif [[ $status != cancelling ]]; then
                    task_statuses[index]='cancelling'
                    kill -TERM $pid >/dev/null 2>&1 # Send SIGTERM
                fi
            else
                task_statuses[index]='executing'
            fi
            return 0 # Task is still running
        fi

        # Task has finished (PID no longer exists)
        if $term_started; then         # If termination has started
            if [[ $status == cancelling ]]; then # If was cancelling
                task_statuses[index]='cancelled'
            else
                task_statuses[index]='failure'
            fi
        elif [[ $status == executing ]]; then # If was executing
            task_statuses[index]='success'
            if ! $interactive; then # If not in interactive mode
                printf '(task:%d ✔)' $index
            fi
        fi
        return 1 # Task is done
    }

    # Process all remaining tasks
    process_tasks() {
        declare -i index
        declare -ar remaining=("${remaining_indexes[@]}") # Copy remaining indexes
        remaining_indexes=() # Clear remaining indexes

        # Process each remaining task
        for index in "${remaining[@]}"; do
            if process_task $index; then # If task is still running or pending
                remaining_indexes+=("$index") # Keep in remaining list
            else
                ((executing_count--)) # Decrement executing count
            fi
        done
    }

    # ---------------------------------------------------------------------
    # Display functions
    # ---------------------------------------------------------------------
    # Print the overall information header
    print_overall_header() {
        printf '%bLogs:\n└── %b%s%b\n\n' "$GRAY" "$CYAN" "${LOGS_DIR}" "$RESET"
    }
    # Print the status icon for a task
    print_task_status_icon() {
        declare use_color=$1
        declare -ir index=$2
        local red='' green='' yellow='' blue=''
        if $use_color; then
            red=$RED
            green=$GREEN
            yellow=$YELLOW
            blue=$BLUE
        fi
        case ${task_statuses[index]} in
        'pending')                printf '%b◷' "$blue" ;;
        'executing'|'cancelling') printf '%b%s' "$green" "${spinner[spinner_index]}" ;;
        'success')                printf '%b✔' "$green" ;;
        'failure')                printf '%b✘' "$red" ;;
        'cancelled')              printf '%b□' "$yellow" ;;
        esac
    }
    # Print the info for a task (task index, PID if available)
    print_task_info() {
        declare -ir index=$1
        declare -ir pid=${task_pids[index]}
        printf 'task:%d' $index 
        ((pid > 0)) && printf ' pid:%d' $pid
    }
    # Print the prefix for a task line (status icon, info)
    print_task_line_prefix() {
        local use_color=$1
        declare -ir index=$2
        local reset='' gray=''
        if $use_color; then
            reset=$RESET
            gray=$GRAY
        fi
        printf '  ' # Initial indent
        print_task_status_icon $use_color $index 
        printf '%b ' "$gray"
        print_task_info $index
        printf ' ─ %b' "$reset"
    }
    # Print a full task line (prefix + command)
    print_task_line() {
        declare -ir index=$1
        if $interactive; then
            declare -r prefix=$(print_task_line_prefix false $index)
            print_task_line_prefix true $index
            print_truncated "${tasks[index]}" $((columns-${#prefix}))
            printf '\033[K' # Clear the rest of the line
        else     
            print_task_line_prefix false $index
            printf '%s' "${tasks[index]}"
        fi
        printf '\n'
    }

    print_loop() {
        if $interactive; then        # Interactive mode display
            for ((i = 0; i < task_count; i++)); do
                print_task_line $i   # Update each task line
            done
            spinner_index=$(((spinner_index + 1) % spinner_count)) # Update spinner
            printf '\033[%sA' $task_count # Move cursor up to beginning of task list
        else                         # Non-interactive mode
            printf '▱'               # Print progress dot
        fi
    }

    # ---------------------------------------------------------------------
    # Summary and reporting functions
    # ---------------------------------------------------------------------
    # Print summary and exit
    exit_with_summary() {
        declare -a success=() cancelled=() # Arrays for successful and cancelled tasks
        declare -i count index

        # Print information about cancelled tasks
        print_cancelled() {
            count=${#cancelled[@]} # Number of cancelled tasks
            if ((count > 0)); then
                printf '  %b%b%sCancelled remaining %d%b:\n' "$YELLOW" "$BOLD" "$1" $count "$RESET"
                for index in "${cancelled[@]}"; do
                    print_task_line $index # Print each cancelled task
                done
                printf '\n'
            fi
        }

        # Print details about a failed task
        print_failure_summary() {
            declare -ir index=$1
            declare -r err_log="${LOGS_DIR}task-${index}/err.log" # Error log path
            
            printf '  %b%bFailed with error%b:\n' "$RED" "$BOLD" "$RESET"
            print_task_line $index # Print the failed task
            printf '\n'
            
            # Use AWK to extract just the last few non-empty lines from the error log
            declare -a log_lines=()
            while IFS= read -r line; do
                log_lines+=("$line")
            done < <(awk -v max_lines="${FAILURE_SUMMARY_MAX_LOG_LINES}" '
            NF { 
                lines[count % max_lines] = $0; # Non-empty lines in circular buffer
                count++;
            } 
            END { 
                if (count > 0) {
                    num_lines = (count < max_lines ? count : max_lines);
                    start = (count < max_lines ? 0 : count - num_lines);
                    for (i = start; i < count; i++) {
                        print lines[i % max_lines];
                    }
                }
            }' "$err_log")

            declare -ir num_lines=${#log_lines[@]}
            if ((num_lines > 0)); then
                printf '  %bLast %d line(s) of %b%s%b\n' "$GRAY" $num_lines "$CYAN" "$err_log" "$RESET"
                for ((i=0; i<num_lines; i++)); do
                    if ((i == num_lines - 1)); then
                        prefix='  └── '
                    else
                        prefix='  │   '
                    fi
                    printf '%b%s%b' "$GRAY" "$prefix" "$RESET"
                    print_truncated "${log_lines[i]}" $((columns-${#prefix}))
                    printf '\n'
                done
                printf '\n'
            fi
        }

        # Collect task results
        for ((i = 0; i < task_count; i++)); do
            case ${task_statuses[i]} in
            'success') success+=("$i") ;; # Add to success array
            'cancelled') cancelled+=("$i") ;; # Add to cancelled array
            esac
        done

        if $interactive; then
            tput ed # Clear to the reset of the screen in interactive mode
        else 
            printf '\n\n' # Print 2 newlines in non-interactive mode
        fi

        # Print successful tasks
        count=${#success[@]} # Number of successful tasks
        if ((count > 0)); then
            if ((count < task_count)); then
                printf '  %b%bCompleted %d successfully%b:\n' "$GREEN" "$BOLD" $count "$RESET"
            else
                printf '  %b%bAll tasks completed successfully%b:\n' "$GREEN" "$BOLD" "$RESET"
            fi
            for index in "${success[@]}"; do
                print_task_line $index # Print each successful task
            done
            printf '\n'
        fi

        # Handle termination conditions
        if ((trapped_code > 0)); then # If terminated by signal
            print_cancelled 'Interrupted! ' >&2
            exit $trapped_code
        elif [[ -f $FAILURE_FILE ]]; then # If any task failed
            print_cancelled 'Failure detected! ' >&2
            print_failure_summary "$(head -1 "$FAILURE_FILE")" >&2
            exit 1
        fi
    }

    # ---------------------------------------------------------------------
    # Main execution
    # ---------------------------------------------------------------------

    # Create log directory and print path
    mkdir -p "${LOGS_DIR}"

    print_overall_header

    if $interactive; then
        # Set up a trap for window resize for interactive mode
        handle_resize() {
            columns=$(tput cols)
        }
        trap handle_resize WINCH
        printf '\e[?25l' # Hide the cursor in interactive mode
    else
        # Print initial task status if not in interactive mode
        for ((i = 0; i < task_count; i++)); do
            print_task_line $i
        done
        printf '\nProgress: '
    fi

    # Main execution loop
    while ((${#remaining_indexes[@]} > 0)); do

        # Check for termination conditions
        if ! $term_started && { ((trapped_code > 0 || SECONDS - start_time > timeout_seconds)) || [[ -f $FAILURE_FILE ]]; }; then
            term_started=true        # Start termination
            term_start_time=$SECONDS # Record when termination started
        fi
        # Check if grace period has expired
        if ! $term_expired && { $term_started && ((SECONDS - term_start_time > term_timeout_seconds)); }; then
            term_expired=true        # Mark grace period as expired
        fi

        process_tasks
        print_loop

        if $interactive; then 
            sleep 0.05
        else
            sleep 1 
        fi
    done

    exit_with_summary
}

main "$@"