clitools/pkg/controller/osimage/write_throttle_linux.go
raw ยท 7914 bytes
//go:build linux
package osimage
import (
"bufio"
"context"
"fmt"
"os"
"path/filepath"
"strconv"
"strings"
"sync"
"time"
"golang.org/x/sys/unix"
)
type adaptiveWriteController struct {
mu sync.Mutex
limiter *rateLimiter
monitor *diskBusyMonitor
sampleInterval time.Duration
nextSampleAt time.Time
minBPS int64
maxBPS int64
busyHighPct float64
busyLowPct float64
}
func newAdaptiveWriteController(targetPath string) (*adaptiveWriteController, error) {
mon, err := newDiskBusyMonitor(targetPath)
if err != nil {
return nil, err
}
now := time.Now()
return &adaptiveWriteController{
limiter: newRateLimiter(defaultInitialWriteBPS, defaultBurstBytes),
monitor: mon,
sampleInterval: defaultSampleInterval,
nextSampleAt: now.Add(defaultSampleInterval),
minBPS: defaultMinWriteBPS,
maxBPS: defaultMaxWriteBPS,
busyHighPct: defaultBusyHighPct,
busyLowPct: defaultBusyLowPct,
}, nil
}
func newNoopAdaptiveWriteController() *adaptiveWriteController {
return &adaptiveWriteController{
limiter: newRateLimiter(0, 0),
sampleInterval: defaultSampleInterval,
}
}
func (c *adaptiveWriteController) Wait(ctx context.Context, n int) error {
if c == nil || c.limiter == nil {
return nil
}
return c.limiter.Wait(ctx, n)
}
func (c *adaptiveWriteController) ObserveWrite(n int) {
c.observe(false)
}
func (c *adaptiveWriteController) ObserveSync() {
c.observe(true)
}
func (c *adaptiveWriteController) observe(afterSync bool) {
if c == nil {
return
}
c.mu.Lock()
defer c.mu.Unlock()
now := time.Now()
if c.monitor == nil || now.Before(c.nextSampleAt) {
return
}
c.nextSampleAt = now.Add(c.sampleInterval)
s, err := c.monitor.Sample(now)
if err != nil {
return
}
cur := c.limiter.Rate()
if cur <= 0 {
cur = c.minBPS
}
switch {
case s.UtilPct >= c.busyHighPct || s.Await >= defaultSlowAwait || afterSync:
// Back off aggressively when the disk is obviously suffering.
next := cur / 2
if next < c.minBPS {
next = c.minBPS
}
c.limiter.SetRate(next)
case s.UtilPct <= c.busyLowPct && s.Await <= defaultFastAwait:
// Recover slowly.
next := cur + (cur / 5) // +20%
if next > c.maxBPS {
next = c.maxBPS
}
c.limiter.SetRate(next)
}
}
type rateLimiter struct {
mu sync.Mutex
rateBPS int64
burst int64
tokens float64
last time.Time
}
func newRateLimiter(rateBPS, burst int64) *rateLimiter {
now := time.Now()
if burst < 0 {
burst = 0
}
return &rateLimiter{
rateBPS: rateBPS,
burst: burst,
tokens: float64(burst),
last: now,
}
}
func (r *rateLimiter) Rate() int64 {
r.mu.Lock()
defer r.mu.Unlock()
return r.rateBPS
}
func (r *rateLimiter) SetRate(rateBPS int64) {
r.mu.Lock()
defer r.mu.Unlock()
r.refillLocked(time.Now())
r.rateBPS = rateBPS
if rateBPS <= 0 {
r.tokens = 0
r.burst = 0
return
}
// Keep burst small and fixed. Do not let burst scale with rate.
r.burst = defaultBurstBytes
if r.tokens > float64(r.burst) {
r.tokens = float64(r.burst)
}
}
func (r *rateLimiter) Wait(ctx context.Context, n int) error {
if n <= 0 {
return nil
}
remaining := n
for remaining > 0 {
r.mu.Lock()
if r.rateBPS <= 0 {
r.mu.Unlock()
return nil
}
now := time.Now()
r.refillLocked(now)
allowed := remaining
if int64(allowed) > r.burst && r.burst > 0 {
allowed = int(r.burst)
}
if allowed <= 0 {
allowed = remaining
}
if r.tokens >= float64(allowed) {
r.tokens -= float64(allowed)
r.mu.Unlock()
remaining -= allowed
continue
}
missing := float64(allowed) - r.tokens
waitDur := time.Duration(missing / float64(r.rateBPS) * float64(time.Second))
if waitDur < 5*time.Millisecond {
waitDur = 5 * time.Millisecond
}
r.mu.Unlock()
timer := time.NewTimer(waitDur)
select {
case <-ctx.Done():
timer.Stop()
return ctx.Err()
case <-timer.C:
}
}
return nil
}
func (r *rateLimiter) refillLocked(now time.Time) {
if r.rateBPS <= 0 {
r.last = now
return
}
elapsed := now.Sub(r.last)
if elapsed <= 0 {
return
}
r.tokens += elapsed.Seconds() * float64(r.rateBPS)
if r.tokens > float64(r.burst) {
r.tokens = float64(r.burst)
}
r.last = now
}
type diskBusySample struct {
UtilPct float64
Await time.Duration
}
type diskBusyMonitor struct {
major int
minor int
lastAt time.Time
lastIOMs uint64
lastWrites uint64
}
func newDiskBusyMonitor(targetPath string) (*diskBusyMonitor, error) {
major, minor, err := resolveWholeDiskMajorMinor(targetPath)
if err != nil {
return nil, err
}
ioMs, writes, err := readDiskStats(major, minor)
if err != nil {
return nil, err
}
return &diskBusyMonitor{
major: major,
minor: minor,
lastAt: time.Now(),
lastIOMs: ioMs,
lastWrites: writes,
}, nil
}
func (m *diskBusyMonitor) Sample(now time.Time) (diskBusySample, error) {
ioMs, writes, err := readDiskStats(m.major, m.minor)
if err != nil {
return diskBusySample{}, err
}
elapsedMs := now.Sub(m.lastAt).Milliseconds()
if elapsedMs <= 0 {
return diskBusySample{}, nil
}
deltaIOMs := int64(ioMs - m.lastIOMs)
deltaWrites := int64(writes - m.lastWrites)
m.lastAt = now
m.lastIOMs = ioMs
m.lastWrites = writes
util := float64(deltaIOMs) * 100 / float64(elapsedMs)
if util < 0 {
util = 0
}
if util > 100 {
util = 100
}
var await time.Duration
if deltaWrites > 0 {
await = time.Duration(deltaIOMs/int64(deltaWrites)) * time.Millisecond
}
return diskBusySample{
UtilPct: util,
Await: await,
}, nil
}
func resolveWholeDiskMajorMinor(targetPath string) (int, int, error) {
var st unix.Stat_t
if err := unix.Stat(targetPath, &st); err != nil {
return 0, 0, fmt.Errorf("stat target %q: %w", targetPath, err)
}
if st.Mode&unix.S_IFMT != unix.S_IFBLK {
return 0, 0, fmt.Errorf("target %q is not a block device", targetPath)
}
major := int(unix.Major(uint64(st.Rdev)))
minor := int(unix.Minor(uint64(st.Rdev)))
sysfsPath := fmt.Sprintf("/sys/dev/block/%d:%d", major, minor)
resolved, err := filepath.EvalSymlinks(sysfsPath)
if err != nil {
return major, minor, nil
}
// Partition path usually looks like .../block/sda/sda3
// Parent whole disk is .../block/sda
parent := filepath.Dir(resolved)
devName := filepath.Base(parent)
ueventPath := filepath.Join(parent, "dev")
data, err := os.ReadFile(ueventPath)
if err != nil {
return major, minor, nil
}
parts := strings.Split(strings.TrimSpace(string(data)), ":")
if len(parts) != 2 {
return major, minor, nil
}
parentMajor, err1 := strconv.Atoi(parts[0])
parentMinor, err2 := strconv.Atoi(parts[1])
if err1 != nil || err2 != nil || devName == "" {
return major, minor, nil
}
return parentMajor, parentMinor, nil
}
func readDiskStats(major, minor int) (ioMs uint64, writesCompleted uint64, err error) {
f, err := os.Open("/proc/diskstats")
if err != nil {
return 0, 0, fmt.Errorf("open /proc/diskstats: %w", err)
}
defer f.Close()
sc := bufio.NewScanner(f)
for sc.Scan() {
line := strings.Fields(sc.Text())
if len(line) < 14 {
continue
}
maj, err := strconv.Atoi(line[0])
if err != nil {
continue
}
min, err := strconv.Atoi(line[1])
if err != nil {
continue
}
if maj != major || min != minor {
continue
}
// writes completed successfully: field 5, index 4
writesCompleted, err = strconv.ParseUint(line[4], 10, 64)
if err != nil {
return 0, 0, fmt.Errorf("parse writes completed for %d:%d: %w", major, minor, err)
}
// time spent doing I/Os (ms): field 13, index 12
ioMs, err = strconv.ParseUint(line[12], 10, 64)
if err != nil {
return 0, 0, fmt.Errorf("parse io_ms for %d:%d: %w", major, minor, err)
}
return ioMs, writesCompleted, nil
}
if err := sc.Err(); err != nil {
return 0, 0, fmt.Errorf("scan /proc/diskstats: %w", err)
}
return 0, 0, fmt.Errorf("device %d:%d not found in /proc/diskstats", major, minor)
}