Skip to content

feat(sandbox): add Firecracker disk I/O rate limiting (bandwidth + IOPS) - #48

Open
rogeroger-yu wants to merge 2 commits into
kvcache-ai:mainfrom
rogeroger-yu:feat/disk-rate-limit-firecracker
Open

feat(sandbox): add Firecracker disk I/O rate limiting (bandwidth + IOPS)#48
rogeroger-yu wants to merge 2 commits into
kvcache-ai:mainfrom
rogeroger-yu:feat/disk-rate-limit-firecracker

Conversation

@rogeroger-yu

@rogeroger-yu rogeroger-yu commented Jul 28, 2026

Copy link
Copy Markdown

Summary

  • Wire Firecracker's per-drive TokenBucket rate limiter into the sandbox lifecycle
  • Configurable via [machine.disk_rate_limit] with bandwidth, IOPS, burst, and refill settings
  • Applied at both fresh-boot (add_drive) and snapshot-resume (PATCH /drives) paths

Closes #46

Changes

  • config/default.toml: Add [machine.disk_rate_limit] section (disabled by default)
  • src/cfg.rs: Add DiskRateLimitConfig struct with config-rs derive
  • src/sandbox/firecracker/config.rs: Plumb config into FirecrackerSandboxConfig
  • src/sandbox/firecracker/instance.rs: Add patch_drive_rate_limiter() method
  • src/sandbox/firecracker/sandbox.rs: Build and apply rate limiter at boot and resume

Test plan

  • Build passes (cargo build --release)
  • Baseline test: sandbox disk I/O without limiting (~1.8 GB/s write)
  • Rate-limited test: sandbox disk I/O with 100 MB/s + 3000 IOPS configured (~12-28 MB/s)
  • Resume path: rate limiter applied correctly after snapshot resume
  • Integration test with cargo test

@github-actions

github-actions Bot commented Jul 28, 2026

Copy link
Copy Markdown
Contributor

🔍 OpenCodeReview found 5 issue(s) in this PR.

  • ✅ Successfully posted inline: 5 comment(s)

Comment thread config/default.toml Outdated
Comment on lines +115 to +119
# bandwidth_bytes_per_sec = 104857600 # 100 MB/s
# bandwidth_burst_bytes = 10485760 # 10 MB burst
# iops = 3000
# iops_burst = 500
# refill_time_ms = 1000

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The *_per_sec naming/documentation is inconsistent with how this configuration is converted to Firecracker token buckets: these values are passed directly as the bucket size, while refill_time_ms controls how often that many tokens are replenished. Thus bandwidth_bytes_per_sec = 104857600 is 100 MB/s only at 1000 ms; setting 500 ms yields 200 MB/s, and iops behaves similarly. Either convert per-second rates into bucket sizes based on refill_time_ms, or rename these options to bucket-size values and document the resulting rate semantics.

Comment thread src/cfg.rs
Comment on lines +264 to +266
/// Sustained disk bandwidth limit in bytes per second (0 = unlimited).
#[config(default = 0u64)]
pub bandwidth_bytes_per_sec: u64,

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This value is described as a per-second rate, but the Firecracker integration passes it directly as the token-bucket size. The effective rate is size / refill_time, so any configured refill_time_ms other than 1000 changes the requested bytes/second (for example, 500 ms doubles it). Either remove/configure a fixed 1000 ms refill period, or convert the per-second value to a bucket size using checked arithmetic before constructing the limiter.

Comment thread src/cfg.rs Outdated
Comment on lines +276 to +278
/// Token bucket refill period in milliseconds.
#[config(default = 1000u64)]
pub refill_time_ms: u64,

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

These public values are later converted from u64 to Firecracker's i64 fields with as, but there is no configuration validation. Values above i64::MAX wrap to negative numbers, and refill_time_ms = 0 creates an invalid token bucket whenever bandwidth or IOPS limiting is configured, causing sandbox startup/restore to fail at the Firecracker API boundary. Validate enabled limiter values (refill_time_ms > 0 and every converted value <= i64::MAX) during AppConfig::validate, or use checked conversions when building the limiter.

Comment thread src/sandbox/firecracker/sandbox.rs Outdated
Comment on lines +67 to +70
let mut bw = firecracker_client::models::TokenBucket::new(
cfg.refill_time_ms as i64,
cfg.bandwidth_bytes_per_sec as i64,
);

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

TokenBucket::size is the number of tokens added per refill_time, not a per-second rate. Passing bandwidth_bytes_per_sec directly only produces the configured rate when refill_time_ms == 1000; for example, 100 MB/s with 500 ms refills becomes 200 MB/s. The IOPS bucket below has the same issue. Scale each configured per-second rate by refill_time_ms / 1000 (with checked arithmetic and a documented rounding policy), or remove the configurable refill period.

Comment thread src/sandbox/firecracker/sandbox.rs Outdated
Comment on lines +68 to +69
cfg.refill_time_ms as i64,
cfg.bandwidth_bytes_per_sec as i64,

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

These values are user-configurable u64s, and as i64 silently wraps values above i64::MAX into negative API fields. The burst and IOPS conversions have the same problem. Validate all fields with i64::try_from and propagate a configuration error; this requires making build_disk_rate_limiter return Result<Option<_>>.

@yingdi-shan
yingdi-shan requested a review from huang-jl July 29, 2026 02:01
Wire Firecracker's per-drive TokenBucket rate limiter into the sandbox
lifecycle. Configurable via [machine.disk_rate_limit] with bandwidth
(bytes/sec), IOPS caps, burst allowances, and refill interval. Applied
at both fresh-boot (add_drive) and snapshot-resume (PATCH /drives)
paths so all sandboxes are governed regardless of launch mode.
@rogeroger-yu
rogeroger-yu force-pushed the feat/disk-rate-limit-firecracker branch from 35f51d9 to ef01c5c Compare July 29, 2026 02:33
Comment thread config/default.toml Outdated
Comment on lines +115 to +119
# bandwidth_bytes_per_sec = 104857600 # 100 MB/s
# bandwidth_burst_bytes = 10485760 # 10 MB burst
# iops = 3000
# iops_burst = 500
# refill_time_ms = 1000

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

These names/comments promise per-second limits, but the implementation passes bandwidth_bytes_per_sec and iops directly as Firecracker token-bucket sizes. The effective rate is bucket size per refill_time_ms, so changing refill_time_ms from 1000 changes the configured per-second rate (for example, 500 ms doubles it). Either remove/configure a fixed 1000 ms refill interval, rename the fields as per-refill bucket sizes, or scale the bucket sizes by refill_time_ms / 1000 with overflow/range validation.

Comment thread src/sandbox/firecracker/sandbox.rs Outdated
Comment on lines +67 to +70
let mut bw = firecracker_client::models::TokenBucket::new(
cfg.refill_time_ms as i64,
cfg.bandwidth_bytes_per_sec as i64,
);

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

TokenBucket::size is the number of tokens replenished per refill_time, but this passes a per-second value unchanged. Therefore any configured refill period other than 1000 ms changes the actual sustained rate (for example, 100 MB/s with 500 ms becomes 200 MB/s). Scale the bandwidth and IOPS bucket sizes by refill_time_ms / 1000, or remove the configurable refill period. Please also use checked arithmetic/conversions because these u64 as i64 casts (and the equivalent IOPS/burst casts below) wrap values above i64::MAX into invalid negative Firecracker fields instead of returning a configuration error.

Suggestion:

Suggested change
let mut bw = firecracker_client::models::TokenBucket::new(
cfg.refill_time_ms as i64,
cfg.bandwidth_bytes_per_sec as i64,
);
let refill_time = i64::try_from(cfg.refill_time_ms)
.context("disk rate-limit refill_time_ms exceeds Firecracker's i64 range")?;
let bandwidth_size = cfg
.bandwidth_bytes_per_sec
.checked_mul(cfg.refill_time_ms)
.and_then(|value| value.checked_div(1000))
.and_then(|value| i64::try_from(value).ok())
.context("disk bandwidth rate limit is outside Firecracker's supported range")?;
let mut bw = firecracker_client::models::TokenBucket::new(refill_time, bandwidth_size);

Comment thread src/sandbox/firecracker/sandbox.rs Outdated
Comment on lines +1508 to +1513
if let Some(rl) = build_disk_rate_limiter(disk_rl_cfg) {
self.fc_instance
.patch_drive_rate_limiter(USER_ROOTFS_DRIVE_ID, rl)
.await
.context("apply disk rate limiter after resume")?;
}

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

When rate limiting is disabled (or both limits are zero), this skips the PATCH entirely. A restored snapshot can already contain the user drive's previous rate limiter, so disabling the current configuration will leave that snapshotted limiter active. The resume path should always reconcile the drive with current configuration, including explicitly patching an empty/default RateLimiter to remove an inherited limiter (assuming Firecracker's empty limiter representation), rather than treating None as “do nothing.”

Comment thread config/default.toml Outdated
# bandwidth_burst_bytes = 10485760 # 10 MB burst
# iops = 3000
# iops_burst = 500
# refill_time_ms = 1000

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

As pointed out in the review, the current configuration exposes too much complexity to users. I'm considering dropping refill_time_ms and hardcoding it to 1000ms instead.

Address review feedback on the Firecracker disk I/O rate limiter:

- Drop the configurable refill_time_ms and pin the token-bucket refill
  period to 1000 ms, so the configured *_per_sec values are the actual
  sustained per-second rates.
- Replace lossy `as i64` casts with checked conversions that surface an
  error instead of silently truncating oversized byte/IOPS values.
- Unify the fresh-boot and resume paths behind a single post-boot PATCH
  helper, and remove the now-unused rate_limiter argument from add_drive.
- Fix resume reconciliation: a restored snapshot inherits its previous
  limiter, and an empty RateLimiter PATCH is a no-op because Firecracker
  treats an absent bucket as "leave unchanged". When the current config
  disables limiting, overwrite both buckets with an effectively-unlimited
  bucket so the inherited throttle is actually cleared.

Add unit tests covering bucket construction, the disabled/zero cases,
checked-conversion overflow, and the unlimited-limiter override.
@rogeroger-yu

Copy link
Copy Markdown
Author

Follow-up: fixed resume reconciliation + e2e validation

Pushed ab2ab1f addressing the review feedback (pinned refill to 1000 ms, checked i64 conversions, unified fresh-boot/resume PATCH path, removed the unused add_drive rate_limiter arg). While running the full e2e — including the resume path — I found the original resume-clear logic was silently ineffective, so this commit also fixes it. Details below.

e2e setup

Cold-start sandbox from a public image (ghcr.io/linuxserver/baseimage-ubuntu:noble), writing to the rate-limited user rootfs drive (vdb, mounted at /) with dd ... oflag=direct. Limiter config: bandwidth_bytes_per_sec = 104857600 (100 MiB/s), burst 10 MiB, iops = 3000. A 2 GB transfer is used for the steady-state number (a 512 MB transfer is skewed by the one-time burst and doesn't reach steady state).

Results

Scenario 2 GB oflag=direct write Interpretation
Fresh boot, limiter enabled 110 MB/s (≈104.9 MiB/s) throttled to the configured rate ✅
Fresh boot, limiter disabled (control) 1.4 GB/s native storage ceiling
Resume under disabled config, old empty-PATCH clear 110 MB/s throttle persisted — bug ❌
Resume under disabled config, fixed clear 1.3 GB/s inherited limiter cleared ✅

The control run (fresh + disabled = 1.4 GB/s) rules out "110 MB/s is just the disk ceiling": the disk can clearly do >10x that, so a resumed sandbox stuck at 110 MB/s is still being throttled.

Root cause

A Firecracker snapshot persists the block device's rate-limiter state, so a resumed VM inherits whatever limiter was active when it was paused. The original code tried to clear this by PATCHing an empty RateLimiter. That is a no-op: Firecracker's PATCH /drives maps an absent token bucket to BucketUpdate::None ("leave unchanged"), so an empty limiter changes nothing and the inherited throttle survives.

I confirmed this directly against the Firecracker API socket of a resumed sandbox:

PATCH /drives/user_rootfs  {"rate_limiter":{}}                       -> HTTP 204, still 110 MB/s
PATCH /drives/user_rootfs  {"rate_limiter":{bandwidth:{size:1e12,refill_time:1}, ops:{...}}}
                                                                     -> HTTP 204, 1.1 GB/s

Fix

Since Firecracker offers no way to remove a limiter via PATCH, when the current node config disables limiting on resume we overwrite both buckets with an effectively-unlimited bucket (size = 1 TiB, refill = 1 ms). That dwarfs any real disk, so throttling no longer bites. Fresh boots are unaffected (their device model starts clean and skip the PATCH). Added a unit test for the override plus the checked-conversion cases.

Comment thread src/cfg.rs
Comment on lines +265 to +275
#[config(default = 0u64)]
pub bandwidth_bytes_per_sec: u64,
/// One-time burst allowance in bytes above the sustained bandwidth.
#[config(default = 0u64)]
pub bandwidth_burst_bytes: u64,
/// Sustained IOPS limit (0 = unlimited).
#[config(default = 0u64)]
pub iops: u64,
/// One-time burst allowance in operations above the sustained IOPS.
#[config(default = 0u64)]
pub iops_burst: u64,

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[bug · medium]
These u64 fields admit values above Firecracker's signed i64 token-bucket range. Such a configuration passes AppConfig::validate() and only fails later whenever a sandbox starts, even though it is statically invalid. Validate all four values against i64::MAX during config loading (and preferably reject a nonzero burst when its corresponding sustained limit is zero) so startup reports the bad configuration immediately.

.or_else(|| Some(DEFAULT_BOOT_ARGS.to_string())),
vcpu_count: config.machine.vcpu_count,
mem_size_mib: config.machine.mem_size_mib,
disk_rate_limit: config.machine.disk_rate_limit.clone(),

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[bug · medium]
This copied per-sandbox setting is never consumed: the launch path applies ConfigManager::global_config().machine.disk_rate_limit instead of FirecrackerSandboxConfig::disk_rate_limit. Consequently, callers of from_app_config_with_user_image using a non-global AppConfig can configure one limit here but launch with another. Pass this field into the rate-limiter application path (or remove it and explicitly require global configuration) so the constructor's configuration is honored.

Comment on lines 1323 to +1324
self.fc_instance.start().await?;
self.apply_disk_rate_limiter(false).await?;

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[security · high]
Install the limiter before starting guest execution. As written, the guest can issue unrestricted disk I/O between start() and this PATCH; if the PATCH fails, startup returns an error while Firecracker remains running without the configured isolation. Fresh drives already support Drive::rate_limiter, so pass the limiter when configuring the user-rootfs drive (or explicitly stop the VM on PATCH failure).

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think the correct way to apply the disk rate limiter for start_fresh is to configure it during configure_microvm -> add_drive(USER_ROOTFS_DRIVE_ID, ...), rather than patching it after the microVM has started.

That said, this would likely require more extensive code changes.

Our development team needs to discuss this internally before finalizing our review, I will follow up soon.

Comment on lines 1536 to +1541
self.fc_instance.resume().await?;

// A restored snapshot may carry a previously configured limiter, so
// reconcile against the node's current config, clearing any inherited
// limiter when disk rate limiting is now disabled.
self.apply_disk_rate_limiter(true).await?;

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[security · high]
Reconcile the restored drive while the VM is still paused, before calling resume(). This ordering allows a snapshot's stale/disabled limiter—or no limiter—to govern disk I/O immediately after resume, and a failed PATCH leaves a running VM even though restore is reported as failed. PATCH the loaded paused device first; if Firecracker cannot accept that state transition, add rollback that stops the VM when reconciliation fails.

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This AI review is correct. Consider moving apply_disk_rate_limiter to after load_snapshot_file but before resume().

Comment on lines +1721 to +1722
let cfg = &ConfigManager::global_config().machine.disk_rate_limit;
let rl = match build_disk_rate_limiter(cfg)? {

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[bug · medium]
This ignores FirecrackerSandboxConfig::disk_rate_limit, so a fresh sandbox created from an explicit AppConfig can be throttled using a different global setting. The newly captured launch-config field is otherwise unused. Pass the fresh launch configuration into this method (while the snapshot-resume path can intentionally use the current global configuration), so the configuration used to construct the sandbox is the one applied.

@huang-jl huang-jl Aug 4, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is also correct. The disk_rate_limit field you added to FirecrackerSandboxConfig is currently unused, it's never read in apply_disk_rate_limiter.

There are two ways to resolve this:

  • Drop the new field entirely and have apply_disk_rate_limiter read from the global configuration.
  • Actually read FirecrackerSandboxConfig.disk_rate_limit in apply_disk_rate_limiter, so it's actually consumed.

I'd prefer the second option.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

feat: support per-sandbox disk I/O rate limiting via Firecracker TokenBucket

2 participants