From 4b6451176d2350ebd58a440857ddfe0fc828bd92 Mon Sep 17 00:00:00 2001
From: Jeongik Cha <jeongik@google.com>
Date: Fri, 27 Oct 2023 18:49:04 +0900
Subject: [PATCH] vsock: try epoll_modify before epoll_register in recv_pkt

in this context, epoll listener can be already registered via other TX
event, so let it try epoll_modify first to avoid 'silent' failure which
possibly drops packets.

Signed-off-by: Jeongik Cha <jeongik@google.com>
---
 vhost-device-vsock/src/vsock_conn.rs | 25 +++++++++++++++++++------
 1 file changed, 19 insertions(+), 6 deletions(-)

diff --git a/vhost-device-vsock/src/vsock_conn.rs b/vhost-device-vsock/src/vsock_conn.rs
index 058c2e1..1304f1e 100644
--- a/vhost-device-vsock/src/vsock_conn.rs
+++ b/vhost-device-vsock/src/vsock_conn.rs
@@ -6,7 +6,7 @@ use std::{
     os::unix::prelude::{AsRawFd, RawFd},
 };
 
-use log::info;
+use log::{error, info};
 use virtio_vsock::packet::{VsockPacket, PKT_HEADER_SIZE};
 use vm_memory::{bitmap::BitmapSlice, Bytes, VolatileSlice};
 
@@ -174,11 +174,22 @@ impl<S: AsRawFd + Read + Write> VsockConnection<S> {
                         pkt.set_op(VSOCK_OP_RW).set_len(read_cnt as u32);
 
                         // Re-register the stream file descriptor for read and write events
-                        VhostUserVsockThread::epoll_register(
+                        if VhostUserVsockThread::epoll_modify(
                             self.epoll_fd,
                             self.stream.as_raw_fd(),
                             epoll::Events::EPOLLIN | epoll::Events::EPOLLOUT,
-                        )?;
+                        )
+                        .is_err()
+                        {
+                            if let Err(e) = VhostUserVsockThread::epoll_register(
+                                self.epoll_fd,
+                                self.stream.as_raw_fd(),
+                                epoll::Events::EPOLLIN | epoll::Events::EPOLLOUT,
+                            ) {
+                                // TODO: let's move this logic out of this func, and handle it properly
+                                error!("epoll_register failed: {:?}, but proceed further.", e);
+                            }
+                        };
                     }
 
                     // Update the rx_cnt with the amount of data in the vsock packet.
@@ -680,13 +691,15 @@ mod tests {
         );
 
         // VSOCK_OP_RW: finite data read from stream/file
-        conn_local.stream.write_all(b"hello").unwrap();
+        let payload = b"hello";
+        conn_local.stream.write_all(payload).unwrap();
         conn_local.rx_queue.enqueue(RxOps::Rw);
         let op_zero_read = conn_local.recv_pkt(&mut pkt);
-        // below error due to epoll add
-        assert!(op_zero_read.is_err());
+        assert!(op_zero_read.is_ok());
         assert_eq!(pkt.op(), VSOCK_OP_RW);
         assert!(!conn_local.rx_queue.pending_rx());
+        assert_eq!(conn_local.rx_cnt, Wrapping(payload.len() as u32));
+        assert_eq!(conn_local.last_fwd_cnt, Wrapping(1024));
         assert_eq!(pkt.len(), 5);
         let buf = &mut [0u8; 5];
         assert!(pkt.data_slice().unwrap().read_slice(buf, 0).is_ok());