#!/bin/sh
# script/loadtest: measure pixad's throughput, latency and peak memory.
#
#   script/loadtest [duration [clients]]    (defaults: 10s and 4)
#
# A benchmark, not a check: script/check does not run it. It needs Docker
# and Go. It builds the image with script/docker and builds vegeta, the
# load tool, from a pinned commit. Each scenario then gets a new pixad
# container and a new origin container (cmd/loadtest-origin, which answers
# every path with the same JPEG), and vegeta sends requests for <duration>
# from <clients> clients at once, each asking for an image resized to
# 400x300 WebP:
#
#   hit   the same image every time, put in the cache first
#   miss  a new source image every time
#   herd  each new source image once per client in a row, so that all
#         clients ask for it at the same time
#
# For each, it prints vegeta's report, pixad's peak resident memory and
# how many requests reached the origin. README.md says how to read them.
set -eu

SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd -P)"
ROOT="$(cd "$SCRIPT_DIR/.." && pwd -P)"

# vegeta v12.13.0, 2026-10-04
VEGETA_COMMIT=4b240c3089fa4aa10816542d64a74294d974211f

# pixad refuses upstream hosts with private or local addresses, so the
# containers share a network in 203.0.113.0/24, a range set aside for
# documentation (RFC 5737) that pixad does not refuse and that is never
# routed on the internet.
SUBNET=203.0.113.0/24

usage() {
    echo "usage: script/loadtest [duration [clients]]" >&2
    exit 2
}

main() {
    duration="${1:-10s}"
    clients="${2:-4}"
    # The duration is a whole number, not zero (vegeta takes 0 to mean no
    # end), followed by ms, s, m or h.
    case "$duration" in
        *ms) number="${duration%ms}" ;;
        *s | *m | *h) number="${duration%?}" ;;
        *) usage ;;
    esac
    case "$number" in
        "" | *[!0-9]*) usage ;;
    esac
    [ "$number" -gt 0 ] || usage
    # The number of clients is a whole number that does not start with 0,
    # which also refuses zero: vegeta reads a leading 0 as octal.
    case "$clients" in
        *[!0-9]* | 0*) usage ;;
    esac

    cd "$ROOT"
    run="pixa-loadtest-$$"
    tmp="$(mktemp -d)"
    trap cleanup EXIT
    trap 'exit 1' HUP INT TERM

    "$SCRIPT_DIR/docker"
    # The image's ID, so a build elsewhere that moves the tag does not
    # change what a later scenario starts.
    image="$(docker image inspect --format '{{.Id}}' \
        "$("$SCRIPT_DIR/projectname")")"

    GOBIN="$tmp" go install "github.com/tsenart/vegeta/v12@$VEGETA_COMMIT"
    # The origin runs in a container, so it is built for the Docker host.
    CGO_ENABLED=0 GOOS=linux \
        GOARCH="$(docker version --format '{{.Server.Arch}}')" \
        go build -o "$tmp/loadtest-origin" ./cmd/loadtest-origin

    docker network create --subnet "$SUBNET" "$run" >/dev/null

    start_containers
    # Put the image the hit scenario asks for in the cache.
    docker exec "$run-pixad" wget -q -O /dev/null \
        "http://localhost:8080/v1/image/origin/hit.jpg/400x300.webp"
    attack hit hit_targets
    stop_containers

    start_containers
    attack miss miss_targets
    stop_containers

    start_containers
    attack herd herd_targets
    stop_containers
}

# start_containers starts a new origin and a new pixad, and waits up to 30
# seconds for pixad's health check to pass.
start_containers() {
    docker run -d --name "$run-origin" \
        --network "$run" --network-alias origin \
        -v "$tmp/loadtest-origin:/usr/local/bin/loadtest-origin:ro" \
        --entrypoint /usr/local/bin/loadtest-origin "$image" >/dev/null
    docker run -d --name "$run-pixad" \
        --network "$run" -p 127.0.0.1::8080 --health-interval=1s \
        -e PIXA_SIGNING_KEY="$(head -c 32 /dev/urandom | base64)" \
        -e PIXA_ALLOWLIST_HOSTS=origin -e PIXA_ALLOW_HTTP=true \
        "$image" >/dev/null

    waited=0
    until [ "$(docker inspect --format '{{.State.Health.Status}}' \
        "$run-pixad")" = healthy ]; do
        if [ "$waited" -ge 30 ]; then
            echo "loadtest: pixad not healthy after 30 seconds; its log:" >&2
            docker logs "$run-pixad" >&2
            exit 1
        fi
        sleep 1
        waited=$((waited + 1))
    done
    pixa="http://$(docker port "$run-pixad" 8080/tcp)"
}

stop_containers() {
    docker rm -f "$run-pixad" "$run-origin" >/dev/null
}

# attack <scenario> <targets>: send the requests <targets> prints and
# report on them.
attack() {
    echo
    echo "== $1: $clients clients for $duration"
    "$2" | "$tmp/vegeta" attack -lazy -rate 0 -workers "$clients" \
        -max-workers "$clients" -duration "$duration" -max-body 0 |
        "$tmp/vegeta" report
    # pixad is process 1 in its container: the entrypoint execs it.
    echo "pixad peak memory (VmHWM):" \
        "$(docker exec "$run-pixad" awk '/^VmHWM:/ { print $2, $3 }' \
            /proc/1/status)"
    echo "requests to the origin:" \
        "$(docker logs "$run-origin" 2>&1 | grep -c ' request ')"
}

# The targets functions print vegeta targets until vegeta stops reading.

hit_targets() {
    while :; do
        echo "GET $pixa/v1/image/origin/hit.jpg/400x300.webp"
    done
}

miss_targets() {
    i=0
    while :; do
        i=$((i + 1))
        echo "GET $pixa/v1/image/origin/miss/$i.jpg/400x300.webp"
    done
}

herd_targets() {
    i=0
    while :; do
        i=$((i + 1))
        n=0
        while [ "$n" -lt "$clients" ]; do
            n=$((n + 1))
            echo "GET $pixa/v1/image/origin/herd/$i.jpg/400x300.webp"
        done
    done
}

cleanup() {
    docker rm -f "$run-pixad" "$run-origin" >/dev/null 2>&1 || :
    docker network rm "$run" >/dev/null 2>&1 || :
    rm -rf "$tmp"
}

main "$@"
