StratoSDK is a framework with a declarative approach similar to Flutter/React, written and designed entirely for Rust.
| 1 | use crate::platform::CapturedFrame; |
| 2 | use image::ImageEncoder; |
| 3 | #[cfg(feature = "integration_tests")] |
| 4 | use instant::Instant; |
| 5 | use std::{ |
| 6 | path::{Path, PathBuf}, |
| 7 | sync::{ |
| 8 | atomic::{AtomicBool, Ordering}, |
| 9 | Arc, Mutex, |
| 10 | }, |
| 11 | }; |
| 12 | |
| 13 | /// Well-known key used to store the `VideoRecorder` inside `StepDataMap`. |
| 14 | pub const VIDEO_RECORDER_KEY: &str = "video_recorder"; |
| 15 | |
| 16 | /// Well-known key prefix for screenshot path requests stored in `StepDataMap`. |
| 17 | pub const SCREENSHOT_PATH_KEY: &str = "pending_screenshot_path"; |
| 18 | |
| 19 | /// Environment variable that enables automatic video recording for all |
| 20 | /// integration test steps. When set, the driver starts recording at the |
| 21 | /// beginning of the test and writes the video on completion. |
| 22 | pub const VIDEO_ENABLED_ENV_VAR: &str = "WARP_INTEGRATION_TEST_VIDEO"; |
| 23 | |
| 24 | /// Environment variable that sets the output directory for video recordings |
| 25 | /// and screenshots. Defaults to `$TMPDIR/warp_integration_video_captures` when |
| 26 | /// unset. |
| 27 | pub const VIDEO_DIR_ENV_VAR: &str = "WARP_INTEGRATION_TEST_VIDEO_DIR"; |
| 28 | |
| 29 | /// A captured frame paired with the wall-clock time it was taken. |
| 30 | pub(super) struct TimestampedFrame { |
| 31 | pub(super) frame: CapturedFrame, |
| 32 | #[cfg(feature = "integration_tests")] |
| 33 | pub(super) captured_at: Instant, |
| 34 | } |
| 35 | |
| 36 | /// Shared state passed to the capture loop task so it can push frames and |
| 37 | /// check whether recording/stopping is requested. All fields use |
| 38 | /// atomics/mutex so they are `Send`. |
| 39 | #[cfg(feature = "integration_tests")] |
| 40 | pub struct CaptureLoopState { |
| 41 | pub(super) recording: Arc<AtomicBool>, |
| 42 | pub(super) stopped: Arc<AtomicBool>, |
| 43 | pub(super) frames: Arc<Mutex<Vec<TimestampedFrame>>>, |
| 44 | } |
| 45 | |
| 46 | /// Records captured frames during integration tests and can produce |
| 47 | /// individual PNGs or an encoded video file. |
| 48 | pub struct VideoRecorder { |
| 49 | /// Whether frames should currently be pushed (shared with the capture loop). |
| 50 | recording: Arc<AtomicBool>, |
| 51 | /// Set to `true` to tell the capture loop to exit. |
| 52 | stopped: Arc<AtomicBool>, |
| 53 | /// Accumulated frames (shared with the capture loop callback). |
| 54 | frames: Arc<Mutex<Vec<TimestampedFrame>>>, |
| 55 | /// Wall-clock time when `start_recording` was called. |
| 56 | #[cfg(feature = "integration_tests")] |
| 57 | recording_start: Option<Instant>, |
| 58 | } |
| 59 | |
| 60 | impl Default for VideoRecorder { |
| 61 | fn default() -> Self { |
| 62 | Self { |
| 63 | recording: Arc::new(AtomicBool::new(false)), |
| 64 | stopped: Arc::new(AtomicBool::new(false)), |
| 65 | frames: Arc::new(Mutex::new(Vec::new())), |
| 66 | #[cfg(feature = "integration_tests")] |
| 67 | recording_start: None, |
| 68 | } |
| 69 | } |
| 70 | } |
| 71 | |
| 72 | impl VideoRecorder { |
| 73 | pub fn new() -> Self { |
| 74 | Self::default() |
| 75 | } |
| 76 | |
| 77 | pub fn start_recording(&mut self) { |
| 78 | #[cfg(feature = "integration_tests")] |
| 79 | { |
| 80 | self.recording_start = Some(Instant::now()); |
| 81 | } |
| 82 | self.recording.store(true, Ordering::Relaxed); |
| 83 | } |
| 84 | |
| 85 | pub fn stop_recording(&mut self) { |
| 86 | self.recording.store(false, Ordering::Relaxed); |
| 87 | } |
| 88 | |
| 89 | /// Signals the capture loop to exit on its next iteration. |
| 90 | pub fn stop_capture_loop(&self) { |
| 91 | self.stopped.store(true, Ordering::Relaxed); |
| 92 | } |
| 93 | |
| 94 | pub fn is_recording(&self) -> bool { |
| 95 | self.recording.load(Ordering::Relaxed) |
| 96 | } |
| 97 | |
| 98 | /// Returns the instant recording started, if recording has been started. |
| 99 | #[cfg(feature = "integration_tests")] |
| 100 | pub fn recording_start(&self) -> Option<Instant> { |
| 101 | self.recording_start |
| 102 | } |
| 103 | |
| 104 | /// Returns a `CaptureLoopState` that the capture loop task can use to |
| 105 | /// push frames and read the recording/stopped flags. |
| 106 | #[cfg(feature = "integration_tests")] |
| 107 | pub fn capture_loop_state(&self) -> CaptureLoopState { |
| 108 | CaptureLoopState { |
| 109 | recording: self.recording.clone(), |
| 110 | stopped: self.stopped.clone(), |
| 111 | frames: self.frames.clone(), |
| 112 | } |
| 113 | } |
| 114 | |
| 115 | /// Returns the number of frames captured so far. |
| 116 | pub fn frame_count(&self) -> usize { |
| 117 | self.frames.lock().map(|g| g.len()).unwrap_or(0) |
| 118 | } |
| 119 | |
| 120 | /// Encodes all captured frames into a video at `output_path`. |
| 121 | /// Falls back to saving individual PNGs if encoding fails. |
| 122 | pub fn finalize( |
| 123 | &mut self, |
| 124 | output_path: &Path, |
| 125 | overlay_log: Option<&super::overlay::OverlayLog>, |
| 126 | ) -> anyhow::Result<()> { |
| 127 | let frames: Vec<TimestampedFrame> = self |
| 128 | .frames |
| 129 | .lock() |
| 130 | .map(|mut g| std::mem::take(&mut *g)) |
| 131 | .unwrap_or_default(); |
| 132 | |
| 133 | if frames.is_empty() { |
| 134 | log::info!("VideoRecorder: no frames captured, nothing to finalize"); |
| 135 | return Ok(()); |
| 136 | } |
| 137 | |
| 138 | if let Some(parent) = output_path.parent() { |
| 139 | std::fs::create_dir_all(parent)?; |
| 140 | } |
| 141 | |
| 142 | #[cfg(feature = "integration_tests")] |
| 143 | match encode_to_mp4(output_path, &frames, overlay_log) { |
| 144 | Ok(()) => { |
| 145 | log::info!( |
| 146 | "VideoRecorder: wrote {} frames to {}", |
| 147 | frames.len(), |
| 148 | output_path.display() |
| 149 | ); |
| 150 | } |
| 151 | Err(e) => { |
| 152 | log::warn!("VideoRecorder: MP4 encoding failed ({e}), falling back to PNGs"); |
| 153 | save_frames_as_pngs(output_path, &frames)?; |
| 154 | } |
| 155 | } |
| 156 | |
| 157 | #[cfg(not(feature = "integration_tests"))] |
| 158 | { |
| 159 | let _ = overlay_log; |
| 160 | save_frames_as_pngs(output_path, &frames)?; |
| 161 | } |
| 162 | |
| 163 | Ok(()) |
| 164 | } |
| 165 | } |
| 166 | |
| 167 | /// Encodes a slice of timestamped frames into an H.264/MP4 file. |
| 168 | /// Runs entirely on the calling thread at test finalization time. |
| 169 | #[cfg(feature = "integration_tests")] |
| 170 | fn encode_to_mp4( |
| 171 | output_path: &Path, |
| 172 | frames: &[TimestampedFrame], |
| 173 | overlay_log: Option<&super::overlay::OverlayLog>, |
| 174 | ) -> anyhow::Result<()> { |
| 175 | use super::overlay::OverlayState; |
| 176 | use minimp4::Mp4Muxer; |
| 177 | use openh264::encoder::Encoder; |
| 178 | use openh264::formats::{RgbSliceU8, YUVBuffer}; |
| 179 | use std::io::Cursor; |
| 180 | |
| 181 | const TARGET_FPS: u32 = 60; |
| 182 | const FRAME_DURATION_MS: u128 = 1000 / TARGET_FPS as u128; |
| 183 | |
| 184 | let width = frames[0].frame.width; |
| 185 | let height = frames[0].frame.height; |
| 186 | |
| 187 | let mut encoder = Encoder::new().map_err(|e| anyhow::anyhow!("openh264 init: {e}"))?; |
| 188 | |
| 189 | let mut overlay_state = OverlayState::new(); |
| 190 | let overlay_events = overlay_log.map(|ol| ol.events()).unwrap_or(&[]); |
| 191 | let overlay_scale = overlay_log.map(|ol| ol.scale_factor()).unwrap_or(2.0); |
| 192 | let has_overlays = !overlay_events.is_empty(); |
| 193 | |
| 194 | let mut h264_buf = Vec::new(); |
| 195 | let mut total_encoded_frames = 0u32; |
| 196 | |
| 197 | for i in 0..frames.len() { |
| 198 | let ts_frame = &frames[i]; |
| 199 | |
| 200 | let rgb_data = if has_overlays { |
| 201 | overlay_state.advance_to(ts_frame.captured_at, overlay_events); |
| 202 | let mut composited = ts_frame.frame.data.clone(); |
| 203 | overlay_state.render_onto( |
| 204 | &mut composited, |
| 205 | width, |
| 206 | height, |
| 207 | ts_frame.captured_at, |
| 208 | overlay_scale, |
| 209 | ); |
| 210 | rgba_to_rgb(&composited) |
| 211 | } else { |
| 212 | rgba_to_rgb(&ts_frame.frame.data) |
| 213 | }; |
| 214 | |
| 215 | let rgb_source = RgbSliceU8::new(&rgb_data, (width as usize, height as usize)); |
| 216 | let yuv = YUVBuffer::from_rgb_source(rgb_source); |
| 217 | |
| 218 | let repeat_count = if i + 1 < frames.len() { |
| 219 | let gap_ms = frames[i + 1] |
| 220 | .captured_at |
| 221 | .duration_since(ts_frame.captured_at) |
| 222 | .as_millis(); |
| 223 | (gap_ms / FRAME_DURATION_MS).max(1) as u32 |
| 224 | } else { |
| 225 | 1 |
| 226 | }; |
| 227 | |
| 228 | for _ in 0..repeat_count { |
| 229 | let bitstream = encoder |
| 230 | .encode(&yuv) |
| 231 | .map_err(|e| anyhow::anyhow!("openh264 encode: {e}"))?; |
| 232 | bitstream.write_vec(&mut h264_buf); |
| 233 | total_encoded_frames += 1; |
| 234 | } |
| 235 | } |
| 236 | |
| 237 | log::info!( |
| 238 | "VideoRecorder: encoded {total_encoded_frames} video frames from {} captured frames", |
| 239 | frames.len() |
| 240 | ); |
| 241 | |
| 242 | let mut mp4_buf = Cursor::new(Vec::new()); |
| 243 | let mut muxer = Mp4Muxer::new(&mut mp4_buf); |
| 244 | muxer.init_video( |
| 245 | width as i32, |
| 246 | height as i32, |
| 247 | false, |
| 248 | "integration test recording", |
| 249 | ); |
| 250 | muxer.write_video_with_fps(&h264_buf, TARGET_FPS); |
| 251 | muxer.close(); |
| 252 | |
| 253 | std::fs::write(output_path, mp4_buf.into_inner())?; |
| 254 | Ok(()) |
| 255 | } |
| 256 | |
| 257 | /// Saves each frame as a PNG into a subdirectory next to `output_path`. |
| 258 | /// Heavy PNG encoding is offloaded to a Tokio blocking thread. |
| 259 | fn save_frames_as_pngs(output_path: &Path, frames: &[TimestampedFrame]) -> anyhow::Result<()> { |
| 260 | let stem = output_path |
| 261 | .file_stem() |
| 262 | .and_then(|s| s.to_str()) |
| 263 | .unwrap_or("frame"); |
| 264 | let dir = output_path |
| 265 | .parent() |
| 266 | .unwrap_or_else(|| Path::new(".")) |
| 267 | .join(format!("{stem}_frames")); |
| 268 | std::fs::create_dir_all(&dir)?; |
| 269 | |
| 270 | for (i, ts_frame) in frames.iter().enumerate() { |
| 271 | let path = dir.join(format!("{stem}_{i:04}.png")); |
| 272 | save_captured_frame_as_png(&ts_frame.frame, &path)?; |
| 273 | } |
| 274 | |
| 275 | log::info!( |
| 276 | "VideoRecorder: saved {} PNGs to {}", |
| 277 | frames.len(), |
| 278 | dir.display() |
| 279 | ); |
| 280 | Ok(()) |
| 281 | } |
| 282 | |
| 283 | /// Background task that drives continuous frame capture at ~60 FPS. |
| 284 | /// |
| 285 | /// This future is `!Send` (it holds a clone of `App` which is `Rc<RefCell<...>>`) and |
| 286 | /// must be spawned on the foreground (main-thread) executor. It interleaves with the |
| 287 | /// step execution loop at every `Timer` yield point. |
| 288 | /// |
| 289 | /// When `state.recording` is `false` the loop sleeps without requesting any captures, |
| 290 | /// so there is zero rendering overhead when recording is not active. When |
| 291 | /// `state.stopped` is set the loop exits cleanly. |
| 292 | #[cfg(feature = "integration_tests")] |
| 293 | pub async fn run_capture_loop(app: crate::App, state: CaptureLoopState) { |
| 294 | use crate::r#async::Timer; |
| 295 | use std::time::Duration; |
| 296 | |
| 297 | loop { |
| 298 | Timer::after(Duration::from_millis(16)).await; |
| 299 | |
| 300 | if state.stopped.load(Ordering::Relaxed) { |
| 301 | break; |
| 302 | } |
| 303 | |
| 304 | if !state.recording.load(Ordering::Relaxed) { |
| 305 | continue; |
| 306 | } |
| 307 | |
| 308 | let window = app.read(|ctx| { |
| 309 | let windowing_state = ctx.windows(); |
| 310 | windowing_state |
| 311 | .active_window() |
| 312 | .and_then(|id| windowing_state.platform_window(id)) |
| 313 | }); |
| 314 | |
| 315 | let Some(window) = window else { |
| 316 | continue; |
| 317 | }; |
| 318 | |
| 319 | let frames = state.frames.clone(); |
| 320 | window |
| 321 | .as_ctx() |
| 322 | .request_frame_capture(Box::new(move |frame| { |
| 323 | let captured_at = Instant::now(); |
| 324 | if let Ok(mut guard) = frames.lock() { |
| 325 | guard.push(TimestampedFrame { frame, captured_at }); |
| 326 | } |
| 327 | })); |
| 328 | window.as_ctx().request_redraw(); |
| 329 | } |
| 330 | } |
| 331 | |
| 332 | #[cfg(feature = "integration_tests")] |
| 333 | fn rgba_to_rgb(rgba: &[u8]) -> Vec<u8> { |
| 334 | let pixel_count = rgba.len() / 4; |
| 335 | let mut rgb = Vec::with_capacity(pixel_count * 3); |
| 336 | for chunk in rgba.chunks_exact(4) { |
| 337 | rgb.push(chunk[0]); |
| 338 | rgb.push(chunk[1]); |
| 339 | rgb.push(chunk[2]); |
| 340 | } |
| 341 | rgb |
| 342 | } |
| 343 | |
| 344 | /// Saves a single `CapturedFrame` to a PNG file at the given path. |
| 345 | pub fn save_captured_frame_as_png(frame: &CapturedFrame, path: &Path) -> anyhow::Result<()> { |
| 346 | let mut frame = frame.clone(); |
| 347 | frame.ensure_rgba(); |
| 348 | if let Some(parent) = path.parent() { |
| 349 | std::fs::create_dir_all(parent)?; |
| 350 | } |
| 351 | |
| 352 | let file = std::fs::File::create(path)?; |
| 353 | let mut writer = std::io::BufWriter::new(file); |
| 354 | |
| 355 | let encoder = image::codecs::png::PngEncoder::new_with_quality( |
| 356 | &mut writer, |
| 357 | image::codecs::png::CompressionType::Fast, |
| 358 | image::codecs::png::FilterType::NoFilter, |
| 359 | ); |
| 360 | |
| 361 | encoder.write_image( |
| 362 | &frame.data, |
| 363 | frame.width, |
| 364 | frame.height, |
| 365 | image::ColorType::Rgba8.into(), |
| 366 | )?; |
| 367 | |
| 368 | Ok(()) |
| 369 | } |
| 370 | |
| 371 | /// Returns the directory where video recordings and screenshots should be |
| 372 | /// written, creating it if necessary. |
| 373 | pub fn output_dir() -> PathBuf { |
| 374 | let dir = std::env::var(VIDEO_DIR_ENV_VAR) |
| 375 | .map(PathBuf::from) |
| 376 | .unwrap_or_else(|_| std::env::temp_dir().join("warp_integration_video_captures")); |
| 377 | std::fs::create_dir_all(&dir).ok(); |
| 378 | dir |
| 379 | } |
| 380 | |
| 381 | /// Helper to retrieve a mutable reference to the recorder from a `StepDataMap`. |
| 382 | pub fn get_recorder_mut( |
| 383 | step_data_map: &mut super::step::StepDataMap, |
| 384 | ) -> Option<&mut VideoRecorder> { |
| 385 | step_data_map.get_mut::<_, VideoRecorder>(VIDEO_RECORDER_KEY) |
| 386 | } |
| 387 | |
| 388 | /// Helper to retrieve a shared reference to the recorder from a `StepDataMap`. |
| 389 | pub fn get_recorder(step_data_map: &super::step::StepDataMap) -> Option<&VideoRecorder> { |
| 390 | step_data_map.get::<_, VideoRecorder>(VIDEO_RECORDER_KEY) |
| 391 | } |
| 392 |