From bac1e56842a67fca28fbba720408d8e6309a75f1 Mon Sep 17 00:00:00 2001
From: DataTriny <datatriny@gmail.com>
Date: Tue, 26 Dec 2023 19:01:01 +0100
Subject: [PATCH] Properly drop the top-level adapter struct

---
 platforms/unix/src/adapter.rs | 20 ++++++++++++++++----
 platforms/unix/src/context.rs | 22 +++++++++++++++++-----
 2 files changed, 33 insertions(+), 9 deletions(-)

diff --git a/platforms/unix/src/adapter.rs b/platforms/unix/src/adapter.rs
index e8c37e53b..0102eac7f 100644
--- a/platforms/unix/src/adapter.rs
+++ b/platforms/unix/src/adapter.rs
@@ -198,8 +198,6 @@ impl TreeChangeHandler for AdapterChangeHandler<'_> {
     }
 }
 
-static NEXT_ADAPTER_ID: AtomicUsize = AtomicUsize::new(0);
-
 pub(crate) struct AdapterImpl {
     id: usize,
     messages: Sender<Message>,
@@ -208,13 +206,13 @@ pub(crate) struct AdapterImpl {
 
 impl AdapterImpl {
     fn new(
+        id: usize,
         initial_state: TreeUpdate,
         is_window_focused: bool,
         root_window_bounds: WindowBounds,
         action_handler: Box<dyn ActionHandler + Send>,
     ) -> Self {
         let tree = Tree::new(initial_state, is_window_focused);
-        let id = NEXT_ADAPTER_ID.fetch_add(1, Ordering::SeqCst);
         let (messages, context) = {
             let mut app_context = AppContext::write();
             let messages = app_context.messages.clone().unwrap();
@@ -430,7 +428,10 @@ impl Drop for AdapterImpl {
 
 pub(crate) type LazyAdapter = Pin<Arc<Lazy<AdapterImpl, Boxed<AdapterImpl>>>>;
 
+static NEXT_ADAPTER_ID: AtomicUsize = AtomicUsize::new(0);
+
 pub struct Adapter {
+    id: usize,
     r#impl: LazyAdapter,
     is_window_focused: Arc<AtomicBool>,
     root_window_bounds: Arc<Mutex<WindowBounds>>,
@@ -442,6 +443,7 @@ impl Adapter {
         source: impl 'static + FnOnce() -> TreeUpdate + Send,
         action_handler: Box<dyn ActionHandler + Send>,
     ) -> Self {
+        let id = NEXT_ADAPTER_ID.fetch_add(1, Ordering::SeqCst);
         let is_window_focused = Arc::new(AtomicBool::new(false));
         let is_window_focused_copy = is_window_focused.clone();
         let root_window_bounds = Arc::new(Mutex::new(Default::default()));
@@ -451,6 +453,7 @@ impl Adapter {
                 let is_window_focused = is_window_focused_copy.load(Ordering::Relaxed);
                 let root_window_bounds = *root_window_bounds_copy.lock().unwrap();
                 AdapterImpl::new(
+                    id,
                     source(),
                     is_window_focused,
                     root_window_bounds,
@@ -460,11 +463,12 @@ impl Adapter {
             .boxed(),
         ));
         let adapter = Self {
+            id,
             r#impl: r#impl.clone(),
             is_window_focused,
             root_window_bounds,
         };
-        block_on(async move { ActivationContext::activate_eventually(r#impl).await });
+        block_on(async move { ActivationContext::activate_eventually(id, r#impl).await });
         adapter
     }
 
@@ -496,6 +500,14 @@ impl Adapter {
     }
 }
 
+impl Drop for Adapter {
+    fn drop(&mut self) {
+        block_on(async {
+            ActivationContext::remove_adapter(self.id).await;
+        })
+    }
+}
+
 pub(crate) enum Message {
     RegisterInterfaces {
         adapter_id: usize,
diff --git a/platforms/unix/src/context.rs b/platforms/unix/src/context.rs
index cde8c7215..794181d79 100644
--- a/platforms/unix/src/context.rs
+++ b/platforms/unix/src/context.rs
@@ -116,7 +116,7 @@ impl AppContext {
 
 pub(crate) struct ActivationContext {
     _task: Task<()>,
-    adapters: Vec<LazyAdapter>,
+    adapters: Vec<(usize, LazyAdapter)>,
 }
 
 static ACTIVATION_CONTEXT: AsyncOnceCell<Arc<AsyncMutex<ActivationContext>>> = AsyncOnceCell::new();
@@ -143,15 +143,27 @@ impl ActivationContext {
             .await
     }
 
-    pub(crate) async fn activate_eventually(adapter: LazyAdapter) {
+    pub(crate) async fn activate_eventually(id: usize, adapter: LazyAdapter) {
         let mut activation_context = ActivationContext::get_or_init().await;
-        activation_context.adapters.push(adapter);
+        activation_context.adapters.push((id, adapter));
         let is_a11y_enabled = AppContext::get_or_init().messages.is_some();
         if is_a11y_enabled {
-            let adapter = activation_context.adapters.last().unwrap();
+            let adapter = &activation_context.adapters.last().unwrap().1;
             adapter.as_ref().await;
         }
     }
+
+    pub(crate) async fn remove_adapter(id: usize) {
+        if let Some(activation_context) = ACTIVATION_CONTEXT.get() {
+            let mut context = activation_context.lock().await;
+            if let Ok(index) = context
+                .adapters
+                .binary_search_by(|adapter| adapter.0.cmp(&id))
+            {
+                context.adapters.remove(index);
+            }
+        }
+    }
 }
 
 async fn listen(session_bus: Connection) -> zbus::Result<()> {
@@ -182,7 +194,7 @@ async fn listen(session_bus: Connection) -> zbus::Result<()> {
                 if atspi_bus.is_some() {
                     if let Some(activation_context) = ACTIVATION_CONTEXT.get() {
                         let activation_context = activation_context.lock().await;
-                        for adapter in &activation_context.adapters {
+                        for (_, adapter) in &activation_context.adapters {
                             adapter.as_ref().await.register_tree().await;
                         }
                     }