gpu-trace-perf 1.8.2

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

use anyhow::{Context, Result, bail};
use log::{error, warn};

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

pub struct AngleTrace {
    file: PathBuf,
    test: String,
    name: String,
    extra_args: Vec<String>,
}

impl AngleTrace {
    pub fn new(file: &Path, test: &str, extra_args: Vec<String>) -> AngleTrace {
        AngleTrace {
            file: file.to_owned(),
            test: test.to_owned(),
            name: format!("angle_trace_tests/{test}"),
            extra_args,
        }
    }

    /// Gets the list of tests from the angle_perf_test binary and returns the AngleTraces for them.
    pub fn enumerate(bin: &Path, extra_args: &[String]) -> Result<Vec<AngleTrace>> {
        let mut command = Command::new(bin);
        command.arg("--list-tests");

        let output = command
            .output()
            .with_context(|| format!("Executing {:?}", &command))?;

        let stdout = String::from_utf8_lossy(&output.stdout);

        if !output.status.success() {
            error!("Enumerating angle tests failed: {}", &stdout);
            bail!("Enumerating angle tests");
        }
        let tests = parse_gtest_test_list(&stdout);

        Ok(tests
            .iter()
            .map(|test| AngleTrace::new(bin, test, extra_args.to_vec()))
            .collect())
    }
}

fn parse_gtest_test_list(output: &str) -> Vec<String> {
    output
        .lines()
        .skip_while(|line| *line != "Tests list:")
        .skip(1)
        .filter(|x| *x != "End tests list.")
        .map(|x| x.to_string())
        .collect()
}

fn angle_result_field<'a>(stdout: &'a str, result: &str, suffix: &str) -> Result<&'a str> {
    stdout
        .lines()
        .find(|x| x.contains(result))
        .with_context(|| format!("finding {result} in {stdout}"))?
        .split_once("= ")
        .with_context(|| format!("finding delimiter for {result} in {stdout}"))?
        .1
        .strip_suffix(suffix)
        .with_context(|| format!("stripping suffix for {suffix} in {stdout}"))
}

fn angle_output_fps(stdout: &str) -> Result<f64> {
    let wall_time = angle_result_field(stdout, "wall_time", " ms")?;
    let wall_time =
        str::parse::<f64>(wall_time).with_context(|| format!("parsing wall time {wall_time}"))?;
    Ok(1000.0 / wall_time)
}

impl TraceTool for AngleTrace {
    fn replay(&self, wrapper: Option<&str>, envs: &[(String, String)]) -> Result<ReplayOutput> {
        let filter = format!("--gtest_filter={}", &self.test);
        let mut angle_command: Vec<&OsStr> = vec![
            self.file.as_os_str(),
            OsStr::new("--no-warmup"),
            OsStr::new("--trials=1"),
            OsStr::new("--fixed-test-time=1"),
        ];
        for arg in &self.extra_args {
            angle_command.push(OsStr::new(arg));
        }
        angle_command.push(OsStr::new("--offscreen"));
        angle_command.push(OsStr::new(&filter));

        let command = replay_command(&angle_command, wrapper, envs);
        let output: ReplayOutput = self.run_replay_command(command);

        if output.exit_code != 0 && output.stdout.contains("Could not load trace.") {
            warn!("Is your build synced with the trace sources you have downloaded?");
        }

        Ok(output)
    }

    fn fps(&self, output: &ReplayOutput) -> Result<f64> {
        angle_output_fps(&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_path = self.output_dir(output_dir)?.unwrap();
        let test_name = self.test.split('/').next_back().unwrap();
        // ANGLE names screenshots angle_{backend}_offscreen_{test}.png, where backend
        // matches the --use-gl= value (default: vulkan).
        let gl_backend = self
            .extra_args
            .iter()
            .find_map(|arg| arg.strip_prefix("--use-gl="))
            .unwrap_or("vulkan");
        let screenshot_name = format!("angle_{gl_backend}_offscreen_{test_name}.png");

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

        // ANGLE looks for trace .so files in the working directory, so run
        // from the real binary's directory (resolving any symlink).  Also
        // resolve the binary to an absolute path so it remains valid when the
        // working directory changes.
        let canonical_file = self
            .file
            .canonicalize()
            .unwrap_or_else(|_| self.file.clone());
        let binary_dir = canonical_file.parent().map(|p| p.to_path_buf());

        // Test names from enumerate() include the suite (e.g. "TraceTest.foo"),
        // but subtest paths only give the test name (e.g. "foo").  Gtest
        // requires the full suite.test form to match.
        let test_filter = if self.test.contains('.') {
            self.test.clone()
        } else {
            format!("TraceTest.{}", &self.test)
        };
        let filter = format!("--gtest_filter={test_filter}");
        for i in 1..=loops {
            let mut angle_command: Vec<&OsStr> = vec![
                canonical_file.as_os_str(),
                OsStr::new("--max-steps-performed=1"), // New flag for playing a single frame of the trace.
                OsStr::new("--one-frame-only"),        // Old flag for compatibility pre-2023.
            ];
            for arg in &self.extra_args {
                angle_command.push(OsStr::new(arg));
            }
            angle_command.extend(&[
                OsStr::new("--save-screenshots"),
                OsStr::new("--offscreen"),
                OsStr::new("--screenshot-dir"),
                output_dir_path.as_os_str(),
                OsStr::new(&filter),
            ]);

            let mut command = replay_command(&angle_command, None, &[]);
            if let Some(ref dir) = binary_dir {
                command.current_dir(dir);
            }
            last_output = Some(self.run_replay_command_with_timeout(command, None, timeout));

            let screenshot_path = output_dir_path.join(&screenshot_name);

            if last_output.as_ref().unwrap().exit_code != 0 {
                // Delete any screenshot file the tool may have written before
                // failing, so it doesn't accumulate in the output directory.
                let _ = std::fs::remove_file(&screenshot_path);
                break;
            }
            let snapshot_filename = format!("snapshot{i:04}.png");
            let snapshot_path = output_dir_path.join(&snapshot_filename);

            std::fs::rename(&screenshot_path, &snapshot_path).with_context(|| {
                format!(
                    "renaming {} to {}",
                    screenshot_path.display(),
                    snapshot_path.display()
                )
            })?;

            files.push(PathBuf::from(&snapshot_filename));
        }

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

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

    #[test]
    fn test_parse_gtest_test_list() {
        let list = "
1 GPUs:
  0 - AMD device id: 0x150E, revision id: 0xD3, system device id: 0x0

Active GPU: 0

Optimus: false
AMD Switchable: false
Mac Switchable: false
Needs EAGL on Mac: false


Tests list:
TraceTest.1945_air_force
TraceTest.20_minutes_till_dawn
TraceTest.age_of_origins_z
End tests list.
";
        assert_eq!(
            parse_gtest_test_list(list),
            vec![
                "TraceTest.1945_air_force".to_string(),
                "TraceTest.20_minutes_till_dawn".to_string(),
                "TraceTest.age_of_origins_z".to_string(),
            ]
        );
    }

    #[test]
    fn test_fps() {
        let stdout = r#"
1 GPUs:
  0 - AMD device id: 0x150E, revision id: 0xD3, system device id: 0x0

Active GPU: 0

Optimus: false
AMD Switchable: false
Mac Switchable: false
Needs EAGL on Mac: false


Note: Google Test filter = TraceTest.1945_air_force
[==========] Running 1 test from 1 test suite.
[----------] Global test environment set-up.
[----------] 1 test from TraceTest
[ RUN      ] TraceTest.1945_air_force
Warning: setpriority failed in StabilizeCPUForBenchmarking. Process will retain default priority: Permission denied
running test name: "TracePerf", backend: "_native", story: "1945_air_force"
*RESULT TracePerf_native.cpu_time: 1945_air_force= 1.3413200000 ms
*RESULT TracePerf_native.wall_time: 1945_air_force= 1.2631850701 ms
RESULT TracePerf_native.trial_steps: 1945_air_force= 200 count
RESULT TracePerf_native.total_steps: 1945_air_force= 200 count
*RESULT TracePerf_native.memory_median: 1945_air_force= 379936000 sizeInBytes
*RESULT TracePerf_native.memory_max: 1945_air_force= 382492000 sizeInBytes
[       OK ] TraceTest.1945_air_force (1059 ms)
[----------] 1 test from TraceTest (1059 ms total)

[----------] Global test environment tear-down
[==========] 1 test from 1 test suite ran. (1059 ms total)
[  PASSED  ] 1 test.
"#;

        assert_approx_eq!(angle_output_fps(stdout).unwrap(), 1000.0 / 1.2631850701);
    }
}