gpu-trace-perf 1.8.2

use std::{
    ffi::OsStr,
    path::{Path, PathBuf},
    process::Command,
    time::Duration,
};

use anyhow::{Context, Result, bail};
use log::{debug, error, warn};
use regex::Regex;
use serde::Deserialize;

use crate::{ReplayOutput, TraceTool, replay_command, snapshot::SnapshotResult};

pub struct ApitraceTrace {
    file: PathBuf,
    name: String,
    is_directx: bool,
    nonloopable: bool,
    extra_args: Vec<String>,
}

impl ApitraceTrace {
    pub fn new(
        root: &Path,
        file: &Path,
        nonloopable: bool,
        extra_args: Vec<String>,
    ) -> ApitraceTrace {
        ApitraceTrace {
            file: file.to_owned(),
            name: crate::relative_test_name(root, file),
            is_directx: apitrace_file_is_directx(file).unwrap_or_else(|e| {
                error!("Failure calling apitrace info, assuming file is GL: {}", e);
                false
            }),
            nonloopable,
            extra_args,
        }
    }
}

pub fn call_apitrace<R: TraceTool>(trace: &R, command: Command) -> Result<ReplayOutput> {
    let output = trace.run_replay_command(command);

    if output.exit_code != 0 && output.stderr.contains("waffle_context_create failed") {
        warn!(
            "apitrace reported waffle_context_create() failed, likely due to trace requiring too new of a GL version"
        );
    }

    Ok(output)
}

fn parse_snapshot_line(line: &str) -> Option<String> {
    if line.starts_with("Wrote ") {
        Some(line.trim_start_matches("Wrote ").to_string())
    } else {
        None
    }
}

impl TraceTool for ApitraceTrace {
    fn replay(&self, wrapper: Option<&str>, envs: &[(String, String)]) -> Result<ReplayOutput> {
        // apitrace replay otherwise assumes x11 (glretrace) for glx traces instead of using waffle
        let mut apitrace_command: Vec<&OsStr> = vec![
            OsStr::new("eglretrace"),
            OsStr::new("--pgpu"),
            OsStr::new("--headless"),
            OsStr::new("--loop=1"), // loop the last frame once so we know that caches are hot, shaders are compiled, etc.
        ];
        for arg in &self.extra_args {
            apitrace_command.push(OsStr::new(arg));
        }
        apitrace_command.push(self.file.as_os_str());

        call_apitrace(self, replay_command(&apitrace_command, wrapper, envs))
    }

    fn fps(&self, output: &ReplayOutput) -> Result<f64> {
        parse_apitrace_pgpu_output(&output.stdout)
    }

    fn name(&self) -> &str {
        &self.name
    }

    fn can_snapshot(&self) -> bool {
        true
    }

    fn snapshot(&self, output_dir: &str, loops: u32, timeout: Duration) -> Result<SnapshotResult> {
        let output_dir = self.output_dir(output_dir)?.unwrap();
        let trace_invocations = if self.nonloopable { loops } else { 1 };

        let mut args = Vec::new();
        if self.is_directx {
            args.push("wine".to_string());
            args.push("d3dretrace.exe".to_string());
        } else {
            args.push("eglretrace".to_string());
        }
        args.push("--headless".to_string());

        if loops > 1 && !self.nonloopable {
            args.push(format!("--loop={}", loops - 1));
            args.push("--call-nos=false".to_string());
        } else {
            let last_frame = apitrace_last_frame(&self.file)?;
            args.push(format!("--snapshot={last_frame}"));
        }
        args.push(format!(
            "--snapshot-prefix={}/snapshot",
            output_dir.display()
        ));
        args.extend(self.extra_args.iter().cloned());
        args.push(self.file.display().to_string());

        let mut last_output = None;
        let mut files = Vec::new();
        let mut remove_files = Vec::new();
        let loops = loops as usize;
        let start_time = std::time::Instant::now();

        for i in 1..=trace_invocations {
            let command = replay_command(&args, None, &[]);
            let output = self.run_replay_command_with_timeout(command, None, timeout);
            if output.exit_code != 0 && output.stderr.contains("waffle_context_create failed") {
                warn!(
                    "apitrace reported waffle_context_create() failed, likely due to trace requiring too new of a GL version"
                );
            }

            let mut run_files: Vec<_> = output
                .stdout
                .lines()
                .filter_map(parse_snapshot_line)
                .map(PathBuf::from)
                .collect();

            last_output = Some(output);

            if last_output.as_ref().unwrap().exit_code != 0 {
                // The tool may have written PNGs before failing; mark them for
                // cleanup so they don't accumulate in the output directory.
                remove_files.append(&mut run_files);
                break;
            }

            if self.nonloopable {
                if run_files.is_empty() {
                    error!("Replay succeeded but did not capture any frames.");
                    break;
                }

                // Save off the last frame of the replay, let the rest get deleted.
                if let Some(file) = run_files.pop() {
                    let to = output_dir.join(format!("frame{i}.png"));
                    std::fs::rename(file, &to)
                        .with_context(|| "Moving looped frame from {file:?} to {to:?}")?;
                    files.push(to);
                }
                remove_files.append(&mut run_files);
            } else if run_files.len() >= loops {
                // Save off all of the looped frames, and delete the rest.
                let first_save = run_files.len() - loops;
                for (i, file) in run_files.into_iter().enumerate() {
                    if i < first_save {
                        remove_files.push(file);
                    } else {
                        files.push(file);
                    }
                }
            } else {
                warn!(
                    "Frame looping on {} didn't appear to generate per-frame snapshots, do you have https://github.com/apitrace/apitrace/pull/968",
                    self.name()
                );
            }
        }

        // Delete .pngs for the frames we're not saving.
        for file in &remove_files {
            std::fs::remove_file(file)
                .unwrap_or_else(|e| error!("Removing unneeded snapshot {file:?}: {e}"));
        }

        for file in files.iter_mut() {
            *file = file
                .strip_prefix(&output_dir)
                .context("getting output dir relative to snapshot")?
                .to_path_buf();
        }

        let last_output = last_output.context("Finding the output for the last command")?;
        Ok(SnapshotResult {
            files,
            output: last_output,
            runtime: start_time.elapsed(),
        })
    }
}

// Returns the fps from the last frame of an apitrace replay --pgpu output
fn parse_apitrace_pgpu_output(output: &str) -> Result<f64> {
    lazy_static! {
        static ref CALL_RE: Regex = Regex::new("^call [0-9]+ -?[0-9]+ ([0-9]+)").unwrap();
    }

    let mut total = 0;
    let mut start_of_frame = true;
    for line in output.lines() {
        if line == "frame_end" {
            start_of_frame = true;
        } else {
            let cap = CALL_RE.captures(line);
            if let Some(cap) = cap {
                if start_of_frame {
                    total = 0;
                    start_of_frame = false;
                }
                match cap[1].parse::<i64>() {
                    Ok(gpu) => {
                        if gpu >= 0 {
                            total += gpu;
                        } else {
                            anyhow::bail!(
                                "apitrace produced GL_TIME_ELAPSED < 0, skipping(gpu hang?)"
                            );
                        }
                    }
                    Err(_) => {
                        anyhow::bail!("failed to parse apitrace's GL_TIME_ELAPSED");
                    }
                }
            }
        }
    }

    if total == 0 {
        anyhow::bail!("No times parsed");
    }

    Ok(1_000_000_000.0 / (total as f64))
}

/// If we're measuring times with u_trace, we don't need timestamps, but we do
/// need to loop the frame because there's a bunch of setup at the start that
/// we're not trying to measure.
pub struct ApitraceUtraceTrace {
    file: PathBuf,
    name: String,
    is_directx: bool,
    extra_args: Vec<String>,
}

impl ApitraceUtraceTrace {
    pub fn new(root: &Path, file: &Path, extra_args: Vec<String>) -> ApitraceUtraceTrace {
        ApitraceUtraceTrace {
            file: file.to_owned(),
            name: crate::relative_test_name(root, file),
            is_directx: apitrace_file_is_directx(file).unwrap_or_else(|e| {
                error!("Failure calling apitrace info, assuming file is GL: {}", e);
                true
            }),
            extra_args,
        }
    }
}

impl TraceTool for ApitraceUtraceTrace {
    fn replay(&self, wrapper: Option<&str>, envs: &[(String, String)]) -> Result<ReplayOutput> {
        // apitrace replay otherwise assumes x11 (glretrace) for glx traces
        // instead of using waffle
        //
        // loops the last frame twice so we know that caches are hot, shaders
        // are compiled, etc., and the middle frame utrace results will be one
        // of the hot ones.
        let mut args: Vec<&OsStr> = if self.is_directx {
            vec![
                OsStr::new("wine"),
                OsStr::new("d3dretrace.exe"),
                OsStr::new("--headless"),
                OsStr::new("--loop=2"),
            ]
        } else {
            vec![
                OsStr::new("eglretrace"),
                OsStr::new("--headless"),
                OsStr::new("--loop=2"),
            ]
        };
        for arg in &self.extra_args {
            args.push(OsStr::new(arg));
        }
        args.push(self.file.as_os_str());

        call_apitrace(self, replay_command(&args, wrapper, envs))
    }

    fn fps(&self, _: &ReplayOutput) -> Result<f64> {
        unreachable!("shouldn't be called");
    }

    fn name(&self) -> &str {
        &self.name
    }
}

#[derive(Deserialize)]
struct ApitraceInfo {
    #[serde(rename = "API")]
    api: String,
}

pub fn apitrace_info_is_directx(input: &str) -> Result<bool> {
    let info = serde_json::from_str::<ApitraceInfo>(input)
        .with_context(|| format!("Parsing apitrace info output:\n{input}"))?;
    Ok(info.api == "DirectX")
}

pub fn apitrace_file_is_directx_native(path: &Path) -> Result<bool> {
    let file = std::fs::File::open(path).with_context(|| format!("opening {}", path.display()))?;
    let reader = apitrace::TraceReader::new(file)
        .with_context(|| format!("opening {} as apitrace trace", path.display()))?;
    match reader.guess_api()? {
        apitrace::call_flags::CallAPI::GL => Ok(false),
        apitrace::call_flags::CallAPI::D3D => Ok(true),
        apitrace::call_flags::CallAPI::Unknown => {
            bail!("Failed to detect API for {}", path.display())
        }
    }
}

pub fn apitrace_file_is_directx(file: &Path) -> Result<bool> {
    debug!("checking directx on {}", file.display());

    match apitrace_file_is_directx_native(file) {
        Ok(dx) => return Ok(dx),
        Err(e) => warn!("Failed to check apitrace file for directx-ness: {e}"),
    }

    let mut command = replay_command(
        &[OsStr::new("apitrace"), OsStr::new("info"), file.as_os_str()],
        None,
        &[],
    );
    let output = command.output().context("Calling apitrae info")?;
    apitrace_info_is_directx(&String::from_utf8_lossy(&output.stdout))
}

pub fn apitrace_last_call_number(input: &str) -> Result<u64> {
    let mut last_line = None;

    for line in input.lines() {
        if !line.trim_end().is_empty() {
            last_line = Some(line);
        }
    }

    let last_line = last_line.context("finding a non-empty line in apitrace dump output")?;
    let result = str::parse::<u64>(last_line.split_once(' ').context("finding space")?.0)
        .context("parsing u64 from last line")?;
    Ok(result)
}

pub fn apitrace_last_frame(file: &Path) -> Result<u64> {
    // TODO: Convert this over to using apitrace-rs (once we flag all the frame-end call sigs)
    debug!("apitrace dumping {}", file.display());
    let args = ["apitrace", "dump", "--calls=frame"];
    let mut command = Command::new(args[0]);
    for arg in &args[1..] {
        command.arg(arg);
    }
    command.arg(file);

    let output = ReplayOutput::from(command.output().context("calling apitrace dump")?);

    apitrace_last_call_number(&output.stdout)
        .with_context(|| format!("Getting last frame in {file:?}: {output}"))
}

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

    #[test]
    fn test_apitrace_parsing() {
        // Actual apitrace output, trimmed down for the testcase.
        let apitrace_input = "
# call no gpu_start gpu_dura cpu_start cpu_dura vsize_start vsize_dura rss_start rss_dura pixels program name
call 44 0 0 0 0 0 0 0 0 0 0 glViewport
call 56 25082334 50166 0 0 0 0 0 0 0 0 glClear
call 81 41719667 0 0 0 0 0 0 0 0 0 glClear
call 176 42206667 472583 0 0 0 0 0 0 0 7 glDrawArrays
frame_end
call 222 0 0 0 0 0 0 0 0 0 4 glClearColor
call 224 45001334 21666 0 0 0 0 0 0 0 4 glClear
call 231 45023750 38000 0 0 0 0 0 0 0 7 glClear
call 239 45062584 519333 0 0 0 0 0 0 0 7 glDrawArrays
frame_end
call 222 0 0 0 0 0 0 0 0 0 4 glClearColor
call 224 47438000 13666 0 0 0 0 0 0 0 4 glClear
call 231 47452417 59583 0 0 0 0 0 0 0 7 glClear
call 239 47512917 579083 0 0 0 0 0 0 0 7 glDrawArrays
frame_end
Rendered 3 frames in 0.0539452 secs, average of 55.612 fps
";
        assert_approx_eq!(
            parse_apitrace_pgpu_output(apitrace_input).unwrap(),
            1.0 / ((13_666 + 59_583 + 579_083) as f64 / 1_000_000_000.0)
        )
    }

    #[test]
    fn test_apitrace_parsing_negatve_start() {
        let apitrace_input = "call 318 -8883437858 156 0 0 0 0 0 0 0 0 glBlitFramebuffer";
        assert_approx_eq!(
            parse_apitrace_pgpu_output(apitrace_input).unwrap(),
            1.0 / (156.0 / 1_000_000_000.0)
        );
    }

    #[test]
    fn test_apitrace_parsing_empty() {
        let apitrace_input = "
# call no gpu_start gpu_dura cpu_start cpu_dura vsize_start vsize_dura rss_start rss_dura pixels program name
call 44 0 0 0 0 0 0 0 0 0 0 glViewport
frame_end
";
        assert!(parse_apitrace_pgpu_output(apitrace_input).is_err());
    }

    #[test]
    fn test_apitrace_directx_info() -> Result<()> {
        assert!(apitrace_info_is_directx(
            r#"
{
  "FileName": "/home/anholt/src/traces-db/unigine/heaven-scene1-low-d3d11.trace-dxgi",
  "ContainerVersion": 6,
  "ContainerType": "Brotli",
  "API": "DirectX",
  "FramesCount": 104,
  "ActualDataSize": 130120914,
  "ContainerSize": 63387386
}
"#
        )?);
        assert!(!apitrace_info_is_directx(
            r#"
{
  "FileName": "/home/anholt/src/traces-db/neverball/neverball-v2.trace",
  "ContainerVersion": 6,
  "ContainerType": "Brotli",
  "API": "OpenGL + GLX/WGL/CGL",
  "FramesCount": 147,
  "ActualDataSize": 21503984,
  "ContainerSize": 1554696
}
"#
        )?);
        Ok(())
    }

    #[test]
    fn test_apitrace_last_call_number() {
        let input = r#"
// process.name = "/usr/bin/glxgears"
1384 glXSwapBuffers(dpy = 0x56060e921f80, drawable = 31457282)

1413 glXSwapBuffers(dpy = 0x56060e921f80, drawable = 31457282)
"#;

        assert_eq!(apitrace_last_call_number(input).unwrap(), 1413);
    }

    #[test]
    fn test_parse_snapshot_line() {
        assert_eq!(
            parse_snapshot_line("Wrote /path/to/snapshot0001.png"),
            Some("/path/to/snapshot0001.png".to_string())
        );
        assert_eq!(
            parse_snapshot_line("Wrote output/dir/frame-123.png"),
            Some("output/dir/frame-123.png".to_string())
        );
        assert_eq!(parse_snapshot_line("Some other output"), None);
        assert_eq!(parse_snapshot_line(""), None);
    }

    #[cfg(feature = "testapitrace")]
    fn test_snapshot(nonloopable: bool) {
        use tempfile::{NamedTempFile, TempDir};
        let loops = 5usize;

        let output_dir = TempDir::new().unwrap();
        let output_dir_str = output_dir.path().to_str().unwrap();
        let trace = include_bytes!("test_data/glmark2-pulsar.trace");
        let trace_file = NamedTempFile::with_suffix(".trace").unwrap();
        std::fs::write(trace_file.path(), trace).unwrap();
        let trace = ApitraceTrace::new(Path::new("/"), trace_file.path(), nonloopable, vec![]);
        let output = trace
            .snapshot(output_dir_str, loops as u32, Duration::MAX)
            .context("snapshotting apitrace")
            .unwrap();

        println!("output: {}", &output.output);

        assert_eq!(output.output.exit_code, 0, "{}", &output.output);

        assert_eq!(
            output.output.cmdline.contains("--loop="),
            !nonloopable,
            "cmdline: {}",
            &output.output.cmdline
        );

        for file in &output.files {
            let path = trace
                .output_dir(output_dir_str)
                .unwrap()
                .unwrap()
                .join(file);
            assert!(
                std::fs::exists(&path).unwrap(),
                "finding reported snapshot png {path:?}"
            );
        }
        for file in walkdir::WalkDir::new(&output_dir).into_iter() {
            println!("output: {file:?}");
        }
        assert_eq!(
            walkdir::WalkDir::new(&output_dir)
                .into_iter()
                .filter_map(|x| x.ok())
                .filter(|x| x.path().extension().is_some_and(|x| x == "png"))
                .count(),
            loops,
        );
    }

    #[cfg(feature = "testapitrace")]
    #[test]
    fn test_loopable_snapshot() {
        test_snapshot(false);
    }

    #[cfg(feature = "testapitrace")]
    #[test]
    fn test_nonloopable_snapshot() {
        test_snapshot(true);
    }
}