From b93ddfce5b2c987babe7018284d4f5ad56c59db7 Mon Sep 17 00:00:00 2001
From: galister <22305755+galister@users.noreply.github.com>
Date: Thu, 22 Feb 2024 02:01:15 +0100
Subject: [PATCH] openxr: fix segfault

---
 src/backend/common.rs           | 10 ++---
 src/backend/openvr/mod.rs       |  2 +-
 src/backend/openxr/mod.rs       | 75 ++++++++++++++++++---------------
 src/backend/openxr/swapchain.rs |  2 +-
 4 files changed, 48 insertions(+), 41 deletions(-)

diff --git a/src/backend/common.rs b/src/backend/common.rs
index 41d9714..7f5b899 100644
--- a/src/backend/common.rs
+++ b/src/backend/common.rs
@@ -90,20 +90,18 @@ where
         }
     }
 
-    pub fn drop_by_selector(&mut self, selector: &OverlaySelector) {
+    pub fn remove_by_selector(&mut self, selector: &OverlaySelector) -> Option<OverlayData<T>> {
         match selector {
-            OverlaySelector::Id(id) => {
-                self.overlays.remove(id);
-            }
+            OverlaySelector::Id(id) => self.overlays.remove(id),
             OverlaySelector::Name(name) => {
                 let id = self
                     .overlays
                     .iter()
                     .find(|(_, o)| *o.state.name == **name)
                     .map(|(id, _)| *id);
-                id.and_then(|id| self.overlays.remove(id));
+                id.and_then(|id| self.overlays.remove(id))
             }
-        };
+        }
     }
 
     pub fn get_by_id(&mut self, id: usize) -> Option<&OverlayData<T>> {
diff --git a/src/backend/openvr/mod.rs b/src/backend/openvr/mod.rs
index 3a3ffcf..85e6b6b 100644
--- a/src/backend/openvr/mod.rs
+++ b/src/backend/openvr/mod.rs
@@ -182,7 +182,7 @@ pub fn openvr_run(running: Arc<AtomicBool>) -> Result<(), BackendError> {
                 TaskType::DropOverlay(sel) => {
                     if let Some(o) = overlays.mut_by_selector(&sel) {
                         o.destroy(&mut overlay_mngr);
-                        overlays.drop_by_selector(&sel);
+                        overlays.remove_by_selector(&sel);
                     }
                 }
             }
diff --git a/src/backend/openxr/mod.rs b/src/backend/openxr/mod.rs
index a394e29..4183d14 100644
--- a/src/backend/openxr/mod.rs
+++ b/src/backend/openxr/mod.rs
@@ -1,7 +1,7 @@
 use std::{
     collections::VecDeque,
     sync::{
-        atomic::{AtomicBool, Ordering},
+        atomic::{AtomicBool, AtomicUsize, Ordering},
         Arc,
     },
     time::Duration,
@@ -34,6 +34,7 @@ mod swapchain;
 
 const VIEW_TYPE: xr::ViewConfigurationType = xr::ViewConfigurationType::PRIMARY_STEREO;
 const VIEW_COUNT: u32 = 2;
+static FRAME_COUNTER: AtomicUsize = AtomicUsize::new(0);
 
 struct XrState {
     instance: xr::Instance,
@@ -68,6 +69,8 @@ pub fn openxr_run(running: Arc<AtomicBool>) -> Result<(), BackendError> {
     notifications.run_dbus();
     notifications.run_udp();
 
+    let mut delete_queue = vec![];
+
     #[cfg(feature = "osc")]
     let mut osc_sender =
         crate::backend::osc::OscSender::new(app_state.session.config.osc_out_port).ok();
@@ -121,6 +124,8 @@ pub fn openxr_run(running: Arc<AtomicBool>) -> Result<(), BackendError> {
     let mut due_tasks = VecDeque::with_capacity(4);
 
     'main_loop: loop {
+        let cur_frame = FRAME_COUNTER.fetch_add(1, Ordering::Relaxed);
+
         if !running.load(Ordering::Relaxed) {
             log::warn!("Received shutdown signal.");
             match xr_state.session.request_exit() {
@@ -184,38 +189,6 @@ pub fn openxr_run(running: Arc<AtomicBool>) -> Result<(), BackendError> {
             continue 'main_loop;
         }
 
-        notifications.submit_pending(&mut app_state);
-
-        app_state.tasks.retrieve_due(&mut due_tasks);
-        while let Some(task) = due_tasks.pop_front() {
-            match task {
-                TaskType::Global(f) => f(&mut app_state),
-                TaskType::Overlay(sel, f) => {
-                    if let Some(o) = overlays.mut_by_selector(&sel) {
-                        f(&mut app_state, &mut o.state);
-                    }
-                }
-                TaskType::CreateOverlay(sel, f) => {
-                    let None = overlays.mut_by_selector(&sel) else {
-                        continue;
-                    };
-
-                    let Some((state, backend)) = f(&mut app_state) else {
-                        continue;
-                    };
-
-                    overlays.add(OverlayData {
-                        state,
-                        backend,
-                        ..Default::default()
-                    });
-                }
-                TaskType::DropOverlay(sel) => {
-                    overlays.drop_by_selector(&sel);
-                }
-            }
-        }
-
         app_state.input_state.pre_update();
         input_source.update(&xr_state, &mut app_state)?;
         app_state.input_state.post_update();
@@ -326,6 +299,42 @@ pub fn openxr_run(running: Arc<AtomicBool>) -> Result<(), BackendError> {
             &frame_ref,
         )?;
 
+        notifications.submit_pending(&mut app_state);
+
+        app_state.tasks.retrieve_due(&mut due_tasks);
+        while let Some(task) = due_tasks.pop_front() {
+            match task {
+                TaskType::Global(f) => f(&mut app_state),
+                TaskType::Overlay(sel, f) => {
+                    if let Some(o) = overlays.mut_by_selector(&sel) {
+                        f(&mut app_state, &mut o.state);
+                    }
+                }
+                TaskType::CreateOverlay(sel, f) => {
+                    let None = overlays.mut_by_selector(&sel) else {
+                        continue;
+                    };
+
+                    let Some((state, backend)) = f(&mut app_state) else {
+                        continue;
+                    };
+
+                    overlays.add(OverlayData {
+                        state,
+                        backend,
+                        ..Default::default()
+                    });
+                }
+                TaskType::DropOverlay(sel) => {
+                    let o = overlays.remove_by_selector(&sel);
+                    // set for deletion after all images are done showing
+                    delete_queue.push((o, cur_frame + 5));
+                }
+            }
+        }
+
+        delete_queue.retain(|(_, frame)| *frame > cur_frame);
+
         app_state.hid_provider.on_new_frame();
 
         let watch = overlays.mut_by_id(watch_id).unwrap(); // want panic
diff --git a/src/backend/openxr/swapchain.rs b/src/backend/openxr/swapchain.rs
index 65bd89e..4ba811e 100644
--- a/src/backend/openxr/swapchain.rs
+++ b/src/backend/openxr/swapchain.rs
@@ -49,7 +49,7 @@ pub(super) fn create_swapchain_render_data(
             let vk_image = vk::Image::from_raw(handle);
             // thanks @yshui
             let raw_image = unsafe {
-                RawImage::from_handle(
+                RawImage::from_handle_borrowed(
                     graphics.device.clone(),
                     vk_image,
                     ImageCreateInfo {