handle macos stream errors with error flag

Author: Brendonovich

src/capturer/engine/mac/mod.rs

@@ -1,5 +1,6 @@
-use std::cmp;
+use std::sync::atomic::AtomicBool;
 use std::sync::mpsc;
+use std::{cmp, sync::Arc};
 
 use pixelformat::get_pts_in_nanoseconds;
 use screencapturekit::{
@@ -31,10 +32,15 @@ mod pixelformat;
 
 pub use pixel_buffer::PixelBuffer;
 
-struct ErrorHandler;
+struct ErrorHandler {
+    error_flag: Arc<AtomicBool>,
+}
+
 impl StreamErrorHandler for ErrorHandler {
     fn on_error(&self) {
-        println!("Error!");
+        eprintln!("Screen capture error occurred.");
+        self.error_flag
+            .store(true, std::sync::atomic::Ordering::Relaxed);
     }
 }
 
@@ -54,7 +60,11 @@ impl StreamOutput for Capturer {
     }
 }
 
-pub fn create_capturer(options: &Options, tx: mpsc::Sender<ChannelItem>) -> SCStream {
+pub fn create_capturer(
+    options: &Options,
+    tx: mpsc::Sender<ChannelItem>,
+    error_flag: Arc<AtomicBool>,
+) -> SCStream {
     // If no target is specified, capture the main display
     let target = options
         .target
@@ -147,7 +157,7 @@ pub fn create_capturer(options: &Options, tx: mpsc::Sender<ChannelItem>) -> SCSt
         ..Default::default()
     };
 
-    let mut stream = SCStream::new(filter, stream_config, ErrorHandler);
+    let mut stream = SCStream::new(filter, stream_config, ErrorHandler { error_flag });
     stream.add_output(Capturer::new(tx), SCStreamOutputType::Screen);
 
     stream

src/capturer/engine/mac/pixel_buffer.rs

@@ -166,8 +166,26 @@ impl<'a> Drop for PixelBufferData<'a> {
 impl RawCapturer<'_> {
     #[cfg(target_os = "macos")]
     pub fn get_next_pixel_buffer(&self) -> Result<PixelBuffer, mpsc::RecvError> {
+        use std::time::Duration;
+
+        let capturer = &self.capturer;
+
         loop {
-            if let Some(frame) = PixelBuffer::new(self.capturer.rx.recv()?) {
+            let error_flag = capturer
+                .engine
+                .error_flag
+                .load(std::sync::atomic::Ordering::Relaxed);
+            if error_flag {
+                return Err(mpsc::RecvError);
+            }
+
+            let res = match capturer.rx.recv_timeout(Duration::from_millis(1)) {
+                Ok(v) => Ok(v),
+                Err(mpsc::RecvTimeoutError::Timeout) => continue,
+                Err(mpsc::RecvTimeoutError::Disconnected) => Err(mpsc::RecvError),
+            }?;
+
+            if let Some(frame) = PixelBuffer::new(res) {
                 return Ok(frame);
             }
         }

src/capturer/engine/mod.rs

@@ -40,8 +40,11 @@ pub fn get_output_frame_size(options: &Options) -> [u32; 2] {
 
 pub struct Engine {
     options: Options,
+
     #[cfg(target_os = "macos")]
     mac: screencapturekit::sc_stream::SCStream,
+    #[cfg(target_os = "macos")]
+    error_flag: std::sync::Arc<std::sync::atomic::AtomicBool>,
 
     #[cfg(target_os = "windows")]
     win: win::WCStream,
@@ -54,9 +57,12 @@ impl Engine {
     pub fn new(options: &Options, tx: mpsc::Sender<ChannelItem>) -> Engine {
         #[cfg(target_os = "macos")]
         {
-            let mac = mac::create_capturer(options, tx);
+            let error_flag = std::sync::Arc::new(std::sync::atomic::AtomicBool::new(false));
+            let mac = mac::create_capturer(options, tx, error_flag.clone());
+
             Engine {
                 mac,
+                error_flag,
                 options: (*options).clone(),
             }
         }