io_uring flush fixes

Author: dobrac

src/vmm/src/device_manager/pci_mngr.rs

@@ -311,7 +311,7 @@ impl<'a> Persist<'a> for PciDevices {
                              snapshotting yet"
                         );
                     } else {
-                        block_dev.prepare_save();
+                        block_dev.prepare_save()?;
                         let device_state = block_dev.save();
                         state.block_devices.push(VirtioDeviceState {
                             device_id: block_dev.id().to_string(),

src/vmm/src/device_manager/persist.rs

@@ -265,7 +265,7 @@ impl<'a> Persist<'a> for MMIODeviceManager {
                              snapshotting yet"
                         );
                     } else {
-                        block.prepare_save();
+                        block.prepare_save()?;
                         let device_state = block.save();
                         states.block_devices.push(VirtioDeviceState {
                             device_id,

src/vmm/src/devices/virtio/block/device.rs

@@ -82,10 +82,10 @@ impl Block {
         }
     }
 
-    pub fn prepare_save(&mut self) {
+    pub fn prepare_save(&mut self) -> Result<(), BlockError> {
         match self {
-            Self::Virtio(b) => b.prepare_save(),
-            Self::VhostUser(b) => b.prepare_save(),
+            Self::Virtio(b) => b.prepare_save().map_err(BlockError::VirtioBackend),
+            Self::VhostUser(_) => Err(BlockError::InvalidBlockBackend),
         }
     }
 

src/vmm/src/devices/virtio/block/virtio/device.rs

@@ -563,22 +563,32 @@ impl VirtioBlock {
         }
     }
 
-    fn drain_and_flush(&mut self, discard: bool) {
-        if let Err(err) = self.disk.file_engine.drain_and_flush(discard) {
-            error!("Failed to drain ops and flush block data: {:?}", err);
-        }
+    fn drain_and_flush(&mut self, discard: bool) -> Result<(), VirtioBlockError> {
+        self.disk
+            .file_engine
+            .drain_and_flush(discard)
+            .map_err(VirtioBlockError::FileEngine)
     }
 
     /// Prepare device for being snapshotted.
-    pub fn prepare_save(&mut self) {
+    pub fn prepare_save(&mut self) -> Result<(), VirtioBlockError> {
         if !self.is_activated() {
-            return;
+            return Ok(());
         }
 
-        self.drain_and_flush(false);
-        if let FileEngine::Async(ref _engine) = self.disk.file_engine {
+        self.drain_and_flush(false)?;
+        let is_async = matches!(self.disk.file_engine, FileEngine::Async(_));
+        if is_async {
             self.process_async_completion_queue();
+            if let FileEngine::Async(ref engine) = self.disk.file_engine {
+                let pending = engine.pending_ops();
+                if pending != 0 {
+                    return Err(VirtioBlockError::PendingAsyncOperations(pending));
+                }
+            }
         }
+
+        Ok(())
     }
 }
 
@@ -682,7 +692,9 @@ impl Drop for VirtioBlock {
                 }
             }
             CacheType::Writeback => {
-                self.drain_and_flush(true);
+                if let Err(err) = self.drain_and_flush(true) {
+                    error!("Failed to drain ops and flush block data: {:?}", err);
+                }
             }
         };
     }
@@ -1673,7 +1685,7 @@ mod tests {
             // Add a batch of flush requests.
             add_flush_requests_batch(&mut block, &vq, 5);
             simulate_queue_event(&mut block, None);
-            block.prepare_save();
+            block.prepare_save().unwrap();
 
             // Check that all the pending flush requests were processed during `prepare_save()`.
             check_flush_requests_batch(5, &vq);

src/vmm/src/devices/virtio/block/virtio/io/async_io.rs

@@ -119,6 +119,10 @@ impl AsyncFileEngine {
         &self.completion_evt
     }
 
+    pub fn pending_ops(&self) -> u32 {
+        self.ring.num_ops()
+    }
+
     pub fn push_read(
         &mut self,
         offset: u64,

src/vmm/src/devices/virtio/block/virtio/io/mod.rs

@@ -172,6 +172,13 @@ impl FileEngine {
             FileEngine::Sync(engine) => engine.flush().map_err(BlockIoError::Sync),
         }
     }
+
+    pub fn pending_async_ops(&self) -> u32 {
+        match self {
+            FileEngine::Async(engine) => engine.pending_ops(),
+            FileEngine::Sync(_) => 0,
+        }
+    }
 }
 
 #[cfg(test)]

src/vmm/src/devices/virtio/block/virtio/mod.rs

@@ -63,4 +63,6 @@ pub enum VirtioBlockError {
     RateLimiter(std::io::Error),
     /// Persistence error: {0}
     Persist(crate::devices::virtio::persist::PersistError),
+    /// Async file engine still has {0} pending operations.
+    PendingAsyncOperations(u32),
 }