From 28c6a3e1e60a18efce1e769ace49e01cc44f0970 Mon Sep 17 00:00:00 2001
From: mars <mars@tebibyte.media>
Date: Sat, 29 Oct 2022 17:07:34 -0600
Subject: [PATCH] Automatically close client's windows on disconnect

---
 apps/magpie/src/ipc.rs | 55 ++++++++++++++++++++++++++++++++----------
 1 file changed, 42 insertions(+), 13 deletions(-)

diff --git a/apps/magpie/src/ipc.rs b/apps/magpie/src/ipc.rs
index 51dbdd1..f464780 100644
--- a/apps/magpie/src/ipc.rs
+++ b/apps/magpie/src/ipc.rs
@@ -1,12 +1,11 @@
-use std::io::Read;
+use std::collections::HashMap;
 use std::ops::{Deref, DerefMut};
 use std::path::{Path, PathBuf};
-use std::str::from_utf8;
 use std::sync::mpsc::{channel, Receiver, Sender};
 use std::sync::Arc;
 use std::time::Duration;
 
-use magpie_types::{CreatePanel, MagpieClientMsg, MagpieServerMsg, ServerMessenger};
+use magpie_types::{CreatePanel, MagpieServerMsg, ServerMessenger};
 use mio::net::{UnixListener, UnixStream};
 use mio::{Events, Interest, Poll, Token, Waker};
 use mio_signals::{Signal, Signals};
@@ -55,6 +54,15 @@ impl DerefMut for Listener {
 
 pub struct IpcData {
     poll: Poll,
+    next_window_id: usize,
+}
+
+impl IpcData {
+    pub fn new_window_id(&mut self) -> usize {
+        let id = self.next_window_id;
+        self.next_window_id = self.next_window_id.wrapping_add(1);
+        id
+    }
 }
 
 pub struct Client {
@@ -62,6 +70,23 @@ pub struct Client {
     window_sender: WindowMessageSender,
     messenger: ServerMessenger<UnixStream>,
     token: Token,
+    id_to_window: HashMap<u32, usize>,
+}
+
+impl Drop for Client {
+    fn drop(&mut self) {
+        println!("Client #{} disconnected", self.token.0);
+        let data = self.data.write();
+        let _ = data
+            .poll
+            .registry()
+            .deregister(&mut self.messenger.transport);
+
+        for (_id, window) in self.id_to_window.drain() {
+            let msg = WindowMessage::CloseWindow { id: window };
+            let _ = self.window_sender.send_event(msg);
+        }
+    }
 }
 
 impl Client {
@@ -72,8 +97,14 @@ impl Client {
             println!("Client #{}: {:?}", self.token.0, msg);
             match msg {
                 MagpieServerMsg::CreatePanel(CreatePanel { id }) => {
-                    let id = id as usize;
-                    let msg = WindowMessage::OpenWindow { id };
+                    let window = self.data.write().new_window_id();
+
+                    if let Some(old_id) = self.id_to_window.insert(id, window) {
+                        let msg = WindowMessage::CloseWindow { id: old_id };
+                        let _ = self.window_sender.send_event(msg);
+                    }
+
+                    let msg = WindowMessage::OpenWindow { id: window };
                     let _ = self.window_sender.send_event(msg);
                 }
                 _ => unimplemented!(),
@@ -130,7 +161,10 @@ impl Ipc {
             sender,
         };
 
-        let data = IpcData { poll };
+        let data = IpcData {
+            poll,
+            next_window_id: 0,
+        };
 
         let ipc = Self {
             data: Arc::new(RwLock::new(data)),
@@ -175,6 +209,7 @@ impl Ipc {
                                 token,
                                 data: self.data.clone(),
                                 window_sender: self.window_sender.clone(),
+                                id_to_window: Default::default(),
                             });
                         }
                         Err(ref err) if err.kind() == std::io::ErrorKind::WouldBlock => break,
@@ -190,13 +225,7 @@ impl Ipc {
             } else if let Some(client) = self.clients.get_mut(event.token().0) {
                 let disconnected = client.on_readable()?;
                 if disconnected {
-                    println!("Client #{} disconnected", event.token().0);
-                    let mut client = self.clients.remove(event.token().0);
-                    self.data
-                        .write()
-                        .poll
-                        .registry()
-                        .deregister(&mut client.messenger.transport)?;
+                    self.clients.remove(event.token().0);
                 }
             } else {
                 panic!("Unrecognized event token: {:?}", event);