fix: resolve TODO to preserve TokenBucket budget during updates

Also remove some stale TODOs and update docs.

Signed-off-by: Aaron Ang <aaron.angyd@gmail.com>

Author: aaron-ang

CHANGELOG.md

@@ -18,6 +18,11 @@ and this project adheres to
 
 ### Fixed
 
+- [#5594](https://github.com/firecracker-microvm/firecracker/pull/5594): Fixed
+  rate limiter bucket updates to preserve budget and one-time burst state
+  (capped at new bucket limits if smaller) instead of resetting to full
+  capacity.
+
 ## [v1.14.0]
 
 ### Added

docs/api_requests/patch-network-interface.md

@@ -63,6 +63,11 @@ found in our [OpenAPI spec](../../src/firecracker/swagger/firecracker.yaml).
 > The data provided for the update is merged with the existing data. In the
 > above example, the RX rate limit is updated, but the TX rate limit remains
 > unchanged.
+>
+> When updating rate limiter parameters, the current budget and one-time burst
+> state are preserved (capped at the new bucket's limits if the new bucket is
+> smaller), ensuring rate limiting continues smoothly without resetting to full
+> capacity.
 
 ## Removing Rate Limiting
 

src/acpi-tables/src/aml.rs

@@ -160,8 +160,6 @@ pub struct Name {
 
 impl Aml for Name {
     fn append_aml_bytes(&self, bytes: &mut Vec<u8>) -> Result<(), AmlError> {
-        // TODO: Refactor this to make more efficient but there are
-        // lifetime/ownership challenges.
         bytes.extend_from_slice(&self.bytes);
         Ok(())
     }

src/vmm/src/devices/virtio/vsock/unix/muxer.rs

@@ -586,10 +586,6 @@ impl VsockMuxer {
 
     /// Allocate a host-side port to be assigned to a new host-initiated connection.
     fn allocate_local_port(&mut self) -> u32 {
-        // TODO: this doesn't seem very space-efficient.
-        // Mybe rewrite this to limit port range and use a bitmap?
-        //
-
         loop {
             self.local_port_last = (self.local_port_last + 1) & !(1 << 31) | (1 << 30);
             if self.local_port_set.insert(self.local_port_last) {

src/vmm/src/rate_limiter/mod.rs

@@ -474,21 +474,42 @@ impl RateLimiter {
         }
     }
 
-    /// Updates the parameters of the token buckets associated with this RateLimiter.
-    // TODO: Please note that, right now, the buckets become full after being updated.
+    /// Updates the parameters of the token buckets associated with this `RateLimiter`.
+    ///
+    /// When updating bucket parameters, the existing state (`budget`, `one_time_burst`, and
+    /// `last_update` timestamp) is preserved from the current bucket. The `budget` and
+    /// `one_time_burst` values are capped at the new bucket's limits if the new limits
+    /// are smaller than the existing values.
     pub fn update_buckets(&mut self, bytes: BucketUpdate, ops: BucketUpdate) {
         match bytes {
             BucketUpdate::Disabled => self.bandwidth = None,
-            BucketUpdate::Update(tb) => self.bandwidth = Some(tb),
+            BucketUpdate::Update(tb) => {
+                let updated = Self::update_bucket(tb, self.bandwidth.as_ref());
+                self.bandwidth = Some(updated);
+            }
             BucketUpdate::None => (),
         };
         match ops {
             BucketUpdate::Disabled => self.ops = None,
-            BucketUpdate::Update(tb) => self.ops = Some(tb),
+            BucketUpdate::Update(tb) => {
+                let updated = Self::update_bucket(tb, self.ops.as_ref());
+                self.ops = Some(updated);
+            }
             BucketUpdate::None => (),
         };
     }
 
+    /// Updates a token bucket with preserved state from an existing bucket.
+    fn update_bucket(mut new_bucket: TokenBucket, existing: Option<&TokenBucket>) -> TokenBucket {
+        if let Some(existing) = existing {
+            new_bucket.budget = std::cmp::min(existing.budget, new_bucket.size);
+            new_bucket.one_time_burst =
+                std::cmp::min(existing.one_time_burst, new_bucket.initial_one_time_burst);
+            new_bucket.last_update = existing.last_update;
+        }
+        new_bucket
+    }
+
     /// Returns an immutable view of the inner bandwidth token bucket.
     pub fn bandwidth(&self) -> Option<&TokenBucket> {
         self.bandwidth.as_ref()
@@ -1167,20 +1188,112 @@ pub(crate) mod tests {
             BucketUpdate::Update(new_ops.clone()),
         );
 
-        // We have manually adjust the last_update field, because it changes when update_buckets()
-        // constructs new buckets (and thus gets a different value for last_update). We do this so
-        // it makes sense to test the following assertions.
-        x.bandwidth.as_mut().unwrap().last_update = new_bw.last_update;
-        x.ops.as_mut().unwrap().last_update = new_ops.last_update;
+        // Verify that new bucket parameters are applied
+        assert_eq!(x.bandwidth.as_ref().unwrap().capacity(), new_bw.capacity());
+        assert_eq!(
+            x.bandwidth.as_ref().unwrap().refill_time_ms(),
+            new_bw.refill_time_ms()
+        );
+        assert_eq!(
+            x.bandwidth.as_ref().unwrap().initial_one_time_burst(),
+            new_bw.initial_one_time_burst()
+        );
+        assert_eq!(x.ops.as_ref().unwrap().capacity(), new_ops.capacity());
+        assert_eq!(
+            x.ops.as_ref().unwrap().refill_time_ms(),
+            new_ops.refill_time_ms()
+        );
+        assert_eq!(
+            x.ops.as_ref().unwrap().initial_one_time_burst(),
+            new_ops.initial_one_time_burst()
+        );
 
-        assert_eq!(x.bandwidth, Some(new_bw));
-        assert_eq!(x.ops, Some(new_ops));
+        // Verify that state is preserved (budget is capped at new bucket size if smaller)
+        assert_eq!(
+            x.bandwidth.as_ref().unwrap().budget(),
+            std::cmp::min(initial_bw.as_ref().unwrap().budget(), new_bw.capacity())
+        );
+        assert_eq!(
+            x.bandwidth.as_ref().unwrap().one_time_burst(),
+            std::cmp::min(
+                initial_bw.as_ref().unwrap().one_time_burst(),
+                new_bw.initial_one_time_burst()
+            )
+        );
+        assert_eq!(
+            x.bandwidth.as_ref().unwrap().get_last_update(),
+            initial_bw.as_ref().unwrap().get_last_update()
+        );
+        assert_eq!(
+            x.ops.as_ref().unwrap().budget(),
+            std::cmp::min(initial_ops.as_ref().unwrap().budget(), new_ops.capacity())
+        );
+        assert_eq!(
+            x.ops.as_ref().unwrap().one_time_burst(),
+            std::cmp::min(
+                initial_ops.as_ref().unwrap().one_time_burst(),
+                new_ops.initial_one_time_burst()
+            )
+        );
+        assert_eq!(
+            x.ops.as_ref().unwrap().get_last_update(),
+            initial_ops.as_ref().unwrap().get_last_update()
+        );
 
         x.update_buckets(BucketUpdate::Disabled, BucketUpdate::Disabled);
         assert_eq!(x.bandwidth, None);
         assert_eq!(x.ops, None);
     }
 
+    #[test]
+    fn test_update_buckets_preserves_budget_after_consumption() {
+        let mut x = RateLimiter::new(1000, 0, 1000, 500, 0, 1000).unwrap();
+
+        // Consume some tokens to reduce budget
+        assert!(x.consume(300, TokenType::Bytes));
+        assert!(x.consume(200, TokenType::Ops));
+
+        // Verify budget was reduced
+        assert_eq!(x.bandwidth.as_ref().unwrap().budget(), 700);
+        assert_eq!(x.ops.as_ref().unwrap().budget(), 300);
+
+        // Save state before update
+        let initial_bw = x.bandwidth.as_ref().unwrap().clone();
+        let initial_ops = x.ops.as_ref().unwrap().clone();
+
+        // Update buckets with new parameters
+        let new_bw = TokenBucket::new(2000, 0, 500).unwrap();
+        let new_ops = TokenBucket::new(800, 0, 500).unwrap();
+        x.update_buckets(
+            BucketUpdate::Update(new_bw.clone()),
+            BucketUpdate::Update(new_ops.clone()),
+        );
+
+        // Verify new parameters are applied
+        assert_eq!(x.bandwidth.as_ref().unwrap().capacity(), new_bw.capacity());
+        assert_eq!(x.ops.as_ref().unwrap().capacity(), new_ops.capacity());
+
+        // Verify budget is preserved (not reset to full)
+        assert_eq!(x.bandwidth.as_ref().unwrap().budget(), initial_bw.budget());
+        assert_eq!(x.ops.as_ref().unwrap().budget(), initial_ops.budget());
+
+        // Test edge case: update with smaller size - budget should be capped
+        let mut x2 = RateLimiter::new(1000, 0, 1000, 500, 0, 1000).unwrap();
+        assert!(x2.consume(300, TokenType::Bytes));
+        assert_eq!(x2.bandwidth.as_ref().unwrap().budget(), 700);
+
+        // Update with smaller bucket size
+        let smaller_bw = TokenBucket::new(500, 0, 1000).unwrap();
+        x2.update_buckets(BucketUpdate::Update(smaller_bw.clone()), BucketUpdate::None);
+
+        // Budget should be capped at new size
+        assert_eq!(x2.bandwidth.as_ref().unwrap().budget(), smaller_bw.budget());
+        assert_eq!(
+            x2.bandwidth.as_ref().unwrap().capacity(),
+            smaller_bw.capacity()
+        );
+    }
+
     #[test]
     fn test_rate_limiter_debug() {
         let l = RateLimiter::new(1, 2, 3, 4, 5, 6).unwrap();

src/vmm/src/utils/net/mac.rs

@@ -84,8 +84,6 @@ impl MacAddr {
     /// * `src` - slice from which to copy MAC address content.
     #[inline]
     pub fn from_bytes_unchecked(src: &[u8]) -> MacAddr {
-        // TODO: using something like std::mem::uninitialized could avoid the extra initialization,
-        // if this ever becomes a performance bottleneck.
         let mut bytes = [0u8; MAC_ADDR_LEN as usize];
         bytes[..].copy_from_slice(src);