#!/bin/bash
# SPDX-License-Identifier: BSD-3-Clause
# Copyright (c) 2026 Open-IOV contributors
# pin_supervisor.sh - supervise the host-side pin tools (pin_gr_stream.py and
# semfix_uvmt.py) so a silent death cannot strand every later CUDA client at
# first cudaMalloc.
#
# Both tools write a heartbeat file; this script checks, every 5 s, that each
# process is alive (exact script path) and that its heartbeat mtime is < 75 s.
# A startup grace (120 s, skipped early once the tool log shows it has done
# first useful work) keeps a slow cold start from being restarted. On failure
# it kills the tool (exact path match) and restarts it, appending RESTART
# lines to the supervisor log. After 10 restarts in one hour for a tool it
# logs GIVEUP and exits 1. Only one semfix_uvmt instance is ever started
# (QEMU's QMP socket accepts a single client).
#
# Liveness policy (do not tighten without re-checking the soak evidence):
#  * the tools refresh their heartbeats from dedicated loops/threads even
#    while idle (pin_gr_stream: 5 s thread; semfix: per loop + around every
#    blocking guest-agent/dmesg call), so 75 s never trips on a healthy
#    client between phases;
#  * a DEAD process is restarted on the next 5 s tick regardless of grace;
#  * dead-pid/state recovery (the kill cascade) is handled INSIDE the tools
#    (in-band #nopid + idle liveness probe + stall watchdog), NOT by this
#    supervisor - a log-based restart here could fire mid-run of a healthy
#    client (the m178 lesson).
#
# Usage: pin_supervisor.sh LOGDIR TAG
# Logs: LOGDIR/12-pin-supervisor-TAG.log (supervisor)
#       LOGDIR/12-pin_gr_stream-TAG.log, LOGDIR/12-semfix-TAG.log (tools)
set -uo pipefail

LOGDIR="${1:?usage: pin_supervisor.sh LOGDIR TAG}"
TAG="${2:?usage: pin_supervisor.sh LOGDIR TAG}"

HERE=$(cd "$(dirname "$0")" && pwd)
PIN="$HERE/pin_gr_stream.py"
SEMFIX="$HERE/semfix_uvmt.py"
PIN_HB="${PIN_GR_HEARTBEAT:-/tmp/opencode/pin_gr_heartbeat}"
SEMFIX_HB="${SEMFIX_HEARTBEAT_FILE:-/tmp/opencode/semfix_heartbeat}"
SUPLOG="$LOGDIR/12-pin-supervisor-$TAG.log"
PIN_LOG="$LOGDIR/12-pin_gr_stream-$TAG.log"
SEMFIX_LOG="$LOGDIR/12-semfix-$TAG.log"
WINDOW=3600
CAP=10
STALE=75
GRACE=120
mkdir -p "$LOGDIR" /tmp/opencode

log() { echo "[$(date +%H:%M:%S)] $*" | tee -a "$SUPLOG"; }

tool_alive() { pgrep -f "python3 $1" >/dev/null 2>&1; }

hb_age() {  # age in seconds of the heartbeat file, 9999 if missing
    local f="$1" m
    m=$(stat -c %Y "$f" 2>/dev/null) || { echo 9999; return; }
    echo $(( $(date +%s) - m ))
}

kill_tool() {
    pkill -f "python3 $1" 2>/dev/null
    sleep 1
    if pgrep -f "python3 $1" >/dev/null 2>&1; then
        pkill -9 -f "python3 $1" 2>/dev/null
        sleep 1
    fi
    return 0
}

start_pin() {
    if tool_alive "$PIN"; then return 0; fi
    touch "$PIN_HB" 2>/dev/null || true
    ( setsid nohup env SEMFIX_COUNTER_RETRY=10 \
        python3 "$PIN" >> "$PIN_LOG" 2>&1 & )
    pin_started=$(date +%s)
    log "START pin_gr_stream"
}

start_semfix() {
    # pgrep exact-count check before start: if an instance already exists
    # (ours or a leftover), never start a second one.
    if tool_alive "$SEMFIX"; then
        log "semfix already running - not starting another"
        return 0
    fi
    touch "$SEMFIX_HB" 2>/dev/null || true
    ( setsid nohup env SEMFIX_SKIP_PB=0 SEMFIX_COUNTER_RETRY=10 \
        python3 "$SEMFIX" 0.5 >> "$SEMFIX_LOG" 2>&1 & )
    semfix_started=$(date +%s)
    log "START semfix_uvmt"
}

prune() {  # $1 now, $2 list -> list with entries older than WINDOW dropped
    local now="$1" t out=""
    for t in $2; do
        [ $(( now - t )) -lt "$WINDOW" ] && out="$out $t"
    done
    printf '%s' "$out"
}

count_list() {  # $1 list -> element count
    local n=0 t
    for t in $1; do n=$(( n + 1 )); done
    echo "$n"
}

restarts_pin=""
restarts_semfix=""
pin_started=$(date +%s)
semfix_started=$(date +%s)

restart() {  # $1 label, $2 script path, $3 start function
    local now list n
    now=$(date +%s)
    if [ "$1" = pin_gr_stream ]; then
        restarts_pin=$(prune "$now" "$restarts_pin")
        n=$(count_list "$restarts_pin")
    else
        restarts_semfix=$(prune "$now" "$restarts_semfix")
        n=$(count_list "$restarts_semfix")
    fi
    if [ "$n" -ge "$CAP" ]; then
        log "GIVEUP $1: $n restarts within ${WINDOW}s"
        exit 1
    fi
    kill_tool "$2"
    if [ "$1" = pin_gr_stream ]; then
        restarts_pin="$restarts_pin $now"
    else
        restarts_semfix="$restarts_semfix $now"
    fi
    log "RESTART $(( n + 1 )) $1 $(date +%Y-%m-%dT%H:%M:%S)"
    "$3"
}

check_tool() {  # $1 label, $2 path, $3 heartbeat file, $4 start function,
                # $5 tool start TS, $6 tool log file, $7 first-work pattern
    local age runtime
    age=$(hb_age "$3")
    if tool_alive "$2" && [ "$age" -le "$STALE" ]; then
        return 0
    fi
    if tool_alive "$2"; then
        # Startup grace: a cold tool can spend its first ~2 min installing the
        # guest scanner or connecting QMP with a lagging heartbeat; only
        # enforce staleness once the grace has passed or the log shows the
        # tool has already done successful work.
        runtime=$(( $(date +%s) - $5 ))
        if [ "$runtime" -lt "$GRACE" ] && \
                ! grep -q -- "$7" "$6" 2>/dev/null; then
            return 0
        fi
        log "$1 heartbeat stale (${age}s) - restarting"
    else
        log "$1 not running (heartbeat age ${age}s) - restarting"
    fi
    restart "$1" "$2" "$4"
}

log "pin_supervisor up (LOGDIR=$LOGDIR TAG=$TAG)"
start_pin
start_semfix
while true; do
    sleep 5
    check_tool pin_gr_stream "$PIN" "$PIN_HB" start_pin "$pin_started" \
        "$PIN_LOG" "listening on vsock"
    check_tool semfix_uvmt "$SEMFIX" "$SEMFIX_HB" start_semfix \
        "$semfix_started" "$SEMFIX_LOG" "semfix_uvmt: watching"
done