Seregon/StratoSDK

StratoSDK is a framework with a declarative approach similar to Flutter/React, written and designed entirely for Rust.

Rust/27.3 KB/No license
crates/strato-ui-core/src/integration/capture_recorder.rs
raw
StratoSDK / crates / strato-ui-core / src / integration / capture_recorder.rs
1use crate::platform::CapturedFrame;
2use image::ImageEncoder;
3use instant::Instant;
4use std::{
5 path::Path,
6 sync::{Arc, Mutex},
7};
8 
9/// The lifecycle state of the capture recorder / capture loop.
10#[derive(Clone, Copy, Debug, PartialEq, Eq)]
11enum RecorderState {
12 Idle,
13 Recording,
14 Stopping,
15}
16 
17/// Well-known key used to store the `CaptureRecorder` inside `StepDataMap`.
18pub const CAPTURE_RECORDER_KEY: &str = "capture_recorder";
19 
20/// Well-known key prefix for screenshot path requests stored in `StepDataMap`.
21pub const SCREENSHOT_PATH_KEY: &str = "pending_screenshot_path";
22 
23/// Environment variable that enables automatic video recording for all
24/// integration test steps. When set, the driver starts recording at the
25/// beginning of the test and writes the video on completion.
26pub const CAPTURE_RECORDING_ENABLED_ENV_VAR: &str = "WARP_INTEGRATION_TEST_VIDEO";
27 
28/// A captured frame paired with the wall-clock time it was taken.
29struct TimestampedFrame {
30 frame: CapturedFrame,
31 #[allow(dead_code)]
32 captured_at: Instant,
33}
34 
35/// Mutable state shared between `CaptureRecorder` and `CaptureLoopState`.
36/// All access is serialised through a single `Mutex`.
37#[allow(dead_code)]
38struct SharedState {
39 recorder_state: RecorderState,
40 raw_frames: Vec<TimestampedFrame>,
41 h264_buf: Vec<u8>,
42 encoded_frame_count: u32,
43 dimensions: Option<(u32, u32)>,
44 encoding_in_progress: bool,
45}
46 
47impl Default for SharedState {
48 fn default() -> Self {
49 Self {
50 recorder_state: RecorderState::Idle,
51 raw_frames: Vec::new(),
52 h264_buf: Vec::new(),
53 encoded_frame_count: 0,
54 dimensions: None,
55 encoding_in_progress: false,
56 }
57 }
58}
59 
60/// Shared handle passed to the capture loop task.
61#[allow(dead_code)]
62pub struct CaptureLoopState(Arc<Mutex<SharedState>>);
63 
64/// Records captured frames during integration tests and can produce
65/// individual PNGs or an encoded capture artifact.
66pub struct CaptureRecorder {
67 inner: Arc<Mutex<SharedState>>,
68 recording_start: Option<Instant>,
69}
70 
71impl Default for CaptureRecorder {
72 fn default() -> Self {
73 Self {
74 inner: Arc::new(Mutex::new(SharedState::default())),
75 recording_start: None,
76 }
77 }
78}
79 
80impl CaptureRecorder {
81 pub fn new() -> Self {
82 Self::default()
83 }
84 
85 pub fn start_recording(&mut self) {
86 self.recording_start = Some(Instant::now());
87 if let Ok(mut s) = self.inner.lock() {
88 s.recorder_state = RecorderState::Recording;
89 }
90 }
91 
92 pub fn stop_recording(&mut self) {
93 if let Ok(mut s) = self.inner.lock() {
94 if s.recorder_state == RecorderState::Recording {
95 s.recorder_state = RecorderState::Idle;
96 }
97 }
98 }
99 
100 /// Signals the capture loop to exit on its next iteration.
101 pub fn stop_capture_loop(&self) {
102 if let Ok(mut s) = self.inner.lock() {
103 s.recorder_state = RecorderState::Stopping;
104 }
105 }
106 
107 pub fn is_recording(&self) -> bool {
108 self.inner
109 .lock()
110 .map(|s| s.recorder_state == RecorderState::Recording)
111 .unwrap_or(false)
112 }
113 
114 pub fn recording_start(&self) -> Option<Instant> {
115 self.recording_start
116 }
117 
118 pub fn capture_loop_state(&self) -> CaptureLoopState {
119 CaptureLoopState(self.inner.clone())
120 }
121 
122 pub fn raw_frame_count(&self) -> usize {
123 self.inner.lock().map(|s| s.raw_frames.len()).unwrap_or(0)
124 }
125 
126 pub fn is_encoding(&self) -> bool {
127 self.inner
128 .lock()
129 .map(|s| s.encoding_in_progress)
130 .unwrap_or(false)
131 }
132 
133 pub fn frame_count(&self) -> usize {
134 self.inner
135 .lock()
136 .map(|s| s.encoded_frame_count as usize)
137 .unwrap_or(0)
138 }
139}
140 
141// ---------------------------------------------------------------------------
142// Feature-gated recording implementation.
143//
144// When `integration_tests` is enabled we have access to the `openh264` and
145// `minimp4` crates and can encode captured frames to H.264 / MP4.
146// Otherwise we provide a no-op capture loop and a PNG-based fallback for
147// `finalize`.
148// ---------------------------------------------------------------------------
149cfg_if::cfg_if! {
150 if #[cfg(feature = "integration_tests")] {
151 impl CaptureRecorder {
152 pub fn finalize(&mut self, output_path: &Path) -> anyhow::Result<()> {
153 if let Some(parent) = output_path.parent() {
154 std::fs::create_dir_all(parent)?;
155 }
156 
157 let (h264_data, total, dims) = {
158 let mut s = self.inner.lock().map_err(|e| anyhow::anyhow!("{e}"))?;
159 (
160 std::mem::take(&mut s.h264_buf),
161 s.encoded_frame_count,
162 s.dimensions,
163 )
164 };
165 
166 if h264_data.is_empty() || dims.is_none() {
167 log::info!("CaptureRecorder: no frames encoded, nothing to finalize");
168 return Ok(());
169 }
170 
171 let (width, height) = dims.expect("dimensions set when h264_data is non-empty");
172 match mux_h264_to_mp4(output_path, &h264_data, width, height) {
173 Ok(()) => {
174 log::info!(
175 "CaptureRecorder: wrote {total} frames to {}",
176 output_path.display()
177 );
178 }
179 Err(e) => {
180 log::warn!("CaptureRecorder: MP4 muxing failed ({e})");
181 return Err(e);
182 }
183 }
184 
185 Ok(())
186 }
187 }
188 
189 pub async fn run_capture_loop(app: crate::App, state: CaptureLoopState) {
190 use crate::r#async::Timer;
191 use std::time::Duration;
192 
193 const CAPTURE_INTERVAL_MS: u64 = 66;
194 const FLUSH_THRESHOLD: usize = 60;
195 const KEEP_RECENT: usize = 15;
196 
197 let inner = &state.0;
198 
199 loop {
200 Timer::after(Duration::from_millis(CAPTURE_INTERVAL_MS)).await;
201 
202 let (current_state, encoding, backlog) = {
203 let s = inner.lock().unwrap_or_else(|e| e.into_inner());
204 (s.recorder_state, s.encoding_in_progress, s.raw_frames.len())
205 };
206 let should_stop = current_state == RecorderState::Stopping;
207 
208 if should_stop && encoding {
209 Timer::after(Duration::from_millis(50)).await;
210 continue;
211 }
212 
213 let should_flush =
214 (backlog >= FLUSH_THRESHOLD || (should_stop && backlog > 0)) && !encoding;
215 
216 if should_flush {
217 let drain_count = if should_stop {
218 backlog
219 } else {
220 backlog.saturating_sub(KEEP_RECENT)
221 };
222 
223 let to_encode: Vec<TimestampedFrame> = {
224 let mut s = inner.lock().unwrap_or_else(|e| e.into_inner());
225 s.raw_frames.drain(..drain_count).collect()
226 };
227 
228 if !to_encode.is_empty() {
229 {
230 let mut s = inner.lock().unwrap_or_else(|e| e.into_inner());
231 s.encoding_in_progress = true;
232 }
233 let inner_clone = inner.clone();
234 std::thread::Builder::new()
235 .name("video-encoder".to_string())
236 .spawn(move || {
237 encode_frame_batch(&to_encode, &inner_clone);
238 })
239 .ok();
240 }
241 }
242 
243 let still_encoding = inner
244 .lock()
245 .map(|s| s.encoding_in_progress)
246 .unwrap_or(false);
247 if should_stop && !still_encoding {
248 break;
249 }
250 
251 if current_state != RecorderState::Recording {
252 continue;
253 }
254 
255 let window = app.read(|ctx| {
256 let windowing_state = ctx.windows();
257 windowing_state
258 .active_window()
259 .and_then(|id| windowing_state.platform_window(id))
260 });
261 
262 let Some(window) = window else {
263 continue;
264 };
265 
266 let inner_clone = inner.clone();
267 window
268 .as_ctx()
269 .request_frame_capture(Box::new(move |frame| {
270 let captured_at = Instant::now();
271 if let Ok(mut s) = inner_clone.lock() {
272 s.raw_frames.push(TimestampedFrame { frame, captured_at });
273 }
274 }));
275 }
276 }
277 
278 fn mux_h264_to_mp4(
279 output_path: &Path,
280 h264_data: &[u8],
281 width: u32,
282 height: u32,
283 ) -> anyhow::Result<()> {
284 use minimp4::Mp4Muxer;
285 use std::io::Cursor;
286 
287 const TARGET_FPS: u32 = 15;
288 
289 let mut mp4_buf = Cursor::new(Vec::new());
290 let mut muxer = Mp4Muxer::new(&mut mp4_buf);
291 muxer.init_video(
292 width as i32,
293 height as i32,
294 false,
295 "integration test recording",
296 );
297 muxer.write_video_with_fps(h264_data, TARGET_FPS);
298 muxer.close();
299 
300 std::fs::write(output_path, mp4_buf.into_inner())?;
301 Ok(())
302 }
303 
304 /// Encodes a batch of raw frames to H.264 on the calling (background)
305 /// thread. Writes results back into `SharedState` under the lock.
306 fn encode_frame_batch(
307 frames: &[TimestampedFrame],
308 inner: &Arc<Mutex<SharedState>>,
309 ) {
310 use openh264::encoder::Encoder;
311 use openh264::formats::{RgbSliceU8, YUVBuffer};
312 
313 const FRAME_DURATION_MS: u128 = 1000 / 15;
314 
315 let mut encoder = match Encoder::new() {
316 Ok(e) => e,
317 Err(e) => {
318 log::error!("CaptureRecorder: failed to create encoder on background thread: {e}");
319 if let Ok(mut s) = inner.lock() {
320 s.encoding_in_progress = false;
321 }
322 return;
323 }
324 };
325 
326 if let Some(first) = frames.first() {
327 let mut s = inner.lock().unwrap_or_else(|e| e.into_inner());
328 if s.dimensions.is_none() {
329 s.dimensions = Some((first.frame.width, first.frame.height));
330 }
331 }
332 
333 let mut prev_captured_at: Option<Instant> = None;
334 let mut batch_h264 = Vec::new();
335 let mut batch_encoded = 0u32;
336 
337 for ts_frame in frames {
338 let width = ts_frame.frame.width;
339 let height = ts_frame.frame.height;
340 let rgb_data = pixel_data_to_rgb(&ts_frame.frame.data, ts_frame.frame.format);
341 let rgb_source = RgbSliceU8::new(&rgb_data, (width as usize, height as usize));
342 let yuv = YUVBuffer::from_rgb_source(rgb_source);
343 
344 let repeat_count = if let Some(prev) = prev_captured_at {
345 let gap_ms = ts_frame.captured_at.duration_since(prev).as_millis();
346 (gap_ms / FRAME_DURATION_MS).max(1) as u32
347 } else {
348 1
349 };
350 prev_captured_at = Some(ts_frame.captured_at);
351 
352 for _ in 0..repeat_count {
353 match encoder.encode(&yuv) {
354 Ok(bitstream) => {
355 bitstream.write_vec(&mut batch_h264);
356 batch_encoded += 1;
357 }
358 Err(e) => {
359 log::error!("CaptureRecorder: encode error: {e}");
360 break;
361 }
362 }
363 }
364 }
365 
366 {
367 let mut s = inner.lock().unwrap_or_else(|e| e.into_inner());
368 if !batch_h264.is_empty() {
369 s.h264_buf.extend_from_slice(&batch_h264);
370 }
371 s.encoded_frame_count += batch_encoded;
372 s.encoding_in_progress = false;
373 }
374 
375 log::info!(
376 "CaptureRecorder: background-encoded {batch_encoded} H.264 frames from {} raw frames",
377 frames.len()
378 );
379 }
380 
381 fn pixel_data_to_rgb(data: &[u8], format: crate::platform::CapturedFrameFormat) -> Vec<u8> {
382 use crate::platform::CapturedFrameFormat;
383 let pixel_count = data.len() / 4;
384 let mut rgb = Vec::with_capacity(pixel_count * 3);
385 for chunk in data.chunks_exact(4) {
386 match format {
387 CapturedFrameFormat::Rgba => {
388 rgb.push(chunk[0]);
389 rgb.push(chunk[1]);
390 rgb.push(chunk[2]);
391 }
392 CapturedFrameFormat::Bgra => {
393 rgb.push(chunk[2]);
394 rgb.push(chunk[1]);
395 rgb.push(chunk[0]);
396 }
397 }
398 }
399 rgb
400 }
401 } else {
402 impl CaptureRecorder {
403 pub fn finalize(&mut self, output_path: &Path) -> anyhow::Result<()> {
404 if let Some(parent) = output_path.parent() {
405 std::fs::create_dir_all(parent)?;
406 }
407 
408 let frames: Vec<TimestampedFrame> = self
409 .inner
410 .lock()
411 .map(|mut s| std::mem::take(&mut s.raw_frames))
412 .unwrap_or_default();
413 if frames.is_empty() {
414 log::info!("CaptureRecorder: no frames captured, nothing to finalize");
415 return Ok(());
416 }
417 save_frames_as_pngs(output_path, &frames)?;
418 
419 Ok(())
420 }
421 }
422 
423 pub async fn run_capture_loop(_app: crate::App, _state: CaptureLoopState) {}
424 
425 fn save_frames_as_pngs(output_path: &Path, frames: &[TimestampedFrame]) -> anyhow::Result<()> {
426 let stem = output_path
427 .file_stem()
428 .and_then(|s| s.to_str())
429 .unwrap_or("frame");
430 let dir = output_path
431 .parent()
432 .unwrap_or_else(|| Path::new("."))
433 .join(format!("{stem}_frames"));
434 std::fs::create_dir_all(&dir)?;
435 
436 for (i, ts_frame) in frames.iter().enumerate() {
437 let path = dir.join(format!("{stem}_{i:04}.png"));
438 save_captured_frame_as_png(&ts_frame.frame, &path)?;
439 }
440 
441 log::info!(
442 "CaptureRecorder: saved {} PNGs to {}",
443 frames.len(),
444 dir.display()
445 );
446 Ok(())
447 }
448 }
449}
450 
451/// Saves a single `CapturedFrame` to a PNG file at the given path.
452pub fn save_captured_frame_as_png(frame: &CapturedFrame, path: &Path) -> anyhow::Result<()> {
453 let mut frame = frame.clone();
454 frame.ensure_rgba();
455 if let Some(parent) = path.parent() {
456 std::fs::create_dir_all(parent)?;
457 }
458 
459 let file = std::fs::File::create(path)?;
460 let mut writer = std::io::BufWriter::new(file);
461 
462 let encoder = image::codecs::png::PngEncoder::new_with_quality(
463 &mut writer,
464 image::codecs::png::CompressionType::Fast,
465 image::codecs::png::FilterType::NoFilter,
466 );
467 
468 encoder.write_image(
469 &frame.data,
470 frame.width,
471 frame.height,
472 image::ColorType::Rgba8.into(),
473 )?;
474 
475 Ok(())
476}
477 
478/// Helper to retrieve a mutable reference to the recorder from a `StepDataMap`.
479pub fn get_capture_recorder_mut(
480 step_data_map: &mut super::step::StepDataMap,
481) -> Option<&mut CaptureRecorder> {
482 step_data_map.get_mut::<_, CaptureRecorder>(CAPTURE_RECORDER_KEY)
483}
484 
485/// Helper to retrieve a shared reference to the recorder from a `StepDataMap`.
486pub fn get_capture_recorder(step_data_map: &super::step::StepDataMap) -> Option<&CaptureRecorder> {
487 step_data_map.get::<_, CaptureRecorder>(CAPTURE_RECORDER_KEY)
488}
489