Repository: firecracker-microvm/firecracker
File: src/vmm/src/devices/virtio/balloon/device.rs (Rust)

Code before:
    pub(crate) fn process_deflate_queue(&mut self) -> Result<(), BalloonError> {
        METRICS.deflate_count.inc();

        let queue = &mut self.queues[DEFLATE_INDEX];
        let mut needs_interrupt = false;

        while let Some(head) = queue.pop().map_err(QueueError::InvalidAvailIdx)? {
            queue.add_used(head.index, 0)?;
            needs_interrupt = true;
        }

        if needs_interrupt {
            self.signal_used_queue()
        } else {
            Ok(())
        }
    }

    pub(crate) fn process_stats_queue(&mut self) -> Result<(), BalloonError> {
        // This is safe since we checked in the event handler that the device is activated.
        let mem = self.device_state.mem().unwrap();
        METRICS.stats_updates_count.inc();

        while let Some(head) = self.queues[STATS_INDEX]
            .pop()
            .map_err(QueueError::InvalidAvailIdx)?
        {
            if let Some(prev_stats_desc) = self.stats_desc_index {
                // We shouldn't ever have an extra buffer if the driver follows
                // the protocol, but return it if we find one.
                error!("balloon: driver is not compliant, more than one stats buffer received");
                self.queues[STATS_INDEX].add_used(prev_stats_desc, 0)?;
            }
            for index in (0..head.len).step_by(SIZE_OF_STAT) {
                // Read the address at position `index`. The only case
                // in which this fails is if there is overflow,
                // in which case this descriptor is malformed,
                // so we ignore the rest of it.
                let addr = head
                    .addr
                    .checked_add(u64::from(index))
                    .ok_or(BalloonError::MalformedDescriptor)?;
                let stat = mem
                    .read_obj::<BalloonStat>(addr)
                    .map_err(|_| BalloonError::MalformedDescriptor)?;
                self.latest_stats.update_with_stat(&stat).map_err(|_| {
                    METRICS.stats_update_fails.inc();
                    BalloonError::MalformedPayload
                })?;
            }


Review comment (bug):
What if we add a variant in the `BalloonError` like:

```Rust
#[derive(Debug, thiserror::Error, displaydoc::Display)]
pub enum BalloonError {
    ...

    /// Error while processing queue: {0}
    Queue(#[from] QueueError])
}
```

so that we can use `?` and avoid all these `.map_err`?

Code after:
    pub(crate) fn process_deflate_queue(&mut self) -> Result<(), BalloonError> {
        METRICS.deflate_count.inc();

        let queue = &mut self.queues[DEFLATE_INDEX];
        let mut needs_interrupt = false;

        while let Some(head) = queue.pop()? {
            queue.add_used(head.index, 0)?;
            needs_interrupt = true;
        }

        if needs_interrupt {
            self.signal_used_queue()
        } else {
            Ok(())
        }
    }

    pub(crate) fn process_stats_queue(&mut self) -> Result<(), BalloonError> {
        // This is safe since we checked in the event handler that the device is activated.
        let mem = self.device_state.mem().unwrap();
        METRICS.stats_updates_count.inc();

        while let Some(head) = self.queues[STATS_INDEX].pop()? {
            if let Some(prev_stats_desc) = self.stats_desc_index {
                // We shouldn't ever have an extra buffer if the driver follows
                // the protocol, but return it if we find one.
                error!("balloon: driver is not compliant, more than one stats buffer received");
                self.queues[STATS_INDEX].add_used(prev_stats_desc, 0)?;
            }
            for index in (0..head.len).step_by(SIZE_OF_STAT) {
                // Read the address at position `index`. The only case
                // in which this fails is if there is overflow,
                // in which case this descriptor is malformed,
                // so we ignore the rest of it.
                let addr = head
                    .addr
                    .checked_add(u64::from(index))
                    .ok_or(BalloonError::MalformedDescriptor)?;
                let stat = mem
                    .read_obj::<BalloonStat>(addr)
                    .map_err(|_| BalloonError::MalformedDescriptor)?;
                self.latest_stats.update_with_stat(&stat).map_err(|_| {
                    METRICS.stats_update_fails.inc();
                    BalloonError::MalformedPayload
                })?;
            }

            self.stats_desc_index = Some(head.index);
        }