Vibe-coded (Gemini 2.5 Pro) script for simple rate-limit testing

rate_limit_tester.sh

#!/bin/bash

# --- Parallel Rate Limit Tester (v11 - Preserve Exit Status) ---
# This version preserves the final progress line on exit, ensuring full
# context is available for failures or user-initiated aborts (Ctrl+C).

# --- Usage ---
# ./rate_limit_tester.sh <URL> [parallel_requests] [max_requests] [min_delay_ms]

# --- Dependencies ---
# 'bc' is required for floating-point math.
# On Debian/Ubuntu: sudo apt-get install bc
# On RHEL/CentOS: sudo yum install bc

# --- Configuration ---
URL="$1"
PARALLEL_REQUESTS=${2:-50}
MAX_REQUESTS=${3:-50000}
MIN_DELAY_MS=${4:-0}

# --- Argument & Dependency Check ---
if ! command -v bc &> /dev/null; then echo "Error: 'bc' not found. Please install it."; exit 1; fi
if [ -z "$URL" ]; then echo "Usage: $0 <URL> [parallel_requests] [max_requests] [min_delay_ms]"; exit 1; fi

# --- State Management & Cleanup ---
RUN_DIR=$(mktemp -d)

# General cleanup function, called on ANY exit.
cleanup() {
    pkill -P $$ &>/dev/null
    rm -rf "$RUN_DIR"
}

# Specific handler for Ctrl+C (SIGINT signal).
abort_handler() {
    echo "" # Move to a new line, preserving the last progress bar
    echo "--- Test Aborted by User (Ctrl+C) ---"
    # Recalculate and print the absolute final status
    local current_time=$(date +%s.%N)
    local elapsed_time=$(echo "$current_time - $start_time" | bc)
    if (( $(echo "$elapsed_time > 0" | bc -l) )); then
        local rate=$(echo "scale=1; $completed_count / $elapsed_time" | bc)
    else
        local rate="0.0"
    fi
    echo "Final Status: Sent: $request_count | Completed: $completed_count | Rate: ${rate}/s"
    exit 130 # Exit with the standard code for Ctrl+C
}

# Set the traps. The EXIT trap is a catch-all for cleanup.
# The SIGINT trap specifically handles the Ctrl+C event.
trap cleanup EXIT
trap abort_handler SIGINT

MIN_DELAY_S=$(echo "scale=3; $MIN_DELAY_MS / 1000" | bc)

# --- Child Request Function ---
make_request() {
    local request_num=$1; local url=$2; local run_dir=$3; local min_delay=$4
    local lock_dir="$run_dir/lock"
    if [ -d "$lock_dir" ]; then return; fi
    local start_time=$(date +%s.%N)
    http_status=$(curl --connect-timeout 5 -o /dev/null -s -w "%{http_code}" "$url")
    local end_time=$(date +%s.%N)
    if (( $(echo "$min_delay > 0" | bc -l) )); then
        local request_duration=$(echo "$end_time - $start_time" | bc)
        local sleep_duration=$(echo "$min_delay - $request_duration" | bc)
        if (( $(echo "$sleep_duration > 0" | bc -l) )); then
            sleep "$sleep_duration"
        fi
    fi
    if [ "$http_status" -ne 200 ]; then
        if mkdir "$lock_dir" 2>/dev/null; then
            echo "Request #$request_num FAILED with status code: $http_status" > "$run_dir/failure_info.txt"
        fi
    fi
}
export -f make_request; export URL; export RUN_DIR; export MIN_DELAY_S

# --- Script Start ---
echo "--- Starting Dynamic Wait Test on: $URL ---"
echo "--- Running up to $PARALLEL_REQUESTS workers ---"
if (( $(echo "$MIN_DELAY_S > 0" | bc -l) )); then
    echo "--- Enforcing minimum request time of ${MIN_DELAY_MS}ms ---"
fi
echo "Press [CTRL+C] to stop."
echo ""

request_count=0
completed_count=0
start_time=$(date +%s.%N)

# --- Main Loop ---
while [ $request_count -lt $MAX_REQUESTS ]; do
    pids=()
    for i in $(seq 1 $PARALLEL_REQUESTS); do
        ((request_count++))
        bash -c 'make_request "$0" "$URL" "$RUN_DIR" "$MIN_DELAY_S"' "$request_count" &
        pids+=($!)
    done

    while [ ${#pids[@]} -gt 0 ]; do
        if [ -f "$RUN_DIR/failure_info.txt" ]; then
            # A failure was detected.
            echo "" # Move to a new line, preserving the final progress status.
            cat "$RUN_DIR/failure_info.txt"
            echo "--- Test Stopped ---"
            exit 0 # This will trigger the 'cleanup' trap.
        fi

        current_time=$(date +%s.%N)
        elapsed_time=$(echo "$current_time - $start_time" | bc)
        if (( $(echo "$elapsed_time > 0" | bc -l) )); then
            rate=$(echo "scale=1; $completed_count / $elapsed_time" | bc)
        else rate="0.0"; fi
        echo -ne "Sent: $request_count | Completed: $completed_count | Active: ${#pids[@]} | Rate: ${rate}/s \r"

        still_running_pids=()
        for pid in "${pids[@]}"; do
            if kill -0 "$pid" 2>/dev/null; then still_running_pids+=("$pid");
            else ((completed_count++)); fi
        done
        pids=("${still_running_pids[@]}")
        sleep 0.02
    done
done

echo "" # Move to a new line after the final progress bar.
echo "--- Test finished after $MAX_REQUESTS requests with no failures detected. ---"