From 99ec30759a146c25d85892923d3ea44b42e2b8c1 Mon Sep 17 00:00:00 2001
From: lnyng <lnyng@meta.com>
Date: Sun, 26 Jan 2025 20:49:52 -0800
Subject: [PATCH] Disable EPOLLOUT if triggered while txbuf is empty

Currently the vsock connection always triggers on EPOLLOUT. When a host
applicationlistens to the UDS, it's almost always writable, and thus
keeps triggering backend events with EPOLLOUT on the connection. As a
result, vhost-device-vsock daemon consumes 100% CPU.

To make it more CPU efficient, it can instead unsubscribe EPOLLOUT
whenever there is no more data to send. It can re-subscribe if the
connection cannot write out all bytes. This change causes the CPU util
to drop to negligible level.

Signed-off-by: Leon Yang <lnyng@meta.com>
---
 vhost-device-vsock/CHANGELOG.md            |  1 +
 vhost-device-vsock/src/vhu_vsock_thread.rs | 13 ++++++++++---
 vhost-device-vsock/src/vsock_conn.rs       | 10 ++++++++++
 3 files changed, 21 insertions(+), 3 deletions(-)

diff --git a/vhost-device-vsock/CHANGELOG.md b/vhost-device-vsock/CHANGELOG.md
index 8006ffd..71a33c2 100644
--- a/vhost-device-vsock/CHANGELOG.md
+++ b/vhost-device-vsock/CHANGELOG.md
@@ -9,6 +9,7 @@
 ### Changed
 
 ### Fixed
+- [#800](https://github.com/rust-vmm/vhost-device/pull/800) Disable EPOLLOUT if triggered while txbuf is empty
 
 ### Deprecated
 
diff --git a/vhost-device-vsock/src/vhu_vsock_thread.rs b/vhost-device-vsock/src/vhu_vsock_thread.rs
index 1db9f09..fca2fce 100644
--- a/vhost-device-vsock/src/vhu_vsock_thread.rs
+++ b/vhost-device-vsock/src/vhu_vsock_thread.rs
@@ -18,9 +18,7 @@ use std::{
     thread,
 };
 
-#[cfg(feature = "backend_vsock")]
-use log::error;
-use log::warn;
+use log::{error, warn};
 use vhost_user_backend::{VringEpollHandler, VringRwLock, VringT};
 use virtio_queue::QueueOwnedT;
 use virtio_vsock::packet::{VsockPacket, PKT_HEADER_SIZE};
@@ -431,6 +429,9 @@ impl VhostUserVsockThread {
                 }
             } else {
                 // Previously connected connection
+
+                // Get epoll fd before getting conn as that takes self mut ref
+                let epoll_fd = self.get_epoll_fd();
                 let key = self.thread_backend.listener_map.get(&fd).unwrap();
                 let conn = self.thread_backend.conn_map.get_mut(key).unwrap();
 
@@ -441,6 +442,12 @@ impl VhostUserVsockThread {
                             if cnt > 0 {
                                 conn.fwd_cnt += Wrapping(cnt as u32);
                                 conn.rx_queue.enqueue(RxOps::CreditUpdate);
+                            } else {
+                                // If no remaining data to flush, try to disable EPOLLOUT
+                                if Self::epoll_modify(epoll_fd, fd, epoll::Events::EPOLLIN).is_err()
+                                {
+                                    error!("Failed to disable EPOLLOUT");
+                                }
                             }
                             self.thread_backend
                                 .backend_rxq
diff --git a/vhost-device-vsock/src/vsock_conn.rs b/vhost-device-vsock/src/vsock_conn.rs
index beb3a01..239203f 100644
--- a/vhost-device-vsock/src/vsock_conn.rs
+++ b/vhost-device-vsock/src/vsock_conn.rs
@@ -345,6 +345,16 @@ impl<S: AsRawFd + ReadVolatile + Write + WriteVolatile + IsHybridVsock> VsockCon
         }
 
         if written_count != buf.len() {
+            // Try to re-enable EPOLLOUT in case it is disabled when txbuf is empty.
+            if VhostUserVsockThread::epoll_modify(
+                self.epoll_fd,
+                self.stream.as_raw_fd(),
+                epoll::Events::EPOLLIN | epoll::Events::EPOLLOUT,
+            )
+            .is_err()
+            {
+                error!("Failed to re-enable EPOLLOUT");
+            }
             return self.tx_buf.push(&buf.offset(written_count).unwrap());
         }