Arch: TestCpu [CPU stress test with journalctl events included in output]

#!/bin/bash

# Version 26

# WARNING: This script will run a CPU stress test on the processor resulting in high temperatures and power draw.
#          Prolonged usage may cause system instability, degradation, or damage.
#          Do not leave the system unattended during use.
#          Do not run this script unless adequate cooling is in place and you understand the risks.
#          You assume all responsibility for using this script.

# Installation:
#     $ sudo cp ./testcpu /usr/local/bin/
#     $ sudo chmod +x /usr/local/bin/testcpu

# ----- Colors -----
BLACK="\033[30m"
RED="\033[31m"
GREEN="\033[32m"
YELLOW="\033[33m"
BLUE="\033[34m"
MAGENTA="\033[35m"
CYAN="\033[36m"
WHITE="\033[37m"
BRIGHT_BLACK="\033[90m"
BRIGHT_RED="\033[91m"
BRIGHT_GREEN="\033[92m"
BRIGHT_YELLOW="\033[93m"
BRIGHT_BLUE="\033[94m"
BRIGHT_MAGENTA="\033[95m"
BRIGHT_CYAN="\033[96m"
BRIGHT_WHITE="\033[97m"

# ----- Styles -----
RESET="\033[0m" # Reset all styles.
BOLD="\033[1m"
HALF_BRIGHT="\033[2m"
ITALIC="\033[3m"
UNDERLINE="\033[4m"
REVERSED="\033[7m" # Reverse foreground/background.

show_help() {
    cat << EOF
Usage: ${0##*/} [OPTIONS] [CPU_LIST]

Options:
    -h, --help           Show this help message and exit.
    -m, --mode MODE      Set test mode: 'async' or 'round' (default: async).
                         Accepts formats: -m async, --mode async, or --mode=async
    -c, --cycles NUM     Set the number of test cycles to run (default: 1).
                         Accepts formats: -c 5, --cycles 5, or --cycles=5
    --layout             Display the Logical/Physical CPU layout table and exit.

CPU_LIST:
    Optional comma-separated list of logical CPUs or CPU ranges to test.
    Examples: 0-3, 1,3,5

Example:
  ${0##*/} --cycles=3 10,11

Note: Terminal history is your friend. Press up arrow if you want to re-run the last command.
EOF
}

show_layout() {
    declare -A seen
    printf "\n%-5s | %-5s | %-10s\n" "CPU" "Core" "Type"
    echo "------------------------"

    while IFS=',' read -r cpu core; do
        if [[ -z ${seen[$core]} ]]; then
            seen[$core]=1
            printf "%-5s | %-5s | %-10s\n" "$cpu" "$core" "Physical"
        else
            printf "%-5s | %-5s | %-10s\n" "$cpu" "$core" "Logical"
        fi
    done < <(lscpu -p=CPU,Core | grep -v '^#')
    exit 0
}

parse_cpu_list() {
    local input=$1
    local -a list=()
    local -a parts
    local part start end i

    IFS=',' read -ra parts <<< "$input"

    for part in "${parts[@]}"; do
        if [[ $part =~ ^([0-9]+)-([0-9]+)$ ]]; then
            start=${BASH_REMATCH[1]}
            end=${BASH_REMATCH[2]}
            if (( start > end )); then
                echo "Invalid CPU range: $part"
                exit 1
            fi
            for ((i=start; i<=end; i++)); do list+=("$i"); done
        elif [[ $part =~ ^[0-9]+$ ]]; then
            list+=("$part")
        else
            echo "Invalid CPU list format: $part"
            exit 1
        fi
    done

    # Properly sort and remove duplicates
    mapfile -t cpus < <(printf "%s\n" "${list[@]}" | sort -nu)
}

# Preserve original arguments before parsing/shifting
original_args=("$@")

# Defaults
run_mode="async"
max_test_cycles=1
cpu_list_input=""

# If no arguments are specified, show help and exit
if [[ "$#" -eq 0 ]]; then
    show_help
    exit 0
fi

# Parse arguments
while [[ "$#" -gt 0 ]]; do
    case $1 in
        -h|--help) show_help; exit 0 ;;
        --layout) show_layout ;;
        -m|--mode) run_mode="${2,,}"; shift ;;
        --mode=*) run_mode="${1#*=}"; run_mode="${run_mode,,}" ;;
        -c|--cycles) max_test_cycles="$2"; shift ;;
        --cycles=*) max_test_cycles="${1#*=}" ;;
        -*) echo "Unknown parameter: $1"; exit 1 ;;
        *) cpu_list_input="$1" ;;
    esac
    shift
done

# Check requirements
if ! command -v stress-ng >/dev/null 2>&1; then
    echo -e "\e[94m::\e[0m stress-ng not found. Install with $ sudo pacman -S stress-ng"
    exit 1
fi

if ! [[ "$max_test_cycles" =~ ^[0-9]+$ ]] || (( max_test_cycles < 1 )); then
    echo "ERROR: Cycles must be a positive integer greater than 0."
    exit 1
fi

# Only request sudo if the script is actually proceeding to run a test
if [[ $EUID -ne 0 ]]; then
    echo "Requesting administrative privileges . . ."
    exec sudo "$0" "${original_args[@]}"
fi

# Get total logical CPUs.
max_cpu=$(( $(nproc) - 1 ))

# Determine CPU indexes to test.
if [[ -n "$cpu_list_input" ]]; then
    parse_cpu_list "$cpu_list_input"
    cpu_warning="logical CPUs: $(IFS=,; echo "${cpus[*]}")"
else
    cpus=($(seq 0 "$max_cpu"))
    cpu_warning="all logical CPUs"
fi

# Validate logical CPUs are within range.
for cpu in "${cpus[@]}"; do
    if (( cpu < 0 || cpu > max_cpu )); then
        echo "ERROR: CPU number out of range: $cpu"
        exit 1
    fi
done

while true; do
    echo
    echo -e "${RED}:: WARNING: This script will run a CPU stress test on $cpu_warning resulting in high temperatures and power draw.${RESET}"
    echo -e "${RED}            Prolonged usage may cause system instability, degradation, or damage.${RESET}"
    echo -e "${RED}            Do not leave the system unattended during use.${RESET}"
    echo -e "${RED}            Do not continue unless adequate cooling is in place and you understand the risks.${RESET}"
    echo -e "${RED}            You assume all responsibility for using this script.${RESET}"
    echo
    read -p "Do you want to continue? [Y/N]: " response
    case "$response" in
        y|Y) break ;;
        n|N) exit 0 ;;
        *) echo ;;
    esac
done
echo

while true; do
    read -p "Do you want to include a RAM test? [Y/N]: " ram_response
    case "$ram_response" in
        y|Y) include_ram_test=true; break ;;
        n|N) include_ram_test=false; break ;;
        *) echo ;;
    esac
done
echo

# Pass cpus array as string to embedded script.
cpus_str="$(IFS=,; echo "${cpus[*]}")"

# Create a temporary file.
stress_script=$(mktemp)

cleanup_main() {
    [[ -n "$stress_script" ]] && rm -f "$stress_script"
}
trap cleanup_main SIGINT SIGTERM SIGHUP SIGQUIT EXIT

# Embedded Script Generation
cat > "$stress_script" <<EOF
#!/bin/bash

blue="\e[94m::\e[0m"
yellow="\e[93m::\e[0m"
red="\e[91m::\e[0m"

color_stress_output() {
    awk -v c_pass="$BRIGHT_GREEN" -v c_fail="$RED" -v c_reset="$RESET" '
    /failed:/ {
        n = split(\$0, a, "failed: ")
        if (n > 1) {
            val = +a[2]
            if (val > 0) {
                print c_fail \$0 c_reset
                next
            } else {
                print c_pass \$0 c_reset
                next
            }
        }
    }
    { print }
    '
}

cleanup() {
    local exit_code=\$?
    [[ -n "\${stress_pid}" && "\${stress_pid}" =~ ^[0-9]+\$ ]] && kill "\${stress_pid}" 2>/dev/null
    [[ -n "\${journal_pid}" ]] && kill "\${journal_pid}" 2>/dev/null
    exit "\$exit_code"
}
trap cleanup SIGINT SIGTERM SIGHUP SIGQUIT EXIT

cpus_str="\$1"
run_mode="\$2"
test_type="\$3"
time_per_test="\$4"
max_test_cycles="\$5"

# Default values.
seconds_between_tests=3
max_ram_usage=90

# Convert HH:MM:SS to seconds.
IFS=: read -r hh mm ss <<< "\$time_per_test"
seconds_per_test=\$((10#\$hh * 3600 + 10#\$mm * 60 + 10#\$ss))

IFS=',' read -ra cpus <<< "\$cpus_str"
num_cpus=\${#cpus[@]}

# Launch interleaved journalctl stream unconditionally
journalctl -fk -o short-iso | sed -e "s/^/$(echo -e "${YELLOW}")[SYS LOG] /" -e "s/\$/$(echo -e "${RESET}")/" &
journal_pid=\$!

cycle_count=1
while true; do
    if (( cycle_count > max_test_cycles )); then exit 0; fi

    if [[ "\$run_mode" == "async" ]]; then
        eta=\$(printf '%02d:%02d:%02d' \$((seconds_per_test / 3600)) \$(((seconds_per_test % 3600) / 60)) \$((seconds_per_test % 60)))
        progress=\$((cycle_count * 100 / max_test_cycles))
        echo -e "\${yellow} Running async stress test on CPUs: \$cpus_str [Cycle \${cycle_count}/\${max_test_cycles}, Progress: \${progress}% ETA: \${eta}]"

        max_ram_per_thread=\$(( max_ram_usage / num_cpus ))

        # 1. Create an empty array to hold the PIDs
        stress_pids=()

        for cpu in "\${cpus[@]}"; do
            taskset -c "\$cpu" stress-ng \
                \$( [[ "\$test_type" == "cpu" ]] && echo "--cpu 1 --cpu-method all" ) \
                \$( [[ "\$test_type" == "ram" ]] && echo "--vm 1 --vm-bytes \${max_ram_per_thread}% --vm-method all --verify --vm-keep" ) \
                \$( [[ "\$test_type" == "all" ]] && echo "--cpu 1 --cpu-method all --vm 1 --vm-bytes \${max_ram_per_thread}% --vm-method all --verify --vm-keep" ) \
                --metrics-brief --timeout "\${seconds_per_test}s" 2>&1 | color_stress_output &

            # 2. Capture the PID of the newly launched stress-ng process
            stress_pids+=(\$!)
        done

        # 3. Wait ONLY for the targeted stress-ng processes to finish
        wait "\${stress_pids[@]}"

        echo -e "\${blue} Cooling down for \${seconds_between_tests} seconds . . ."
        sleep "\${seconds_between_tests}"

    else
        for idx in "\${!cpus[@]}"; do
            cpu="\${cpus[\$idx]}"
            time_remaining=\$(( (num_cpus - idx - 1) * seconds_per_test ))
            progress=\$(( idx * 100 / num_cpus ))
            eta=\$(printf '%02d:%02d:%02d' \$((time_remaining / 3600)) \$(((time_remaining % 3600) / 60)) \$((time_remaining % 60)))

            echo -e "\${yellow} Testing CPU \${cpu} [Cycle \${cycle_count}/\${max_test_cycles}, Progress: \${progress}% ETA: \${eta}]"

            taskset -c "\$cpu" stress-ng \
                \$( [[ "\$test_type" == "cpu" ]] && echo "--cpu 1 --cpu-method all" ) \
                \$( [[ "\$test_type" == "ram" ]] && echo "--vm 1 --vm-bytes \$max_ram_usage% --vm-method all --verify --vm-keep" ) \
                \$( [[ "\$test_type" == "all" ]] && echo "--cpu 1 --cpu-method all --vm 1 --vm-bytes \$max_ram_usage% --vm-method all --verify --vm-keep" ) \
                --metrics-brief --timeout "\${seconds_per_test}s" 2>&1 | color_stress_output &

            stress_pid=\$!
            wait "\$stress_pid"
            unset stress_pid

            echo -e "\${blue} Cooling down for \${seconds_between_tests} seconds . . ."
            sleep "\${seconds_between_tests}"
        done
    fi
    echo -e "\${red} Completed \${cycle_count} of \${max_test_cycles} test cycles . . ."
    ((cycle_count++))
done

[[ -n "\${journal_pid}" ]] && kill "\${journal_pid}" 2>/dev/null
EOF

chmod +x "$stress_script"

if [[ "$include_ram_test" == "true" ]]; then test_type="all"; else test_type="cpu"; fi
time_per_test="00:00:15"

# Execute normally in foreground
"$stress_script" "$cpus_str" "$run_mode" "$test_type" "$time_per_test" "$max_test_cycles"