#!/usr/bin/env bash
set -euo pipefail

MD_DEV="${MD_DEV:-/dev/md/md777}"
STATE_DIR="${STATE_DIR:-/tmp/repro-md-noio-leak}"
MD_MEMBERS_STR="${MD_MEMBERS:-/dev/sdb /dev/sdc}"
WORKERS="${WORKERS:-32}"
ITERATIONS="${ITERATIONS:-200000}"
PF_MEMALLOC_NOIO=$((0x00080000))

need_root()
{
    [ "$(id -u)" -eq 0 ] || {
        echo "run as root" >&2
        exit 1
    }
}

need_cmd()
{
    command -v "$1" >/dev/null 2>&1 || {
        echo "$1 not found" >&2
        exit 1
    }
}

cleanup()
{
    set +e
    if [ -n "${MD_REAL:-}" ] && [ -e "/sys/block/$MD_REAL" ]; then
        mdadm -S "$MD_DEV" >/dev/null 2>&1 || true
    fi
    rm -rf "$STATE_DIR"
}
trap cleanup EXIT

need_root
for cmd in mdadm readlink bash findmnt lsblk; do
    need_cmd "$cmd"
done

mkdir -p "$STATE_DIR"

is_root_device()
{
    local dev="$1"
    local root_src root_real dev_real pkname

    root_src=$(findmnt -n -o SOURCE / || true)
    [ -n "$root_src" ] || return 1

    root_real=$(readlink -f "$root_src" 2>/dev/null || true)
    dev_real=$(readlink -f "$dev" 2>/dev/null || true)

    [ -n "$root_real" ] || return 1
    [ -n "$dev_real" ] || return 1

    if [ "$root_real" = "$dev_real" ]; then
        return 0
    fi

    pkname=$(lsblk -no PKNAME "$root_real" 2>/dev/null | head -n1 || true)
    [ -n "$pkname" ] || return 1

    [ "/dev/$pkname" = "$dev_real" ]
}

check_member_device()
{
    local dev="$1"
    local holders slave

    [ -b "$dev" ] || {
        echo "member $dev is not a block device" >&2
        exit 1
    }

    if findmnt -rn -S "$dev" >/dev/null 2>&1; then
        echo "member $dev is mounted" >&2
        exit 1
    fi

    if is_root_device "$dev"; then
        echo "member $dev looks like the root disk/device" >&2
        exit 1
    fi

    if [ -r "/sys/block/$(basename "$dev")/md/dev" ]; then
        echo "member $dev is already attached to an md array" >&2
        exit 1
    fi

    holders=$(find "/sys/block/$(basename "$dev")/holders" -mindepth 1 -maxdepth 1 2>/dev/null | head -n1 || true)
    if [ -n "$holders" ]; then
        echo "member $dev has active holders" >&2
        exit 1
    fi

    slave=$(find "/sys/block/$(basename "$dev")/slaves" -mindepth 1 -maxdepth 1 2>/dev/null | head -n1 || true)
    if [ -n "$slave" ]; then
        echo "member $dev is layered on another block device" >&2
        exit 1
    fi
}

read -r -a MD_MEMBERS <<< "$MD_MEMBERS_STR"
[ "${#MD_MEMBERS[@]}" -eq 2 ] || {
    echo "set MD_MEMBERS to exactly two block devices, e.g. '/dev/sdb /dev/sdc'" >&2
    exit 1
}

for dev in "${MD_MEMBERS[@]}"; do
    check_member_device "$dev"
done

mdadm -S "$MD_DEV" >/dev/null 2>&1 || true
mdadm --stop --scan >/dev/null 2>&1 || true

for dev in "${MD_MEMBERS[@]}"; do
    mdadm --zero-superblock --force "$dev" >/dev/null
done

mdadm -CR "$MD_DEV" -l1 -n2 --metadata=1.2 "${MD_MEMBERS[@]}" --assume-clean
MD_REAL=$(basename "$(readlink -f "$MD_DEV")")
MD_SYS="/sys/block/$MD_REAL/md"

[ -w "$MD_SYS/suspend_lo" ] || {
    echo "missing $MD_SYS/suspend_lo" >&2
    exit 1
}
[ -w "$MD_SYS/suspend_hi" ] || {
    echo "missing $MD_SYS/suspend_hi" >&2
    exit 1
}

worker_script='
set -euo pipefail

id="$1"
attr="$2"
value="$3"
iters="$4"
pf_noio="$5"
sys="$6"

read_self_flags() {
    local pid comm state ppid pgrp session tty_nr tpgid flags rest
    read -r pid comm state ppid pgrp session tty_nr tpgid flags rest < /proc/self/stat
    printf "%s\n" "$flags"
}

baseline=$(read_self_flags)
if (( baseline & pf_noio )); then
    echo "worker=$id pid=$$ baseline already has PF_MEMALLOC_NOIO flags=$baseline" >&2
    exit 2
fi

for ((i = 0; i < iters; i++)); do
    printf "%s\n" "$value" > "$sys/$attr"

    flags=$(read_self_flags)
    if (( flags & pf_noio )); then
        echo
        echo "REPRODUCED"
        echo "worker=$id pid=$$ attr=$attr iter=$i flags=$flags"
        echo "proc_stat: $(cat /proc/self/stat)"
        exit 99
    fi
done

exit 0
'

pids=()
for ((i = 0; i < WORKERS; i++)); do
    if (( i % 2 == 0 )); then
        attr="suspend_lo"
        value="0"
    else
        attr="suspend_hi"
        value="0"
    fi

    bash -c "$worker_script" _ \
        "$i" "$attr" "$value" "$ITERATIONS" "$PF_MEMALLOC_NOIO" "$MD_SYS" &
    pids+=("$!")
done

found=0
for pid in "${pids[@]}"; do
    if wait "$pid"; then
        :
    else
        rc=$?
        if [ "$rc" -eq 99 ]; then
            found=1
            break
        else
            echo "worker pid=$pid exited rc=$rc" >&2
        fi
    fi
done

if [ "$found" -eq 1 ]; then
    echo
    echo "SUCCESS: observed leaked PF_MEMALLOC_NOIO in a writer task"
    for pid in "${pids[@]}"; do
        kill "$pid" >/dev/null 2>&1 || true
    done
    exit 0
fi

echo
echo "NOT REPRODUCED in this run"
echo "Try increasing WORKERS and ITERATIONS"
exit 1