#!/bin/sh
# SPDX-License-Identifier: BSD-3-Clause
# Copyright (c) 2026 Open-IOV contributors
# guest_kprobes.sh - (re)install the UVM trace probes used by semfix_watch.py.
# Runs in the guest as root; kprobes do not survive a guest boot, so this must
# run after every VM restart.
#
#   uvmt  uvm_gpu_tracking_semaphore_update_completed_value  sem=%di
#   uvch  __uvm_push_begin_acquire_on_channel_with_info      ch=%di
#   uvbp  uvm_pushbuffer_begin_push                          pb=%di
#   uvpbf uvm_pushbuffer_create                              pb=%si
#
# Tries the plain symbol first (module symbols usually resolve), then the
# module:symbol spelling, then without the fetch argument (fetchargs on some
# symbols are rejected).

TR=/sys/kernel/tracing
[ -d "$TR" ] || TR=/sys/kernel/debug/tracing
# A fresh boot may leave tracefs unmounted even though the mount point exists
# (observed: mkdir of events/kprobes/enable fails with "Directory
# nonexistent" and every probe add fails silently). Mount it if needed.
if [ ! -d "$TR/events" ]; then
    mount -t tracefs tracefs "$TR" 2>/dev/null || \
        mount -t debugfs debugfs "$TR" 2>/dev/null || true
fi

# The UVM symbols only exist after nvidia_uvm is loaded. It loads lazily with
# the first CUDA client, but the probes must be armed BEFORE that client's
# first CE push - otherwise no #sem reaches the host, the push page never
# translates and the client stalls in cudaMalloc forever. modprobe it and
# wait (bounded) for the tracking-symbol to appear.
modprobe nvidia_uvm 2>/dev/null || true
i=0
while [ $i -lt 30 ]; do
    grep -qw uvm_gpu_tracking_semaphore_update_completed_value /proc/kallsyms 2>/dev/null && break
    i=$((i+1))
    sleep 1
done

# An enabled kprobe event cannot be deleted (EBUSY); disable the group first.
echo 0 > "$TR/events/kprobes/enable" 2>/dev/null || true

add_probe() {
    name=$1; sym=$2; fmt=$3
    echo "-:kprobes/$name" >> "$TR/kprobe_events" 2>/dev/null
    echo "-:$name" >> "$TR/kprobe_events" 2>/dev/null
    if ! grep -qw "$sym" /proc/kallsyms 2>/dev/null; then
        echo "FAIL $name: symbol $sym not in kallsyms"
        return 1
    fi
    if echo "p:$name $sym $fmt" >> "$TR/kprobe_events" 2>/dev/null; then
        if grep -qE "p:[^ ]*[ /]$name( |$)" "$TR/kprobe_events" 2>/dev/null; then
            echo "ok $name $sym"
            return 0
        fi
    fi
    echo "-:kprobes/$name" >> "$TR/kprobe_events" 2>/dev/null
    echo "-:$name" >> "$TR/kprobe_events" 2>/dev/null
    if echo "p:$name nvidia_uvm:$sym $fmt" >> "$TR/kprobe_events" 2>/dev/null; then
        if grep -qE "p:[^ ]*[ /]$name( |$)" "$TR/kprobe_events" 2>/dev/null; then
            echo "ok-module $name $sym"
            return 0
        fi
    fi
    echo "-:kprobes/$name" >> "$TR/kprobe_events" 2>/dev/null
    echo "-:$name" >> "$TR/kprobe_events" 2>/dev/null
    if echo "p:$name $sym" >> "$TR/kprobe_events" 2>/dev/null; then
        if grep -qE "p:[^ ]*[ /]$name( |$)" "$TR/kprobe_events" 2>/dev/null; then
            echo "ok-bare $name $sym"
            return 0
        fi
    fi
    echo "FAIL $name $sym"
    return 1
}

add_probe uvmt uvm_gpu_tracking_semaphore_update_completed_value 'sem=%di val=+16(%di):u32'
add_probe uvch __uvm_push_begin_acquire_on_channel_with_info 'ch=%di'
add_probe uvbp uvm_pushbuffer_begin_push 'pb=%di'
add_probe uvpbf uvm_pushbuffer_create 'pb=%si'
# uvm_tracker_entry_t { uvm_channel_t *channel; NvU64 value; } - the exact
# (channel, value) the client's tracker wait is blocked on.
add_probe uvtw uvm_tracker_wait_for_entry 'chan=+0(%di):u64 value=+8(%di):u64'
# Completion check: entry with (sem, value) and its boolean result tell
# whether the client can see the release the host applied.
add_probe uvsv uvm_gpu_tracking_semaphore_is_value_completed 'sem=%di value=%si'
echo "-:kprobes/uvsvr" >> "$TR/kprobe_events" 2>/dev/null
echo "-:uvsvr" >> "$TR/kprobe_events" 2>/dev/null
echo 'r:uvsvr uvm_gpu_tracking_semaphore_is_value_completed $retval' >> "$TR/kprobe_events" 2>/dev/null
# RPC wait tracing: _issueRpcAndWait is notrace in 570.144; probe the
# send/poll transport instead. An rpcp entry whose return is ~4s later is
# the timed-out RPC that fails RmInitAdapter (observed: first init aborts,
# second succeeds).
echo "-:kprobes/rpcs" >> "$TR/kprobe_events" 2>/dev/null
echo "-:rpcs" >> "$TR/kprobe_events" 2>/dev/null
echo 'p:rpcs _rpcSendMessage_VGPUGSP rpc=%si' >> "$TR/kprobe_events" 2>/dev/null
echo "-:kprobes/rpcp" >> "$TR/kprobe_events" 2>/dev/null
echo "-:rpcp" >> "$TR/kprobe_events" 2>/dev/null
echo 'p:rpcp rpcRecvPoll_IMPL gpu=%di' >> "$TR/kprobe_events" 2>/dev/null
echo "-:kprobes/rpcpr" >> "$TR/kprobe_events" 2>/dev/null
echo "-:rpcpr" >> "$TR/kprobe_events" 2>/dev/null
echo 'r:rpcpr rpcRecvPoll_IMPL $retval' >> "$TR/kprobe_events" 2>/dev/null
# Control-cache set tracing: retval 0x3b (INVALID_PARAMETER) names the cmd/size
# whose reply breaks the guest's cache bookkeeping.
echo "-:kprobes/ccset" >> "$TR/kprobe_events" 2>/dev/null
echo "-:ccset" >> "$TR/kprobe_events" 2>/dev/null
echo 'p:ccset rmapiControlCacheSet hCli=%di hObj=%si cmd=%dx size=%r8' >> "$TR/kprobe_events" 2>/dev/null
echo "-:kprobes/ccsetr" >> "$TR/kprobe_events" 2>/dev/null
echo "-:ccsetr" >> "$TR/kprobe_events" 2>/dev/null
echo 'r:ccsetr rmapiControlCacheSet $retval' >> "$TR/kprobe_events" 2>/dev/null

echo 1 > "$TR/events/kprobes/enable" 2>/dev/null || true
# The client's spin floods the trace with uvmt/uvsv events; a one-shot uvbp
# (first push-buffer begin) can be overwritten before the pin scanner reads
# it. Enlarge the ring and clear it so the first events survive.
echo 65536 > "$TR/buffer_size_kb" 2>/dev/null || true
echo > "$TR/trace" 2>/dev/null || true
echo "--- installed probes ---"
grep -E 'uvmt|uvch|uvbp|uvpbf|uvtw|uvsv' "$TR/kprobe_events" 2>/dev/null || echo "(none)"

# Arming check: the host recovery pipeline (#sem -> pushbuffer pin) depends on
# these probes; make success/failure explicit for the bring-up scripts.
if grep -q "uvmt" "$TR/kprobe_events" 2>/dev/null; then
    echo "KP_ARMED"
else
    echo "KP_FAIL uvmt missing"
fi