feat: pmem: add rate-limiter for virtio-pmem device

Add rate-limiter support to the virtio-pmem device to allow users to
configure limits of the I/O bandwidth generated by the `msync` call in
the device which could be triggered by the guest FLUSH requests.

Signed-off-by: Egor Lazarchuk <yegorlz@amazon.co.uk>

Author: ShadowCurse

src/firecracker/src/api_server/parsed_request.rs

@@ -24,7 +24,7 @@ use super::request::machine_configuration::{
 use super::request::metrics::parse_put_metrics;
 use super::request::mmds::{parse_get_mmds, parse_patch_mmds, parse_put_mmds};
 use super::request::net::{parse_patch_net, parse_put_net};
-use super::request::pmem::parse_put_pmem;
+use super::request::pmem::{parse_patch_pmem, parse_put_pmem};
 use super::request::snapshot::{parse_patch_vm_state, parse_put_snapshot};
 use super::request::version::parse_get_version;
 use super::request::vsock::parse_put_vsock;
@@ -120,6 +120,7 @@ impl TryFrom<&Request> for ParsedRequest {
             (Method::Patch, "network-interfaces", Some(body)) => {
                 parse_patch_net(body, path_tokens.next())
             }
+            (Method::Patch, "pmem", Some(body)) => parse_patch_pmem(body, path_tokens.next()),
             (Method::Patch, "vm", Some(body)) => parse_patch_vm_state(body),
             (Method::Patch, "hotplug", Some(body)) if path_tokens.next() == Some("memory") => {
                 parse_patch_memory_hotplug(body)

src/firecracker/src/api_server/request/pmem.rs

@@ -3,7 +3,7 @@
 
 use vmm::logger::{IncMetric, METRICS};
 use vmm::rpc_interface::VmmAction;
-use vmm::vmm_config::pmem::PmemConfig;
+use vmm::vmm_config::pmem::{PmemConfig, PmemDeviceUpdateConfig};
 
 use super::super::parsed_request::{ParsedRequest, RequestError, checked_id};
 use super::{Body, StatusCode};
@@ -37,6 +37,36 @@ pub(crate) fn parse_put_pmem(
     }
 }
 
+pub(crate) fn parse_patch_pmem(
+    body: &Body,
+    id_from_path: Option<&str>,
+) -> Result<ParsedRequest, RequestError> {
+    METRICS.patch_api_requests.pmem_count.inc();
+    let id = if let Some(id) = id_from_path {
+        checked_id(id)?
+    } else {
+        METRICS.patch_api_requests.pmem_fails.inc();
+        return Err(RequestError::EmptyID);
+    };
+
+    let update_cfg =
+        serde_json::from_slice::<PmemDeviceUpdateConfig>(body.raw()).inspect_err(|_| {
+            METRICS.patch_api_requests.pmem_fails.inc();
+        })?;
+
+    if id == update_cfg.id {
+        Ok(ParsedRequest::new_sync(VmmAction::UpdatePmemDevice(
+            update_cfg,
+        )))
+    } else {
+        METRICS.patch_api_requests.pmem_fails.inc();
+        Err(RequestError::Generic(
+            StatusCode::BadRequest,
+            "The id from the path does not match the id from the body!".to_string(),
+        ))
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -60,7 +90,8 @@ mod tests {
             "id": "1000",
             "path_on_host": "dummy",
             "root_device": true,
-            "read_only": true
+            "read_only": true,
+            "rate_limiter": {}
         }"#;
         let r = vmm_action_from_request(parse_put_pmem(&Body::new(body), Some("1000")).unwrap());
 
@@ -69,7 +100,35 @@ mod tests {
             path_on_host: "dummy".to_string(),
             root_device: true,
             read_only: true,
+            rate_limiter: Some(Default::default()),
         };
         assert_eq!(r, VmmAction::InsertPmemDevice(expected_config));
     }
+
+    #[test]
+    fn test_parse_patch_pmem_request() {
+        parse_put_pmem(&Body::new("invalid_payload"), None).unwrap_err();
+        parse_put_pmem(&Body::new("invalid_payload"), Some("id")).unwrap_err();
+
+        let body = r#"{
+            "id": "bar",
+        }"#;
+        parse_put_pmem(&Body::new(body), Some("1")).unwrap_err();
+        let body = r#"{
+            "foo": "1",
+        }"#;
+        parse_put_pmem(&Body::new(body), Some("1")).unwrap_err();
+
+        let body = r#"{
+            "id": "1000",
+            "rate_limiter": {}
+        }"#;
+        let r = vmm_action_from_request(parse_patch_pmem(&Body::new(body), Some("1000")).unwrap());
+
+        let expected_config = PmemDeviceUpdateConfig {
+            id: "1000".to_string(),
+            rate_limiter: Some(Default::default()),
+        };
+        assert_eq!(r, VmmAction::UpdatePmemDevice(expected_config));
+    }
 }

src/firecracker/swagger/firecracker.yaml

@@ -370,6 +370,35 @@ paths:
           description: Internal server error.
           schema:
             $ref: "#/definitions/Error"
+    patch:
+      summary: Updates the rate limiter of a pmem device. Post-boot only.
+      description:
+        Updates the rate limiter applied to the pmem device with the ID specified
+        by the id path parameter.
+      operationId: patchGuestPmemByID
+      parameters:
+        - name: id
+          in: path
+          description: The id of the guest pmem device
+          required: true
+          type: string
+        - name: body
+          in: body
+          description: Pmem rate limiter properties
+          required: true
+          schema:
+            $ref: "#/definitions/PartialPmem"
+      responses:
+        204:
+          description: Pmem device updated
+        400:
+          description: Pmem device cannot be updated due to bad input
+          schema:
+            $ref: "#/definitions/Error"
+        default:
+          description: Internal server error.
+          schema:
+            $ref: "#/definitions/Error"
 
   /logger:
     put:
@@ -1255,6 +1284,8 @@ definitions:
         type: boolean
         description:
           Flag to map backing file in read-only mode.
+      rate_limiter:
+        $ref: "#/definitions/RateLimiter"
 
   Error:
     type: object
@@ -1538,6 +1569,19 @@ definitions:
       tx_rate_limiter:
         $ref: "#/definitions/RateLimiter"
 
+  PartialPmem:
+    type: object
+    description:
+      Defines a partial pmem device structure, used to update the rate limiter
+      for that device, after microvm start.
+    required:
+      - id
+    properties:
+      id:
+        type: string
+      rate_limiter:
+        $ref: "#/definitions/RateLimiter"
+
   RateLimiter:
     type: object
     description:

src/vmm/src/builder.rs

@@ -1311,6 +1311,7 @@ pub(crate) mod tests {
             path_on_host: "".into(),
             root_device: true,
             read_only: true,
+            ..Default::default()
         }];
         let mut vmm = default_vmm();
         let mut cmdline = default_kernel_cmdline();

src/vmm/src/device_manager/pci_mngr.rs

@@ -705,6 +705,7 @@ mod tests {
                 path_on_host: "".into(),
                 root_device: true,
                 read_only: true,
+                ..Default::default()
             }];
             _pmem_files =
                 insert_pmem_devices(&mut vmm, &mut cmdline, &mut event_manager, pmem_configs);
@@ -812,7 +813,8 @@ mod tests {
       "id": "pmem",
       "path_on_host": "{}",
       "root_device": true,
-      "read_only": true
+      "read_only": true,
+      "rate_limiter": null
     }}
   ],
   "memory-hotplug": {{

src/vmm/src/device_manager/persist.rs

@@ -738,6 +738,7 @@ mod tests {
                 path_on_host: "".into(),
                 root_device: true,
                 read_only: true,
+                ..Default::default()
             }];
             _pmem_files =
                 insert_pmem_devices(&mut vmm, &mut cmdline, &mut event_manager, pmem_configs);
@@ -842,7 +843,8 @@ mod tests {
       "id": "pmem",
       "path_on_host": "{}",
       "root_device": true,
-      "read_only": true
+      "read_only": true,
+      "rate_limiter": null
     }}
   ],
   "memory-hotplug": {{

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

@@ -22,7 +22,9 @@ use crate::devices::virtio::pmem::metrics::{PmemMetrics, PmemMetricsPerDevice};
 use crate::devices::virtio::queue::{DescriptorChain, InvalidAvailIdx, Queue, QueueError};
 use crate::devices::virtio::transport::{VirtioInterrupt, VirtioInterruptType};
 use crate::logger::{IncMetric, error, info};
+use crate::rate_limiter::{BucketUpdate, RateLimiter, TokenType};
 use crate::utils::{align_up, u64_to_usize};
+use crate::vmm_config::RateLimiterConfig;
 use crate::vmm_config::pmem::PmemConfig;
 use crate::vstate::memory::{ByteValued, Bytes, GuestMemoryMmap, GuestMmapRegion};
 use crate::vstate::vm::VmError;
@@ -58,6 +60,8 @@ pub enum PmemError {
     Queue(#[from] QueueError),
     /// Error during obtaining the descriptor from the queue: {0}
     QueuePop(#[from] InvalidAvailIdx),
+    /// Error creating rate limiter: {0}
+    RateLimiter(std::io::Error),
 }
 
 const VIRTIO_PMEM_REQ_TYPE_FLUSH: u32 = 0;
@@ -94,6 +98,7 @@ pub struct Pmem {
     pub file_len: u64,
     pub mmap_ptr: u64,
     pub metrics: Arc<PmemMetrics>,
+    pub rate_limiter: RateLimiter,
 
     pub config: PmemConfig,
 }
@@ -125,6 +130,13 @@ impl Pmem {
         let (file, file_len, mmap_ptr, mmap_len) =
             Self::mmap_backing_file(&config.path_on_host, config.read_only)?;
 
+        let rate_limiter = config
+            .rate_limiter
+            .map(RateLimiterConfig::try_into)
+            .transpose()
+            .map_err(PmemError::RateLimiter)?
+            .unwrap_or_default();
+
         Ok(Self {
             avail_features: 1u64 << VIRTIO_F_VERSION_1,
             acked_features: 0u64,
@@ -140,6 +152,7 @@ impl Pmem {
             file_len,
             mmap_ptr,
             metrics: PmemMetricsPerDevice::alloc(config.id.clone()),
+            rate_limiter,
             config,
         })
     }
@@ -254,6 +267,26 @@ impl Pmem {
         // This is safe since we checked in the event handler that the device is activated.
         let active_state = self.device_state.active_state().unwrap();
 
+        if self.queues[0].is_empty() {
+            return Ok(());
+        }
+
+        // There is only 1 type of request pmem supports, so we can consume
+        // rate-limiter before even looking at the queue. This is still valid
+        // even if the queue will not have any valid requests since it indicate
+        // broken guest driver and rate-limiting should still apply for such case.
+        // Rate limit: consume 1 op and file_len bytes for the coalesced msync.
+        // If the rate limiter is blocked, defer notification until the timer fires.
+        if !self.rate_limiter.consume(1, TokenType::Ops) {
+            self.metrics.rate_limiter_throttled_events.inc();
+            return Ok(());
+        }
+        if !self.rate_limiter.consume(self.file_len, TokenType::Bytes) {
+            self.rate_limiter.manual_replenish(1, TokenType::Ops);
+            self.metrics.rate_limiter_throttled_events.inc();
+            return Ok(());
+        }
+
         let mut cached_result = None;
         while let Some(head) = self.queues[0].pop()? {
             let add_result = match self.process_chain(head, &mut cached_result) {
@@ -270,8 +303,8 @@ impl Pmem {
                 break;
             }
         }
-        self.queues[0].advance_used_ring_idx();
 
+        self.queues[0].advance_used_ring_idx();
         if self.queues[0].prepare_kick() {
             active_state
                 .interrupt
@@ -347,6 +380,11 @@ impl Pmem {
         Ok(())
     }
 
+    /// Updates the parameters for the rate limiter.
+    pub fn update_rate_limiter(&mut self, bytes: BucketUpdate, ops: BucketUpdate) {
+        self.rate_limiter.update_buckets(bytes, ops);
+    }
+
     pub fn process_queue(&mut self) {
         self.metrics.queue_event_count.inc();
         if let Err(err) = self.queue_events[0].read() {
@@ -355,6 +393,25 @@ impl Pmem {
             return;
         }
 
+        if self.rate_limiter.is_blocked() {
+            self.metrics.rate_limiter_throttled_events.inc();
+            return;
+        }
+
+        self.handle_queue().unwrap_or_else(|err| {
+            error!("pmem: {err:?}");
+            self.metrics.event_fails.inc();
+        });
+    }
+
+    pub fn process_rate_limiter_event(&mut self) {
+        self.metrics.rate_limiter_event_count.inc();
+        if let Err(err) = self.rate_limiter.event_handler() {
+            error!("pmem: Failed to get rate-limiter event: {err:?}");
+            self.metrics.event_fails.inc();
+            return;
+        }
+
         self.handle_queue().unwrap_or_else(|err| {
             error!("pmem: {err:?}");
             self.metrics.event_fails.inc();
@@ -458,6 +515,7 @@ mod tests {
             path_on_host: "not_a_path".into(),
             root_device: true,
             read_only: false,
+            ..Default::default()
         };
         assert!(matches!(
             Pmem::new(config).unwrap_err(),
@@ -471,6 +529,7 @@ mod tests {
             path_on_host: dummy_path.clone(),
             root_device: true,
             read_only: false,
+            ..Default::default()
         };
         assert!(matches!(
             Pmem::new(config).unwrap_err(),
@@ -483,6 +542,7 @@ mod tests {
             path_on_host: dummy_path,
             root_device: true,
             read_only: false,
+            ..Default::default()
         };
         Pmem::new(config).unwrap();
     }
@@ -497,6 +557,7 @@ mod tests {
             path_on_host: dummy_path,
             root_device: true,
             read_only: false,
+            ..Default::default()
         };
         let mut pmem = Pmem::new(config).unwrap();