fix: ipc error-handling

Signed-off-by: usamoi <usamoi@outlook.com>
tensorchord · Dec 19, 2023 · 6a64c2f · 6a64c2f
1 parent a1c530c
commit 6a64c2f
Show file tree

Hide file tree

Showing 3 changed files with 29 additions and 12 deletions.
diff --git a/src/ipc/client/mod.rs b/src/ipc/client/mod.rs
@@ -217,10 +217,9 @@ impl<T: ClientLike> Drop for ClientGuard<T> {
         let mut x = CLIENT.borrow_mut();
         match *x {
             Status::Borrowed => {
-                if T::RESET {
-                    unsafe {
-                        *x = Status::Reset(ManuallyDrop::take(&mut self.0).to_socket());
-                    }
+                let socket = unsafe { ManuallyDrop::take(&mut self.0).to_socket() };
+                if T::RESET && socket.test() {
+                    *x = Status::Reset(socket);
                 } else {
                     *x = Status::Lost;
                 }

diff --git a/src/ipc/mod.rs b/src/ipc/mod.rs
@@ -34,11 +34,17 @@ pub fn listen_mmap() -> impl Iterator<Item = RpcHandler> {
 }
 
 pub fn connect_unix() -> self::transport::ClientSocket {
-    self::transport::ClientSocket::Unix(self::transport::unix::connect())
+    self::transport::ClientSocket::Unix {
+        ok: true,
+        socket: self::transport::unix::connect(),
+    }
 }
 
 pub fn connect_mmap() -> self::transport::ClientSocket {
-    self::transport::ClientSocket::Mmap(self::transport::mmap::connect())
+    self::transport::ClientSocket::Mmap {
+        ok: true,
+        socket: self::transport::mmap::connect(),
+    }
 }
 
 pub fn init() {

diff --git a/src/ipc/transport/mod.rs b/src/ipc/transport/mod.rs
@@ -11,8 +11,8 @@ pub enum ServerSocket {
 }
 
 pub enum ClientSocket {
-    Unix(unix::Socket),
-    Mmap(mmap::Socket),
+    Unix { ok: bool, socket: unix::Socket },
+    Mmap { ok: bool, socket: mmap::Socket },
 }
 
 impl ServerSocket {
@@ -48,19 +48,31 @@ impl ClientSocket {
     pub fn send<T: Serialize>(&mut self, packet: T) -> Result<(), IpcError> {
         let buffer = bincode::serialize(&packet).expect("Failed to serialize");
         match self {
-            Self::Unix(x) => x.send(&buffer),
-            Self::Mmap(x) => x.send(&buffer),
+            Self::Unix { ok, socket } => socket.send(&buffer).inspect(|_| *ok = false),
+            Self::Mmap { ok, socket } => socket.send(&buffer).inspect(|_| *ok = false),
         }
     }
     pub fn recv<T: for<'a> Deserialize<'a>>(&mut self) -> Result<T, FriendlyError> {
         let buffer = match self {
-            Self::Unix(x) => x.recv().map_err(|_| FriendlyError::Ipc)?,
-            Self::Mmap(x) => x.recv().map_err(|_| FriendlyError::Ipc)?,
+            Self::Unix { ok, socket } => socket
+                .recv()
+                .inspect(|_| *ok = false)
+                .map_err(|_| FriendlyError::Ipc)?,
+            Self::Mmap { ok, socket } => socket
+                .recv()
+                .inspect(|_| *ok = false)
+                .map_err(|_| FriendlyError::Ipc)?,
         };
         match buffer[0] {
             0u8 => Ok(bincode::deserialize(&buffer[1..]).expect("Failed to deserialize.")),
             1u8 => Err(bincode::deserialize(&buffer[1..]).expect("Failed to deserialize.")),
             _ => unreachable!(),
         }
     }
+    pub fn test(&self) -> bool {
+        match self {
+            ClientSocket::Unix { ok, .. } => *ok,
+            ClientSocket::Mmap { ok, .. } => *ok,
+        }
+    }
 }