virtio-devices, vmm: add I/O rate limiter on block device

This patch is based on the 'rate_limiter' module from firecracker[1]. To
simplify dependencies, we reply on 'vmm-sys-util::TimerFd' instead of
the `timerfd` crate.

[1]https://github.com/firecracker-microvm/firecracker/tree/master/src/rate_limiter

Fixes: #1285

Signed-off-by: Bo Chen <chen.bo@intel.com>

virtio-devices/src/block.rs

@@ -11,8 +11,10 @@
 use super::Error as DeviceError;
 use super::{
     ActivateError, ActivateResult, EpollHelper, EpollHelperError, EpollHelperHandler, Queue,
-    VirtioCommon, VirtioDevice, VirtioDeviceType, VirtioInterruptType, EPOLL_HELPER_EVENT_LAST,
+    RateLimiterConfig, VirtioCommon, VirtioDevice, VirtioDeviceType, VirtioInterruptType,
+    EPOLL_HELPER_EVENT_LAST,
 };
+use crate::rate_limiter::{RateLimiter, TokenType};
 use crate::seccomp_filters::{get_seccomp_filter, Thread};
 use crate::VirtioInterrupt;
 use anyhow::anyhow;
@@ -21,7 +23,6 @@ use block_util::{
     RequestType, VirtioBlockConfig,
 };
 use seccomp::{SeccompAction, SeccompFilter};
-use std::collections::HashMap;
 use std::io;
 use std::num::Wrapping;
 use std::os::unix::io::AsRawFd;
@@ -30,6 +31,7 @@ use std::result;
 use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
 use std::sync::{Arc, Barrier};
 use std::thread;
+use std::{collections::HashMap, convert::TryInto};
 use virtio_bindings::bindings::virtio_blk::*;
 use vm_memory::{
     ByteValued, Bytes, GuestAddress, GuestAddressSpace, GuestMemoryAtomic, GuestMemoryError,
@@ -48,6 +50,8 @@ pub const SECTOR_SIZE: u64 = 0x01 << SECTOR_SHIFT;
 const QUEUE_AVAIL_EVENT: u16 = EPOLL_HELPER_EVENT_LAST + 1;
 // New completed tasks are pending on the completion ring.
 const COMPLETION_EVENT: u16 = EPOLL_HELPER_EVENT_LAST + 2;
+// New 'wake up' event from the rate limiter
+const RATE_LIMITER_EVENT: u16 = EPOLL_HELPER_EVENT_LAST + 3;
 
 #[derive(Debug)]
 pub enum Error {
@@ -104,6 +108,7 @@ struct BlockEpollHandler {
     counters: BlockCounters,
     queue_evt: EventFd,
     request_list: HashMap<u16, Request>,
+    rate_limiter: Option<RateLimiter>,
 }
 
 impl BlockEpollHandler {
@@ -116,6 +121,38 @@ impl BlockEpollHandler {
 
         for avail_desc in queue.iter(&mem) {
             let mut request = Request::parse(&avail_desc, &mem).map_err(Error::RequestParsing)?;
+
+            if let Some(rate_limiter) = &mut self.rate_limiter {
+                // If limiter.consume() fails it means there is no more TokenType::Ops
+                // budget and rate limiting is in effect.
+                if !rate_limiter.consume(1, TokenType::Ops) {
+                    // Stop processing the queue and return this descriptor chain to the
+                    // avail ring, for later processing.
+                    queue.go_to_previous_position();
+                    break;
+                }
+                // Exercise the rate limiter only if this request is of data transfer type.
+                if request.request_type == RequestType::In
+                    || request.request_type == RequestType::Out
+                {
+                    let mut bytes = Wrapping(0);
+                    for (_, data_len) in &request.data_descriptors {
+                        bytes += Wrapping(*data_len as u64);
+                    }
+
+                    // If limiter.consume() fails it means there is no more TokenType::Bytes
+                    // budget and rate limiting is in effect.
+                    if !rate_limiter.consume(bytes.0, TokenType::Bytes) {
+                        // Revert the OPS consume().
+                        rate_limiter.manual_replenish(1, TokenType::Ops);
+                        // Stop processing the queue and return this descriptor chain to the
+                        // avail ring, for later processing.
+                        queue.go_to_previous_position();
+                        break;
+                    }
+                };
+            }
+
             request.set_writeback(self.writeback.load(Ordering::Acquire));
 
             if request
@@ -242,6 +279,9 @@ impl BlockEpollHandler {
         let mut helper = EpollHelper::new(&self.kill_evt, &self.pause_evt)?;
         helper.add_event(self.queue_evt.as_raw_fd(), QUEUE_AVAIL_EVENT)?;
         helper.add_event(self.disk_image.notifier().as_raw_fd(), COMPLETION_EVENT)?;
+        if let Some(rate_limiter) = &self.rate_limiter {
+            helper.add_event(rate_limiter.as_raw_fd(), RATE_LIMITER_EVENT)?;
+        }
         helper.run(paused, paused_sync, self)?;
 
         Ok(())
@@ -258,18 +298,24 @@ impl EpollHelperHandler for BlockEpollHandler {
                     return true;
                 }
 
-                match self.process_queue_submit() {
-                    Ok(needs_notification) => {
-                        if needs_notification {
-                            if let Err(e) = self.signal_used_queue() {
-                                error!("Failed to signal used queue: {:?}", e);
-                                return true;
+                let rate_limit_reached =
+                    self.rate_limiter.as_ref().map_or(false, |r| r.is_blocked());
+
+                // Process the queue only when the rate limit is not reached
+                if !rate_limit_reached {
+                    match self.process_queue_submit() {
+                        Ok(needs_notification) => {
+                            if needs_notification {
+                                if let Err(e) = self.signal_used_queue() {
+                                    error!("Failed to signal used queue: {:?}", e);
+                                    return true;
+                                }
                             }
                         }
-                    }
-                    Err(e) => {
-                        error!("Failed to process queue (submit): {:?}", e);
-                        return true;
+                        Err(e) => {
+                            error!("Failed to process queue (submit): {:?}", e);
+                            return true;
+                        }
                     }
                 }
             }
@@ -294,6 +340,31 @@ impl EpollHelperHandler for BlockEpollHandler {
                     }
                 }
             }
+            RATE_LIMITER_EVENT => {
+                if let Some(rate_limiter) = &mut self.rate_limiter {
+                    // Upon rate limiter event, call the rate limiter handler
+                    // and restart processing the queue.
+                    if rate_limiter.event_handler().is_ok() {
+                        match self.process_queue_submit() {
+                            Ok(needs_notification) => {
+                                if needs_notification {
+                                    if let Err(e) = self.signal_used_queue() {
+                                        error!("Failed to signal used queue: {:?}", e);
+                                        return true;
+                                    }
+                                }
+                            }
+                            Err(e) => {
+                                error!("Failed to process queue (submit): {:?}", e);
+                                return true;
+                            }
+                        }
+                    }
+                } else {
+                    error!("Unexpected 'RATE_LIMITER_EVENT' when rate_limiter is not enabled.");
+                    return true;
+                }
+            }
             _ => {
                 error!("Unexpected event: {}", ev_type);
                 return true;
@@ -314,6 +385,7 @@ pub struct Block {
     writeback: Arc<AtomicBool>,
     counters: BlockCounters,
     seccomp_action: SeccompAction,
+    rate_limiter_config: Option<RateLimiterConfig>,
 }
 
 #[derive(Serialize, Deserialize)]
@@ -337,6 +409,7 @@ impl Block {
         num_queues: usize,
         queue_size: u16,
         seccomp_action: SeccompAction,
+        rate_limiter_config: Option<RateLimiterConfig>,
     ) -> io::Result<Self> {
         let disk_size = disk_image.size().map_err(|e| {
             io::Error::new(
@@ -393,6 +466,7 @@ impl Block {
             writeback: Arc::new(AtomicBool::new(true)),
             counters: BlockCounters::default(),
             seccomp_action,
+            rate_limiter_config,
         })
     }
 
@@ -521,6 +595,12 @@ impl VirtioDevice for Block {
                     ActivateError::BadActivate
                 })?;
 
+            let rate_limiter: Option<RateLimiter> = self
+                .rate_limiter_config
+                .map(RateLimiterConfig::try_into)
+                .transpose()
+                .map_err(ActivateError::CreateRateLimiter)?;
+
             let mut handler = BlockEpollHandler {
                 queue,
                 mem: mem.clone(),
@@ -540,6 +620,7 @@ impl VirtioDevice for Block {
                 counters: self.counters.clone(),
                 queue_evt,
                 request_list: HashMap::with_capacity(queue_size.into()),
+                rate_limiter,
             };
 
             let paused = self.common.paused.clone();