avoid potential panics on shutdown

Fixes amqp-rs#358 Signed-off-by: Marc-Antoine Perennou <[email protected]>
romatthe · Jan 30, 2022 · 6d6c27b · 6d6c27b
1 parent 636a903
commit 6d6c27b
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 8 deletions.
diff --git a/async-lapin/src/lib.rs b/async-lapin/src/lib.rs
@@ -135,11 +135,13 @@ async fn heartbeat(heartbeat: Heartbeat) {
 }
 
 async fn poll_read(socket: Arc<Async<TcpStreamWrapper>>, socket_state: SocketStateHandle) {
-    socket.readable().await.unwrap();
-    socket_state.send(SocketEvent::Readable);
+    if socket.readable().await.is_ok() {
+        socket_state.send(SocketEvent::Readable);
+    }
 }
 
 async fn poll_write(socket: Arc<Async<TcpStreamWrapper>>, socket_state: SocketStateHandle) {
-    socket.writable().await.unwrap();
-    socket_state.send(SocketEvent::Writable);
+    if socket.writable().await.is_ok() {
+        socket_state.send(SocketEvent::Writable);
+    }
 }
diff --git a/tokio/src/lib.rs b/tokio/src/lib.rs
@@ -169,12 +169,16 @@ mod unix {
     }
 
     async fn poll_read(socket: Arc<AsyncFd<TcpStreamWrapper>>, socket_state: SocketStateHandle) {
-        socket.readable().await.unwrap().clear_ready();
-        socket_state.send(SocketEvent::Readable);
+        if let Ok(mut events) = socket.readable().await {
+            events.clear_ready();
+            socket_state.send(SocketEvent::Readable);
+        }
     }
 
     async fn poll_write(socket: Arc<AsyncFd<TcpStreamWrapper>>, socket_state: SocketStateHandle) {
-        socket.writable().await.unwrap().clear_ready();
-        socket_state.send(SocketEvent::Writable);
+        if let Ok(mut events) = socket.writable().await {
+            events.clear_ready();
+            socket_state.send(SocketEvent::Writable);
+        }
     }
 }