WayVR: Process cleanup support, Refactoring

src/backend/wayvr/client.rs

@@ -9,13 +9,13 @@ use smithay::{
 
 use super::{
     comp::{self},
-    display,
+    display, process,
 };
 
 pub struct WayVRClient {
     pub client: wayland_server::Client,
     pub display_handle: display::DisplayHandle,
-    pub pid: i32,
+    pub pid: u32,
 }
 
 pub struct WayVRManager {
@@ -28,6 +28,8 @@ pub struct WayVRManager {
     display: wayland_server::Display<comp::Application>,
     listener: wayland_server::ListeningSocket,
 
+    toplevel_surf_count: u32, // for logging purposes
+
     pub clients: Vec<WayVRClient>,
 }
 
@@ -67,6 +69,7 @@ impl WayVRManager {
             wayland_env,
             serial_counter: SerialCounter::new(),
             clients: Vec::new(),
+            toplevel_surf_count: 0,
         })
     }
 
@@ -74,6 +77,7 @@ impl WayVRManager {
         &mut self,
         stream: UnixStream,
         displays: &mut display::DisplayVec,
+        processes: &mut process::ProcessVec,
     ) -> anyhow::Result<()> {
         let client = self
             .display
@@ -84,28 +88,34 @@ impl WayVRManager {
         let creds = client.get_credentials(&self.display.handle())?;
         let auth_key = get_display_auth_from_pid(creds.pid)?;
 
-        for (idx, cell) in displays.vec.iter().enumerate() {
+        // Find suitable auth key from the process list
-            if let Some(cell) = &cell {
+        for process in processes.vec.iter().flatten() {
-                let display = &cell.obj;
+            let process = &process.obj;
-                if display.auth_key_matches(auth_key.as_str()) {
-                    let display_handle = display::DisplayVec::get_handle(cell, idx);
 
+            // Find process with matching auth key
+            if process.auth_key.as_str() == auth_key {
+                // Check if display handle is valid
+                if displays.get(&process.display_handle).is_some() {
+                    // Add client
                     self.clients.push(WayVRClient {
                         client,
-                        display_handle,
+                        display_handle: process.display_handle,
-                        pid: creds.pid,
+                        pid: creds.pid as u32,
                     });
                     return Ok(());
                 }
             }
         }
-
         anyhow::bail!("Process auth key is invalid or selected display is non-existent");
     }
 
-    fn accept_connections(&mut self, displays: &mut display::DisplayVec) -> anyhow::Result<()> {
+    fn accept_connections(
+        &mut self,
+        displays: &mut display::DisplayVec,
+        processes: &mut process::ProcessVec,
+    ) -> anyhow::Result<()> {
         if let Some(stream) = self.listener.accept()? {
-            if let Err(e) = self.accept_connection(stream, displays) {
+            if let Err(e) = self.accept_connection(stream, displays, processes) {
                 log::error!("Failed to accept connection: {}", e);
             }
         }
@@ -113,14 +123,24 @@ impl WayVRManager {
         Ok(())
     }
 
-    pub fn tick_wayland(&mut self, displays: &mut display::DisplayVec) -> anyhow::Result<()> {
+    pub fn tick_wayland(
-        if let Err(e) = self.accept_connections(displays) {
+        &mut self,
+        displays: &mut display::DisplayVec,
+        processes: &mut process::ProcessVec,
+    ) -> anyhow::Result<()> {
+        if let Err(e) = self.accept_connections(displays, processes) {
             log::error!("accept_connections failed: {}", e);
         }
 
         self.display.dispatch_clients(&mut self.state)?;
         self.display.flush_clients()?;
 
+        let surf_count = self.state.xdg_shell.toplevel_surfaces().len() as u32;
+        if surf_count != self.toplevel_surf_count {
+            self.toplevel_surf_count = surf_count;
+            log::info!("Toplevel surface count changed: {}", surf_count);
+        }
+
         Ok(())
     }
 

src/backend/wayvr/comp.rs

@@ -27,6 +27,7 @@ use wayland_server::protocol::wl_surface::WlSurface;
 use wayland_server::Client;
 
 use super::event_queue::SyncEventQueue;
+use super::WayVRTask;
 
 pub struct Application {
     pub compositor: compositor::CompositorState,
@@ -35,7 +36,7 @@ pub struct Application {
     pub shm: ShmState,
     pub data_device: DataDeviceState,
 
-    pub queue_new_toplevel: SyncEventQueue<(ClientId, ToplevelSurface)>,
+    pub wayvr_tasks: SyncEventQueue<WayVRTask>,
 }
 
 impl compositor::CompositorHandler for Application {
@@ -125,7 +126,8 @@ impl XdgShellHandler for Application {
 
     fn new_toplevel(&mut self, surface: ToplevelSurface) {
         if let Some(client) = surface.wl_surface().client() {
-            self.queue_new_toplevel.send((client.id(), surface.clone()));
+            self.wayvr_tasks
+                .send(WayVRTask::NewToplevel(client.id(), surface.clone()));
         }
         surface.with_pending_state(|state| {
             state.states.set(xdg_toplevel::State::Activated);

src/backend/wayvr/display.rs

@@ -18,7 +18,8 @@ use smithay::{
 use crate::{backend::overlay::OverlayID, gen_id};
 
 use super::{
-    client::WayVRManager, comp::send_frames_surface_tree, egl_data, smithay_wrapper, window,
+    client::WayVRManager, comp::send_frames_surface_tree, egl_data, event_queue::SyncEventQueue,
+    process, smithay_wrapper, window,
 };
 
 fn generate_auth_key() -> String {
@@ -26,22 +27,21 @@ fn generate_auth_key() -> String {
     uuid.to_string()
 }
 
-struct Process {
-    auth_key: String,
-    child: std::process::Child,
-}
-
-impl Drop for Process {
-    fn drop(&mut self) {
-        let _dont_care = self.child.kill();
-    }
-}
-
 struct DisplayWindow {
-    handle: window::WindowHandle,
+    window_handle: window::WindowHandle,
+    process_handle: process::ProcessHandle,
     toplevel: ToplevelSurface,
 }
 
+pub struct SpawnProcessResult {
+    pub auth_key: String,
+    pub child: std::process::Child,
+}
+
+pub enum DisplayTask {
+    ProcessCleanup(process::ProcessHandle),
+}
+
 pub struct Display {
     // Display info stuff
     pub width: u32,
@@ -59,7 +59,7 @@ pub struct Display {
     egl_data: Rc<egl_data::EGLData>,
     pub dmabuf_data: egl_data::DMAbufData,
 
-    processes: Vec<Process>,
+    pub tasks: SyncEventQueue<DisplayTask>,
 }
 
 impl Drop for Display {
@@ -113,25 +113,22 @@ impl Display {
             egl_image,
             gles_texture,
             wayland_env,
-            processes: Vec::new(),
             visible: true,
             overlay_id: None,
+            tasks: SyncEventQueue::new(),
         })
     }
 
-    pub fn auth_key_matches(&self, auth_key: &str) -> bool {
+    pub fn add_window(
-        for process in &self.processes {
+        &mut self,
-            if process.auth_key.as_str() == auth_key {
+        window_handle: window::WindowHandle,
-                return true;
+        process_handle: process::ProcessHandle,
-            }
+        toplevel: &ToplevelSurface,
-        }
+    ) {
-        false
-    }
-
-    pub fn add_window(&mut self, window_handle: window::WindowHandle, toplevel: &ToplevelSurface) {
         log::debug!("Attaching toplevel surface into display");
         self.displayed_windows.push(DisplayWindow {
-            handle: window_handle,
+            window_handle,
+            process_handle,
             toplevel: toplevel.clone(),
         });
         self.reposition_windows();
@@ -141,7 +138,7 @@ impl Display {
         let window_count = self.displayed_windows.len();
 
         for (i, win) in self.displayed_windows.iter_mut().enumerate() {
-            if let Some(window) = self.wm.borrow_mut().windows.get_mut(&win.handle) {
+            if let Some(window) = self.wm.borrow_mut().windows.get_mut(&win.window_handle) {
                 let d_cur = i as f32 / window_count as f32;
                 let d_next = (i + 1) as f32 / window_count as f32;
 
@@ -154,6 +151,24 @@ impl Display {
         }
     }
 
+    pub fn tick(&mut self) {
+        while let Some(task) = self.tasks.read() {
+            match task {
+                DisplayTask::ProcessCleanup(process_handle) => {
+                    self.displayed_windows
+                        .retain(|win| win.process_handle != process_handle);
+                    log::info!(
+                        "Cleanup finished for display \"{}\". Current window count: {}",
+                        self.name,
+                        self.displayed_windows.len()
+                    );
+
+                    self.reposition_windows();
+                }
+            }
+        }
+    }
+
     pub fn tick_render(&self, renderer: &mut GlesRenderer, time_ms: u64) -> anyhow::Result<()> {
         renderer.bind(self.gles_texture.clone())?;
 
@@ -166,7 +181,7 @@ impl Display {
             .iter()
             .flat_map(|display_window| {
                 let wm = self.wm.borrow_mut();
-                if let Some(window) = wm.windows.get(&display_window.handle) {
+                if let Some(window) = wm.windows.get(&display_window.window_handle) {
                     render_elements_from_surface_tree(
                         renderer,
                         display_window.toplevel.wl_surface(),
@@ -207,13 +222,13 @@ impl Display {
         let wm = self.wm.borrow();
 
         for cell in self.displayed_windows.iter() {
-            if let Some(window) = wm.windows.get(&cell.handle) {
+            if let Some(window) = wm.windows.get(&cell.window_handle) {
                 if (cursor_x as i32) >= window.pos_x
                     && (cursor_x as i32) < window.pos_x + window.size_x as i32
                     && (cursor_y as i32) >= window.pos_y
                     && (cursor_y as i32) < window.pos_y + window.size_y as i32
                 {
-                    return Some(cell.handle);
+                    return Some(cell.window_handle);
                 }
             }
         }
@@ -335,7 +350,7 @@ impl Display {
         exec_path: &str,
         args: &[&str],
         env: &[(&str, &str)],
-    ) -> anyhow::Result<()> {
+    ) -> anyhow::Result<SpawnProcessResult> {
         log::info!("Spawning subprocess with exec path \"{}\"", exec_path);
 
         let auth_key = generate_auth_key();
@@ -349,9 +364,7 @@ impl Display {
         }
 
         match cmd.spawn() {
-            Ok(child) => {
+            Ok(child) => Ok(SpawnProcessResult { auth_key, child }),
-                self.processes.push(Process { child, auth_key });
-            }
             Err(e) => {
                 anyhow::bail!(
 					"Failed to launch process with path \"{}\": {}. Make sure your exec path exists.",
@@ -360,8 +373,6 @@ impl Display {
 				);
             }
         }
-
-        Ok(())
     }
 }
 

src/backend/wayvr/egl_ex.rs

@@ -1,3 +1,5 @@
+#![allow(clippy::all)]
+
 //eglExportDMABUFImageMESA
 pub type PFNEGLEXPORTDMABUFIMAGEMESAPROC = Option<
     unsafe extern "C" fn(

src/backend/wayvr/handle.rs

@@ -51,7 +51,7 @@ macro_rules! gen_id {
                 }
             }
 
-			pub fn iter(&self, callback: &mut dyn FnMut($handle_name, &$instance_name)) {
+            pub fn iter(&self, callback: &dyn Fn($handle_name, &$instance_name)) {
                 for (idx, opt_cell) in self.vec.iter().enumerate() {
                     if let Some(cell) = opt_cell {
                         let handle = $container_name::get_handle(&cell, idx);
@@ -60,6 +60,18 @@ macro_rules! gen_id {
                 }
             }
 
+            pub fn iter_mut(
+                &mut self,
+                callback: &mut dyn FnMut($handle_name, &mut $instance_name),
+            ) {
+                for (idx, opt_cell) in self.vec.iter_mut().enumerate() {
+                    if let Some(cell) = opt_cell {
+                        let handle = $container_name::get_handle(&cell, idx);
+                        callback(handle, &mut cell.obj);
+                    }
+                }
+            }
+
             pub fn get_handle(cell: &$cell_name, idx: usize) -> $handle_name {
                 $handle_name {
                     idx: idx as u32,

src/backend/wayvr/mod.rs

@@ -5,6 +5,7 @@ pub mod egl_data;
 mod egl_ex;
 mod event_queue;
 mod handle;
+mod process;
 mod smithay_wrapper;
 mod time;
 mod window;
@@ -14,6 +15,8 @@ use std::{cell::RefCell, rc::Rc};
 use comp::Application;
 use display::DisplayVec;
 use event_queue::SyncEventQueue;
+use process::ProcessVec;
+use smallvec::SmallVec;
 use smithay::{
     backend::renderer::gles::GlesRenderer,
     input::SeatState,
@@ -27,6 +30,9 @@ use smithay::{
 };
 use time::get_millis;
 
+const STR_INVALID_HANDLE_DISP: &str = "Invalid display handle";
+const STR_INVALID_HANDLE_PROCESS: &str = "Invalid process handle";
+
 #[derive(Clone)]
 pub struct WaylandEnv {
     pub display_num: u32,
@@ -39,6 +45,12 @@ impl WaylandEnv {
     }
 }
 
+#[derive(Clone)]
+pub enum WayVRTask {
+    NewToplevel(ClientId, ToplevelSurface),
+    ProcessTerminationRequest(process::ProcessHandle),
+}
+
 #[allow(dead_code)]
 pub struct WayVR {
     time_start: u64,
@@ -47,8 +59,9 @@ pub struct WayVR {
     manager: client::WayVRManager,
     wm: Rc<RefCell<window::WindowManager>>,
     egl_data: Rc<egl_data::EGLData>,
+    pub processes: process::ProcessVec,
 
-    queue_new_toplevel: SyncEventQueue<(ClientId, ToplevelSurface)>,
+    tasks: SyncEventQueue<WayVRTask>,
 }
 
 pub enum MouseIndex {
@@ -72,7 +85,7 @@ impl WayVR {
         let seat_keyboard = seat.add_keyboard(Default::default(), 100, 100)?;
         let seat_pointer = seat.add_pointer();
 
-        let queue_new_toplevel = SyncEventQueue::new();
+        let tasks = SyncEventQueue::new();
 
         let state = Application {
             compositor,
@@ -80,7 +93,7 @@ impl WayVR {
             seat_state,
             shm,
             data_device,
-            queue_new_toplevel: queue_new_toplevel.clone(),
+            wayvr_tasks: tasks.clone(),
         };
 
         let time_start = get_millis();
@@ -94,19 +107,19 @@ impl WayVR {
             time_start,
             manager: client::WayVRManager::new(state, display, seat_keyboard, seat_pointer)?,
             displays: DisplayVec::new(),
+            processes: ProcessVec::new(),
             egl_data: Rc::new(egl_data),
             wm: Rc::new(RefCell::new(window::WindowManager::new())),
-            queue_new_toplevel,
+            tasks,
         })
     }
 
     pub fn tick_display(&mut self, display: display::DisplayHandle) -> anyhow::Result<()> {
         // millis since the start of wayvr
-
         let display = self
             .displays
             .get(&display)
-            .ok_or(anyhow::anyhow!("Invalid display handle"))?;
+            .ok_or(anyhow::anyhow!(STR_INVALID_HANDLE_DISP))?;
 
         let time_ms = get_millis() - self.time_start;
 
@@ -121,25 +134,68 @@ impl WayVR {
     }
 
     pub fn tick_events(&mut self) -> anyhow::Result<()> {
-        // Attach newly created toplevel surfaces to displayes
+        // Tick all child processes
-        while let Some((client_id, toplevel)) = self.queue_new_toplevel.read() {
+        let mut to_remove: SmallVec<[(process::ProcessHandle, display::DisplayHandle); 2]> =
+            SmallVec::new();
+        self.processes.iter_mut(&mut |handle, process| {
+            if !process.is_running() {
+                to_remove.push((handle, process.display_handle));
+            }
+        });
+
+        for (p_handle, disp_handle) in to_remove {
+            self.processes.remove(&p_handle);
+
+            if let Some(display) = self.displays.get(&disp_handle) {
+                display
+                    .tasks
+                    .send(display::DisplayTask::ProcessCleanup(p_handle));
+            }
+        }
+
+        for display in self.displays.vec.iter_mut().flatten() {
+            display.obj.tick();
+        }
+
+        while let Some(task) = self.tasks.read() {
+            match task {
+                WayVRTask::NewToplevel(client_id, toplevel) => {
+                    // Attach newly created toplevel surfaces to displays
                     for client in &self.manager.clients {
                         if client.client.id() == client_id {
                             let window_handle = self.wm.borrow_mut().create_window(&toplevel);
 
-                    if let Some(display) = self.displays.get_mut(&client.display_handle) {
+                            if let Some(process_handle) =
-                        display.add_window(window_handle, &toplevel);
+                                process::find_by_pid(&self.processes, client.pid)
+                            {
+                                if let Some(display) = self.displays.get_mut(&client.display_handle)
+                                {
+                                    display.add_window(window_handle, process_handle, &toplevel);
                                 } else {
                                     // This shouldn't happen, scream if it does
                                     log::error!("Could not attach window handle into display");
                                 }
+                            } else {
+                                log::error!(
+                                    "Failed to find process by PID {}. It was probably spawned externally.",
+                                    client.pid
+                                );
+                            }
 
                             break;
                         }
                     }
                 }
+                WayVRTask::ProcessTerminationRequest(process_handle) => {
+                    if let Some(process) = self.processes.get_mut(&process_handle) {
+                        process.terminate();
+                    }
+                }
+            }
+        }
 
-        self.manager.tick_wayland(&mut self.displays)
+        self.manager
+            .tick_wayland(&mut self.displays, &mut self.processes)
     }
 
     pub fn tick_finish(&mut self) -> anyhow::Result<()> {
@@ -222,16 +278,59 @@ impl WayVR {
         self.displays.remove(&handle);
     }
 
+    // Check if process with given arguments already exists
+    pub fn process_query(
+        &self,
+        display_handle: display::DisplayHandle,
+        exec_path: &str,
+        args: &[&str],
+        _env: &[(&str, &str)],
+    ) -> Option<process::ProcessHandle> {
+        for (idx, cell) in self.processes.vec.iter().enumerate() {
+            if let Some(cell) = &cell {
+                let process = &cell.obj;
+                if process.display_handle != display_handle
+                    || process.exec_path != exec_path
+                    || process.args != args
+                {
+                    continue;
+                }
+
+                return Some(process::ProcessVec::get_handle(cell, idx));
+            }
+        }
+
+        None
+    }
+
+    pub fn terminate_process(&mut self, process_handle: process::ProcessHandle) {
+        self.tasks
+            .send(WayVRTask::ProcessTerminationRequest(process_handle));
+    }
+
     pub fn spawn_process(
         &mut self,
-        display: display::DisplayHandle,
+        display_handle: display::DisplayHandle,
         exec_path: &str,
         args: &[&str],
         env: &[(&str, &str)],
-    ) -> anyhow::Result<()> {
+    ) -> anyhow::Result<process::ProcessHandle> {
-        if let Some(display) = self.displays.get_mut(&display) {
+        let display = self
-            display.spawn_process(exec_path, args, env)?
+            .displays
-        }
+            .get_mut(&display_handle)
-        Ok(())
+            .ok_or(anyhow::anyhow!(STR_INVALID_HANDLE_DISP))?;
+
+        let res = display.spawn_process(exec_path, args, env)?;
+        Ok(self.processes.add(process::Process {
+            auth_key: res.auth_key,
+            child: res.child,
+            display_handle,
+            exec_path: String::from(exec_path),
+            args: args.iter().map(|x| String::from(*x)).collect(),
+            env: env
+                .iter()
+                .map(|(a, b)| (String::from(*a), String::from(*b)))
+                .collect(),
+        }))
     }
 }

src/backend/wayvr/process.rs

@@ -0,0 +1,57 @@
+use crate::gen_id;
+
+use super::display;
+
+pub struct Process {
+    pub auth_key: String,
+    pub child: std::process::Child,
+    pub display_handle: display::DisplayHandle,
+
+    pub exec_path: String,
+    pub args: Vec<String>,
+    pub env: Vec<(String, String)>,
+}
+
+impl Drop for Process {
+    fn drop(&mut self) {
+        log::info!(
+            "Sending SIGTERM (graceful exit) to process {}",
+            self.exec_path.as_str()
+        );
+        self.terminate();
+    }
+}
+
+impl Process {
+    pub fn is_running(&mut self) -> bool {
+        match self.child.try_wait() {
+            Ok(Some(_exit_status)) => false,
+            Ok(None) => true,
+            Err(e) => {
+                // this shouldn't happen
+                log::error!("Child::try_wait failed: {}", e);
+                false
+            }
+        }
+    }
+
+    pub fn terminate(&mut self) {
+        unsafe {
+            // Gracefully stop process
+            libc::kill(self.child.id() as i32, libc::SIGTERM);
+        }
+    }
+}
+
+gen_id!(ProcessVec, Process, ProcessCell, ProcessHandle);
+
+pub fn find_by_pid(processes: &ProcessVec, pid: u32) -> Option<ProcessHandle> {
+    for (idx, cell) in processes.vec.iter().enumerate() {
+        if let Some(cell) = cell {
+            if cell.obj.child.id() == pid {
+                return Some(ProcessVec::get_handle(cell, idx));
+            }
+        }
+    }
+    None
+}

src/overlays/wayvr.rs

@@ -218,7 +218,7 @@ impl OverlayRenderer for WayVRRenderer {
 }
 
 #[allow(dead_code)]
-pub fn create_wayvr<O>(
+pub fn create_wayvr_display_overlay<O>(
     app: &mut state::AppState,
     display_width: u32,
     display_height: u32,
@@ -285,7 +285,7 @@ fn action_app_click<O>(
 where
     O: Default,
 {
-    use crate::overlays::wayvr::create_wayvr;
+    use crate::overlays::wayvr::create_wayvr_display_overlay;
 
     let mut created_overlay: Option<OverlayData<O>> = None;
 
@@ -322,7 +322,7 @@ where
                 &app_entry.target_display,
             )?;
 
-            let overlay = create_wayvr::<O>(
+            let overlay = create_wayvr_display_overlay::<O>(
                 app,
                 conf_display.width,
                 conf_display.height,
@@ -354,7 +354,16 @@ where
             vec![]
         };
 
-        wayvr.spawn_process(disp_handle, &app_entry.exec, &args_vec, &env_vec)?
+        // Terminate existing process if required
+        if let Some(process_handle) =
+            wayvr.process_query(disp_handle, &app_entry.exec, &args_vec, &env_vec)
+        {
+            // Terminate process
+            wayvr.terminate_process(process_handle);
+        } else {
+            // Spawn process
+            wayvr.spawn_process(disp_handle, &app_entry.exec, &args_vec, &env_vec)?;
+        }
     }
 
     Ok(created_overlay)