crates/nornir-trace/src/xray.rs
//! **LLVM XRay, basic ("naive") mode — read natively.**
//!
//! How a run is traced (the runner's side, in the runa toolchain image):
//!
//! ```text
//! RUSTFLAGS="-Z instrument-xray=always -C link-arg=-fxray-instrument …" (RUSTFLAGS)
//! XRAY_OPTIONS="patch_premain=true xray_mode=xray-basic xray_logfile_base=<dir>/"
//! XRAY_BASIC_OPTIONS="func_duration_threshold_us=0" (keep short calls)
//! ```
//!
//! `instrument-xray=always` puts sleds in EVERY function the crate graph compiles
//! (the default skips functions under 200 machine instructions); the basic
//! runtime drops calls shorter than 5 µs unless `func_duration_threshold_us=0`.
//! A function inlined away has no frame and so no sled: the trace is the run as
//! the optimiser left it.
//!
//! Each traced process writes one log, `<base><program>.<random>`. Its records
//! name a function by **id**: the 1-based index of the function in the program's
//! `xray_instr_map` section (one 32-byte sled per entry/exit point, grouped per
//! function). This module reads that section from the program's ELF, names each
//! function by the symbol table, and turns the records into [`TraceEvent`]s.
use std::collections::HashMap;
use anyhow::{bail, Context, Result};
use object::{Object, ObjectSection, ObjectSymbol};
use crate::{EventKind, TraceEvent, TraceRun, TracedBinary, TracedFn};
/// The RUSTFLAGS that instrument every function and link the XRay runtime (clang
/// drives the link; it adds compiler-rt's `libclang_rt.xray*` for
/// `-fxray-instrument`). The caller adds its own linker flags (RUSTFLAGS
/// replaces a config's `rustflags`).
pub const RUSTFLAGS: &str = "-Z instrument-xray=always -C link-arg=-fxray-instrument";
/// `XRAY_OPTIONS` for a run whose logs go below `log_base` (a directory with a
/// trailing `/`, or a file-name prefix).
pub fn xray_options(log_base: &str) -> String {
format!("patch_premain=true xray_mode=xray-basic verbosity=0 xray_logfile_base={log_base}")
}
/// `XRAY_BASIC_OPTIONS`: keep every call, however short (the default drops < 5 µs).
pub const BASIC_OPTIONS: &str = "func_duration_threshold_us=0";
/// The program a basic-mode log came from: `<base><program>.<random>` → `program`.
pub fn log_program(file_name: &str, base_prefix: &str) -> Option<String> {
let rest = file_name.strip_prefix(base_prefix)?;
let (prog, rand) = rest.rsplit_once('.')?;
(!prog.is_empty() && !rand.is_empty()).then(|| prog.to_string())
}
/// A basic-mode log's header.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct Header {
pub version: u16,
/// 0 = basic ("naive") log; 1 = flight-data-recorder (not read here).
pub kind: u16,
pub constant_tsc: bool,
pub nonstop_tsc: bool,
/// TSC ticks per second.
pub cycle_frequency: u64,
}
/// One function record of a basic-mode log.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct Record {
pub cpu: u8,
/// 0 entry, 1 exit, 2 tail exit, 3 entry with args (4/5: custom/typed events).
pub kind: u8,
pub func_id: i32,
pub tsc: u64,
pub tid: u32,
pub pid: u32,
}
const HEADER_LEN: usize = 32;
const RECORD_LEN: usize = 32;
fn u16_at(b: &[u8], o: usize) -> u16 {
u16::from_le_bytes([b[o], b[o + 1]])
}
fn u32_at(b: &[u8], o: usize) -> u32 {
u32::from_le_bytes(b[o..o + 4].try_into().unwrap())
}
fn u64_at(b: &[u8], o: usize) -> u64 {
u64::from_le_bytes(b[o..o + 8].try_into().unwrap())
}
/// **Parse a basic-mode XRay log** (versions 1–3): the header, then the function
/// records in file order (per-thread buffers are flushed one after another, so
/// file order is not time order). Argument-payload records are skipped. A torn
/// last record (the process died mid-write) is dropped.
pub fn parse_basic_log(b: &[u8]) -> Result<(Header, Vec<Record>)> {
if b.len() < HEADER_LEN {
bail!("an XRay log is at least a 32-byte header; this one is {} bytes", b.len());
}
let bits = u32_at(b, 4);
let h = Header {
version: u16_at(b, 0),
kind: u16_at(b, 2),
constant_tsc: bits & 1 != 0,
nonstop_tsc: bits & 2 != 0,
cycle_frequency: u64_at(b, 8),
};
if h.kind != 0 {
bail!("XRay log type {} is not a basic-mode log (xray_mode=xray-basic writes type 0)", h.kind);
}
if !(1..=3).contains(&h.version) {
bail!("XRay basic log version {} is not read here (1..=3)", h.version);
}
let mut out = Vec::with_capacity((b.len() - HEADER_LEN) / RECORD_LEN);
let mut o = HEADER_LEN;
while o + RECORD_LEN <= b.len() {
let r = &b[o..o + RECORD_LEN];
o += RECORD_LEN;
// RecordType: 0 = a function record, 1 = an argument payload
if u16_at(r, 0) != 0 {
continue;
}
out.push(Record {
cpu: r[2],
kind: r[3],
func_id: i32::from_le_bytes(r[4..8].try_into().unwrap()),
tsc: u64_at(r, 8),
tid: u32_at(r, 16),
pid: if h.version >= 3 { u32_at(r, 20) } else { 0 },
});
}
Ok((h, out))
}
/// **Function addresses by XRay id** from an `xray_instr_map` section's bytes
/// (`section_addr` = the section's virtual address): `out[id - 1]` is function
/// `id`'s entry address. Sleds are 32 bytes (`address`, `function`, kind,
/// always-instrument, version); version ≥ 2 sleds hold their two addresses
/// relative to the field itself (position-independent), older ones absolute.
/// A new function starts wherever a sled's function address changes — the
/// numbering llvm-xray and the runtime's `xray_fn_idx` agree on.
pub fn sled_functions(section: &[u8], section_addr: u64) -> Result<Vec<u64>> {
const SLED: usize = 32;
if section.len() % SLED != 0 {
bail!("xray_instr_map is {} bytes, not a whole number of 32-byte sleds", section.len());
}
let mut fns: Vec<u64> = Vec::new();
for (i, s) in section.chunks_exact(SLED).enumerate() {
let version = s[18];
let raw_fn = u64_at(s, 8);
let function = if version >= 2 {
// relative to the `function` field's own address
section_addr
.wrapping_add((i * SLED + 8) as u64)
.wrapping_add(raw_fn)
} else {
raw_fn
};
if fns.last() != Some(&function) {
fns.push(function);
}
}
Ok(fns)
}
/// [`sled_functions`] of a program's ELF bytes.
pub fn function_addresses(elf: &[u8]) -> Result<Vec<u64>> {
let file = object::File::parse(elf).context("parse the traced program as an object file")?;
let sec = file
.section_by_name("xray_instr_map")
.context("the program has no xray_instr_map section: it was not built with -Z instrument-xray")?;
let data = sec.data().context("read xray_instr_map")?;
sled_functions(data, sec.address())
}
/// **Name each address** by the program's symbol table (`.symtab`, else the
/// dynamic one): the demangled symbol, its hash dropped, and — when the program
/// carries DWARF line tables — the source file and line of its first
/// instruction. An address no symbol starts at is named `0x<hex>`.
pub fn symbolize(elf: &[u8], addrs: &[u64]) -> Result<Vec<TracedFn>> {
let file = object::File::parse(elf).context("parse the traced program as an object file")?;
let mut by_addr: HashMap<u64, String> = HashMap::new();
for s in file.symbols() {
if s.address() != 0 && s.is_definition() {
if let Ok(name) = s.name() {
by_addr.entry(s.address()).or_insert_with(|| name.to_string());
}
}
}
if by_addr.is_empty() {
for s in file.dynamic_symbols() {
if s.address() != 0 && s.is_definition() {
if let Ok(name) = s.name() {
by_addr.entry(s.address()).or_insert_with(|| name.to_string());
}
}
}
}
let lines = Lines::of(&file);
Ok(addrs
.iter()
.map(|a| {
let (file, line) = lines.as_ref().map(|l| l.at(*a)).unwrap_or_default();
TracedFn {
symbol: match by_addr.get(a) {
Some(n) => demangle(n),
None => format!("{a:#x}"),
},
file,
line,
}
})
.collect())
}
/// **Where each named function's code is** (for a sampled profile, which names
/// functions but not addresses): demangled symbol → its address in the
/// program's symbol table → the DWARF line of that address. `("", 0)` when the
/// program has no such symbol or no line tables.
pub fn locate_symbols(elf: &[u8], symbols: &[String]) -> Result<Vec<(String, u32)>> {
let file = object::File::parse(elf).context("parse the profiled program as an object file")?;
let mut by_name: HashMap<String, u64> = HashMap::new();
for s in file.symbols() {
if s.address() != 0 && s.is_definition() {
if let Ok(name) = s.name() {
by_name.entry(demangle(name)).or_insert(s.address());
}
}
}
let lines = Lines::of(&file);
Ok(symbols
.iter()
.map(|sym| match (by_name.get(sym), lines.as_ref()) {
(Some(a), Some(l)) => l.at(*a),
_ => (String::new(), 0),
})
.collect())
}
/// The program's DWARF line tables (`addr2line` over gimli), when it has any.
struct Lines<'a> {
ctx: addr2line::Context<addr2line::gimli::EndianSlice<'a, addr2line::gimli::RunTimeEndian>>,
}
impl<'a> Lines<'a> {
fn of(file: &object::File<'a>) -> Option<Self> {
use addr2line::gimli;
file.section_by_name(".debug_info")?;
let endian = if file.is_little_endian() { gimli::RunTimeEndian::Little } else { gimli::RunTimeEndian::Big };
let dwarf = gimli::Dwarf::load(|id| -> std::result::Result<_, gimli::Error> {
// a compressed section reads as empty: its lines are then unknown
let data = file
.section_by_name(id.name())
.and_then(|s| s.data().ok())
.unwrap_or(&[]);
Ok(gimli::EndianSlice::new(data, endian))
})
.ok()?;
let ctx = addr2line::Context::from_dwarf(dwarf).ok()?;
Some(Self { ctx })
}
fn at(&self, addr: u64) -> (String, u32) {
match self.ctx.find_location(addr) {
Ok(Some(loc)) => (loc.file.unwrap_or("").to_string(), loc.line.unwrap_or(0)),
_ => (String::new(), 0),
}
}
}
/// A mangled symbol, demangled without its hash (`xt::middle`); a C symbol as is.
pub fn demangle(mangled: &str) -> String {
match rustc_demangle::try_demangle(mangled) {
Ok(d) => format!("{d:#}"),
Err(_) => mangled.to_string(),
}
}
/// One traced process: its log's bytes, its program's ELF bytes and name.
pub struct LogInput<'a> {
pub log: &'a [u8],
pub elf: &'a [u8],
pub program: String,
/// The program's SHA-256 (hex), when the caller computed it.
pub sha256: String,
}
/// **A run from one or more traced processes**: every entry/exit record named by
/// its program's symbols, merged in time order (TSC → ns by each log's cycle
/// frequency; all logs of one machine share the TSC), sequence-numbered, and cut
/// at `max_events` (the earliest kept; [`TraceRun::truncated`] says so).
pub fn run_from_logs(inputs: &[LogInput<'_>], max_events: usize) -> Result<TraceRun> {
let mut run = TraceRun::default();
let mut fn_index: HashMap<String, u32> = HashMap::new();
// (abs ns, thread, func, kind, file order for a stable tie)
let mut all: Vec<(u128, u32, u32, EventKind, usize)> = Vec::new();
let mut order = 0usize;
for inp in inputs {
let (h, recs) = parse_basic_log(inp.log)
.with_context(|| format!("read the XRay log of {}", inp.program))?;
let addrs = function_addresses(inp.elf).with_context(|| format!("read {}", inp.program))?;
let names = symbolize(inp.elf, &addrs)?;
let freq = if h.cycle_frequency == 0 { 1_000_000_000 } else { h.cycle_frequency } as u128;
// id (1-based) → run function index, interned lazily
let mut local: HashMap<i32, u32> = HashMap::new();
let mut n = 0u64;
for r in &recs {
let kind = match r.kind {
0 | 3 => EventKind::Enter,
1 | 2 => EventKind::Exit,
_ => continue,
};
let Some(named) = usize::try_from(r.func_id).ok().and_then(|id| id.checked_sub(1)).and_then(|i| names.get(i))
else {
continue;
};
let f = *local.entry(r.func_id).or_insert_with(|| {
*fn_index.entry(named.symbol.clone()).or_insert_with(|| {
run.functions.push(named.clone());
(run.functions.len() - 1) as u32
})
});
all.push((r.tsc as u128 * 1_000_000_000 / freq, r.tid, f, kind, order));
order += 1;
n += 1;
}
run.binaries.push(TracedBinary {
name: inp.program.clone(),
sha256: inp.sha256.clone(),
instrumented: addrs.len() as u32,
records: n,
});
}
all.sort_by_key(|&(t, _, _, _, o)| (t, o));
run.total_events = all.len() as u64;
if all.len() > max_events {
all.truncate(max_events);
run.truncated = true;
}
let t0 = all.first().map(|e| e.0).unwrap_or(0);
run.events = all
.into_iter()
.enumerate()
.map(|(i, (t, thread, func, kind, _))| TraceEvent {
thread,
seq_no: i as u64,
t_ns: (t - t0) as u64,
func,
kind,
})
.collect();
Ok(run)
}
#[cfg(test)]
mod tests {
use super::*;
fn header(version: u16, freq: u64) -> Vec<u8> {
let mut h = vec![0u8; 32];
h[0..2].copy_from_slice(&version.to_le_bytes());
h[4..8].copy_from_slice(&3u32.to_le_bytes());
h[8..16].copy_from_slice(&freq.to_le_bytes());
h
}
fn record(rtype: u16, kind: u8, func: i32, tsc: u64, tid: u32) -> Vec<u8> {
let mut r = vec![0u8; 32];
r[0..2].copy_from_slice(&rtype.to_le_bytes());
r[2] = 7;
r[3] = kind;
r[4..8].copy_from_slice(&func.to_le_bytes());
r[8..16].copy_from_slice(&tsc.to_le_bytes());
r[16..20].copy_from_slice(&tid.to_le_bytes());
r[20..24].copy_from_slice(&42u32.to_le_bytes());
r
}
#[test]
fn a_basic_log_parses_its_header_and_function_records_and_skips_payloads() {
let mut b = header(3, 2_000_000_000);
b.extend(record(0, 0, 5, 100, 9));
b.extend(record(1, 0, 5, 0, 9)); // an argument payload
b.extend(record(0, 2, 5, 140, 9));
b.extend([1, 2, 3]); // a torn tail
let (h, recs) = parse_basic_log(&b).unwrap();
assert_eq!((h.version, h.kind, h.constant_tsc, h.nonstop_tsc, h.cycle_frequency), (3, 0, true, true, 2_000_000_000));
assert_eq!(recs.len(), 2);
assert_eq!((recs[0].func_id, recs[0].kind, recs[0].tsc, recs[0].tid, recs[0].pid, recs[0].cpu), (5, 0, 100, 9, 42, 7));
assert_eq!(recs[1].kind, 2);
let mut fdr = header(1, 1);
fdr[2] = 1;
assert!(parse_basic_log(&fdr).unwrap_err().to_string().contains("not a basic-mode"));
assert!(parse_basic_log(&[0; 8]).is_err());
}
#[test]
fn relative_sleds_number_functions_where_the_function_address_changes() {
// three sleds at section address 0x1000: f@0x5000 (entry, exit), g@0x6000
let mut sec = Vec::new();
for (i, (addr, func)) in [(0x5000u64, 0x5000u64), (0x5010, 0x5000), (0x6000, 0x6000)].iter().enumerate() {
let at = 0x1000 + (i * 32) as u64;
let mut s = vec![0u8; 32];
s[0..8].copy_from_slice(&addr.wrapping_sub(at).to_le_bytes());
s[8..16].copy_from_slice(&func.wrapping_sub(at + 8).to_le_bytes());
s[18] = 2;
sec.extend(s);
}
assert_eq!(sled_functions(&sec, 0x1000).unwrap(), vec![0x5000, 0x6000]);
assert!(sled_functions(&sec[..40], 0x1000).is_err());
}
#[test]
fn a_log_name_says_its_program() {
assert_eq!(log_program("xray-log.code_room-5a22.0BbtI0", "xray-log.").as_deref(), Some("code_room-5a22"));
assert_eq!(log_program("other", "xray-log."), None);
}
#[test]
fn demangling_drops_the_hash() {
assert_eq!(demangle("_ZN2xt6middle17h0123456789abcdefE"), "xt::middle");
assert_eq!(demangle("memcpy"), "memcpy");
}
/// **A real XRay run, in order.** `NORNIR_TRACE_FIXTURE` names a directory
/// holding `prog` (an XRay-instrumented test binary of the `xt` crate:
/// `top` calls `middle(3)` and `middle(2)`, `middle` calls `leaf` per item)
/// and `log` (its basic-mode log). The runa image makes it; dwarves'
/// `a_real_trace_step_*` test runs the whole road.
#[test]
#[ignore = "needs NORNIR_TRACE_FIXTURE (an XRay-built test binary + its log)"]
fn a_real_xray_run_of_a_small_crate_is_its_calls_in_order() {
let dir = std::path::PathBuf::from(std::env::var("NORNIR_TRACE_FIXTURE").unwrap());
let elf = std::fs::read(dir.join("prog")).unwrap();
let log = std::fs::read(dir.join("log")).unwrap();
let run = run_from_logs(&[LogInput { log: &log, elf: &elf, program: "prog".into(), sha256: String::new() }], usize::MAX).unwrap();
let named: Vec<(String, EventKind)> = run
.events
.iter()
.map(|e| (run.functions[e.func as usize].symbol.clone(), e.kind))
.filter(|(f, _)| f.starts_with("xt::") && !f.starts_with("xt::tests") && f != "xt::main")
.collect();
use EventKind::*;
let mut want = vec![("xt::top".to_string(), Enter)];
for n in [3, 2] {
want.push(("xt::middle".into(), Enter));
for _ in 0..n {
want.push(("xt::leaf".into(), Enter));
want.push(("xt::leaf".into(), Exit));
}
want.push(("xt::middle".into(), Exit));
}
want.push(("xt::top".into(), Exit));
assert_eq!(named, want);
assert!(run.events.windows(2).all(|w| w[0].t_ns <= w[1].t_ns && w[1].seq_no == w[0].seq_no + 1));
// DWARF places each function on its source line
let leaf = run.functions.iter().find(|f| f.symbol == "xt::leaf").unwrap();
assert!(leaf.file.ends_with("src/lib.rs") && leaf.line == 1, "{leaf:?}");
}
}