disk_blockdevice/

lib.rs

// Copyright (c) Microsoft Corporation.
// Licensed under the MIT License.

#![expect(missing_docs)]
#![cfg(target_os = "linux")]

//! Implements the [`DiskIo`] trait for virtual disks backed by a raw block
//! device.

// UNSAFETY: Issuing IOs and calling ioctls.
#![expect(unsafe_code)]

mod ioctl;
mod nvme;
pub mod resolver;

use anyhow::Context;
use blocking::unblock;
use disk_backend::DiskError;
use disk_backend::DiskIo;
use disk_backend::UnmapBehavior;
use disk_backend::pr::PersistentReservation;
use disk_backend::pr::ReservationCapabilities;
use disk_backend::pr::ReservationReport;
use disk_backend::pr::ReservationType;
use fs_err::PathExt;
use guestmem::MemoryRead;
use guestmem::MemoryWrite;
use inspect::Inspect;
use io_uring::opcode;
use io_uring::types;
use nvme::check_nvme_status;
use nvme_spec::nvm;
use pal::unix::affinity;
use pal_async::driver::Driver;
use scsi_buffers::BounceBuffer;
use scsi_buffers::BounceBufferTracker;
use scsi_buffers::RequestBuffers;
use std::fmt::Debug;
use std::fs;
use std::os::unix::io::AsRawFd;
use std::os::unix::prelude::FileTypeExt;
use std::os::unix::prelude::MetadataExt;
use std::path::Path;
use std::path::PathBuf;
use std::str::FromStr;
use std::sync::Arc;
use std::sync::atomic::AtomicU64;
use std::sync::atomic::Ordering;
use thiserror::Error;
use uevent::CallbackHandle;
use uevent::UeventListener;

/// Opens a file for use with [`BlockDevice`] or
/// [`disk_backend_resources::BlockDeviceDiskHandle`].
pub fn open_file_for_block(
    path: &Path,
    read_only: bool,
    direct: bool,
) -> std::io::Result<fs::File> {
    use std::os::unix::prelude::*;

    tracing::debug!(?path, read_only, direct, "open_file_for_block");
    let mut opts = fs::OpenOptions::new();
    opts.read(true).write(!read_only);
    if direct {
        opts.custom_flags(libc::O_DIRECT);
    }
    opts.open(path)
}

/// A bounce buffer that may or may not be tracked by a
/// [`BounceBufferTracker`].
enum MaybeBounceBuffer<'a> {
    Tracked(scsi_buffers::TrackedBounceBuffer<'a>),
    Untracked(BounceBuffer),
}

impl MaybeBounceBuffer<'_> {
    fn io_vecs(&self) -> &[scsi_buffers::IoBuffer<'_>] {
        match self {
            Self::Tracked(t) => t.buffer.io_vecs(),
            Self::Untracked(b) => b.io_vecs(),
        }
    }

    fn as_mut_bytes(&mut self) -> &mut [u8] {
        match self {
            Self::Tracked(t) => t.buffer.as_mut_bytes(),
            Self::Untracked(b) => b.as_mut_bytes(),
        }
    }
}

/// A storvsp disk backed by a raw block device.
#[derive(Inspect)]
#[inspect(extra = "BlockDevice::inspect_extra")]
pub struct BlockDevice {
    file: Arc<fs::File>,
    sector_size: u32,
    physical_sector_size: u32,
    sector_shift: u32,
    sector_count: AtomicU64,
    optimal_unmap_sectors: u32,
    read_only: bool,
    #[inspect(skip)]
    driver: Box<dyn Driver>,
    #[inspect(flatten)]
    device_type: DeviceType,
    supports_pr: bool,
    supports_fua: bool,
    #[inspect(skip)]
    _uevent_filter: Option<CallbackHandle>,
    resize_epoch: Arc<ResizeEpoch>,
    resized_acked: AtomicU64,
    #[inspect(skip)]
    bounce_buffer_tracker: Option<Arc<BounceBufferTracker>>,
    always_bounce: bool,
}

#[derive(Inspect, Debug, Default)]
#[inspect(transparent)]
struct ResizeEpoch {
    epoch: AtomicU64,
    #[inspect(skip)]
    event: event_listener::Event,
}

#[derive(Debug, Copy, Clone, Inspect)]
#[inspect(tag = "device_type")]
enum DeviceType {
    File {
        sector_count: u64,
    },
    UnknownBlock,
    NVMe {
        ns_id: u32,
        rescap: nvm::ReservationCapabilities,
    },
}

impl BlockDevice {
    fn inspect_extra(&self, resp: &mut inspect::Response<'_>) {
        match self.device_type {
            DeviceType::NVMe { .. } => {
                resp.field_mut_with("interrupt_aggregation", |new_value| {
                    self.inspect_interrupt_coalescing(new_value)
                });
            }
            DeviceType::UnknownBlock => {}
            DeviceType::File { .. } => {}
        }
    }

    fn inspect_interrupt_coalescing(&self, new_value: Option<&str>) -> anyhow::Result<String> {
        let coalescing = if let Some(new_value) = new_value {
            let coalescing = (|| {
                let (threshold, time) = new_value.split_once(' ')?;
                Some(
                    nvme::InterruptCoalescing::new()
                        .with_aggregation_threshold(threshold.parse().ok()?)
                        .with_aggregation_time(time.parse().ok()?),
                )
            })()
            .context("expected `<aggregation_threshold> <aggregation_time>`")?;
            nvme::nvme_set_features_interrupt_coalescing(&self.file, coalescing)?;
            coalescing
        } else if let Ok(coalescing) = nvme::nvme_get_features_interrupt_coalescing(&self.file) {
            coalescing
        } else {
            return Ok("not supported".into());
        };
        Ok(format!(
            "{} {}",
            coalescing.aggregation_threshold(),
            coalescing.aggregation_time()
        ))
    }
}

/// New device error
#[derive(Debug, Error)]
pub enum NewDeviceError {
    #[error("block device ioctl error")]
    IoctlError(#[from] DiskError),
    #[error("failed to read device metadata")]
    DeviceMetadata(#[source] anyhow::Error),
    #[error("invalid file type, not a file or block device")]
    InvalidFileType,
    #[error("invalid disk size {0:#x}")]
    InvalidDiskSize(u64),
    #[error("driver does not support io-uring")]
    NoIoUring,
}

impl BlockDevice {
    /// Constructs a new `BlockDevice` backed by the specified file.
    ///
    /// # Arguments
    /// * `file` - The backing device opened for raw access.
    /// * `read_only` - Indicates whether the device is opened for read-only access.
    /// * `driver` - The async driver to use for issuing IOs (must support io-uring).
    /// * `always_bounce` - Whether to always use bounce buffers for IOs, even for those that are aligned.
    pub async fn new(
        file: fs::File,
        read_only: bool,
        driver: impl Driver,
        uevent_listener: Option<&UeventListener>,
        bounce_buffer_tracker: Option<Arc<BounceBufferTracker>>,
        always_bounce: bool,
    ) -> Result<BlockDevice, NewDeviceError> {
        if !driver.io_uring_probe(opcode::Read::CODE) {
            return Err(NewDeviceError::NoIoUring);
        }
        assert!(driver.io_uring_probe(opcode::Write::CODE));
        assert!(driver.io_uring_probe(opcode::Readv::CODE));
        assert!(driver.io_uring_probe(opcode::Writev::CODE));
        assert!(driver.io_uring_probe(opcode::Fsync::CODE));

        let metadata = file.metadata().map_err(DiskError::Io)?;

        let mut uevent_filter = None;
        let resize_epoch = Arc::new(ResizeEpoch::default());

        let devmeta = if metadata.file_type().is_block_device() {
            let rdev = metadata.rdev();
            let (major, minor) = (libc::major(rdev), libc::minor(rdev));

            // Register for resize events.
            if let Some(uevent_listener) = uevent_listener {
                let resize_epoch = resize_epoch.clone();
                uevent_filter = Some(
                    uevent_listener
                        .add_block_resize_callback(major, minor, {
                            move || {
                                tracing::info!(major, minor, "disk resized");
                                resize_epoch.epoch.fetch_add(1, Ordering::SeqCst);
                                resize_epoch.event.notify(usize::MAX);
                            }
                        })
                        .await,
                );
            }

            DeviceMetadata::from_block_device(&file, major, minor)
                .map_err(NewDeviceError::DeviceMetadata)?
        } else if metadata.file_type().is_file() {
            // Discard on a file is serviced via `fallocate(PUNCH_HOLE)`, issued
            // asynchronously through io-uring. Only offer it when the file is
            // writable and the ring supports the FALLOCATE opcode (Linux 5.6+).
            let allow_discard = !read_only && driver.io_uring_probe(opcode::Fallocate::CODE);
            DeviceMetadata::from_file(&file, &metadata, allow_discard)
                .map_err(NewDeviceError::DeviceMetadata)?
        } else {
            return Err(NewDeviceError::InvalidFileType);
        };

        let sector_size = devmeta.logical_block_size;
        let sector_shift = sector_size.trailing_zeros();
        let physical_sector_size = devmeta.physical_block_size.max(sector_size);
        let sector_count = devmeta.disk_size >> sector_shift;
        let unmap_granularity = devmeta.discard_granularity >> sector_shift;
        let file = Arc::new(file);
        let device = BlockDevice {
            file,
            sector_size,
            physical_sector_size,
            sector_shift: sector_size.trailing_zeros(),
            sector_count: sector_count.into(),
            optimal_unmap_sectors: unmap_granularity,
            read_only,
            driver: Box::new(driver),
            device_type: devmeta.device_type,
            supports_pr: devmeta.supports_pr,
            supports_fua: devmeta.fua,
            _uevent_filter: uevent_filter,
            resize_epoch,
            resized_acked: 0.into(),
            bounce_buffer_tracker,
            always_bounce,
        };

        Ok(device)
    }

    /// Use a box to avoid embedding a large `TrackedBounceBuffer` directly in
    /// the calling future.
    async fn acquire_bounce_buffer(&self, size: usize) -> Box<MaybeBounceBuffer<'_>> {
        Box::new(if let Some(tracker) = &self.bounce_buffer_tracker {
            MaybeBounceBuffer::Tracked(
                tracker
                    .acquire_bounce_buffers(size, affinity::get_cpu_number() as usize)
                    .await,
            )
        } else {
            MaybeBounceBuffer::Untracked(BounceBuffer::new(size))
        })
    }

    fn handle_resize(&self) {
        if let Err(err) = self.handle_resize_inner() {
            tracing::error!(
                error = &err as &dyn std::error::Error,
                "failed to update disk size"
            );
        }
    }

    fn handle_resize_inner(&self) -> std::io::Result<()> {
        let mut acked = self.resized_acked.load(Ordering::SeqCst);
        loop {
            let epoch = self.resize_epoch.epoch.load(Ordering::SeqCst);
            if acked == epoch {
                break Ok(());
            }

            let size_in_bytes = ioctl::query_block_device_size_in_bytes(&self.file)?;

            let new_sector_count = size_in_bytes / self.sector_size as u64;
            let original_sector_count = self.sector_count.load(Ordering::SeqCst);

            tracing::debug!(original_sector_count, new_sector_count, "resize");
            if original_sector_count != new_sector_count {
                tracing::info!(
                    original_sector_count,
                    new_sector_count,
                    "Disk size updating..."
                );
                self.sector_count.store(new_sector_count, Ordering::SeqCst);
            }

            acked = self
                .resized_acked
                .compare_exchange(acked, epoch, Ordering::SeqCst, Ordering::SeqCst)
                .unwrap_or_else(|x| x);
        }
    }

    fn map_io_error(&self, err: std::io::Error) -> DiskError {
        if !matches!(self.device_type, DeviceType::File { .. }) {
            match err.raw_os_error() {
                Some(libc::EBADE) => return DiskError::ReservationConflict,
                Some(libc::ENOSPC) => return DiskError::IllegalBlock,
                _ => {}
            }
        }
        DiskError::Io(err)
    }
}

/// Probes whether the filesystem backing `file` supports hole punching via
/// `fallocate(FALLOC_FL_PUNCH_HOLE)`, used to service guest discard/unmap.
///
/// This is non-destructive: it punches a one-byte hole at the end of the file
/// with `FALLOC_FL_KEEP_SIZE`, so the file size and existing data are
/// unchanged. Filesystems that support the operation return success;
/// unsupported ones return `EOPNOTSUPP` (or `ENOSYS` on ancient kernels).
fn probe_file_punch_hole(file: &fs::File, file_size: u64) -> bool {
    // SAFETY: FFI call with a valid fd owned by `file`. Punching past the end
    // of the file with `FALLOC_FL_KEEP_SIZE` touches no existing data.
    let ret = unsafe {
        libc::fallocate(
            file.as_raw_fd(),
            libc::FALLOC_FL_PUNCH_HOLE | libc::FALLOC_FL_KEEP_SIZE,
            file_size as libc::off_t,
            1,
        )
    };
    if ret == 0 {
        true
    } else {
        let err = std::io::Error::last_os_error();
        tracing::debug!(
            error = &err as &dyn std::error::Error,
            "file does not support punch-hole discard; unmap will be a no-op"
        );
        false
    }
}

struct DeviceMetadata {
    device_type: DeviceType,
    disk_size: u64,
    logical_block_size: u32,
    physical_block_size: u32,
    discard_granularity: u32,
    supports_pr: bool,
    fua: bool,
}

fn block_device_queue_path(devpath: &Path) -> anyhow::Result<PathBuf> {
    let device = if devpath
        .join("partition")
        .try_exists()
        .context("failed to inspect sysfs partition state")?
    {
        let partition =
            fs_err::canonicalize(devpath).context("failed to resolve sysfs partition path")?;
        partition
            .parent()
            .context("sysfs partition path has no parent device")?
            .to_path_buf()
    } else {
        devpath.to_path_buf()
    };
    Ok(device.join("queue"))
}

impl DeviceMetadata {
    fn from_block_device(file: &fs::File, major: u32, minor: u32) -> anyhow::Result<Self> {
        // Ensure the sysfs path exists.
        let devpath = PathBuf::from(format!("/sys/dev/block/{major}:{minor}"));
        devpath
            .fs_err_metadata()
            .context("could not find sysfs path for block device")?;

        let mut supports_pr = false;

        // Check for NVMe by looking for the namespace ID.
        let device_type = match fs_err::read_to_string(devpath.join("nsid")) {
            Ok(ns_id) => {
                let ns_id = ns_id
                    .trim()
                    .parse()
                    .context("failed to parse NVMe namespace ID")?;

                let rescap = nvme::nvme_identify_namespace_data(file, ns_id)?.rescap;
                let oncs = nvme::nvme_identify_controller_data(file)?.oncs;
                tracing::debug!(rescap = ?rescap, oncs = ?oncs, "get identify data");
                supports_pr = oncs.reservations() && u8::from(rescap) != 0;
                Some(DeviceType::NVMe { ns_id, rescap })
            }
            Err(err) if err.kind() == std::io::ErrorKind::NotFound => None,
            Err(err) => Err(err).context("failed to read NVMe namespace ID")?,
        };

        // Fall back to unknown.
        let device_type = device_type.unwrap_or(DeviceType::UnknownBlock);

        fn read_val<T: FromStr>(devpath: &Path, path: &str, msg: &str) -> anyhow::Result<T>
        where
            T::Err: 'static + std::error::Error + Send + Sync,
        {
            fs_err::read_to_string(devpath.join(path))
                .with_context(|| format!("failed to read {msg}"))?
                .trim()
                .parse()
                .with_context(|| format!("failed to parse {msg}"))
        }

        let queuepath = block_device_queue_path(&devpath)?;
        let logical_block_size = read_val(&queuepath, "logical_block_size", "sector size")?;
        let physical_block_size =
            read_val(&queuepath, "physical_block_size", "physical sector size")?;

        // sys/dev/block/*/*/size shows the size in 512-byte
        // sectors irrespective of the block device
        let disk_size = read_val::<u64>(&devpath, "size", "disk size")? * 512;
        let discard_granularity =
            read_val(&queuepath, "discard_granularity", "discard granularity")?;

        let fua = read_val::<u8>(&queuepath, "fua", "fua")? != 0;

        Self {
            device_type,
            disk_size,
            logical_block_size,
            physical_block_size,
            discard_granularity,
            supports_pr,
            fua,
        }
        .validate()
    }

    fn from_file(
        file: &fs::File,
        metadata: &fs::Metadata,
        allow_discard: bool,
    ) -> anyhow::Result<Self> {
        let logical_block_size = 512;
        let physical_block_size = metadata.blksize() as u32;
        // Advertise discard support only when the backing filesystem can
        // actually punch holes. Use the filesystem block size as the optimal
        // unmap granularity, since `fallocate(PUNCH_HOLE)` only deallocates
        // whole blocks.
        let discard_granularity = if allow_discard && probe_file_punch_hole(file, metadata.size()) {
            physical_block_size
        } else {
            0
        };
        Self {
            device_type: DeviceType::File {
                sector_count: metadata.len() / logical_block_size as u64,
            },
            disk_size: metadata.size(),
            logical_block_size,
            physical_block_size,
            discard_granularity,
            supports_pr: false,
            fua: false,
        }
        .validate()
    }

    fn validate(self) -> anyhow::Result<Self> {
        let Self {
            device_type: _,
            disk_size,
            logical_block_size,
            physical_block_size,
            discard_granularity,
            supports_pr: _,
            fua: _,
        } = self;
        if logical_block_size < 512 || !logical_block_size.is_power_of_two() {
            anyhow::bail!("invalid sector size {logical_block_size}");
        }
        if !physical_block_size.is_power_of_two() {
            anyhow::bail!("invalid physical sector size {physical_block_size}");
        }
        if disk_size % logical_block_size as u64 != 0 {
            anyhow::bail!("invalid disk size {disk_size:#x}");
        }
        if discard_granularity % logical_block_size != 0 {
            anyhow::bail!("invalid discard granularity {discard_granularity}");
        }
        Ok(self)
    }
}

impl DiskIo for BlockDevice {
    fn disk_type(&self) -> &str {
        "block_device"
    }

    fn sector_count(&self) -> u64 {
        if self.resize_epoch.epoch.load(Ordering::Relaxed)
            != self.resized_acked.load(Ordering::Relaxed)
        {
            self.handle_resize();
        }
        self.sector_count.load(Ordering::Relaxed)
    }

    fn sector_size(&self) -> u32 {
        self.sector_size
    }

    fn disk_id(&self) -> Option<[u8; 16]> {
        None
    }

    fn physical_sector_size(&self) -> u32 {
        self.physical_sector_size
    }

    fn is_fua_respected(&self) -> bool {
        self.supports_fua
    }

    fn is_read_only(&self) -> bool {
        self.read_only
    }

    fn pr(&self) -> Option<&dyn PersistentReservation> {
        if self.supports_pr { Some(self) } else { None }
    }

    async fn eject(&self) -> Result<(), DiskError> {
        let file = self.file.clone();
        unblock(move || {
            ioctl::lockdoor(&file, false)?;
            ioctl::eject(&file)
        })
        .await
        .map_err(|err| self.map_io_error(err))?;
        Ok(())
    }

    async fn read_vectored(
        &self,
        buffers: &RequestBuffers<'_>,
        sector: u64,
    ) -> Result<(), DiskError> {
        let io_size = buffers.len();
        tracing::trace!(sector, io_size, "read_vectored");

        let mut bounce_buffer = None;
        let locked;
        // Memory behind an emulated IOMMU cannot be locked for zero-copy IO, so
        // fall back to bounce buffering in that case.
        let should_bounce = self.always_bounce
            || !buffers.is_aligned(self.sector_size() as usize)
            || !buffers.guest_memory().supports_kernel_lock();
        let io_vecs = if !should_bounce {
            locked = buffers.lock(true)?;
            locked.io_vecs()
        } else {
            tracing::trace!("bounce buffering IO");

            bounce_buffer
                .insert(self.acquire_bounce_buffer(buffers.len()).await)
                .io_vecs()
        };

        // SAFETY: `io_vecs` and the underlying locked pages are locals
        // in this `async fn`--they are part of the same state machine as
        // the returned future and will not be freed before it completes
        // or is dropped (which aborts).
        let bytes_read = unsafe {
            self.driver.io_uring_submit(
                opcode::Readv::new(
                    types::Fd(self.file.as_raw_fd()),
                    io_vecs.as_ptr().cast(),
                    io_vecs.len() as u32,
                )
                .offset((sector * self.sector_size() as u64) as _)
                .build(),
            )
        }
        .await
        .map_err(|err| self.map_io_error(err))?;
        tracing::trace!(bytes_read, "read_vectored");
        if bytes_read != io_size as i32 {
            return Err(DiskError::IllegalBlock);
        }

        if let Some(mut bounce_buffer) = bounce_buffer {
            buffers.writer().write(bounce_buffer.as_mut_bytes())?;
        }
        Ok(())
    }

    async fn write_vectored(
        &self,
        buffers: &RequestBuffers<'_>,
        sector: u64,
        fua: bool,
    ) -> Result<(), DiskError> {
        let io_size = buffers.len();
        tracing::trace!(sector, io_size, "write_vectored");

        // Ensure the write doesn't extend the file.
        if let DeviceType::File { sector_count } = self.device_type {
            if sector + (io_size as u64 >> self.sector_shift) > sector_count {
                return Err(DiskError::IllegalBlock);
            }
        }

        let mut bounce_buffer;
        let locked;
        // Memory behind an emulated IOMMU cannot be locked for zero-copy IO, so
        // fall back to bounce buffering in that case.
        let should_bounce = self.always_bounce
            || !buffers.is_aligned(self.sector_size() as usize)
            || !buffers.guest_memory().supports_kernel_lock();
        let io_vecs = if !should_bounce {
            locked = buffers.lock(false)?;
            locked.io_vecs()
        } else {
            tracing::trace!("bounce buffering IO");
            bounce_buffer = self.acquire_bounce_buffer(buffers.len()).await;
            buffers.reader().read(bounce_buffer.as_mut_bytes())?;
            bounce_buffer.io_vecs()
        };

        // SAFETY: `io_vecs` and the underlying locked pages are locals
        // in this `async fn`--they are part of the same state machine as
        // the returned future and will not be freed before it completes
        // or is dropped (which aborts).
        let bytes_written = unsafe {
            self.driver.io_uring_submit(
                opcode::Writev::new(
                    types::Fd(self.file.as_raw_fd()),
                    io_vecs.as_ptr().cast::<libc::iovec>(),
                    io_vecs.len() as _,
                )
                .offset((sector * self.sector_size() as u64) as _)
                .rw_flags(if fua { libc::RWF_DSYNC } else { 0 })
                .build(),
            )
        }
        .await
        .map_err(|err| self.map_io_error(err))?;
        tracing::trace!(bytes_written, "write_vectored");
        if bytes_written != io_size as i32 {
            return Err(DiskError::IllegalBlock);
        }

        Ok(())
    }

    async fn sync_cache(&self) -> Result<(), DiskError> {
        // SAFETY: No data buffers.
        unsafe {
            self.driver
                .io_uring_submit(opcode::Fsync::new(types::Fd(self.file.as_raw_fd())).build())
        }
        .await
        .map_err(|err| self.map_io_error(err))?;
        Ok(())
    }

    async fn wait_resize(&self, sector_count: u64) -> u64 {
        loop {
            let listen = self.resize_epoch.event.listen();
            let current = self.sector_count();
            if current != sector_count {
                break current;
            }
            listen.await;
        }
    }

    async fn unmap(
        &self,
        sector_offset: u64,
        sector_count: u64,
        _block_level_only: bool,
    ) -> Result<(), DiskError> {
        // Reject out-of-range requests up front: `fallocate` silently succeeds
        // past the end of the file, so the bounds cannot be inferred from its
        // result (as they can for reads and writes).
        if sector_offset
            .checked_add(sector_count)
            .is_none_or(|end| end > self.sector_count())
        {
            return Err(DiskError::IllegalBlock);
        }

        // When unmap is unsupported (e.g. a filesystem that cannot punch holes,
        // or a block device that reports no discard granularity), report success
        // and treat unmap as a no-op.
        if self.optimal_unmap_sectors == 0 {
            return Ok(());
        }

        let file_offset = sector_offset << self.sector_shift;
        let length = sector_count << self.sector_shift;

        match self.device_type {
            // Files can't service BLKDISCARD; punch a hole via io-uring instead,
            // which deallocates the range and makes it read back as zero. Every
            // error must be propagated: `unmap_behavior` reports
            // `UnmapBehavior::Zeroes`, so the guest may skip zeroing this range
            // itself. Swallowing a failure would leave stale, non-zero data
            // while telling the guest it reads as zero.
            DeviceType::File { .. } => {
                // SAFETY: fallocate references no data buffers.
                unsafe {
                    self.driver.io_uring_submit(
                        opcode::Fallocate::new(types::Fd(self.file.as_raw_fd()), length)
                            .offset(file_offset)
                            .mode(libc::FALLOC_FL_PUNCH_HOLE | libc::FALLOC_FL_KEEP_SIZE)
                            .build(),
                    )
                }
                .await
                .map_err(|err| self.map_io_error(err))?;
                Ok(())
            }
            // Block devices discard via the BLKDISCARD ioctl on a blocking thread.
            _ => {
                let file = self.file.clone();
                match unblock(move || ioctl::discard(&file, file_offset, length)).await {
                    Ok(()) => Ok(()),
                    // The device advertised discard support but can't actually
                    // service the ioctl; treat unmap as a best-effort no-op.
                    Err(err)
                        if matches!(
                            err.raw_os_error(),
                            Some(libc::ENOTTY | libc::EOPNOTSUPP | libc::ENOSYS)
                        ) =>
                    {
                        tracing::debug!(
                            error = &err as &dyn std::error::Error,
                            "discard not supported; ignoring"
                        );
                        Ok(())
                    }
                    Err(err) => Err(self.map_io_error(err)),
                }
            }
        }
    }

    fn unmap_behavior(&self) -> UnmapBehavior {
        if self.optimal_unmap_sectors == 0 {
            UnmapBehavior::Ignored
        } else if matches!(self.device_type, DeviceType::File { .. }) {
            // `fallocate(PUNCH_HOLE)` deterministically zeroes the range: reads
            // of a punched hole return zeroes.
            UnmapBehavior::Zeroes
        } else {
            UnmapBehavior::Unspecified
        }
    }

    fn optimal_unmap_sectors(&self) -> u32 {
        self.optimal_unmap_sectors
    }
}

#[async_trait::async_trait]
impl PersistentReservation for BlockDevice {
    fn capabilities(&self) -> ReservationCapabilities {
        match &self.device_type {
            &DeviceType::NVMe { rescap, .. } => {
                nvme_common::from_nvme_reservation_capabilities(rescap)
            }
            DeviceType::File { .. } | DeviceType::UnknownBlock => unreachable!(),
        }
    }

    async fn report(&self) -> Result<ReservationReport, DiskError> {
        assert!(matches!(self.device_type, DeviceType::NVMe { .. }));
        self.nvme_persistent_reservation_report()
            .await
            .map_err(|err| self.map_io_error(err))
    }

    async fn register(
        &self,
        current_key: Option<u64>,
        new_key: u64,
        ptpl: Option<bool>,
    ) -> Result<(), DiskError> {
        assert!(matches!(self.device_type, DeviceType::NVMe { .. }));

        // The Linux kernel interface to register does not allow ptpl to be
        // configured. We could manually issue an NVMe command, but this code
        // path is not really used anyway.
        if ptpl == Some(false) {
            tracing::warn!("ignoring guest request to disable persist through power loss");
        }

        let file = self.file.clone();
        unblock(move || {
            ioctl::pr_register(
                &file,
                current_key.unwrap_or(0),
                new_key,
                if current_key.is_none() {
                    ioctl::PR_FL_IGNORE_KEY
                } else {
                    0
                },
            )
        })
        .await
        .and_then(check_nvme_status)
        .map_err(|err| self.map_io_error(err))?;
        Ok(())
    }

    async fn reserve(&self, key: u64, reservation_type: ReservationType) -> Result<(), DiskError> {
        assert!(matches!(self.device_type, DeviceType::NVMe { .. }));
        let file = self.file.clone();
        unblock(move || ioctl::pr_reserve(&file, reservation_type, key))
            .await
            .and_then(check_nvme_status)
            .map_err(|err| self.map_io_error(err))?;
        Ok(())
    }

    async fn release(&self, key: u64, reservation_type: ReservationType) -> Result<(), DiskError> {
        assert!(matches!(self.device_type, DeviceType::NVMe { .. }));
        let file = self.file.clone();
        unblock(move || ioctl::pr_release(&file, reservation_type, key))
            .await
            .and_then(check_nvme_status)
            .map_err(|err| self.map_io_error(err))?;
        Ok(())
    }

    async fn clear(&self, key: u64) -> Result<(), DiskError> {
        assert!(matches!(self.device_type, DeviceType::NVMe { .. }));
        let file = self.file.clone();
        unblock(move || ioctl::pr_clear(&file, key))
            .await
            .and_then(check_nvme_status)
            .map_err(|err| self.map_io_error(err))?;
        Ok(())
    }

    async fn preempt(
        &self,
        current_key: u64,
        preempt_key: u64,
        reservation_type: ReservationType,
        abort: bool,
    ) -> Result<(), DiskError> {
        assert!(matches!(self.device_type, DeviceType::NVMe { .. }));
        let file = self.file.clone();
        unblock(move || {
            ioctl::pr_preempt(&file, reservation_type, current_key, preempt_key, abort)
        })
        .await
        .and_then(check_nvme_status)
        .map_err(|err| self.map_io_error(err))?;
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use futures::executor::block_on;
    use guestmem::GuestMemory;
    use hvdef::HV_PAGE_SIZE;
    use hvdef::HV_PAGE_SIZE_USIZE;
    use once_cell::sync::OnceCell;
    use pal_async::async_test;
    use pal_uring::IoUringPool;
    use pal_uring::PoolClient;
    use scsi_buffers::OwnedRequestBuffers;
    use test_with_tracing::test;

    #[test]
    fn whole_device_queue_path_is_unchanged() {
        let root = tempfile::tempdir().unwrap();
        let device = root.path().join("nvme0n1");
        fs::create_dir_all(device.join("queue")).unwrap();
        fs::write(device.join("queue/logical_block_size"), "512\n").unwrap();
        let alias = root.path().join("259:0");
        std::os::unix::fs::symlink(&device, &alias).unwrap();

        let queue = block_device_queue_path(&alias).unwrap();
        assert_eq!(queue, alias.join("queue"));
        assert_eq!(
            fs::read_to_string(queue.join("logical_block_size")).unwrap(),
            "512\n"
        );
    }

    #[test]
    fn partition_queue_path_uses_parent_device() {
        let root = tempfile::tempdir().unwrap();
        let device = root.path().join("nvme0n1");
        let partition = device.join("nvme0n1p2");
        fs::create_dir_all(&partition).unwrap();
        fs::create_dir(device.join("queue")).unwrap();
        fs::write(partition.join("partition"), "2\n").unwrap();
        fs::write(device.join("queue/logical_block_size"), "512\n").unwrap();
        fs::write(device.join("queue/physical_block_size"), "4096\n").unwrap();
        let alias = root.path().join("259:2");
        std::os::unix::fs::symlink(&partition, &alias).unwrap();

        let queue = block_device_queue_path(&alias).unwrap();
        assert_eq!(queue, device.join("queue"));
        assert_eq!(
            fs::read_to_string(queue.join("logical_block_size")).unwrap(),
            "512\n"
        );
        assert_eq!(
            fs::read_to_string(queue.join("physical_block_size")).unwrap(),
            "4096\n"
        );
        assert!(!alias.join("queue").exists());
    }

    #[test]
    fn partition_queue_path_does_not_fall_back_to_child_queue() {
        let root = tempfile::tempdir().unwrap();
        let device = root.path().join("nvme0n1");
        let partition = device.join("nvme0n1p2");
        fs::create_dir_all(partition.join("queue")).unwrap();
        fs::write(partition.join("partition"), "2\n").unwrap();
        fs::write(partition.join("queue/logical_block_size"), "4096\n").unwrap();

        let queue = block_device_queue_path(&partition).unwrap();
        assert_eq!(queue, device.join("queue"));
        assert!(fs::read_to_string(queue.join("logical_block_size")).is_err());
    }

    fn is_buggy_kernel() -> bool {
        // 5.13 kernels seem to have a bug with io_uring where tests hang.
        let output = String::from_utf8(
            std::process::Command::new("uname")
                .arg("-r")
                .output()
                .unwrap()
                .stdout,
        )
        .unwrap();

        output.contains("5.13")
    }

    fn new_block_device_from(file: fs::File) -> Result<BlockDevice, NewDeviceError> {
        // TODO: switch to std::sync::OnceLock once `get_or_try_init` is stable
        static POOL: OnceCell<PoolClient> = OnceCell::new();

        let client = POOL
            .get_or_try_init(|| {
                let pool = IoUringPool::new("test", 16)?;
                let client = pool.client().clone();
                std::thread::spawn(|| pool.run());
                Ok(client)
            })
            .map_err(|err| NewDeviceError::IoctlError(DiskError::Io(err)))?;

        block_on(BlockDevice::new(
            file,
            false,
            client.initiator().clone(),
            None,
            None,
            false,
        ))
    }

    fn new_block_device() -> Result<BlockDevice, NewDeviceError> {
        let test_file = tempfile::tempfile().unwrap();
        test_file.set_len(1024 * 64).unwrap();
        new_block_device_from(test_file)
    }

    /// A loop device over a temporary file, detached when dropped.
    ///
    /// Attaching one requires `CAP_SYS_ADMIN`, which is why the test using this
    /// is ignored by default.
    struct LoopDevice {
        path: String,
        _backing: tempfile::NamedTempFile,
    }

    impl LoopDevice {
        fn new(size: u64) -> Self {
            let backing = tempfile::NamedTempFile::new().unwrap();
            backing.as_file().set_len(size).unwrap();
            let output = std::process::Command::new("losetup")
                .arg("--find")
                .arg("--show")
                .arg(backing.path())
                .output()
                .expect("losetup should be installed");
            assert!(
                output.status.success(),
                "losetup failed (needs root): {}",
                String::from_utf8_lossy(&output.stderr)
            );
            Self {
                path: String::from_utf8(output.stdout).unwrap().trim().to_owned(),
                _backing: backing,
            }
        }
    }

    impl Drop for LoopDevice {
        fn drop(&mut self) {
            // Detach on the way out even if the test panicked, so that loop
            // devices are not leaked.
            let _ = std::process::Command::new("losetup")
                .arg("-d")
                .arg(&self.path)
                .status();
        }
    }

    macro_rules! get_block_device_or_skip {
        () => {
            match new_block_device() {
                Ok(pool) => {
                    if is_buggy_kernel() {
                        println!("Test case skipped (buggy kernel version)");
                        return;
                    }

                    pool
                }
                Err(NewDeviceError::IoctlError(DiskError::Io(err)))
                    if err.raw_os_error() == Some(libc::ENOSYS) =>
                {
                    println!("Test case skipped (no IO-Uring support)");
                    return;
                }
                Err(err) => panic!("{}", err),
            }
        };
    }

    /// The shared sector-range conformance suite.
    ///
    /// This exercises the `DeviceType::File` path, where the kernel does not
    /// enforce bounds and the backend's own checks are the only thing stopping
    /// a write past the end from extending the file. See
    /// [`sector_range_conformance_loop_device`] for the other path.
    #[async_test]
    async fn sector_range_conformance() {
        let disk = disk_backend::Disk::new(get_block_device_or_skip!()).unwrap();
        storage_tests::sector_range::test_disk_sector_range_conformance(&disk).await;
    }

    /// The same suite over a real block device.
    ///
    /// This is the case the tempfile path cannot cover. For a block device the
    /// kernel enforces the bounds and this backend delegates to it, rather than
    /// checking the range itself, so it is the delegation that is under test
    /// here.
    ///
    /// Attaching a loop device needs `CAP_SYS_ADMIN`, so this is ignored by
    /// default. To run it:
    ///
    /// ```text
    /// sudo -E $(which cargo) nextest run -p disk_blockdevice \
    ///     --run-ignored all -E 'test(loop_device)'
    /// ```
    #[async_test]
    #[ignore = "needs root to attach a loop device"]
    async fn sector_range_conformance_loop_device() {
        let loop_device = LoopDevice::new(1024 * 1024);
        let file = fs::OpenOptions::new()
            .read(true)
            .write(true)
            .open(&loop_device.path)
            .unwrap();
        let disk = disk_backend::Disk::new(new_block_device_from(file).unwrap()).unwrap();
        storage_tests::sector_range::test_disk_sector_range_conformance(&disk).await;
    }

    async fn run_async_disk_io(fua: bool) {
        let disk = get_block_device_or_skip!();

        let test_guest_mem = GuestMemory::allocate(0x8000);
        test_guest_mem
            .write_at(0, &(0..0x8000).map(|x| x as u8).collect::<Vec<_>>())
            .unwrap();

        let write_buffers = OwnedRequestBuffers::new(&[3, 2, 1, 0]);
        disk.write_vectored(&write_buffers.buffer(&test_guest_mem), 0, fua)
            .await
            .unwrap();

        if !fua {
            disk.sync_cache().await.unwrap();
        }

        let read_buffers = OwnedRequestBuffers::new(&[7, 6, 5, 4]);
        disk.read_vectored(&read_buffers.buffer(&test_guest_mem), 0)
            .await
            .unwrap();

        let mut source = vec![0u8; 4 * HV_PAGE_SIZE_USIZE];
        test_guest_mem.read_at(0, &mut source).unwrap();
        let mut target = vec![0u8; 4 * HV_PAGE_SIZE_USIZE];
        test_guest_mem
            .read_at(4 * HV_PAGE_SIZE, &mut target)
            .unwrap();
        assert_eq!(source, target);
    }

    #[async_test]
    async fn test_async_disk_io() {
        run_async_disk_io(false).await;
    }

    #[async_test]
    async fn test_async_disk_io_fua() {
        run_async_disk_io(true).await;
    }

    async fn run_async_disk_io_unaligned(fua: bool) {
        let disk = get_block_device_or_skip!();

        let test_guest_mem = GuestMemory::allocate(0x8000);
        test_guest_mem
            .write_at(0, &(0..0x8000).map(|x| x as u8).collect::<Vec<_>>())
            .unwrap();

        let write_buffers =
            OwnedRequestBuffers::new_unaligned(&[0, 1, 2, 3], 512, 3 * HV_PAGE_SIZE_USIZE);

        disk.write_vectored(&write_buffers.buffer(&test_guest_mem), 0, fua)
            .await
            .unwrap();

        if !fua {
            disk.sync_cache().await.unwrap();
        }

        let read_buffers =
            OwnedRequestBuffers::new_unaligned(&[4, 5, 6, 7], 512, 3 * HV_PAGE_SIZE_USIZE);
        disk.read_vectored(&read_buffers.buffer(&test_guest_mem), 0)
            .await
            .unwrap();

        let mut source = vec![0u8; 3 * HV_PAGE_SIZE_USIZE];
        test_guest_mem.read_at(512, &mut source).unwrap();
        let mut target = vec![0u8; 3 * HV_PAGE_SIZE_USIZE];
        test_guest_mem
            .read_at(4 * HV_PAGE_SIZE + 512, &mut target)
            .unwrap();
        assert_eq!(source, target);
    }

    #[async_test]
    async fn test_async_disk_io_unaligned() {
        run_async_disk_io_unaligned(false).await;
    }

    #[async_test]
    async fn test_async_disk_io_unaligned_fua() {
        run_async_disk_io_unaligned(true).await;
    }

    #[async_test]
    async fn test_illegal_lba() {
        let disk = get_block_device_or_skip!();
        let gm = GuestMemory::allocate(512);
        match disk
            .write_vectored(
                &OwnedRequestBuffers::linear(0, 512, true).buffer(&gm),
                i64::MAX as u64 / 512,
                false,
            )
            .await
        {
            Err(DiskError::IllegalBlock) => {}
            r => panic!("unexpected result: {:?}", r),
        }
    }

    #[async_test]
    async fn test_unmap_file_punch_hole() {
        let disk = get_block_device_or_skip!();

        // A writable file-backed device should advertise discard support via
        // punch-hole and report that unmapped ranges read back as zero.
        if disk.optimal_unmap_sectors() == 0 {
            println!("Test case skipped (filesystem does not support punch hole)");
            return;
        }
        assert_eq!(disk.unmap_behavior(), UnmapBehavior::Zeroes);

        let gm = GuestMemory::allocate(0x2000);
        gm.write_at(0, &vec![0xcdu8; 0x2000]).unwrap();

        // Write two pages of non-zero data.
        disk.write_vectored(
            &OwnedRequestBuffers::linear(0, 0x2000, false).buffer(&gm),
            0,
            false,
        )
        .await
        .unwrap();
        disk.sync_cache().await.unwrap();

        // Unmap the first page and read both pages back.
        disk.unmap(0, 0x1000 >> disk.sector_shift, false)
            .await
            .unwrap();

        let read = OwnedRequestBuffers::linear(0, 0x2000, true);
        disk.read_vectored(&read.buffer(&gm), 0).await.unwrap();

        let mut buf = vec![0xffu8; 0x2000];
        gm.read_at(0, &mut buf).unwrap();
        // Punched range reads as zero; the untouched page keeps its data.
        assert!(
            buf[..0x1000].iter().all(|&b| b == 0),
            "unmapped range not zeroed"
        );
        assert!(
            buf[0x1000..].iter().all(|&b| b == 0xcd),
            "second page corrupted"
        );
    }
}