#!/bin/bash
# stuck_client_capture.sh - one-shot capture for a stuck CUDA client in the
# vGPU guest (run on the HOST; talks to the guest over ssh_guest.sh).
#
# The in-guest pin scanner consumes trace_pipe continuously, so kprobe reads
# return entries=0 while it runs. This tool:
#   1. pauses the scanner (kill -STOP),
#   2. reads the retained kprobe trace tail for the client's UVM waits
#      (uvtw/uvsv/uvmt) - the tracker semaphore kptr and the value waited,
#   3. walks the tracker object (tracker_walk.py) for its expected payload,
#   4. dumps the semaphore-pool page (pool_page.py) for applied values,
#   5. resumes the scanner.
# Usage: stuck_client_capture.sh <client-comm> [pool-pa-hex]
#   e.g. stuck_client_capture.sh victory_matmul 0x107fd8000
set -u
REPO="$(cd "$(dirname "$0")/../.." && pwd)"
SSH="$REPO/scripts/ssh_guest.sh"
CLIENT="${1:?usage: stuck_client_capture.sh <client-comm> [pool-pa]}"
POOL="${2:-0x107fd8000}"

echo "=== pushing walk helpers to the guest ==="
"$SSH" "cat > /tmp/tracker_walk.py" < "$REPO/tools/guest-diag/tracker_walk.py"
"$SSH" "cat > /tmp/pool_page.py" < "$REPO/tools/guest-diag/pool_page.py"

echo "=== pausing scanner ==="
timeout 30 "$SSH" "echo ubuntu | sudo -S sh -c 'pkill -STOP -f guest_pin_stream.py; echo STOPPED'" 2>/dev/null | tail -1

echo "=== kprobe trace tail (client $CLIENT) ==="
timeout 60 "$SSH" "echo ubuntu | sudo -S sh -c 'tail -c 300000 /sys/kernel/tracing/trace | grep -a \"\$CLIENT\" | grep -aE \"uvtw|uvsv|uvmt\" | tail -8'" 2>/dev/null | tail -10

echo "=== tracker walk (kptr from the uvsv line above) ==="
echo "    (run: sudo python3 /tmp/tracker_walk.py 0x<sem-kptr>)"

echo "=== pool page $POOL ==="
timeout 90 "$SSH" "echo ubuntu | sudo -S python3 /tmp/poolpage.py $POOL" 2>/dev/null | tail -12

echo "=== resuming scanner ==="
timeout 30 "$SSH" "echo ubuntu | sudo -S sh -c 'pkill -CONT -f guest_pin_stream.py; echo RESUMED'" 2>/dev/null | tail -1
echo "=== done ==="