use std::collections::HashMap;
use std::sync::{Arc, Mutex};
use arc_swap::ArcSwapOption;
use tokio::sync::{broadcast, watch};
use tokio::task::JoinHandle;
use tokio::time::MissedTickBehavior;
use std::time::SystemTime;
use crate::error::{Error, Result};
use crate::event::{Event, EventKind};
use visual_cortex_capture::{CaptureError, Frame, FrameSource, Rate, Target};
const BROADCAST_CAPACITY: usize = 256;
const INITIAL_CAPTURE_RATE_HZ: f64 = 1.0;
const SESSION_WATCHER: &str = "$session";
#[derive(Clone, Debug)]
pub struct Session {
pub(crate) inner: Arc<Inner>,
}
#[derive(Debug)]
pub(crate) struct Inner {
pub(crate) slot: Arc<ArcSwapOption<Frame>>,
pub(crate) events: broadcast::Sender<Event>,
pub(crate) rate_tx: watch::Sender<Rate>,
pub(crate) watchers: Mutex<HashMap<String, WatcherEntry>>,
pub(crate) capture_task: Mutex<Option<JoinHandle<()>>>,
}
#[derive(Debug)]
pub(crate) struct WatcherEntry {
pub(crate) rate: Rate,
pub(crate) task: Option<JoinHandle<()>>,
}
impl Drop for Inner {
fn drop(&mut self) {
if let Some(task) = self.capture_task.lock().unwrap().take() {
task.abort();
}
for (_, mut entry) in self.watchers.lock().unwrap().drain() {
if let Some(task) = entry.task.take() {
task.abort();
}
}
}
}
impl Session {
pub fn builder() -> SessionBuilder {
SessionBuilder {
source: None,
target: None,
reattach: true,
}
}
pub fn watch(&self, name: &str) -> crate::watcher::WatcherBuilder<'_> {
crate::watcher::WatcherBuilder::new(self, name)
}
pub fn unsubscribe(&self, name: &str) -> bool {
let entry = self.inner.watchers.lock().unwrap().remove(name);
match entry {
Some(mut entry) => {
if let Some(task) = entry.task.take() {
task.abort();
}
recompute_capture_rate(&self.inner);
true
}
None => false,
}
}
}
pub struct SessionBuilder {
source: Option<Box<dyn FrameSource>>,
target: Option<Target>,
reattach: bool,
}
impl SessionBuilder {
pub fn source(mut self, source: impl FrameSource) -> Self {
self.source = Some(Box::new(source));
self
}
pub fn target(mut self, target: Target) -> Self {
self.target = Some(target);
self
}
pub fn reattach(mut self, reattach: bool) -> Self {
self.reattach = reattach;
self
}
pub async fn build(self) -> Result<Session> {
let source = self.resolve_source()?;
let slot = Arc::new(ArcSwapOption::from(None));
let (events, _) = broadcast::channel(BROADCAST_CAPACITY);
let (rate_tx, rate_rx) = watch::channel(Rate::hz(INITIAL_CAPTURE_RATE_HZ));
let capture_task =
tokio::spawn(capture_loop(source, slot.clone(), rate_rx, events.clone()));
Ok(Session {
inner: Arc::new(Inner {
slot,
events,
rate_tx,
watchers: Mutex::new(HashMap::new()),
capture_task: Mutex::new(Some(capture_task)),
}),
})
}
fn resolve_source(self) -> Result<Box<dyn FrameSource>> {
match (self.source, self.target) {
(Some(_), Some(_)) => Err(Error::AmbiguousSource),
(Some(source), None) => Ok(source),
(None, Some(target)) => build_target_source(&target, self.reattach),
(None, None) => Err(Error::NoSource),
}
}
}
#[cfg(target_os = "macos")]
fn build_target_source(target: &Target, reattach: bool) -> Result<Box<dyn FrameSource>> {
Ok(Box::new(visual_cortex_capture::ScapSource::new(
target, reattach,
)?))
}
#[cfg(not(target_os = "macos"))]
fn build_target_source(_target: &Target, _reattach: bool) -> Result<Box<dyn FrameSource>> {
Err(Error::UnsupportedPlatform)
}
pub(crate) fn recompute_capture_rate(inner: &Inner) {
let fastest = {
let watchers = inner.watchers.lock().unwrap();
watchers
.values()
.map(|entry| entry.rate)
.min_by_key(|rate| rate.period())
.unwrap_or_else(|| Rate::hz(INITIAL_CAPTURE_RATE_HZ))
};
if fastest != *inner.rate_tx.borrow() {
let _ = inner.rate_tx.send(fastest);
}
}
fn send_session_event(events: &broadcast::Sender<Event>, kind: EventKind) {
let _ = events.send(Event {
watcher: Arc::from(SESSION_WATCHER),
timestamp: SystemTime::now(),
kind,
});
}
async fn capture_loop(
mut source: Box<dyn FrameSource>,
slot: Arc<ArcSwapOption<Frame>>,
mut rate_rx: watch::Receiver<Rate>,
events: broadcast::Sender<Event>,
) {
let mut interval = tokio::time::interval(rate_rx.borrow().period());
interval.set_missed_tick_behavior(MissedTickBehavior::Skip);
let mut target_present = true;
loop {
tokio::select! {
_ = interval.tick() => {
match source.capture_frame().await {
Ok(frame) => {
if !target_present {
target_present = true;
send_session_event(&events, EventKind::TargetReattached);
}
slot.store(Some(frame));
}
Err(CaptureError::TargetLost) | Err(CaptureError::TargetNotFound(_)) => {
if target_present {
target_present = false;
send_session_event(&events, EventKind::TargetLost);
}
}
Err(e) => tracing::warn!("capture error: {e}"),
}
}
changed = rate_rx.changed() => {
if changed.is_err() {
break; }
interval = tokio::time::interval(rate_rx.borrow().period());
interval.set_missed_tick_behavior(MissedTickBehavior::Skip);
}
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use visual_cortex_capture::FakeSource;
#[tokio::test]
async fn build_without_source_errors() {
let err = Session::builder().build().await.unwrap_err();
assert!(matches!(err, Error::NoSource));
}
#[tokio::test(start_paused = true)]
async fn capture_loop_fills_the_frame_slot() {
let session = Session::builder()
.source(FakeSource::new([Frame::solid(4, 4, [1, 2, 3, 255])]))
.build()
.await
.unwrap();
tokio::time::sleep(std::time::Duration::from_secs(3)).await;
let frame = session
.inner
.slot
.load_full()
.expect("slot should hold a frame");
assert_eq!(frame.data()[0], 1);
}
use visual_cortex_vision::Brightness;
fn one_frame_session_source() -> FakeSource {
FakeSource::new([Frame::solid(4, 4, [0, 0, 0, 255])])
}
#[tokio::test(start_paused = true)]
async fn capture_rate_tracks_fastest_watcher_and_recovers_on_unsubscribe() {
let session = Session::builder()
.source(one_frame_session_source())
.build()
.await
.unwrap();
let _slow = session
.watch("slow")
.rate(Rate::hz(2.0))
.detector(Brightness)
.subscribe()
.unwrap();
assert_eq!(*session.inner.rate_tx.borrow(), Rate::hz(2.0));
let _fast = session
.watch("fast")
.rate(Rate::hz(10.0))
.detector(Brightness)
.subscribe()
.unwrap();
assert_eq!(*session.inner.rate_tx.borrow(), Rate::hz(10.0));
assert!(session.unsubscribe("fast"));
assert_eq!(*session.inner.rate_tx.borrow(), Rate::hz(2.0));
assert!(session.unsubscribe("slow"));
assert_eq!(
*session.inner.rate_tx.borrow(),
Rate::hz(INITIAL_CAPTURE_RATE_HZ)
);
}
#[tokio::test]
async fn build_rejects_source_and_target_together() {
use visual_cortex_capture::Target;
let err = Session::builder()
.source(FakeSource::new([Frame::solid(2, 2, [0, 0, 0, 255])]))
.target(Target::display(0))
.build()
.await
.unwrap_err();
assert!(matches!(err, Error::AmbiguousSource));
}
#[tokio::test(start_paused = true)]
async fn unsubscribe_unknown_is_false_and_frees_the_name() {
let session = Session::builder()
.source(one_frame_session_source())
.build()
.await
.unwrap();
assert!(!session.unsubscribe("ghost"));
let _w = session.watch("w").detector(Brightness).subscribe().unwrap();
let dup = session.watch("w").detector(Brightness).subscribe();
assert!(matches!(dup, Err(Error::DuplicateWatcher(n)) if n == "w"));
assert!(session.unsubscribe("w"));
let _w2 = session.watch("w").detector(Brightness).subscribe().unwrap();
}
}