Sync with upstream(6d0f7f6) am: a30504f39c am: 43039a076c am: 68e77424ac

Original change: https://android-review.googlesource.com/c/platform/external/rust/crates/vhost-device-vsock/+/2812533

Change-Id: Ia8afd6a517eb1bfca584d67cd47f25b58a43b908
Signed-off-by: Automerger Merge Worker <android-build-automerger-merge-worker@system.gserviceaccount.com>

3 files changed, 31 insertions, 24 deletions

diff --git a/src/main.rs b/src/main.rs
index 3f785a2..4325e64 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -234,13 +234,13 @@ pub(crate) fn start_backend_server(
         )
         .map_err(BackendError::CouldNotCreateDaemon)?;
 
-        let mut vring_workers = daemon.get_epoll_handlers();
+        let mut epoll_handlers = daemon.get_epoll_handlers();
 
         for thread in backend.threads.iter() {
             thread
                 .lock()
                 .unwrap()
-                .set_vring_worker(Some(vring_workers.remove(0)));
+                .register_listeners(epoll_handlers.remove(0));
         }
 
         daemon.start(listener).unwrap();
diff --git a/src/vhu_vsock_thread.rs b/src/vhu_vsock_thread.rs
index fcefc4a..850ad0c 100644
--- a/src/vhu_vsock_thread.rs
+++ b/src/vhu_vsock_thread.rs
@@ -63,8 +63,6 @@ pub(crate) struct VhostUserVsockThread {
     host_sock_path: String,
     /// Listener listening for new connections on the host.
     host_listener: UnixListener,
-    /// Instance of VringWorker.
-    vring_worker: Option<Arc<VringEpollHandler<ArcVhostBknd, VringRwLock, ()>>>,
     /// epoll fd to which new host connections are added.
     epoll_file: File,
     /// VsockThreadBackend instance.
@@ -151,7 +149,6 @@ impl VhostUserVsockThread {
             host_sock: host_sock.as_raw_fd(),
             host_sock_path: uds_path,
             host_listener: host_sock,
-            vring_worker: None,
             epoll_file,
             thread_backend,
             guest_cid,
@@ -242,20 +239,15 @@ impl VhostUserVsockThread {
         self.epoll_file.as_raw_fd()
     }
 
-    /// Set self's VringWorker.
-    pub fn set_vring_worker(
+    /// Register our listeners in the VringEpollHandler
+    pub fn register_listeners(
         &mut self,
-        vring_worker: Option<Arc<VringEpollHandler<ArcVhostBknd, VringRwLock, ()>>>,
+        epoll_handler: Arc<VringEpollHandler<ArcVhostBknd, VringRwLock, ()>>,
     ) {
-        self.vring_worker = vring_worker;
-        self.vring_worker
-            .as_ref()
-            .unwrap()
+        epoll_handler
             .register_listener(self.get_epoll_fd(), EventSet::IN, u64::from(BACKEND_EVENT))
             .unwrap();
-        self.vring_worker
-            .as_ref()
-            .unwrap()
+        epoll_handler
             .register_listener(
                 self.sibling_event_fd.as_raw_fd(),
                 EventSet::IN,
diff --git a/src/vsock_conn.rs b/src/vsock_conn.rs
index 058c2e1..0a766df 100644
--- a/src/vsock_conn.rs
+++ b/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.
@@ -253,12 +264,14 @@ impl<S: AsRawFd + Read + Write> VsockConnection<S> {
                 )
                 .is_err()
                 {
-                    VhostUserVsockThread::epoll_register(
+                    if let Err(e) = VhostUserVsockThread::epoll_register(
                         self.epoll_fd,
                         self.stream.as_raw_fd(),
                         epoll::Events::EPOLLIN | epoll::Events::EPOLLOUT,
-                    )
-                    .unwrap();
+                    ) {
+                        // TODO: let's move this logic out of this func, and handle it properly
+                        error!("epoll_register failed: {:?}, but proceed further.", e);
+                    }
                 };
             }
             VSOCK_OP_CREDIT_REQUEST => {
@@ -680,13 +693,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());