#!/bin/bash
# inject_counters.sh - register the running CUDA client's per-channel GPU
# completion counters with the host sema service (the PART 45/46 manual
# procedure, made durable). The host then tracks each GR channel's consumed
# GPFIFO progress in the counter the client polls.
#
# Usage: tools/guest-diag/inject_counters.sh [guest-pid] [page-va] [count]
#   guest-pid  default: pgrep -x cdt in the guest
#   page-va    default: 0x20400f000 (observed counter page VA)
#   count      default: 20 (one per GR channel)
#
# Env:
#   WATCH=<seconds>  keep watching for new GR channels and register them,
#                    instead of a single pass (default 0 = single pass).
#   MAP=1            also enqueue the sema_map_add pin per channel (default
#                    on). The pin teaches the host the counter page's VA->PA,
#                    which is what lets the client's own release pushes into
#                    the counter page apply; without it the host logs
#                    "release miss ... (no translation)" and the counters
#                    never move. The sysfs write is async now (queue), so
#                    this is cheap.
#   MAP=0            skip the pins (counters only advance from the host's
#                    fallback put publishing).
#   GRS=0x..,0x..    override the channel handle list.
#
# Handle order: GR channels (eng=0) from the host channel dump, descending,
# are zipped with ascending counter index (the observed 1:1 layout; the
# last-created channel owns counter 0).
set -uo pipefail

HERE="$(cd "$(dirname "$0")" && pwd)"
VMEXEC="$HERE/../../scripts/vgpu_vm.py"
PARAM=/sys/module/open_iov_vfio/parameters
DONE=/tmp/inject_counters.done

PID="${1:-}"
PAGE_VA="${2:-0x20400f000}"
COUNT="${3:-20}"
WATCH="${WATCH:-0}"
MAP="${MAP:-1}"

if [ -z "$PID" ]; then
    PID=$(timeout 20 "$VMEXEC" exec "pgrep -x cdt" 2>/dev/null \
        | grep -E '^[0-9]+$' | head -1 || true)
fi
if [ -z "$PID" ]; then
    echo "ERROR: no cdt pid in the guest (pass one: $0 <pid> [page-va] [count])" >&2
    exit 1
fi

# Push the enumeration tool (the guest's /tmp is wiped on boot).
b64=$(base64 -w0 "$HERE/sem_counters.py")
timeout 30 "$VMEXEC" exec "echo '$b64' | base64 -d > /tmp/sem_counters.py" >/dev/null

# Only consider chan_dump lines produced after this script's own trigger:
# dmesg keeps the previous client's dump lines (and handles) forever, and the
# ring buffer wraps so a line-count baseline is not a stable cursor. Write a
# unique marker to /dev/kmsg before each dump and parse only lines after it.
DMESG_MARK="sema-inject-marker-$$-$(date +%s)"

all_grs() {
    echo "$DMESG_MARK" | sudo tee /dev/kmsg >/dev/null
    echo 1 | sudo tee "$PARAM/chan_dump" >/dev/null
    sleep 3
    sudo dmesg | awk -v m="$DMESG_MARK" \
        '{ l[NR] = $0 } index($0, m) { last = NR } END { for (i = last + 1; i <= NR; i++) print l[i] }' \
        | grep -o 'sema-chan-dump: h=0x[0-9a-f]* eng=0' \
        | grep -o '0x[0-9a-f]*' | sort -u | paste -sd, || true
}

register_handles() {
    local grs="$1"
    [ -z "$grs" ] && return 0

    # Counter index order: observed DESCENDING by channel handle (the client's
    # last-created channel owns counter 0). ORDER=asc restores creation order.
    if [ "${ORDER:-desc}" = "desc" ]; then
        grs=$(echo "$grs" | tr ',' '\n' | sort -r | paste -sd,)
    fi

    local pairs
    pairs=$(timeout 30 "$VMEXEC" exec \
        "python3 /tmp/sem_counters.py $PID $PAGE_VA --count=$COUNT --pairs $grs" \
        2>/dev/null | grep -E '^0x[0-9a-f]+,0x[0-9a-f]+,[0-9]+$' || true)
    if [ -z "$pairs" ]; then
        echo "WARN: enumeration returned no pairs for: $grs" >&2
        return 1
    fi

    local h va pa
    while IFS=, read -r h va pa; do
        # sema_counter is a direct state write (no service-thread round trip),
        # so this loop completes in milliseconds.
        if printf '%s,%s' "$h" "$pa" | sudo tee "$PARAM/sema_counter" >/dev/null; then
            echo "$h" >> "$DONE"
            echo "registered h=$h counter va=$va pa=$pa"
        else
            echo "WARN: sema_counter write failed for h=$h pa=$pa" >&2
        fi
        if [ "$MAP" = "1" ]; then
            printf '%s,%s' "$va" "$pa" | sudo tee "$PARAM/sema_map_add" >/dev/null \
                || echo "WARN: sema_map_add timed out for h=$h (pin skipped)" >&2
        fi
    done <<< "$pairs"
    return 0
}

: > "$DONE"

if [ -n "${GRS:-}" ]; then
    register_handles "$GRS"
    exit 0
fi

GRS=$(all_grs)

if [ "$WATCH" = "0" ]; then
    if [ -z "$GRS" ]; then
        echo "ERROR: no GR channel handles (chan_dump empty; is a VM attached?)" >&2
        exit 1
    fi
    register_handles "$GRS"
    exit 0
fi

# Watch mode: register newly appearing channels until the deadline passes and
# no new handles show up for a few consecutive polls.
deadline=$(( $(date +%s) + WATCH ))
quiet=0
while [ "$(date +%s)" -lt "$deadline" ]; do
    GRS=$(all_grs)
    new=""
    for h in $(echo "$GRS" | tr ',' ' '); do
        grep -qx "$h" "$DONE" 2>/dev/null || new="$new,$h"
    done
    new="${new#,}"
    if [ -n "$new" ]; then
        register_handles "$new" || true
        quiet=0
    else
        quiet=$((quiet + 1))
        [ "$quiet" -ge 30 ] && break
    fi
    sleep 2
done
echo "watch done: $(wc -l < "$DONE") channel(s) registered"