feat: add rootfs usage rate limits and observability (#2349)

Author: dobrac

packages/orchestrator/pkg/sandbox/fc/client.go

@@ -205,19 +205,20 @@ func (c *apiClient) setBootSource(ctx context.Context, kernelArgs string, kernel
 	return err
 }
 
-func (c *apiClient) setRootfsDrive(ctx context.Context, rootfsPath string, ioEngine *string) error {
-	rootfs := "rootfs"
+func (c *apiClient) setRootfsDrive(ctx context.Context, rootfsPath string, ioEngine *string, rateLimiter *models.RateLimiter) error {
+	driveID := rootfsDriveID
 
 	isRootDevice := true
 	driversConfig := operations.PutGuestDriveByIDParams{
 		Context: ctx,
-		DriveID: rootfs,
+		DriveID: driveID,
 		Body: &models.Drive{
-			DriveID:      &rootfs,
+			DriveID:      &driveID,
 			PathOnHost:   rootfsPath,
 			IsRootDevice: &isRootDevice,
 			IsReadOnly:   false,
 			IoEngine:     ioEngine,
+			RateLimiter:  rateLimiter,
 		},
 	}
 
@@ -248,10 +249,10 @@ func buildTokenBucket(b TokenBucketConfig) *models.TokenBucket {
 	return bucket
 }
 
-// buildTxRateLimiter constructs a Firecracker RateLimiter from a TxRateLimiterConfig.
+// buildRateLimiter constructs a Firecracker RateLimiter from a RateLimiterConfig.
 // Either bucket is omitted when its BucketSize is < 0.
 // Returns nil only when both buckets are disabled.
-func buildTxRateLimiter(config TxRateLimiterConfig) *models.RateLimiter {
+func buildRateLimiter(config RateLimiterConfig) *models.RateLimiter {
 	ops := buildTokenBucket(config.Ops)
 	bw := buildTokenBucket(config.Bandwidth)
 
@@ -266,8 +267,8 @@ func buildTxRateLimiter(config TxRateLimiterConfig) *models.RateLimiter {
 // Both buckets are disabled when their BucketSize < 0; if all are disabled an empty
 // RateLimiter is sent to reset any limit persisted in a snapshot.
 // This always sends a PATCH so snapshot-persisted limits are overwritten.
-func (c *apiClient) setTxRateLimit(ctx context.Context, ifaceID string, config TxRateLimiterConfig) error {
-	limiter := buildTxRateLimiter(config)
+func (c *apiClient) setTxRateLimit(ctx context.Context, ifaceID string, config RateLimiterConfig) error {
+	limiter := buildRateLimiter(config)
 	if limiter == nil {
 		limiter = &models.RateLimiter{} // empty = reset
 	}
@@ -289,6 +290,33 @@ func (c *apiClient) setTxRateLimit(ctx context.Context, ifaceID string, config T
 	return nil
 }
 
+// setDriveRateLimit applies or clears a Firecracker VMM-level block device rate limit.
+// Both buckets are disabled when their BucketSize < 0; if all are disabled an empty
+// RateLimiter is sent to reset any limit persisted in a snapshot.
+// This always sends a PATCH so snapshot-persisted limits are overwritten.
+func (c *apiClient) setDriveRateLimit(ctx context.Context, driveID string, config RateLimiterConfig) error {
+	limiter := buildRateLimiter(config)
+	if limiter == nil {
+		limiter = &models.RateLimiter{} // empty = reset
+	}
+
+	params := operations.PatchGuestDriveByIDParams{
+		Context: ctx,
+		DriveID: driveID,
+		Body: &models.PartialDrive{
+			DriveID:     &driveID,
+			RateLimiter: limiter,
+		},
+	}
+
+	_, err := c.client.Operations.PatchGuestDriveByID(&params)
+	if err != nil {
+		return fmt.Errorf("error setting drive rate limit: %w", err)
+	}
+
+	return nil
+}
+
 func (c *apiClient) setNetworkInterface(ctx context.Context, ifaceID string, tapName string, tapMac string, txRateLimiter *models.RateLimiter) error {
 	networkConfig := operations.PutGuestNetworkInterfaceByIDParams{
 		Context: ctx,

packages/orchestrator/pkg/sandbox/fc/config.go

@@ -19,6 +19,8 @@ const (
 
 	SandboxRootfsFile = "rootfs.ext4"
 
+	rootfsDriveID = "rootfs"
+
 	entropyBytesSize    int64 = 1024 // 1 KB
 	entropyRefillTime   int64 = 100
 	entropyOneTimeBurst int64 = 0

packages/orchestrator/pkg/sandbox/fc/fc_metrics.go

@@ -35,6 +35,8 @@ var (
 	directionKey = attribute.Key("direction")
 	attrTX       = metric.WithAttributes(directionKey.String("tx"))
 	attrRX       = metric.WithAttributes(directionKey.String("rx"))
+	attrRead     = metric.WithAttributes(directionKey.String("read"))
+	attrWrite    = metric.WithAttributes(directionKey.String("write"))
 
 	// Counters — global totals, no sandbox_id to avoid high cardinality.
 	fcNetFails         = utils.Must(telemetry.GetCounter(fcMeter, telemetry.SandboxFCNetFails))
@@ -49,12 +51,23 @@ var (
 	// TX-only: no RX equivalent in Firecracker metrics.
 	fcNetRateLimiterEventCount = utils.Must(telemetry.GetHistogram(fcMeter, telemetry.SandboxFCNetRateLimiterEventCount))
 	fcNetRemainingReqs         = utils.Must(telemetry.GetHistogram(fcMeter, telemetry.SandboxFCNetRemainingReqs))
+
+	// Block counters.
+	fcBlockFails         = utils.Must(telemetry.GetCounter(fcMeter, telemetry.SandboxFCBlockFails))
+	fcBlockNoAvailBuffer = utils.Must(telemetry.GetCounter(fcMeter, telemetry.SandboxFCBlockNoAvailBuffer))
+
+	// Block histograms.
+	fcBlockBytes                 = utils.Must(telemetry.GetHistogram(fcMeter, telemetry.SandboxFCBlockBytes))
+	fcBlockCount                 = utils.Must(telemetry.GetHistogram(fcMeter, telemetry.SandboxFCBlockCount))
+	fcBlockRateLimiterThrottled  = utils.Must(telemetry.GetHistogram(fcMeter, telemetry.SandboxFCBlockRateLimiterThrottled))
+	fcBlockRateLimiterEventCount = utils.Must(telemetry.GetHistogram(fcMeter, telemetry.SandboxFCBlockRateLimiterEventCount))
+	fcBlockIOEngineThrottled     = utils.Must(telemetry.GetHistogram(fcMeter, telemetry.SandboxFCBlockIOEngineThrottled))
+	fcBlockRemainingReqs         = utils.Must(telemetry.GetHistogram(fcMeter, telemetry.SandboxFCBlockRemainingReqs))
 )
 
-// firecrackerNetMetrics holds the Firecracker net metrics fields we care about.
-// Firecracker serializes SharedIncMetric fields as per-flush deltas (not cumulative totals):
-// each JSON line contains the increment since the previous flush.
-// Flush interval defaults to 60 s; additional flushes are triggered by FlushMetrics API calls.
+// firecrackerNetMetrics is a subset of Firecracker's NetDeviceMetrics we export via OTEL.
+// Full metric list: https://github.com/firecracker-microvm/firecracker/blob/main/docs/metrics.md
+// Values are per-flush deltas; flush defaults to 60 s, additional flushes via FlushMetrics API.
 type firecrackerNetMetrics struct {
 	// TX
 	TxBytesCount            uint64 `json:"tx_bytes_count"`
@@ -76,13 +89,31 @@ type firecrackerNetMetrics struct {
 	TapReadFails           uint64 `json:"tap_read_fails"`
 }
 
+// firecrackerBlockMetrics is a subset of Firecracker's BlockDeviceMetrics we export via OTEL.
+// Full metric list: https://github.com/firecracker-microvm/firecracker/blob/main/docs/metrics.md
+// Values are per-flush deltas. The aggregate "block" key sums over all drives; we only have one (rootfs).
+type firecrackerBlockMetrics struct {
+	ReadBytes                  uint64 `json:"read_bytes"`
+	WriteBytes                 uint64 `json:"write_bytes"`
+	ReadCount                  uint64 `json:"read_count"`
+	WriteCount                 uint64 `json:"write_count"`
+	RateLimiterThrottledEvents uint64 `json:"rate_limiter_throttled_events"`
+	RateLimiterEventCount      uint64 `json:"rate_limiter_event_count"`
+	IOEngineThrottledEvents    uint64 `json:"io_engine_throttled_events"`
+	NoAvailBuffer              uint64 `json:"no_avail_buffer"`
+	ExecuteFails               uint64 `json:"execute_fails"`
+	EventFails                 uint64 `json:"event_fails"`
+	RemainingReqsCount         uint64 `json:"remaining_reqs_count"`
+}
+
 // firecrackerMetrics is the top-level structure of one Firecracker metrics JSON line.
 type firecrackerMetrics struct {
-	Net firecrackerNetMetrics `json:"net"`
+	Net   firecrackerNetMetrics   `json:"net"`
+	Block firecrackerBlockMetrics `json:"block"`
 }
 
 // startMetricsReader opens the metrics FIFO and starts a goroutine that reads
-// Firecracker metrics lines and exports net device metrics via OTEL.
+// Firecracker metrics lines and exports metrics via OTEL.
 // It must be called before setMetrics so that the FIFO is open for reading
 // before Firecracker opens the write end in response to PUT /metrics.
 func (p *Process) startMetricsReader(ctx context.Context) {
@@ -204,6 +235,31 @@ func (p *Process) startMetricsReader(ctx context.Context) {
 			if n.TapReadFails > 0 {
 				fcNetTapIOFails.Add(ctx, int64(n.TapReadFails), attrRX)
 			}
+
+			// Block histograms — values are already per-flush deltas from Firecracker.
+			b := &m.Block
+
+			fcBlockBytes.Record(ctx, int64(b.ReadBytes), attrRead)
+			fcBlockBytes.Record(ctx, int64(b.WriteBytes), attrWrite)
+			fcBlockCount.Record(ctx, int64(b.ReadCount), attrRead)
+			fcBlockCount.Record(ctx, int64(b.WriteCount), attrWrite)
+			fcBlockRateLimiterEventCount.Record(ctx, int64(b.RateLimiterEventCount))
+			fcBlockRemainingReqs.Record(ctx, int64(b.RemainingReqsCount))
+
+			if b.RateLimiterThrottledEvents > 0 {
+				fcBlockRateLimiterThrottled.Record(ctx, int64(b.RateLimiterThrottledEvents))
+			}
+			if b.IOEngineThrottledEvents > 0 {
+				fcBlockIOEngineThrottled.Record(ctx, int64(b.IOEngineThrottledEvents))
+			}
+
+			// Block global error/event counters.
+			if b.ExecuteFails > 0 || b.EventFails > 0 {
+				fcBlockFails.Add(ctx, int64(b.ExecuteFails)+int64(b.EventFails))
+			}
+			if b.NoAvailBuffer > 0 {
+				fcBlockNoAvailBuffer.Add(ctx, int64(b.NoAvailBuffer))
+			}
 		}
 
 		if err := scanner.Err(); err != nil {

packages/orchestrator/pkg/sandbox/fc/process.go

@@ -108,9 +108,9 @@ type TokenBucketConfig struct {
 	RefillTimeMs int64
 }
 
-// TxRateLimiterConfig holds TX rate limit parameters for a VM's network interface.
+// RateLimiterConfig holds rate limit parameters for a Firecracker device (network or block).
 // Mirrors the Firecracker RateLimiter structure: two independent token buckets.
-type TxRateLimiterConfig struct {
+type RateLimiterConfig struct {
 	Ops       TokenBucketConfig // packets; effective rate = BucketSize * 1000 / RefillTimeMs ops/s
 	Bandwidth TokenBucketConfig // bytes;   effective rate = BucketSize * 1000 / RefillTimeMs bytes/s
 }
@@ -300,7 +300,8 @@ func (p *Process) Create(
 	memoryMB int64,
 	hugePages bool,
 	options ProcessOptions,
-	txRateLimit TxRateLimiterConfig,
+	txRateLimit RateLimiterConfig,
+	driveRateLimit RateLimiterConfig,
 	cgroupFD int,
 ) error {
 	ctx, childSpan := tracer.Start(ctx, "create-fc")
@@ -405,15 +406,15 @@ func (p *Process) Create(
 	}
 	telemetry.ReportEvent(ctx, "symlinked rootfs")
 
-	err = p.client.setRootfsDrive(ctx, p.rootfsPath, options.IoEngine)
+	err = p.client.setRootfsDrive(ctx, p.rootfsPath, options.IoEngine, buildRateLimiter(driveRateLimit))
 	if err != nil {
 		fcStopErr := p.Stop(ctx)
 
 		return errors.Join(fmt.Errorf("error setting fc drivers config: %w", err), fcStopErr)
 	}
 	telemetry.ReportEvent(ctx, "set fc drivers config")
 
-	err = p.client.setNetworkInterface(ctx, p.slot.VpeerName(), p.slot.TapName(), p.slot.TapMAC(), buildTxRateLimiter(txRateLimit))
+	err = p.client.setNetworkInterface(ctx, p.slot.VpeerName(), p.slot.TapName(), p.slot.TapMAC(), buildRateLimiter(txRateLimit))
 	if err != nil {
 		fcStopErr := p.Stop(ctx)
 
@@ -457,7 +458,8 @@ func (p *Process) Resume(
 	uffdReady chan struct{},
 	accessToken *string,
 	cgroupFD int,
-	txRateLimit TxRateLimiterConfig,
+	txRateLimit RateLimiterConfig,
+	driveRateLimit RateLimiterConfig,
 ) error {
 	ctx, span := tracer.Start(ctx, "resume-fc")
 	defer span.End()
@@ -551,15 +553,22 @@ func (p *Process) Resume(
 		return errors.Join(fmt.Errorf("error loading snapshot: %w", err), fcStopErr)
 	}
 
-	// Always apply/reset the TX rate limit before resuming so any rate limit
-	// persisted in the snapshot is overwritten by the current config.
+	// Always apply/reset rate limits before resuming so any limits
+	// persisted in the snapshot are overwritten by the current config.
 	if setErr := p.client.setTxRateLimit(ctx, p.slot.VpeerName(), txRateLimit); setErr != nil {
 		fcStopErr := p.Stop(ctx)
 
 		return errors.Join(fmt.Errorf("error setting TX rate limit: %w", setErr), fcStopErr)
 	}
 	telemetry.ReportEvent(ctx, "configured tx rate limit")
 
+	if setErr := p.client.setDriveRateLimit(ctx, rootfsDriveID, driveRateLimit); setErr != nil {
+		fcStopErr := p.Stop(ctx)
+
+		return errors.Join(fmt.Errorf("error setting drive rate limit: %w", setErr), fcStopErr)
+	}
+	telemetry.ReportEvent(ctx, "configured drive rate limit")
+
 	err = p.client.resumeVM(ctx)
 	if err != nil {
 		fcStopErr := p.Stop(ctx)

packages/orchestrator/pkg/sandbox/sandbox.go

@@ -430,6 +430,7 @@ func (f *Factory) CreateSandbox(
 	}
 
 	throttleConfig := featureflags.GetTCPFirewallEgressThrottleConfig(ctx, f.featureFlags)
+	driveThrottleConfig := featureflags.GetBlockDriveThrottleConfig(ctx, f.featureFlags)
 
 	telemetry.ReportEvent(ctx, "created fc client")
 
@@ -500,10 +501,14 @@ func (f *Factory) CreateSandbox(
 		config.RamMB,
 		config.HugePages,
 		processOptions,
-		fc.TxRateLimiterConfig{
+		fc.RateLimiterConfig{
 			Ops:       fc.TokenBucketConfig(throttleConfig.Ops),
 			Bandwidth: fc.TokenBucketConfig(throttleConfig.Bandwidth),
 		},
+		fc.RateLimiterConfig{
+			Ops:       fc.TokenBucketConfig(driveThrottleConfig.Ops),
+			Bandwidth: fc.TokenBucketConfig(driveThrottleConfig.Bandwidth),
+		},
 		cgroupFD,
 	)
 	if err != nil {
@@ -756,6 +761,7 @@ func (f *Factory) ResumeSandbox(
 	}
 
 	resumeThrottleConfig := featureflags.GetTCPFirewallEgressThrottleConfig(ctx, f.featureFlags)
+	resumeDriveThrottleConfig := featureflags.GetBlockDriveThrottleConfig(ctx, f.featureFlags)
 
 	telemetry.ReportEvent(ctx, "created FC process")
 
@@ -861,10 +867,14 @@ func (f *Factory) ResumeSandbox(
 		fcUffd.Ready(),
 		config.Envd.AccessToken,
 		cgroupFD,
-		fc.TxRateLimiterConfig{
+		fc.RateLimiterConfig{
 			Ops:       fc.TokenBucketConfig(resumeThrottleConfig.Ops),
 			Bandwidth: fc.TokenBucketConfig(resumeThrottleConfig.Bandwidth),
 		},
+		fc.RateLimiterConfig{
+			Ops:       fc.TokenBucketConfig(resumeDriveThrottleConfig.Ops),
+			Bandwidth: fc.TokenBucketConfig(resumeDriveThrottleConfig.Bandwidth),
+		},
 	)
 
 	if fcStartErr != nil {

packages/shared/pkg/featureflags/flags.go

@@ -352,10 +352,8 @@ type TCPFirewallEgressThrottleConfigValue struct {
 	Bandwidth TokenBucketConfig
 }
 
-// GetTCPFirewallEgressThrottleConfig fetches and parses the TCPFirewallEgressThrottleConfig flag.
-func GetTCPFirewallEgressThrottleConfig(ctx context.Context, ff *Client) TCPFirewallEgressThrottleConfigValue {
-	value := ff.JSONFlag(ctx, TCPFirewallEgressThrottleConfig)
-
+// parseThrottleBuckets parses "ops" and "bandwidth" token bucket configs from a JSON flag value.
+func parseThrottleBuckets(value ldvalue.Value) (ops, bandwidth TokenBucketConfig) {
 	parseBucket := func(key string) TokenBucketConfig {
 		b := value.GetByKey(key)
 		if b.IsNull() {
@@ -375,8 +373,45 @@ func GetTCPFirewallEgressThrottleConfig(ctx context.Context, ff *Client) TCPFire
 		}
 	}
 
+	return parseBucket("ops"), parseBucket("bandwidth")
+}
+
+// GetTCPFirewallEgressThrottleConfig fetches and parses the TCPFirewallEgressThrottleConfig flag.
+func GetTCPFirewallEgressThrottleConfig(ctx context.Context, ff *Client) TCPFirewallEgressThrottleConfigValue {
+	value := ff.JSONFlag(ctx, TCPFirewallEgressThrottleConfig)
+	ops, bw := parseThrottleBuckets(value)
+
 	return TCPFirewallEgressThrottleConfigValue{
-		Ops:       parseBucket("ops"),
-		Bandwidth: parseBucket("bandwidth"),
+		Ops:       ops,
+		Bandwidth: bw,
+	}
+}
+
+// BlockDriveThrottleConfig controls per-sandbox block device (disk) throttling via Firecracker's
+// VMM-level token bucket rate limiters on the rootfs drive.
+// Structure mirrors the Firecracker RateLimiter API: two independent token buckets.
+// Set bucketSize to -1 to disable a bucket.
+//
+// Ops bucket (IOPS):       effective rate = ops.bucketSize * 1000 / ops.refillTimeMs ops/s.
+// Bandwidth bucket (bytes): effective rate = bandwidth.bucketSize * 1000 / bandwidth.refillTimeMs bytes/s.
+var BlockDriveThrottleConfig = newJSONFlag("block-drive-throttle-config", ldvalue.FromJSONMarshal(map[string]any{
+	"ops":       map[string]any{"bucketSize": -1, "oneTimeBurst": 0, "refillTimeMs": 1000},
+	"bandwidth": map[string]any{"bucketSize": -1, "oneTimeBurst": 0, "refillTimeMs": 1000},
+}))
+
+// BlockDriveThrottleConfigValue holds the parsed values of BlockDriveThrottleConfig.
+type BlockDriveThrottleConfigValue struct {
+	Ops       TokenBucketConfig
+	Bandwidth TokenBucketConfig
+}
+
+// GetBlockDriveThrottleConfig fetches and parses the BlockDriveThrottleConfig flag.
+func GetBlockDriveThrottleConfig(ctx context.Context, ff *Client) BlockDriveThrottleConfigValue {
+	value := ff.JSONFlag(ctx, BlockDriveThrottleConfig)
+	ops, bw := parseThrottleBuckets(value)
+
+	return BlockDriveThrottleConfigValue{
+		Ops:       ops,
+		Bandwidth: bw,
 	}
 }