From 5269356089ad6280c0e9fac5d41dbfe3ebcf2437 Mon Sep 17 00:00:00 2001
From: Ian Douglas Scott <idscott@system76.com>
Date: Wed, 24 Jan 2024 13:58:12 -0800
Subject: [PATCH] Use `Weak<Capture>` in session state

---
 src/wayland/capture.rs    | 16 +++++++++++-----
 src/wayland/screencopy.rs | 12 ++++--------
 2 files changed, 15 insertions(+), 13 deletions(-)

diff --git a/src/wayland/capture.rs b/src/wayland/capture.rs
index f2842c7..961610a 100644
--- a/src/wayland/capture.rs
+++ b/src/wayland/capture.rs
@@ -11,7 +11,7 @@ use cosmic::cctk;
 
 use std::sync::{
     atomic::{AtomicBool, Ordering},
-    Arc, Mutex,
+    Arc, Mutex, Weak,
 };
 
 use super::{AppData, Buffer};
@@ -49,10 +49,16 @@ impl Capture {
         }
     }
 
+    // Returns `None` if capture is no longer active
+    // (or if `session` wasn't created with `SessionData`)
     pub fn for_session(
         session: &zcosmic_screencopy_session_v1::ZcosmicScreencopySessionV1,
-    ) -> Option<&Arc<Self>> {
-        Some(&session.data::<SessionData>()?.capture)
+    ) -> Option<Arc<Self>> {
+        session
+            .data::<SessionData>()?
+            .capture
+            .upgrade()
+            .filter(|c| c.running())
     }
 
     pub fn running(&self) -> bool {
@@ -79,7 +85,7 @@ impl Capture {
 
         let udata = SessionData {
             session_data: Default::default(),
-            capture: self.clone(),
+            capture: Arc::downgrade(self),
         };
         match &self.source {
             CaptureSource::Toplevel(toplevel) => {
@@ -105,7 +111,7 @@ impl Capture {
 
 struct SessionData {
     session_data: ScreencopySessionData,
-    capture: Arc<Capture>,
+    capture: Weak<Capture>,
 }
 
 impl ScreencopySessionDataExt for SessionData {
diff --git a/src/wayland/screencopy.rs b/src/wayland/screencopy.rs
index 1626887..589563f 100644
--- a/src/wayland/screencopy.rs
+++ b/src/wayland/screencopy.rs
@@ -19,11 +19,10 @@ impl ScreencopyHandler for AppData {
         session: &zcosmic_screencopy_session_v1::ZcosmicScreencopySessionV1,
         buffer_infos: &[BufferInfo],
     ) {
-        let capture = Capture::for_session(session).unwrap();
-        if !capture.running() {
+        let Some(capture) = Capture::for_session(session) else {
             session.destroy();
             return;
-        }
+        };
 
         let mut buffer = capture.buffer.lock().unwrap();
         // Create new buffer if none, or different format
@@ -53,11 +52,10 @@ impl ScreencopyHandler for AppData {
         _qh: &QueueHandle<Self>,
         session: &zcosmic_screencopy_session_v1::ZcosmicScreencopySessionV1,
     ) {
-        let capture = Capture::for_session(session).unwrap();
-        if !capture.running() {
+        let Some(capture) = Capture::for_session(session) else {
             session.destroy();
             return;
-        }
+        };
 
         let mut buffer = capture.buffer.lock().unwrap();
         if buffer.is_none() {
@@ -95,8 +93,6 @@ impl ScreencopyHandler for AppData {
         println!("Failed");
         if let Some(capture) = Capture::for_session(session) {
             capture.cancel();
-        } else {
-            println!("Capture not found?")
         }
         session.destroy();
     }