Gpu socket final2

.buildkite/rust-vmm-ci-tests.json

@@ -22,63 +22,63 @@
     },
     {
       "test_name": "unittests-gnu-all-with-xen",
-      "command": "cargo test --workspace --no-default-features --features test-utils,vhost-vsock,vhost-kern,vhost-vdpa,vhost-net,vhost-user,vhost-user-frontend,vhost-user-backend,xen",
+      "command": "cargo test --workspace --no-default-features --features test-utils,vhost-vsock,vhost-kern,vhost-vdpa,vhost-net,vhost-user,vhost-user-frontend,gpu-socket,vhost-user-backend,xen",
       "platform": [
         "x86_64",
         "aarch64"
       ]
     },
     {
       "test_name": "unittests-gnu-all-without-xen",
-      "command": "cargo test --workspace --no-default-features --features test-utils,vhost-vsock,vhost-kern,vhost-vdpa,vhost-net,vhost-user,vhost-user-frontend,vhost-user-backend,postcopy",
+      "command": "cargo test --workspace --no-default-features --features test-utils,vhost-vsock,vhost-kern,vhost-vdpa,vhost-net,vhost-user,vhost-user-frontend,gpu-socket,vhost-user-backend,postcopy",
       "platform": [
         "x86_64",
         "aarch64"
       ]
     },
     {
       "test_name": "unittests-musl-all-with-xen",
-      "command": "cargo test --workspace --target {target_platform}-unknown-linux-musl --no-default-features --features test-utils,vhost-vsock,vhost-kern,vhost-vdpa,vhost-net,vhost-user,vhost-user-frontend,vhost-user-backend,xen",
+      "command": "cargo test --workspace --target {target_platform}-unknown-linux-musl --no-default-features --features test-utils,vhost-vsock,vhost-kern,vhost-vdpa,vhost-net,vhost-user,vhost-user-frontend,gpu-socket,vhost-user-backend,xen",
       "platform": [
         "x86_64",
         "aarch64"
       ]
     },
     {
       "test_name": "unittests-musl-all-without-xen",
-      "command": "cargo test --workspace --target {target_platform}-unknown-linux-musl --no-default-features --features test-utils,vhost-vsock,vhost-kern,vhost-vdpa,vhost-net,vhost-user,vhost-user-frontend,vhost-user-backend,postcopy",
+      "command": "cargo test --workspace --target {target_platform}-unknown-linux-musl --no-default-features --features test-utils,vhost-vsock,vhost-kern,vhost-vdpa,vhost-net,vhost-user,vhost-user-frontend,gpu-socket,vhost-user-backend,postcopy",
       "platform": [
         "x86_64",
         "aarch64"
       ]
     },
     {
       "test_name": "clippy-all-with-xen",
-      "command": "cargo clippy --workspace --bins --examples --benches --all-targets --no-default-features --features test-utils,vhost-vsock,vhost-kern,vhost-vdpa,vhost-net,vhost-user,vhost-user-frontend,vhost-user-backend,xen -- -D warnings -D clippy::undocumented_unsafe_blocks",
+      "command": "cargo clippy --workspace --bins --examples --benches --all-targets --no-default-features --features test-utils,vhost-vsock,vhost-kern,vhost-vdpa,vhost-net,vhost-user,vhost-user-frontend,gpu-socket,vhost-user-backend,xen -- -D warnings -D clippy::undocumented_unsafe_blocks",
       "platform": [
         "x86_64",
         "aarch64"
       ]
     },
     {
       "test_name": "clippy-all-without-xen",
-      "command": "cargo clippy --workspace --bins --examples --benches --all-targets --no-default-features --features test-utils,vhost-vsock,vhost-kern,vhost-vdpa,vhost-net,vhost-user,vhost-user-frontend,vhost-user-backend,postcopy -- -D warnings -D clippy::undocumented_unsafe_blocks",
+      "command": "cargo clippy --workspace --bins --examples --benches --all-targets --no-default-features --features test-utils,vhost-vsock,vhost-kern,vhost-vdpa,vhost-net,vhost-user,vhost-user-frontend,gpu-socket,vhost-user-backend,postcopy -- -D warnings -D clippy::undocumented_unsafe_blocks",
       "platform": [
         "x86_64",
         "aarch64"
       ]
     },
     {
       "test_name": "check-warnings-all-with-xen",
-      "command": "RUSTFLAGS=\"-D warnings\" cargo check --all-targets --workspace --no-default-features --features test-utils,vhost-vsock,vhost-kern,vhost-vdpa,vhost-net,vhost-user,vhost-user-frontend,vhost-user-backend,xen",
+      "command": "RUSTFLAGS=\"-D warnings\" cargo check --all-targets --workspace --no-default-features --features test-utils,vhost-vsock,vhost-kern,vhost-vdpa,vhost-net,vhost-user,vhost-user-frontend,gpu-socket,vhost-user-backend,xen",
       "platform": [
         "x86_64",
         "aarch64"
       ]
     },
     {
       "test_name": "check-warnings-all-without-xen",
-      "command": "RUSTFLAGS=\"-D warnings\" cargo check --all-targets --workspace --no-default-features --features test-utils,vhost-vsock,vhost-kern,vhost-vdpa,vhost-net,vhost-user,vhost-user-frontend,vhost-user-backend,postcopy",
+      "command": "RUSTFLAGS=\"-D warnings\" cargo check --all-targets --workspace --no-default-features --features test-utils,vhost-vsock,vhost-kern,vhost-vdpa,vhost-net,vhost-user,vhost-user-frontend,gpu-socket,vhost-user-backend,postcopy",
       "platform": [
         "x86_64",
         "aarch64"

vhost-user-backend/CHANGELOG.md

@@ -2,7 +2,7 @@
 ## [Unreleased]
 
 ### Added
-
+- [[#239]](https://github.com/rust-vmm/vhost/pull/239) Add support for `VHOST_USER_GPU_SET_SOCKET`
 ### Changed
 
 ### Fixed

vhost-user-backend/Cargo.toml

@@ -11,6 +11,7 @@ license = "Apache-2.0"
 [features]
 xen = ["vm-memory/xen", "vhost/xen"]
 postcopy = ["vhost/postcopy", "userfaultfd"]
+gpu-socket = ["vhost/gpu-socket"]
 
 [dependencies]
 libc = "0.2.39"

vhost-user-backend/src/backend.rs

@@ -31,6 +31,9 @@ use vm_memory::bitmap::Bitmap;
 use vmm_sys_util::epoll::EventSet;
 use vmm_sys_util::eventfd::EventFd;
 
+#[cfg(feature = "gpu-socket")]
+use vhost::vhost_user::GpuBackend;
+
 use super::vring::VringT;
 use super::GM;
 
@@ -84,6 +87,9 @@ pub trait VhostUserBackend: Send + Sync {
     /// function.
     fn set_backend_req_fd(&self, _backend: Backend) {}
 
+    #[cfg(feature = "gpu-socket")]
+    fn set_gpu_socket(&self, _gpu_backend: GpuBackend) {}
+
     /// Get the map to map queue index to worker thread index.
     ///
     /// A return value of [2, 2, 4] means: the first two queues will be handled by worker thread 0,
@@ -194,6 +200,9 @@ pub trait VhostUserBackendMut: Send + Sync {
     /// function.
     fn set_backend_req_fd(&mut self, _backend: Backend) {}
 
+    #[cfg(feature = "gpu-socket")]
+    fn set_gpu_socket(&mut self, _gpu_backend: GpuBackend) {}
+
     /// Get the map to map queue index to worker thread index.
     ///
     /// A return value of [2, 2, 4] means: the first two queues will be handled by worker thread 0,
@@ -299,6 +308,11 @@ impl<T: VhostUserBackend> VhostUserBackend for Arc<T> {
         self.deref().set_backend_req_fd(backend)
     }
 
+    #[cfg(feature = "gpu-socket")]
+    fn set_gpu_socket(&self, gpu_backend: GpuBackend) {
+        self.deref().set_gpu_socket(gpu_backend)
+    }
+
     fn queues_per_thread(&self) -> Vec<u64> {
         self.deref().queues_per_thread()
     }
@@ -376,6 +390,11 @@ impl<T: VhostUserBackendMut> VhostUserBackend for Mutex<T> {
         self.lock().unwrap().set_backend_req_fd(backend)
     }
 
+    #[cfg(feature = "gpu-socket")]
+    fn set_gpu_socket(&self, gpu_backend: GpuBackend) {
+        self.lock().unwrap().set_gpu_socket(gpu_backend)
+    }
+
     fn queues_per_thread(&self) -> Vec<u64> {
         self.lock().unwrap().queues_per_thread()
     }
@@ -456,6 +475,11 @@ impl<T: VhostUserBackendMut> VhostUserBackend for RwLock<T> {
         self.write().unwrap().set_backend_req_fd(backend)
     }
 
+    #[cfg(feature = "gpu-socket")]
+    fn set_gpu_socket(&self, gpu_backend: GpuBackend) {
+        self.write().unwrap().set_gpu_socket(gpu_backend)
+    }
+
     fn queues_per_thread(&self) -> Vec<u64> {
         self.read().unwrap().queues_per_thread()
     }
@@ -576,6 +600,9 @@ pub mod tests {
 
         fn set_backend_req_fd(&mut self, _backend: Backend) {}
 
+        #[cfg(feature = "gpu-socket")]
+        fn set_gpu_socket(&mut self, _gpu_backend: GpuBackend) {}
+
         fn queues_per_thread(&self) -> Vec<u64> {
             vec![1, 1]
         }

vhost-user-backend/src/handler.rs

@@ -21,9 +21,12 @@ use vhost::vhost_user::message::{
     VhostUserMemoryRegion, VhostUserProtocolFeatures, VhostUserSingleMemoryRegion,
     VhostUserVirtioFeatures, VhostUserVringAddrFlags, VhostUserVringState,
 };
+#[cfg(feature = "gpu-socket")]
+use vhost::vhost_user::GpuBackend;
 use vhost::vhost_user::{
     Backend, Error as VhostUserError, Result as VhostUserResult, VhostUserBackendReqHandlerMut,
 };
+
 use virtio_bindings::bindings::virtio_ring::VIRTIO_RING_F_EVENT_IDX;
 use virtio_queue::{Error as VirtQueError, QueueT};
 use vm_memory::mmap::NewBitmap;
@@ -549,6 +552,11 @@ where
         self.backend.set_backend_req_fd(backend);
     }
 
+    #[cfg(feature = "gpu-socket")]
+    fn set_gpu_socket(&mut self, gpu_backend: GpuBackend) {
+        self.backend.set_gpu_socket(gpu_backend);
+    }
+
     fn get_inflight_fd(
         &mut self,
         _inflight: &vhost::vhost_user::message::VhostUserInflight,

vhost/CHANGELOG.md

@@ -2,7 +2,7 @@
 ## [Unreleased]
 
 ### Added
-
+- [[#239]](https://github.com/rust-vmm/vhost/pull/239) Add support for `VHOST_USER_GPU_SET_SOCKET`
 ### Changed
 
 ### Fixed

vhost/Cargo.toml

@@ -23,6 +23,7 @@ vhost-net = ["vhost-kern"]
 vhost-user = []
 vhost-user-frontend = ["vhost-user"]
 vhost-user-backend = ["vhost-user"]
+gpu-socket = ["vhost-user", "vhost-user-backend"]
 xen = ["vm-memory/xen"]
 postcopy = []
 

vhost/src/vhost_user/backend.rs

@@ -32,7 +32,7 @@ impl<S: VhostUserBackendReqHandler> BackendListener<S> {
     pub fn accept(&mut self) -> Result<Option<BackendReqHandler<S>>> {
         if let Some(fd) = self.listener.accept()? {
             return Ok(Some(BackendReqHandler::new(
-                Endpoint::<FrontendReq>::from_stream(fd),
+                Endpoint::<VhostUserMsgHeader<FrontendReq>>::from_stream(fd),
                 self.backend.take().unwrap(),
             )));
         }

vhost/src/vhost_user/backend_req.rs

@@ -14,7 +14,7 @@ use super::{Error, HandlerResult, Result, VhostUserFrontendReqHandler};
 use vm_memory::ByteValued;
 
 struct BackendInternal {
-    sock: Endpoint<BackendReq>,
+    sock: Endpoint<VhostUserMsgHeader<BackendReq>>,
 
     // Protocol feature VHOST_USER_PROTOCOL_F_REPLY_ACK has been negotiated.
     reply_ack_negotiated: bool,
@@ -83,7 +83,7 @@ pub struct Backend {
 }
 
 impl Backend {
-    fn new(ep: Endpoint<BackendReq>) -> Self {
+    fn new(ep: Endpoint<VhostUserMsgHeader<BackendReq>>) -> Self {
         Backend {
             node: Arc::new(Mutex::new(BackendInternal {
                 sock: ep,
@@ -110,7 +110,9 @@ impl Backend {
 
     /// Create a new instance from a `UnixStream` object.
     pub fn from_stream(sock: UnixStream) -> Self {
-        Self::new(Endpoint::<BackendReq>::from_stream(sock))
+        Self::new(Endpoint::<VhostUserMsgHeader<BackendReq>>::from_stream(
+            sock,
+        ))
     }
 
     /// Set the negotiation state of the `VHOST_USER_PROTOCOL_F_REPLY_ACK` protocol feature.
@@ -175,7 +177,7 @@ mod tests {
     fn test_backend_req_recv_negative() {
         let (p1, p2) = UnixStream::pair().unwrap();
         let backend = Backend::from_stream(p1);
-        let mut frontend = Endpoint::<BackendReq>::from_stream(p2);
+        let mut frontend = Endpoint::<VhostUserMsgHeader<BackendReq>>::from_stream(p2);
 
         let len = mem::size_of::<VhostUserFSBackendMsg>();
         let mut hdr = VhostUserMsgHeader::new(

vhost/src/vhost_user/backend_req_handler.rs

@@ -12,6 +12,8 @@ use vm_memory::ByteValued;
 
 use super::backend_req::Backend;
 use super::connection::Endpoint;
+#[cfg(feature = "gpu-socket")]
+use super::gpu_backend_req::GpuBackend;
 use super::message::*;
 use super::{take_single_file, Error, Result};
 
@@ -65,6 +67,8 @@ pub trait VhostUserBackendReqHandler {
     fn get_config(&self, offset: u32, size: u32, flags: VhostUserConfigFlags) -> Result<Vec<u8>>;
     fn set_config(&self, offset: u32, buf: &[u8], flags: VhostUserConfigFlags) -> Result<()>;
     fn set_backend_req_fd(&self, _backend: Backend) {}
+    #[cfg(feature = "gpu-socket")]
+    fn set_gpu_socket(&self, _gpu_backend: GpuBackend) {}
     fn get_inflight_fd(&self, inflight: &VhostUserInflight) -> Result<(VhostUserInflight, File)>;
     fn set_inflight_fd(&self, inflight: &VhostUserInflight, file: File) -> Result<()>;
     fn get_max_mem_slots(&self) -> Result<u64>;
@@ -124,6 +128,8 @@ pub trait VhostUserBackendReqHandlerMut {
     ) -> Result<Vec<u8>>;
     fn set_config(&mut self, offset: u32, buf: &[u8], flags: VhostUserConfigFlags) -> Result<()>;
     fn set_backend_req_fd(&mut self, _backend: Backend) {}
+    #[cfg(feature = "gpu-socket")]
+    fn set_gpu_socket(&mut self, _gpu_backend: GpuBackend) {}
     fn get_inflight_fd(
         &mut self,
         inflight: &VhostUserInflight,
@@ -235,6 +241,11 @@ impl<T: VhostUserBackendReqHandlerMut> VhostUserBackendReqHandler for Mutex<T> {
         self.lock().unwrap().set_backend_req_fd(backend)
     }
 
+    #[cfg(feature = "gpu-socket")]
+    fn set_gpu_socket(&self, gpu_backend: GpuBackend) {
+        self.lock().unwrap().set_gpu_socket(gpu_backend);
+    }
+
     fn get_inflight_fd(&self, inflight: &VhostUserInflight) -> Result<(VhostUserInflight, File)> {
         self.lock().unwrap().get_inflight_fd(inflight)
     }
@@ -302,7 +313,7 @@ impl<T: VhostUserBackendReqHandlerMut> VhostUserBackendReqHandler for Mutex<T> {
 /// [BackendReqHandler]: struct.BackendReqHandler.html
 pub struct BackendReqHandler<S: VhostUserBackendReqHandler> {
     // underlying Unix domain socket for communication
-    main_sock: Endpoint<FrontendReq>,
+    main_sock: Endpoint<VhostUserMsgHeader<FrontendReq>>,
     // the vhost-user backend device object
     backend: Arc<S>,
 
@@ -319,7 +330,10 @@ pub struct BackendReqHandler<S: VhostUserBackendReqHandler> {
 
 impl<S: VhostUserBackendReqHandler> BackendReqHandler<S> {
     /// Create a vhost-user backend endpoint.
-    pub(super) fn new(main_sock: Endpoint<FrontendReq>, backend: Arc<S>) -> Self {
+    pub(super) fn new(
+        main_sock: Endpoint<VhostUserMsgHeader<FrontendReq>>,
+        backend: Arc<S>,
+    ) -> Self {
         BackendReqHandler {
             main_sock,
             backend,
@@ -359,7 +373,10 @@ impl<S: VhostUserBackendReqHandler> BackendReqHandler<S> {
     /// * - `path` - path of Unix domain socket listener to connect to
     /// * - `backend` - handler for requests from the frontend to the backend
     pub fn connect(path: &str, backend: Arc<S>) -> Result<Self> {
-        Ok(Self::new(Endpoint::<FrontendReq>::connect(path)?, backend))
+        Ok(Self::new(
+            Endpoint::<VhostUserMsgHeader<FrontendReq>>::connect(path)?,
+            backend,
+        ))
     }
 
     /// Mark endpoint as failed with specified error code.
@@ -554,6 +571,11 @@ impl<S: VhostUserBackendReqHandler> BackendReqHandler<S> {
                 let res = self.backend.set_inflight_fd(&msg, file);
                 self.send_ack_message(&hdr, res)?;
             }
+            #[cfg(feature = "gpu-socket")]
+            Ok(FrontendReq::GPU_SET_SOCKET) => {
+                let res = self.set_gpu_socket(files);
+                self.send_ack_message(&hdr, res)?;
+            }
             Ok(FrontendReq::GET_MAX_MEM_SLOTS) => {
                 self.check_proto_feature(VhostUserProtocolFeatures::CONFIGURE_MEM_SLOTS)?;
                 self.check_request_size(&hdr, size, 0)?;
@@ -790,6 +812,18 @@ impl<S: VhostUserBackendReqHandler> BackendReqHandler<S> {
         Ok(())
     }
 
+    #[cfg(feature = "gpu-socket")]
+    fn set_gpu_socket(&mut self, files: Option<Vec<File>>) -> Result<()> {
+        let file = take_single_file(files).ok_or(Error::InvalidMessage)?;
+        // SAFETY: Safe because we have ownership of the files that were
+        // checked when received. We have to trust that they are Unix sockets
+        // since we have no way to check this. If not, it will fail later.
+        let sock = unsafe { UnixStream::from_raw_fd(file.into_raw_fd()) };
+        let gpu_backend = GpuBackend::from_stream(sock);
+        self.backend.set_gpu_socket(gpu_backend);
+        Ok(())
+    }
+
     fn handle_vring_fd_request(
         &mut self,
         buf: &[u8],
@@ -859,7 +893,8 @@ impl<S: VhostUserBackendReqHandler> BackendReqHandler<S> {
                 | FrontendReq::SET_BACKEND_REQ_FD
                 | FrontendReq::SET_INFLIGHT_FD
                 | FrontendReq::ADD_MEM_REG
-                | FrontendReq::SET_DEVICE_STATE_FD,
+                | FrontendReq::SET_DEVICE_STATE_FD
+                | FrontendReq::GPU_SET_SOCKET,
             ) => Ok(()),
             _ if files.is_some() => Err(Error::InvalidMessage),
             _ => Ok(()),
@@ -965,7 +1000,7 @@ mod tests {
     #[test]
     fn test_backend_req_handler_new() {
         let (p1, _p2) = UnixStream::pair().unwrap();
-        let endpoint = Endpoint::<FrontendReq>::from_stream(p1);
+        let endpoint = Endpoint::<VhostUserMsgHeader<FrontendReq>>::from_stream(p1);
         let backend = Arc::new(Mutex::new(DummyBackendReqHandler::new()));
         let mut handler = BackendReqHandler::new(endpoint, backend);