#!/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 service then publishes each GR channel's
# consumed GPFIFO count into the client's counter, which is the completion
# signal its trackers poll.
#
# 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)
#
# Handle order: GR channels (eng=0) from the host channel dump, ascending, are
# zipped with ascending counter index (the observed 1:1 layout). Override the
# list with GRS="0x...,0x...,...".
set -euo pipefail

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

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

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

# Channel handles: GR channels from the live channel dump, ascending.
if [ -z "${GRS:-}" ]; then
    echo 1 | sudo tee "$PARAM/chan_dump" >/dev/null
    sleep 1
    GRS=$(sudo dmesg | grep -o 'sema-chan-dump: h=0x[0-9a-f]* eng=0' \
        | grep -o '0x[0-9a-f]*' | sort -u | paste -sd,)
fi
if [ -z "$GRS" ]; then
    echo "ERROR: no GR channel handles (chan_dump empty; is a VM attached?)" >&2
    exit 1
fi

# Counter index order: observed DESCENDING by channel handle (the client's
# last-created channel owns counter 0; verified live: trackers waiting on
# counter idx2/idx3 belong to the channels 2nd/3rd from the top). ORDER=asc
# restores creation order.
if [ "${ORDER:-desc}" = "desc" ]; then
    GRS=$(echo "$GRS" | tr ',' '\n' | sort -r | paste -sd,)
fi

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]+$')
if [ -z "$PAIRS" ]; then
    echo "ERROR: enumeration returned no pairs (pagemap access?)" >&2
    exit 1
fi

while IFS=, read -r h va pa; do
    # printf: sysfs writes carry a trailing newline that kstrtoull tolerates
    # for both fields; use it anyway (no newline is also fine).
    printf '%s,%s' "$h" "$pa" | sudo tee "$PARAM/sema_counter" >/dev/null
    # Pin the counter VA->PA too: a push that releases the counter then
    # resolves and applies through the normal release path, and pins survive
    # the periodic map rebuilds.
    printf '%s,%s' "$va" "$pa" | sudo tee "$PARAM/sema_map_add" >/dev/null
    echo "registered h=$h counter va=$va pa=$pa"
done <<< "$PAIRS"