chore: Extract reusable device-related code and add small improvements (#6060)

There was quite a bit of code related to devices in guest_vm_runner that
would be useful in other codepaths so this PR extracts the code out to a
separate crate.

I've also found that when we attach a disk image to a loop device, there
is a race condition since finding a free loop device and attaching the
image is not atomic. This PR adds a fix.

---------

Co-authored-by: IDX GitHub Automation <infra+github-automation@dfinity.org>

Author: frankdavid

Cargo.lock

@@ -144,6 +144,7 @@ members = [
     "rs/ic_os/config_types",
     "rs/ic_os/config_types/compatibility_tests",
     "rs/ic_os/deterministic_ips",
+    "rs/ic_os/device",
     "rs/ic_os/dev_test_tools/launch-single-vm",
     "rs/ic_os/dev_test_tools/setupos-disable-checks",
     "rs/ic_os/dev_test_tools/setupos-image-config",
@@ -551,6 +552,7 @@ curve25519-dalek = { version = "4.1.3", features = [
 ] }
 der = { version = "0.7", default-features = false }
 derive-new = "0.7.0"
+devicemapper = "0.34"
 dfx-core = { version = "0.1.4" }
 ed25519-dalek = { version = "2.1.1", features = [
     "std",
@@ -636,6 +638,7 @@ libc = "0.2.158"
 libflate = "2.1.0"
 libnss = "0.5.0"
 local-ip-address = "0.5.6"
+loopdev-3 = "0.5"
 macaddr = "1.0"
 memmap2 = "0.9.5"
 minicbor = { version = "0.19.1", features = ["alloc", "derive"] }

Cargo.toml

@@ -0,0 +1,46 @@
+load("@rules_rust//rust:defs.bzl", "rust_library", "rust_test")
+
+package(default_visibility = ["//rs/ic_os:__subpackages__"])
+
+DEPENDENCIES = [
+    # Keep sorted.
+    "//rs/ic_os/build_tools/partition_tools",
+    "@crate_index//:anyhow",
+    "@crate_index//:devicemapper",
+    "@crate_index//:gpt",
+    "@crate_index//:loopdev-3",
+    "@crate_index//:nix",
+    "@crate_index//:sys-mount",
+    "@crate_index//:tempfile",
+    "@crate_index//:tokio",
+    "@crate_index//:uuid",
+]
+
+MACRO_DEPENDENCIES = [
+    # Keep sorted.
+    "@crate_index//:async-trait",
+]
+
+DEV_DEPENDENCIES = [
+    # Keep sorted.
+]
+
+rust_library(
+    name = "device",
+    srcs = glob(["src/**/*.rs"]),
+    crate_name = "ic_device",
+    proc_macro_deps = MACRO_DEPENDENCIES,
+    target_compatible_with = [
+        "@platforms//os:linux",
+    ],
+    deps = DEPENDENCIES,
+)
+
+rust_test(
+    name = "device_test",
+    crate = ":device",
+    target_compatible_with = [
+        "@platforms//os:linux",
+    ],
+    deps = DEV_DEPENDENCIES,
+)

rs/ic_os/device/BUILD.bazel

@@ -0,0 +1,19 @@
+[package]
+name = "ic_device"
+version = "0.1.0"
+edition = "2021"
+
+[dependencies]
+anyhow = { workspace = true }
+async-trait = { workspace = true }
+devicemapper = { workspace = true }
+gpt = { workspace = true }
+loopdev-3 = { workspace = true }
+nix = { workspace = true }
+partition_tools = { path = "../build_tools/partition_tools" }
+tempfile = { workspace = true }
+tokio = { workspace = true }
+uuid = { workspace = true }
+
+[target.'cfg(target_os = "linux")'.dependencies]
+sys-mount = { workspace = true }

rs/ic_os/device/Cargo.toml

@@ -1,3 +1,4 @@
+use crate::io::retry_if_busy;
 use anyhow::{ensure, Context, Result};
 use devicemapper::{
     devnode_to_devno, Bytes, DevId, Device, DmName, DmOptions, LinearDevTargetParams,
@@ -13,6 +14,7 @@ use std::path::{Path, PathBuf};
 use std::sync::Arc;
 use tempfile::{NamedTempFile, TempPath};
 
+#[allow(clippy::len_without_is_empty)]
 pub trait DeviceTrait: Send + Sync {
     fn len(&self) -> Sectors;
     fn device(&self) -> Device;
@@ -158,10 +160,8 @@ impl TempDevice {
 
         let temp_path = temp_file.into_temp_path();
 
-        let loop_device = LoopDeviceWrapper(loopdev::LoopControl::open()?.next_free()?);
-        loop_device
-            .attach_file(&temp_path)
-            .context("Temp loopback attach_file failed")?;
+        let loop_device = LoopDeviceWrapper::attach_to_next_free(&temp_path)
+            .context("Temp loopback creation failed")?;
 
         let minor = loop_device
             .minor()
@@ -192,6 +192,22 @@ impl DeviceTrait for TempDevice {
 /// Wrapper around a loop device that automatically detaches it when dropped.
 pub struct LoopDeviceWrapper(pub LoopDevice);
 
+impl LoopDeviceWrapper {
+    /// Opens a loop device and attaches it to the specified file.
+    pub fn attach_to_next_free(path: &Path) -> Result<Self> {
+        // next_free() will return the same loop device until a file is attached to it, so
+        // it can happen that a parallel process gets the same loop device and attaches a file
+        // to it before we do. In this case attach_file will fail with ResourceBusy.
+        // We solve this by retrying the operation.
+        retry_if_busy(|| {
+            let loop_device = Self(loopdev::LoopControl::open()?.next_free()?);
+            loop_device.attach_file(path)?;
+            Ok(loop_device)
+        })
+        .context("Failed to attach loop device")
+    }
+}
+
 impl Drop for LoopDeviceWrapper {
     fn drop(&mut self) {
         #[cfg(debug_assertions)]
@@ -230,7 +246,7 @@ impl Deref for LoopDeviceWrapper {
 /// dependencies, these are other devices that the [MappedDevice] depends on and should be cleaned
 /// up when the [MappedDevice] is dropped.
 pub struct MappedDevice {
-    name: &'static str,
+    name: String,
     path: PathBuf,
     len: Sectors,
     device_mapper: Arc<DM>,
@@ -286,7 +302,7 @@ impl MappedDevice {
     /// `source` must stay valid for the lifetime of the snapshot.
     pub fn create_snapshot(
         device_mapper: Arc<DM>,
-        name: &'static str,
+        name: &str,
         source: Box<dyn DeviceTrait>,
         copy_on_write: Box<dyn DeviceTrait>,
     ) -> Result<MappedDevice> {
@@ -307,7 +323,7 @@ impl MappedDevice {
 
     fn create(
         dm: Arc<DM>,
-        name: &'static str,
+        name: &str,
         table: &[(u64, u64, String, String)],
         len: Sectors,
         dependencies: Vec<Box<dyn DeviceTrait>>,
@@ -319,7 +335,7 @@ impl MappedDevice {
         // Wrap the device right away by creating a MappedDevice so it gets detached in the
         // MappedDevice Drop impl if there is an error later.
         let mapped_device = MappedDevice {
-            name,
+            name: name.to_string(),
             path: format!("/dev/mapper/{name}").into(),
             device: device.device(),
             len,
@@ -351,7 +367,7 @@ impl Drop for MappedDevice {
         debug_assert_valid(self);
 
         if let Err(err) = self.device_mapper.device_remove(
-            &DevId::Name(DmName::new(self.name).unwrap()),
+            &DevId::Name(DmName::new(&self.name).unwrap()),
             DmOptions::default(),
         ) {
             debug_panic(&format!("Failed to remove device mapper device: {err}"));

rs/ic_os/os_tools/guest_vm_runner/src/device_mapping.rs renamed to rs/ic_os/device/src/device_mapping.rs

@@ -0,0 +1,48 @@
+use std::io;
+
+/// Retries a function, returning its result if it succeeds, or retrying if it fails with
+/// ResourceBusy.
+pub fn retry_if_busy<T>(mut f: impl FnMut() -> io::Result<T>) -> io::Result<T> {
+    const MAX_ATTEMPTS: i32 = 10;
+    let mut attempts = MAX_ATTEMPTS;
+    loop {
+        match f() {
+            Ok(res) => return Ok(res),
+            Err(e) => {
+                if e.kind() == io::ErrorKind::ResourceBusy && attempts > 0 {
+                    attempts -= 1;
+                } else {
+                    return Err(e);
+                }
+            }
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_retry_if_busy() {
+        let mut attempts = 0;
+        let result = retry_if_busy(|| {
+            attempts += 1;
+            if attempts < 3 {
+                Err(io::Error::new(io::ErrorKind::ResourceBusy, "busy"))
+            } else {
+                Ok("success")
+            }
+        });
+
+        assert_eq!(result.unwrap(), "success");
+        assert_eq!(attempts, 3);
+    }
+
+    #[test]
+    fn test_retry_if_busy_failure() {
+        let result: Result<(), _> =
+            retry_if_busy(|| Err(io::Error::new(io::ErrorKind::Other, "fail")));
+        assert!(result.is_err());
+    }
+}

rs/ic_os/device/src/io.rs

@@ -0,0 +1,3 @@
+pub mod device_mapping;
+mod io;
+pub mod mount;

rs/ic_os/device/src/lib.rs

@@ -1,3 +1,4 @@
+use crate::io::retry_if_busy;
 use anyhow::{Context, Error, Result};
 use async_trait::async_trait;
 use gpt::GptDisk;
@@ -151,6 +152,7 @@ impl PartitionProvider for GptPartitionProvider {
         self.mounter
             .mount_range(self.device.clone(), offset_bytes, len_bytes, options)
             .await
+            .context("mount_range failed")
     }
 }
 
@@ -188,12 +190,15 @@ impl Mounter for LoopDeviceMounter {
             let tempdir = TempDir::new()?;
             let mount_point = tempdir.path();
             Ok::<LoopDeviceMount, Error>(LoopDeviceMount {
-                mount: Mount::builder()
-                    .fstype(FilesystemType::Manual(options.file_system.as_str()))
-                    .loopback_offset(offset_bytes)
-                    .flags(MountFlags::empty())
-                    .explicit_loopback()
-                    .mount_autodrop(device, mount_point, UnmountFlags::empty())?,
+                mount: retry_if_busy(|| {
+                    Mount::builder()
+                        .fstype(FilesystemType::Manual(options.file_system.as_str()))
+                        .loopback_offset(offset_bytes)
+                        .flags(MountFlags::empty())
+                        .explicit_loopback()
+                        .mount_autodrop(&device, mount_point, UnmountFlags::empty())
+                })
+                .context("Failed to create mount")?,
                 _tempdir: tempdir,
             })
         })
@@ -202,8 +207,6 @@ impl Mounter for LoopDeviceMounter {
     }
 }
 
-#[cfg(test)]
-#[allow(dead_code)]
 pub mod testing {
     use super::*;
     use anyhow::Context;
@@ -264,7 +267,7 @@ pub mod testing {
     ///
     /// The mounter "remembers" the extracted partitions, so extracting the same device/offset/len
     /// always returns the same directory.
-    #[derive(Clone)]
+    #[derive(Clone, Default)]
     pub struct ExtractingFilesystemMounter {
         #[allow(clippy::disallowed_types)] // Using tokio Mutex in testing is fine.
         mounts: Arc<tokio::sync::Mutex<PartitionMap>>,
@@ -301,13 +304,6 @@ pub mod testing {
     }
 
     impl ExtractingFilesystemMounter {
-        #[allow(dead_code)]
-        pub fn new() -> Self {
-            Self {
-                mounts: Default::default(),
-            }
-        }
-
         async fn extract_partition_to_tempdir(
             &self,
             device: PathBuf,

rs/ic_os/os_tools/guest_vm_runner/src/mount.rs renamed to rs/ic_os/device/src/mount.rs

@@ -8,6 +8,7 @@ DEPENDENCIES = [
     ":build_script",
     "//rs/ic_os/config_types",
     "//rs/ic_os/deterministic_ips",
+    "//rs/ic_os/device",
     "//rs/ic_os/grub",
     "//rs/ic_os/metrics_tool",
     "//rs/ic_os/sev",
@@ -16,11 +17,9 @@ DEPENDENCIES = [
     "@crate_index//:clap",
     "@crate_index//:devicemapper",
     "@crate_index//:gpt",
-    "@crate_index//:loopdev-3",
     "@crate_index//:macaddr",
     "@crate_index//:nix",
     "@crate_index//:regex",
-    "@crate_index//:sys-mount",
     "@crate_index//:systemd",
     "@crate_index//:tempfile",
     "@crate_index//:thiserror",
@@ -38,9 +37,7 @@ DEV_DEPENDENCIES = [
     "@crate_index//:url",
 ]
 
-MACRO_DEPENDENCIES = [
-    "@crate_index//:async-trait",
-]
+MACRO_DEPENDENCIES = []
 
 cargo_build_script(
     name = "build_script",