Try to make UFFD handlers more robust

src/firecracker/examples/uffd/uffd_utils.rs

@@ -52,21 +52,45 @@ pub struct UffdHandler {
 }
 
 impl UffdHandler {
-    pub fn from_unix_stream(stream: &UnixStream, backing_buffer: *const u8, size: usize) -> Self {
-        let mut message_buf = vec![0u8; 1024];
-        let (bytes_read, file) = stream
-            .recv_with_fd(&mut message_buf[..])
-            .expect("Cannot read from a stream");
-        message_buf.resize(bytes_read, 0);
+    fn get_mappings_and_file(stream: &UnixStream) -> (String, File) {
+        // Sometimes, reading from the stream succeeds but we don't receive any
+        // UFFD descriptor. We don't really have a good understanding why this is
+        // happening, but let's try to be a bit more robust and retry a few times
+        // before we declare defeat.
+        for _ in 1..=5 {
+            let mut message_buf = vec![0u8; 1024];
+            let (bytes_read, file) = match stream.recv_with_fd(&mut message_buf[..]) {
+                Ok(res) => res,
+                Err(err) => {
+                    println!("Could not receive message from stream: {err}");
+                    continue;
+                }
+            };
+            message_buf.resize(bytes_read, 0);
+
+            // We do not expect to receive non-UTF-8 data from Firecracker, so this is probably
+            // an error we can't recover from. Just immediately abort
+            let body = String::from_utf8(message_buf.clone()).unwrap_or_else(|_| {
+                panic!(
+                    "Received body is not a utf-8 valid string. Raw bytes received: \
+                     {message_buf:#?}"
+                )
+            });
+            let file = match file {
+                Some(file) => file,
+                None => {
+                    println!("Did not receive Uffd from UDS. Received body: {body}");
+                    continue;
+                }
+            };
+            return (body, file);
+        }
 
-        let body = String::from_utf8(message_buf.clone()).unwrap_or_else(|_| {
-            panic!(
-                "Received body is not a utf-8 valid string. Raw bytes received: {message_buf:#?}"
-            )
-        });
-        let file =
-            file.unwrap_or_else(|| panic!("Did not receive Uffd from UDS. Received body: {body}"));
+        panic!("Could not get UFFD and mappings after 5 retries");
+    }
 
+    pub fn from_unix_stream(stream: &UnixStream, backing_buffer: *const u8, size: usize) -> Self {
+        let (body, file) = Self::get_mappings_and_file(stream);
         let mappings =
             serde_json::from_str::<Vec<GuestRegionUffdMapping>>(&body).unwrap_or_else(|_| {
                 panic!("Cannot deserialize memory mappings. Received body: {body}")