1use anyhow::{Context, Result, bail};
4use std::collections::HashMap;
5use std::path::{Path, PathBuf};
6use std::process::Command;
7
8pub const DEFAULT_REF_IMAGE: &str = "rlx-tsac-ref";
9pub const DEFAULT_REF_PLATFORM: &str = "linux/amd64";
10
11#[derive(Debug, Clone)]
12pub struct DockerRefOptions {
13 pub image: String,
14 pub platform: String,
15}
16
17impl Default for DockerRefOptions {
18 fn default() -> Self {
19 Self {
20 image: std::env::var("RLX_TSAC_REF_IMAGE").unwrap_or_else(|_| DEFAULT_REF_IMAGE.into()),
21 platform: std::env::var("RLX_TSAC_DOCKER_PLATFORM")
22 .unwrap_or_else(|_| DEFAULT_REF_PLATFORM.into()),
23 }
24 }
25}
26
27#[derive(Debug, Clone)]
28pub struct RefRoundtrip {
29 pub engine: String,
30 pub encode_ms: f64,
31 pub decode_ms: f64,
32 pub output_bytes: u64,
33 pub tsac_path: PathBuf,
34 pub wav_path: PathBuf,
35}
36
37pub fn docker_ref_available(image: &str) -> bool {
38 Command::new("docker")
39 .args(["image", "inspect", image])
40 .stdout(std::process::Stdio::null())
41 .stderr(std::process::Stdio::null())
42 .status()
43 .map(|s| s.success())
44 .unwrap_or(false)
45}
46
47pub fn run_docker_ref_roundtrip(
48 opts: &DockerRefOptions,
49 engine: &str,
50 in_wav: &Path,
51 work_dir: &Path,
52 quality: u8,
53 fast: bool,
54) -> Result<RefRoundtrip> {
55 if !docker_ref_available(&opts.image) {
56 bail!(
57 "Docker image `{}` not found — build with: bash crates/rlx-tsac/docker/run.sh build",
58 opts.image
59 );
60 }
61 if !in_wav.is_file() {
62 bail!("input wav missing: {}", in_wav.display());
63 }
64 std::fs::create_dir_all(work_dir).with_context(|| format!("create {}", work_dir.display()))?;
65
66 let container_wav = format!("/data/{}", in_wav.file_name().unwrap().to_string_lossy());
67 let host_in = work_dir.join(in_wav.file_name().unwrap());
68 if host_in != in_wav {
69 std::fs::copy(in_wav, &host_in)
70 .with_context(|| format!("copy {} -> {}", in_wav.display(), host_in.display()))?;
71 }
72
73 let fast_flag = if fast { "1" } else { "0" };
74 let output = Command::new("docker")
75 .args([
76 "run",
77 "--rm",
78 "--platform",
79 &opts.platform,
80 "-v",
81 &format!("{}:/data", work_dir.display()),
82 &opts.image,
83 engine,
84 "/data",
85 &container_wav,
86 &quality.to_string(),
87 fast_flag,
88 ])
89 .output()
90 .context("docker run reference bench")?;
91
92 if !output.status.success() {
93 let stderr = String::from_utf8_lossy(&output.stderr);
94 let stdout = String::from_utf8_lossy(&output.stdout);
95 bail!(
96 "docker reference bench failed (status {})\nstdout:\n{stdout}\nstderr:\n{stderr}",
97 output.status
98 );
99 }
100
101 let stdout = String::from_utf8_lossy(&output.stdout);
102 let kv = parse_kv_lines(&stdout);
103 let encode_ms = kv
104 .get("ENCODE_MS")
105 .and_then(|s| s.parse().ok())
106 .context("missing ENCODE_MS from docker bench")?;
107 let decode_ms = kv
108 .get("DECODE_MS")
109 .and_then(|s| s.parse().ok())
110 .context("missing DECODE_MS from docker bench")?;
111 let output_bytes = kv
112 .get("BYTES")
113 .and_then(|s| s.parse().ok())
114 .context("missing BYTES from docker bench")?;
115 let tsac_path = kv
116 .get("TSAC_PATH")
117 .map(PathBuf::from)
118 .context("missing TSAC_PATH from docker bench")?;
119 let wav_path = kv
120 .get("WAV_PATH")
121 .map(PathBuf::from)
122 .context("missing WAV_PATH from docker bench")?;
123
124 let tsac_path = remap_container_path(&tsac_path, work_dir);
125 let wav_path = remap_container_path(&wav_path, work_dir);
126
127 Ok(RefRoundtrip {
128 engine: engine.to_string(),
129 encode_ms,
130 decode_ms,
131 output_bytes,
132 tsac_path,
133 wav_path,
134 })
135}
136
137fn remap_container_path(container: &Path, work_dir: &Path) -> PathBuf {
138 container
139 .strip_prefix("/data")
140 .map(|rel| work_dir.join(rel))
141 .unwrap_or_else(|_| work_dir.join(container.file_name().unwrap()))
142}
143
144fn parse_kv_lines(stdout: &str) -> HashMap<String, String> {
145 let mut map = HashMap::new();
146 for line in stdout.lines() {
147 if let Some((k, v)) = line.split_once('=') {
148 map.insert(k.trim().to_string(), v.trim().to_string());
149 }
150 }
151 map
152}