1mod capture;
2mod config;
3mod image_copy;
4
5pub use config::{HeadlessRenderConfig, HeadlessRenderError};
6
7use super::runtime::{asset_root, configure_shared_app};
8use super::{ViewerCommand, ViewerConfig, ViewerSessionHandle};
9use ak_core::{Structure, Trajectory};
10use bevy::app::ScheduleRunnerPlugin;
11use bevy::prelude::*;
12use bevy::transform::TransformSystems;
13use bevy::window::ExitCondition;
14use bevy::winit::WinitPlugin;
15use bevy_panorbit_camera::{PanOrbitCameraPlugin, PanOrbitCameraSystemSet};
16use std::panic::{AssertUnwindSafe, catch_unwind};
17use std::sync::{
18 Arc, Mutex,
19 atomic::{AtomicBool, Ordering},
20 mpsc,
21};
22use std::thread;
23use std::time::Duration;
24
25use capture::{
26 CaptureResult, CaptureSettings, CaptureState, HeadlessTargetSpec, SessionDriverState,
27 queue_capture, save_capture, setup_headless_target,
28};
29use image_copy::setup_image_copy;
30
31pub fn export_image(
32 trajectory: Trajectory,
33 config: ViewerConfig,
34 export: HeadlessRenderConfig,
35) -> Result<(), HeadlessRenderError> {
36 run_headless(trajectory, config, export, None, None)
37}
38
39pub fn export_structure_image(
40 structure: Structure,
41 config: ViewerConfig,
42 export: HeadlessRenderConfig,
43) -> Result<(), HeadlessRenderError> {
44 export_image(Trajectory::new(vec![structure]), config, export)
45}
46
47pub fn export_image_with_session<F>(
48 trajectory: Trajectory,
49 config: ViewerConfig,
50 export: HeadlessRenderConfig,
51 driver: F,
52) -> Result<(), HeadlessRenderError>
53where
54 F: FnOnce(ViewerSessionHandle) + Send + 'static,
55{
56 let (sender, receiver) = mpsc::channel();
57 let handle = ViewerSessionHandle::new(sender);
58 let driver_complete = Arc::new(AtomicBool::new(false));
59 let driver_complete_thread = driver_complete.clone();
60 let driver_thread = thread::Builder::new()
61 .name("ak-headless-session-driver".to_string())
62 .spawn(move || {
63 driver(handle);
64 driver_complete_thread.store(true, Ordering::Release);
65 })
66 .map_err(|err| HeadlessRenderError::new(format!("failed to spawn driver: {err}")))?;
67
68 let result = run_headless(
69 trajectory,
70 config,
71 export,
72 Some(receiver),
73 Some(driver_complete),
74 );
75 let _ = driver_thread.join();
76 result
77}
78
79pub fn export_prepared_image(
80 trajectory: Trajectory,
81 config: ViewerConfig,
82 export: HeadlessRenderConfig,
83 receiver: mpsc::Receiver<ViewerCommand>,
84) -> Result<(), HeadlessRenderError> {
85 run_headless(trajectory, config, export, Some(receiver), None)
86}
87
88fn run_headless(
89 trajectory: Trajectory,
90 mut config: ViewerConfig,
91 export: HeadlessRenderConfig,
92 receiver: Option<mpsc::Receiver<ViewerCommand>>,
93 session_driver_complete: Option<Arc<AtomicBool>>,
94) -> Result<(), HeadlessRenderError> {
95 config.render.show_ui = false;
96 config.render.show_orientation_widget = false;
97
98 let result = CaptureResult {
99 value: Arc::new(Mutex::new(None)),
100 };
101 let render_result = catch_unwind(AssertUnwindSafe(|| -> Result<(), HeadlessRenderError> {
102 let mut app = build_app(
103 trajectory,
104 config,
105 export,
106 result.clone(),
107 receiver,
108 session_driver_complete,
109 )?;
110 app.run();
111 result.take().unwrap_or_else(|| {
112 Err(HeadlessRenderError::new(
113 "headless render exited without producing a capture result",
114 ))
115 })
116 }));
117
118 match render_result {
119 Ok(result) => result,
120 Err(payload) => {
121 let message = if let Some(message) = payload.downcast_ref::<String>() {
122 message.clone()
123 } else if let Some(message) = payload.downcast_ref::<&'static str>() {
124 (*message).to_string()
125 } else {
126 "headless render panicked".to_string()
127 };
128 Err(HeadlessRenderError::new(message))
129 }
130 }
131}
132
133fn build_app(
134 trajectory: Trajectory,
135 config: ViewerConfig,
136 export: HeadlessRenderConfig,
137 result: CaptureResult,
138 receiver: Option<mpsc::Receiver<ViewerCommand>>,
139 session_driver_complete: Option<Arc<AtomicBool>>,
140) -> Result<App, HeadlessRenderError> {
141 if export.width == 0 || export.height == 0 {
142 return Err(HeadlessRenderError::new(
143 "headless render image size must be non-zero",
144 ));
145 }
146
147 let mut app = App::new();
148 app.add_plugins(
149 DefaultPlugins
150 .build()
151 .disable::<bevy::log::LogPlugin>()
152 .set(ImagePlugin::default_nearest())
153 .set(AssetPlugin {
154 file_path: asset_root(),
155 ..default()
156 })
157 .set(WindowPlugin {
158 primary_window: None,
159 exit_condition: ExitCondition::DontExit,
160 ..default()
161 })
162 .disable::<WinitPlugin>(),
163 )
164 .add_plugins(PanOrbitCameraPlugin)
165 .add_plugins(ScheduleRunnerPlugin::run_loop(Duration::from_secs_f64(
166 1.0 / 60.0,
167 )));
168
169 app.insert_resource(HeadlessTargetSpec {
170 width: export.width,
171 height: export.height,
172 })
173 .insert_resource(CaptureSettings {
174 path: export.path,
175 stable_frames: export.stable_frames.max(1),
176 })
177 .insert_resource(CaptureState {
178 preroll_remaining: export.preroll_frames,
179 stable_frames: 0,
180 requested: false,
181 request_delay_remaining: 0,
182 })
183 .insert_resource(result);
184
185 if let Some(complete) = session_driver_complete {
186 app.insert_resource(SessionDriverState { complete });
187 }
188
189 configure_shared_app(
190 &mut app,
191 trajectory,
192 config,
193 receiver,
194 Arc::new(Mutex::new(super::session::ViewerSnapshot::default())),
195 );
196 app.add_systems(
197 Startup,
198 setup_headless_target.before(super::systems::setup_camera),
199 );
200 app.add_systems(
201 PostUpdate,
202 queue_capture
203 .after(PanOrbitCameraSystemSet)
204 .after(TransformSystems::Propagate),
205 );
206 setup_image_copy(&mut app);
207 app.add_systems(PostUpdate, save_capture.after(queue_capture));
208
209 Ok(app)
210}
211
212#[cfg(test)]
213mod tests {
214 use super::*;
215 use crate::{BondList, Face, FaceList};
216 use image::GenericImageView;
217 use std::path::{Path, PathBuf};
218 use std::sync::{
219 Mutex,
220 atomic::{AtomicU64, Ordering},
221 };
222
223 static UNIQUE_ID: AtomicU64 = AtomicU64::new(0);
224 static HEADLESS_TEST_LOCK: Mutex<()> = Mutex::new(());
225
226 fn fixture_structure() -> Structure {
227 Structure::new(
228 vec![
229 [0.0, 0.0, 0.0],
230 [0.0, 0.0, 0.74],
231 [1.0, 0.0, 0.0],
232 [0.0, 1.0, 0.0],
233 ],
234 vec![8, 1, 1, 1],
235 [[8.0, 0.0, 0.0], [0.0, 8.0, 0.0], [0.0, 0.0, 8.0]],
236 [false, false, false],
237 )
238 }
239
240 fn temp_png(name: &str) -> PathBuf {
241 let id = UNIQUE_ID.fetch_add(1, Ordering::Relaxed);
242 std::env::temp_dir().join(format!("ak-headless-{name}-{id}.png"))
243 }
244
245 fn test_render_config(path: &Path) -> HeadlessRenderConfig {
246 let mut config = HeadlessRenderConfig::new(path, 320, 240);
247 if std::env::var_os("CI").is_some() {
248 config.preroll_frames = 48;
249 config.stable_frames = 12;
250 }
251 config
252 }
253
254 fn should_run_headless_render_tests() -> bool {
255 if std::env::var_os("ATOMIC_KERNELS_RUN_RUST_HEADLESS_TESTS").is_some() {
256 return true;
257 }
258 std::env::var_os("CI").is_none()
259 }
260
261 fn image_has_content(path: &Path, width: u32, height: u32) -> bool {
262 let image = image::open(path).expect("saved image should be readable");
263 assert_eq!(image.dimensions(), (width, height));
264 let rgba = image.to_rgba8();
265 let mut unique = std::collections::BTreeSet::new();
266 for pixel in rgba.pixels() {
267 unique.insert(pixel.0);
268 if unique.len() > 8 {
269 break;
270 }
271 }
272 unique.len() > 1
273 }
274
275 fn assert_render_succeeds<F>(name: &str, mut render_once: F)
276 where
277 F: FnMut(&Path) -> Result<(), HeadlessRenderError>,
278 {
279 let attempts = if std::env::var_os("CI").is_some() {
280 5
281 } else {
282 1
283 };
284 for attempt in 0..attempts {
285 let path = temp_png(name);
286 match render_once(&path) {
287 Ok(()) => {
288 if image_has_content(&path, 320, 240) {
289 let _ = std::fs::remove_file(path);
290 return;
291 }
292 }
293 Err(err) if err.to_string().contains("Unable to find a GPU") => return,
294 Err(err) => panic!("headless export should succeed: {err}"),
295 }
296 let _ = std::fs::remove_file(&path);
297 if attempt + 1 == attempts {
298 panic!("rendered image should not be a flat color");
299 }
300 }
301 }
302
303 #[test]
304 fn exports_default_scene_to_png() {
305 if !should_run_headless_render_tests() {
306 return;
307 }
308 let _guard = HEADLESS_TEST_LOCK.lock().unwrap();
309 let config = ViewerConfig::default();
310 assert_render_succeeds("default", |path| {
311 export_structure_image(
312 fixture_structure(),
313 config.clone(),
314 test_render_config(path),
315 )
316 });
317 }
318
319 #[test]
320 fn exports_scripted_scene_with_shared_session_commands() {
321 if !should_run_headless_render_tests() {
322 return;
323 }
324 let _guard = HEADLESS_TEST_LOCK.lock().unwrap();
325 let mut config = ViewerConfig::default();
326 config.render.show_axes = false;
327 assert_render_succeeds("scripted", |path| {
328 export_image_with_session(
329 Trajectory::new(vec![fixture_structure()]),
330 config.clone(),
331 test_render_config(path),
332 |session| {
333 session
334 .set_bonds(BondList::new([(0, 1), (0, 2), (0, 3)]), Some(0))
335 .unwrap();
336 session
337 .set_faces(
338 FaceList::new([Face::new([1, 2, 3], [0.2, 0.6, 0.9, 0.35]).unwrap()]),
339 Some(0),
340 )
341 .unwrap();
342 session
343 .set_render_style(
344 super::super::RenderStyle::BallAndStick(
345 super::super::BallAndStickStyle {
346 atom_scale: 0.5,
347 bond_radius: 0.06,
348 bond_color: [0.5, 0.5, 0.5, 1.0],
349 bond_scope: super::super::BondScope::TouchSelection,
350 },
351 ),
352 vec![true, true, true, true],
353 Some(0),
354 false,
355 )
356 .unwrap();
357 session.frame_all().unwrap();
358 },
359 )
360 });
361 }
362}