use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;
use std::time::{Duration, Instant};
use async_trait::async_trait;
use base64::engine::general_purpose::STANDARD as BASE64;
use base64::Engine as _;
use chromiumoxide::cdp::browser_protocol::emulation::{
ClearDeviceMetricsOverrideParams, SetDeviceMetricsOverrideParams,
};
use chromiumoxide::cdp::browser_protocol::page::{
EventScreencastFrame, ScreencastFrameAckParams, StartScreencastFormat, StartScreencastParams,
StopScreencastParams,
};
use chromiumoxide::Page;
use futures::{Stream, StreamExt};
use tokio::sync::{mpsc, Mutex};
use crate::backend::BrowserError;
use crate::models::Viewport;
const CDP_TEARDOWN_TIMEOUT: Duration = Duration::from_secs(2);
#[derive(Debug, Clone)]
pub struct ScreencastFrame {
pub jpeg: Arc<[u8]>,
pub viewport: Viewport,
pub captured_at: f64,
}
pub type FrameReceiver = mpsc::Receiver<ScreencastFrame>;
pub const FRAME_CHANNEL_CAP: usize = 16;
#[async_trait]
trait FrameAck: Send + Sync {
async fn ack(&self, session_id: i64);
}
#[async_trait]
impl FrameAck for Page {
async fn ack(&self, session_id: i64) {
let _ = self
.execute(ScreencastFrameAckParams::new(session_id))
.await;
}
}
struct RawFrame {
session_id: i64,
data_base64: String,
}
pub struct ScreencastPump {
page: Page,
viewport: Viewport,
quality: i64,
every_nth_frame: i64,
subscribers: Arc<Mutex<Vec<mpsc::Sender<ScreencastFrame>>>>,
running: Mutex<Option<RunningCapture>>,
metrics_overridden: AtomicBool,
}
struct RunningCapture {
stop: Arc<AtomicBool>,
task: tokio::task::JoinHandle<()>,
started: Instant,
}
fn reusable_start(running: &Option<RunningCapture>) -> Option<Instant> {
running
.as_ref()
.filter(|r| !r.task.is_finished())
.map(|r| r.started)
}
impl ScreencastPump {
pub async fn attach(
page: &Page,
quality: i64,
every_nth_frame: i64,
max_width: u32,
max_height: u32,
) -> Result<Self, BrowserError> {
page.execute(
SetDeviceMetricsOverrideParams::builder()
.width(max_width as i64)
.height(max_height as i64)
.device_scale_factor(1.0)
.mobile(false)
.build()
.map_err(|e| {
BrowserError::ScreenshotFailed(format!("device metrics params: {e}"))
})?,
)
.await
.map_err(|e| BrowserError::ScreenshotFailed(format!("setDeviceMetricsOverride: {e}")))?;
Ok(Self {
page: page.clone(),
viewport: Viewport {
width: max_width,
height: max_height,
device_pixel_ratio: 1.0,
},
quality: quality.clamp(1, 100),
every_nth_frame: every_nth_frame.max(1),
subscribers: Arc::new(Mutex::new(Vec::new())),
running: Mutex::new(None),
metrics_overridden: AtomicBool::new(true),
})
}
pub async fn subscribe(&self) -> Result<(FrameReceiver, Instant), BrowserError> {
let (tx, rx) = mpsc::channel(FRAME_CHANNEL_CAP);
self.subscribers.lock().await.push(tx);
let started = match self.ensure_running().await {
Ok(started) => started,
Err(e) => {
drop(rx);
self.subscribers.lock().await.retain(|s| !s.is_closed());
return Err(e);
}
};
Ok((rx, started))
}
async fn ensure_running(&self) -> Result<Instant, BrowserError> {
let mut running = self.running.lock().await;
if let Some(started) = reusable_start(&running) {
return Ok(started);
}
let events = self
.page
.event_listener::<EventScreencastFrame>()
.await
.map_err(|e| BrowserError::ScreenshotFailed(format!("screencast listener: {e}")))?
.map(|frame| RawFrame {
session_id: frame.session_id,
data_base64: AsRef::<str>::as_ref(&frame.data).to_string(),
})
.boxed();
self.page
.execute(
StartScreencastParams::builder()
.format(StartScreencastFormat::Jpeg)
.quality(self.quality)
.every_nth_frame(self.every_nth_frame)
.max_width(self.viewport.width as i64)
.max_height(self.viewport.height as i64)
.build(),
)
.await
.map_err(|e| BrowserError::ScreenshotFailed(format!("startScreencast: {e}")))?;
let stop = Arc::new(AtomicBool::new(false));
let started = Instant::now();
let task = {
let stop = Arc::clone(&stop);
let subscribers = Arc::clone(&self.subscribers);
let page = self.page.clone();
let viewport = self.viewport;
tokio::spawn(async move {
run(events, &page, &stop, started, viewport, &subscribers).await;
})
};
*running = Some(RunningCapture {
stop,
task,
started,
});
Ok(started)
}
pub async fn stop(&self) {
self.clear_metrics_override().await;
let mut running_guard = self.running.lock().await;
let Some(running) = running_guard.take() else {
return;
};
running.stop.store(true, Ordering::SeqCst);
let _ = tokio::time::timeout(
CDP_TEARDOWN_TIMEOUT,
self.page.execute(StopScreencastParams::default()),
)
.await;
running.task.abort();
}
async fn clear_metrics_override(&self) {
if !self.metrics_overridden.swap(false, Ordering::SeqCst) {
return;
}
let _ = tokio::time::timeout(
CDP_TEARDOWN_TIMEOUT,
self.page
.execute(ClearDeviceMetricsOverrideParams::default()),
)
.await;
}
}
impl Drop for ScreencastPump {
fn drop(&mut self) {
let overridden = self.metrics_overridden.swap(false, Ordering::SeqCst);
let running = self.running.get_mut().take();
if running.is_none() && !overridden {
return;
}
if let Some(running) = &running {
running.stop.store(true, Ordering::SeqCst);
running.task.abort();
}
let page = self.page.clone();
let had_capture = running.is_some();
if let Ok(handle) = tokio::runtime::Handle::try_current() {
handle.spawn(async move {
if had_capture {
let _ = page.execute(StopScreencastParams::default()).await;
}
if overridden {
let _ = page
.execute(ClearDeviceMetricsOverrideParams::default())
.await;
}
});
}
}
}
async fn run<A: FrameAck>(
mut events: impl Stream<Item = RawFrame> + Unpin,
acker: &A,
stop: &AtomicBool,
started: Instant,
viewport: Viewport,
subscribers: &Mutex<Vec<mpsc::Sender<ScreencastFrame>>>,
) {
while let Some(raw) = events.next().await {
if stop.load(Ordering::SeqCst) {
break;
}
acker.ack(raw.session_id).await;
let Ok(bytes) = BASE64.decode(&raw.data_base64) else {
continue;
};
publish(
subscribers,
ScreencastFrame {
jpeg: bytes.into(),
viewport,
captured_at: started.elapsed().as_secs_f64(),
},
)
.await;
}
}
async fn publish(subscribers: &Mutex<Vec<mpsc::Sender<ScreencastFrame>>>, frame: ScreencastFrame) {
let mut subs = subscribers.lock().await;
subs.retain(|tx| match tx.try_send(frame.clone()) {
Ok(()) => true,
Err(mpsc::error::TrySendError::Full(_)) => true,
Err(mpsc::error::TrySendError::Closed(_)) => false,
});
}
#[cfg(test)]
mod tests {
use super::*;
use futures::stream;
use std::sync::atomic::AtomicUsize;
struct CountingAck(AtomicUsize);
#[async_trait]
impl FrameAck for CountingAck {
async fn ack(&self, _session_id: i64) {
self.0.fetch_add(1, Ordering::SeqCst);
}
}
fn raw_frames(n: usize) -> Vec<RawFrame> {
(0..n)
.map(|i| RawFrame {
session_id: i as i64,
data_base64: BASE64.encode(format!("frame-{i}")),
})
.collect()
}
fn test_viewport() -> Viewport {
Viewport {
width: 1920,
height: 1080,
device_pixel_ratio: 1.0,
}
}
async fn drain(mut rx: FrameReceiver) -> Vec<ScreencastFrame> {
let mut out = Vec::new();
while let Ok(frame) = rx.try_recv() {
out.push(frame);
}
out
}
#[tokio::test]
async fn a_consumer_that_falls_behind_loses_frames_but_keeps_its_subscription() {
let subscribers: Mutex<Vec<mpsc::Sender<ScreencastFrame>>> = Mutex::new(Vec::new());
let (tx, mut rx) = mpsc::channel(FRAME_CHANNEL_CAP);
subscribers.lock().await.push(tx);
for i in 0..(FRAME_CHANNEL_CAP * 2) {
publish(
&subscribers,
ScreencastFrame {
jpeg: vec![i as u8].into(),
viewport: test_viewport(),
captured_at: i as f64,
},
)
.await;
}
assert_eq!(
subscribers.lock().await.len(),
1,
"a consumer that is behind is still a consumer — only a CLOSED channel deregisters"
);
let mut buffered = 0usize;
while rx.try_recv().is_ok() {
buffered += 1;
}
assert_eq!(
buffered, FRAME_CHANNEL_CAP,
"memory is bounded by the cap, not by how long the consumer stays slow"
);
}
#[tokio::test]
async fn every_subscriber_receives_every_frame() {
let acker = CountingAck(AtomicUsize::new(0));
let stop = AtomicBool::new(false);
let subscribers = Mutex::new(Vec::new());
let (tx_a, rx_a) = mpsc::channel(FRAME_CHANNEL_CAP);
let (tx_b, rx_b) = mpsc::channel(FRAME_CHANNEL_CAP);
subscribers.lock().await.push(tx_a);
subscribers.lock().await.push(tx_b);
let frames = raw_frames(3);
run(
stream::iter(frames),
&acker,
&stop,
Instant::now(),
test_viewport(),
&subscribers,
)
.await;
let a = drain(rx_a).await;
let b = drain(rx_b).await;
assert_eq!(a.len(), 3, "subscriber A should see all 3 frames");
assert_eq!(b.len(), 3, "subscriber B should see all 3 frames");
for (got, i) in a.iter().zip(0..) {
assert_eq!(
got.jpeg,
Arc::<[u8]>::from(format!("frame-{i}").into_bytes())
);
assert_eq!(got.viewport.width, 1920);
}
assert_eq!(
a.iter().map(|f| &f.jpeg).collect::<Vec<_>>(),
b.iter().map(|f| &f.jpeg).collect::<Vec<_>>(),
"both subscribers must see identical frame content"
);
for (left, right) in a.iter().zip(&b) {
assert!(
Arc::ptr_eq(&left.jpeg, &right.jpeg),
"fan-out must share the JPEG allocation instead of deep-cloning it"
);
}
}
#[tokio::test]
async fn ack_count_is_independent_of_subscriber_count() {
for subscriber_count in [0usize, 1, 3] {
let acker = CountingAck(AtomicUsize::new(0));
let stop = AtomicBool::new(false);
let subscribers = Mutex::new(Vec::new());
let mut receivers = Vec::new();
for _ in 0..subscriber_count {
let (tx, rx) = mpsc::channel(FRAME_CHANNEL_CAP);
subscribers.lock().await.push(tx);
receivers.push(rx);
}
run(
stream::iter(raw_frames(5)),
&acker,
&stop,
Instant::now(),
test_viewport(),
&subscribers,
)
.await;
assert_eq!(
acker.0.load(Ordering::SeqCst),
5,
"5 incoming frames must ack exactly 5 times with {subscriber_count} subscribers"
);
for rx in receivers {
assert_eq!(drain(rx).await.len(), 5);
}
}
}
#[tokio::test]
async fn zero_subscribers_means_zero_frames_delivered() {
let acker = CountingAck(AtomicUsize::new(0));
let stop = AtomicBool::new(false);
let subscribers: Mutex<Vec<mpsc::Sender<ScreencastFrame>>> = Mutex::new(Vec::new());
run(
stream::iter(raw_frames(4)),
&acker,
&stop,
Instant::now(),
test_viewport(),
&subscribers,
)
.await;
assert!(
subscribers.lock().await.is_empty(),
"no subscriber ever registered, so none should exist after the run"
);
}
#[tokio::test]
async fn dropped_subscriber_is_pruned_not_fatal() {
let acker = CountingAck(AtomicUsize::new(0));
let stop = AtomicBool::new(false);
let subscribers = Mutex::new(Vec::new());
let (tx_survivor, rx_survivor) = mpsc::channel(FRAME_CHANNEL_CAP);
let (tx_dropped, rx_dropped) = mpsc::channel(FRAME_CHANNEL_CAP);
subscribers.lock().await.push(tx_survivor);
subscribers.lock().await.push(tx_dropped);
drop(rx_dropped);
run(
stream::iter(raw_frames(2)),
&acker,
&stop,
Instant::now(),
test_viewport(),
&subscribers,
)
.await;
assert_eq!(drain(rx_survivor).await.len(), 2);
assert_eq!(
subscribers.lock().await.len(),
1,
"the dropped receiver's sender should have been pruned"
);
}
#[tokio::test]
async fn stop_flag_halts_processing() {
let acker = CountingAck(AtomicUsize::new(0));
let stop = AtomicBool::new(true);
let subscribers = Mutex::new(Vec::new());
run(
stream::iter(raw_frames(3)),
&acker,
&stop,
Instant::now(),
test_viewport(),
&subscribers,
)
.await;
assert_eq!(acker.0.load(Ordering::SeqCst), 0);
}
async fn dead_running_capture() -> RunningCapture {
let task = tokio::spawn(async {});
for _ in 0..1000 {
if task.is_finished() {
break;
}
tokio::task::yield_now().await;
}
assert!(task.is_finished(), "dummy task never finished");
RunningCapture {
stop: Arc::new(AtomicBool::new(false)),
task,
started: Instant::now(),
}
}
fn alive_running_capture() -> RunningCapture {
RunningCapture {
stop: Arc::new(AtomicBool::new(false)),
task: tokio::spawn(std::future::pending::<()>()),
started: Instant::now(),
}
}
#[tokio::test]
async fn reusable_start_is_none_with_no_capture_yet() {
assert!(reusable_start(&None).is_none());
}
#[tokio::test]
async fn reusable_start_reuses_a_live_capture() {
let capture = alive_running_capture();
let started = capture.started;
let running = Some(capture);
assert_eq!(reusable_start(&running), Some(started));
if let Some(r) = running {
r.task.abort();
}
}
#[tokio::test]
async fn reusable_start_discards_a_dead_capture() {
let running = Some(dead_running_capture().await);
assert_eq!(
reusable_start(&running),
None,
"a finished capture task must not be treated as still running"
);
}
}