benilla-octo/crates/benilla-ui/src/script/raise_prof.rs

//! Raises an addon's own `pcall` swallows (`WOW_UI_RAISES=<secs>`): how many a frame, what they
//! cost, and which line raised.
//!
//! **Why this needs an instrument of its own.** A raise out of a handler or a `C_Timer` callback
//! reaches the host: [`super::event::protected_call`] returns it, [`UiScript::take_errors`] hands
//! it over and the app logs one `ui_script: …` line a frame. A raise an addon wraps in its own
//! `pcall` reaches nothing — `pcall` catches it before the host is involved — so the error log of
//! a session paying for a raise every frame is indistinguishable from the log of a session paying
//! for none. The only symptom is a slow frame, attributed to whichever system ticked the VM.
//!
//! That is not a hypothetical shape. pfUI alone carries 25 `pcall(<binding>, …)` sites, several
//! inside an `OnUpdate`: `pcall(UnitXP, "distanceBetween", …)`, `pcall(GetCVar, cvar)`,
//! `pcall(UnitName, unit .. "target")`. Each is a deliberate feature probe — the addon is asking
//! "does this client have it?" — and each costs the whole of whatever a failed binding call costs.
//! That price used to be 6.95 ms apiece on a loaded interface (`third_party/mlua-sys/BENILLA.md`);
//! it is 2.8 us now, which is why the count is worth having even so — five a frame still names an
//! addon asking a question this client answers with a raise.
//!
//! A bench shape is not a substitute: `addon_harness::tick_bench` run with this armed over the
//! real AddOns folder reports `host raises 0.00/f` in every window it has, grouped and solo. The
//! shapes that raise need a world the bench does not seed, which is the whole reason the counter
//! has to work in a live session.
//!
//! Off, this module is one `OnceLock` read at VM construction and nothing else; the wrappers are
//! not installed and `pcall` is the stock one.

use std::collections::HashMap;
use std::time::Instant;

use mlua::{Lua, Table, Value};

use super::UiScript;

/// How many message rows the periodic report prints before rolling the rest into a `… N more`.
const REPORT_ROWS: usize = 8;

/// The report period from `WOW_UI_RAISES=<secs>`, read once; unparseable or non-positive is off.
fn env_period() -> Option<f32> {
    static PERIOD: std::sync::OnceLock<Option<f32>> = std::sync::OnceLock::new();
    *PERIOD.get_or_init(|| {
        std::env::var("WOW_UI_RAISES")
            .ok()
            .and_then(|v| v.parse::<f32>().ok())
            .filter(|secs| *secs > 0.0)
    })
}

/// One raise message's accumulation over the current window.
#[derive(Default)]
struct Slot {
    count: u32,
    us: f64,
}

/// The counter's state, one per VM. Its own `app_data` slot, not the model's: the notes are taken
/// from inside a running `pcall`, where a model borrow may be held by whatever the `pcall` called.
#[derive(Default)]
pub(crate) struct RaiseProf {
    /// Whether this VM counts. Set by [`UiScript::profile_raises`] or the environment.
    on: bool,
    /// Whether the Lua wrappers are in place; arming twice must not stack them.
    wrapped: bool,
    /// Seconds between reports; `0.0` records and never prints.
    period: f32,
    /// The epoch the Lua side's marks are measured from, so a mark is a plain number.
    epoch: Option<Instant>,
    /// Wall seconds and frames since the last report, the divisor of every printed rate.
    window: f32,
    frames: u32,
    /// Failures whose error value was a host userdata — a Rust binding's raise, which is the one
    /// that pays mlua's traceback.
    host_raises: u32,
    host_us: f64,
    /// Failures whose error value was a Lua string: an ordinary Lua error, which costs nothing
    /// beyond the unwind. Counted so a report can say the raises are *not* the host's.
    lua_errors: u32,
    lua_us: f64,
    rows: HashMap<String, Slot>,
}

/// The Lua half: `pcall` and `xpcall` replaced by wrappers that time the call and, when it fails,
/// hand the host the error value's kind and text.
///
/// The chunk is a builder taking the two host functions as named parameters rather than reading
/// them out of `_G`, so nothing new is published as a global and an addon cannot reach them.
///
/// Varargs are 5.0's: `...` is not an expression in this dialect (the client's parser has no
/// `TK_DOTS` arm in `simpleexp`, `third_party/lua-src/BENILLA.md`), so a vararg function reads its
/// arguments from `arg` and `arg.n` — which is exact, where `#arg` would stop at the first nil a
/// caller passed. `collect` captures the protected call's own returns the same way, so a `pcall`
/// that returned `true, nil, 5` still returns three values through the wrapper.
const WRAP: &str = r#"
    return function(mark, note)
      local pcall_raw, xpcall_raw = pcall, xpcall
      local type, tostring, unpack = type, tostring, unpack
      local function collect(...) return arg end
      -- A Rust binding's raise arrives as a host userdata (mlua's `WrappedFailure`); a Lua
      -- `error()` arrives as whatever was thrown, a string in every corpus case. `tostring` on the
      -- userdata runs mlua's `__tostring`, which is the message plus the traceback already built.
      local function seen(t0, e)
        if type(e) == "userdata" then note(t0, tostring(e)) else note(t0, false) end
      end
      local function counted_pcall(...)
        local t0 = mark()
        local r = collect(pcall_raw(arg[1], unpack(arg, 2, arg.n)))
        if not r[1] then seen(t0, r[2]) end
        return unpack(r, 1, r.n)
      end
      -- `xpcall`'s handler runs at the raise site, before the unwind, so the error value has to be
      -- read there: by the time `xpcall` returns, `h` has already replaced it with its own answer.
      local function counted_xpcall(f, h)
        local t0 = mark()
        local r = collect(xpcall_raw(f, function(e) seen(t0, e) return h(e) end))
        return unpack(r, 1, r.n)
      end
      return counted_pcall, counted_xpcall
    end
"#;

/// Put the wrappers in place. Idempotent: a second arm is a no-op, so the wrappers can never wrap
/// each other and double-count.
fn wrap(lua: &Lua) -> mlua::Result<()> {
    match lua.app_data_mut::<RaiseProf>() {
        Some(mut prof) if !prof.wrapped => {
            prof.wrapped = true;
            prof.epoch = Some(Instant::now());
        }
        _ => return Ok(()),
    }
    let mark = lua.create_function(|lua, ()| {
        let prof = lua
            .app_data_ref::<RaiseProf>()
            .expect("raise prof app_data");
        Ok(prof.epoch.map_or(0.0, |e| e.elapsed().as_secs_f64() * 1e6))
    })?;
    let note = lua.create_function(|lua, (t0, err): (f64, Value)| {
        let mut prof = lua
            .app_data_mut::<RaiseProf>()
            .expect("raise prof app_data");
        if !prof.on {
            return Ok(());
        }
        let us = (prof.epoch.map_or(0.0, |e| e.elapsed().as_secs_f64() * 1e6) - t0).max(0.0);
        match err {
            Value::String(s) => {
                prof.host_raises += 1;
                prof.host_us += us;
                let key = message_key(&s.to_string_lossy());
                let slot = prof.rows.entry(key).or_default();
                slot.count += 1;
                slot.us += us;
            }
            // `false`: a plain Lua error, counted but not keyed — its message is the addon's own
            // business and costs nothing to raise.
            _ => {
                prof.lua_errors += 1;
                prof.lua_us += us;
            }
        }
        Ok(())
    })?;
    let builder: mlua::Function = lua
        .load(WRAP)
        .set_name("=[benilla raise counter]")
        .set_mode(mlua::ChunkMode::Text)
        .eval()?;
    let (counted_pcall, counted_xpcall): (mlua::Function, mlua::Function) =
        builder.call((mark, note))?;
    let globals: Table = lua.globals();
    globals.set("pcall", counted_pcall)?;
    globals.set("xpcall", counted_xpcall)?;
    Ok(())
}

/// A raise's report key: the message, then the first *addon* frame of the traceback mlua already
/// built, which is the line that made the call.
///
/// mlua formats a `CallbackError` as the cause, a newline, then the traceback, so line 1 is the
/// message. The frames above the raise are the binding (`[C]: in function 'UnitName'`) and the
/// `pcall` that caught it, which every raise of this message shares, and then — because this
/// counter works by standing in for `pcall` — [`WRAP`]'s own two frames. All of those are skipped:
/// the first frame left is the addon's, which is the only one a report can act on.
fn message_key(text: &str) -> String {
    let mut lines = text.lines();
    let head = lines.next().unwrap_or_default().trim();
    let site = lines
        .map(str::trim)
        .find(|l| {
            // `[C]`: a C function, i.e. the binding or the stock `pcall`. `[benilla `: one of the
            // host's own chunks, this counter's wrapper above all.
            !l.starts_with("[C]") && !l.starts_with("[benilla ") && looks_like_a_source_frame(l)
        })
        .unwrap_or("<no addon frame>");
    format!("{head}  ·  {site}")
}

/// Whether a traceback line names a source position: `<src>:<line>:` with digits after the last
/// colon pair. Cheap and tolerant — a key that misreads a frame is a worse label, never a panic.
fn looks_like_a_source_frame(line: &str) -> bool {
    line.split(':')
        .any(|part| !part.is_empty() && part.bytes().all(|b| b.is_ascii_digit()))
}

/// One raise message's cost over the counter's current window.
#[derive(Clone, Debug, PartialEq)]
pub struct RaiseRow {
    /// The raise's message and the line that called the binding, as one label.
    pub message: String,
    pub count: u32,
    /// Microseconds spent inside the `pcall`s that ended in this message.
    pub us: f64,
}

impl UiScript {
    /// Arm or disarm swallowed-raise counting for this VM, reporting every `period_secs` (`0.0`:
    /// no report, read [`UiScript::raise_profile`]). Arming replaces `pcall` and `xpcall` with
    /// wrappers, so call it before the code under measurement captures them into locals.
    pub fn profile_raises(&self, on: bool, period_secs: f32) {
        {
            let Some(mut prof) = self.lua.app_data_mut::<RaiseProf>() else {
                return;
            };
            prof.on = on;
            prof.period = period_secs;
            prof.window = 0.0;
            prof.frames = 0;
            prof.host_raises = 0;
            prof.host_us = 0.0;
            prof.lua_errors = 0;
            prof.lua_us = 0.0;
            prof.rows.clear();
        }
        if on {
            // The wrappers stay in place once armed: taking them back out would restore whatever
            // `pcall` was at the time, which by then may be another addon's own hook.
            let _ = wrap(&self.lua);
        }
    }

    /// This VM's swallowed host raises over the current window, costliest first.
    pub fn raise_profile(&self) -> Vec<RaiseRow> {
        let Some(prof) = self.lua.app_data_ref::<RaiseProf>() else {
            return Vec::new();
        };
        let mut rows: Vec<RaiseRow> = prof
            .rows
            .iter()
            .map(|(message, slot)| RaiseRow {
                message: message.clone(),
                count: slot.count,
                us: slot.us,
            })
            .collect();
        rows.sort_by(|a, b| {
            b.us.total_cmp(&a.us)
                .then_with(|| a.message.cmp(&b.message))
        });
        rows
    }

    /// Count a ticked frame and print the report when due, at the end of [`UiScript::tick`].
    pub(super) fn report_raise_profile(&mut self, elapsed: f32) {
        let due = match self.lua.app_data_mut::<RaiseProf>() {
            Some(mut prof) if prof.on => {
                prof.frames += 1;
                prof.window += elapsed;
                prof.period > 0.0 && prof.window >= prof.period
            }
            _ => false,
        };
        if !due {
            return;
        }
        let rows = self.raise_profile();
        let mut prof = self
            .lua
            .app_data_mut::<RaiseProf>()
            .expect("raise prof app_data — checked above");
        print_report(
            &rows,
            prof.window,
            prof.frames,
            (prof.host_raises, prof.host_us),
            (prof.lua_errors, prof.lua_us),
        );
        prof.window = 0.0;
        prof.frames = 0;
        prof.host_raises = 0;
        prof.host_us = 0.0;
        prof.lua_errors = 0;
        prof.lua_us = 0.0;
        prof.rows.clear();
    }
}

/// The `[ui-raises]` block: the window, the two kinds of caught failure per frame, then the
/// heaviest messages with the line that called the binding.
fn print_report(rows: &[RaiseRow], window: f32, frames: u32, host: (u32, f64), lua: (u32, f64)) {
    let per_frame = f64::from(frames.max(1));
    eprintln!(
        "[ui-raises] {window:.2}s · {frames} frames · host raises {:.2}/f ({:.2} ms/f) · \
         lua errors {:.2}/f ({:.2} ms/f)",
        f64::from(host.0) / per_frame,
        host.1 / per_frame / 1000.0,
        f64::from(lua.0) / per_frame,
        lua.1 / per_frame / 1000.0,
    );
    if rows.is_empty() {
        eprintln!("[ui-raises] no `pcall` caught a host raise this window");
        return;
    }
    eprintln!("[ui-raises]  count/f    ms/f  message  ·  raised from");
    for row in rows.iter().take(REPORT_ROWS) {
        eprintln!(
            "[ui-raises] {:8.2} {:7.2}  {}",
            f64::from(row.count) / per_frame,
            row.us / per_frame / 1000.0,
            row.message,
        );
    }
    if let Some(rest) = rows.len().checked_sub(REPORT_ROWS).filter(|n| *n > 0) {
        let tail: f64 = rows.iter().skip(REPORT_ROWS).map(|r| r.us).sum();
        eprintln!(
            "[ui-raises]          {:7.2}  … {rest} more",
            tail / per_frame / 1000.0
        );
    }
}

/// Install the counter's `app_data` slot, armed from the environment. Unconditional, like the
/// handler profiler's: mlua refuses to insert app data while any is borrowed, so the slot must
/// exist before the VM runs.
///
/// Called last in [`UiScript::new`], after [`super::event::install`], so the handler trampoline's
/// captured `xpcall` is the stock one. Handler and `C_Timer` raises already reach the error log;
/// this counter is for the ones that do not.
pub(super) fn install(lua: &Lua) -> mlua::Result<()> {
    let period = env_period();
    lua.set_app_data(RaiseProf {
        on: period.is_some(),
        period: period.unwrap_or_default(),
        ..RaiseProf::default()
    });
    if period.is_some() {
        eprintln!(
            "[ui-raises] armed — counting the host raises addon `pcall`s swallow, every {:.1}s \
             (WOW_UI_RAISES)",
            period.unwrap_or_default()
        );
        wrap(lua)?;
    }
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;

    /// A silent instrument is worse than none: this is the test that the counter sees a raise
    /// `pcall` swallowed, which is the whole class it exists for and the one class no log shows.
    #[test]
    fn the_counter_sees_a_binding_raise_that_pcall_swallowed() {
        let s = UiScript::new().expect("VM");
        s.profile_raises(true, 0.0);
        // `UnitName`'s argument goes through `binding_abi::string_arg`: a table raises
        // `Usage: UnitName("unit")`, and the `pcall` means nothing else ever hears about it.
        s.run(
            r#"
            local f = CreateFrame("Frame")
            for _ = 1, 3 do pcall(UnitName, f) end
            "#,
        )
        .expect("the pcall swallows the raise, so the chunk itself succeeds");
        assert!(
            s.errors().is_empty(),
            "a swallowed raise must reach no error channel — that is why it needs counting: {:?}",
            s.errors()
        );
        let rows = s.raise_profile();
        assert_eq!(rows.len(), 1, "one message, three times: {rows:?}");
        assert_eq!(rows[0].count, 3);
        assert!(
            rows[0].message.contains(r#"Usage: UnitName("unit")"#),
            "the row names the raise: {}",
            rows[0].message
        );
        // And the frame it names is the caller's, not the counter's own wrapper: a report that
        // pointed at `[benilla raise counter]` would name the instrument in every row.
        assert!(
            !rows[0].message.contains("[benilla "),
            "the row must name the addon frame, not a host chunk: {}",
            rows[0].message
        );
    }

    /// An ordinary Lua error is counted apart from a binding's raise, because only the binding's
    /// pays for a traceback: a report that lumped them would point at the wrong half.
    #[test]
    fn a_plain_lua_error_is_counted_apart_from_a_binding_raise() {
        let s = UiScript::new().expect("VM");
        s.profile_raises(true, 0.0);
        s.run("for _ = 1, 4 do pcall(function() error(\"mine\") end) end")
            .expect("chunk");
        assert!(
            s.raise_profile().is_empty(),
            "a Lua error is not a host raise: {:?}",
            s.raise_profile()
        );
        let prof = s
            .lua()
            .app_data_ref::<RaiseProf>()
            .expect("raise prof app_data");
        assert_eq!(prof.lua_errors, 4);
        assert_eq!(prof.host_raises, 0);
    }

    /// The wrapper must be transparent: `pcall`'s return arity and its embedded nils are what the
    /// corpus reads, and a wrapper that collapsed them would break addons only while armed.
    #[test]
    fn the_wrapper_returns_exactly_what_pcall_returned() {
        let s = UiScript::new().expect("VM");
        s.profile_raises(true, 0.0);
        let shape: String = s
            .eval(
                r##"
                local function f(a, b) return nil, a, nil, b end
                local t = { pcall(f, 7, nil) }
                -- `ok` then f's four returns: true, nil, 7, nil, nil — five values, three nil.
                return tostring(select("#", pcall(f, 7, nil))) .. "/" .. tostring(t[1])
                  .. "/" .. tostring(t[3])
                "##,
            )
            .expect("eval");
        assert_eq!(
            shape, "5/true/7",
            "the wrapper must forward all five values, nils included"
        );
    }

    /// Arming twice must not stack a wrapper on a wrapper: a doubled count would read as twice
    /// the raises, and the cost would be attributed twice over.
    #[test]
    fn arming_twice_does_not_double_count() {
        let s = UiScript::new().expect("VM");
        s.profile_raises(true, 0.0);
        s.profile_raises(true, 0.0);
        s.run("local f = CreateFrame(\"Frame\") pcall(UnitName, f)")
            .expect("chunk");
        let rows = s.raise_profile();
        assert_eq!(rows.len(), 1, "{rows:?}");
        assert_eq!(rows[0].count, 1, "one raise, counted once");
    }
}