fancy-dns/internal/config/size.go

package config

import (
	"fmt"
	"math"
	"strconv"
	"strings"

	"gopkg.in/yaml.v3"
)

// Size is a byte count that unmarshals from a string such as "512MiB" or
// "2GiB". A bare integer is bytes. It exists so an operator can write a memory
// budget in the same units everywhere else, instead of a raw byte count that
// has to be converted by hand.
type Size int64

// Longest suffix first, so "MIB" is not shadowed by "B" and "M" does not eat
// the "iB" of "MiB".
var sizeUnits = []struct {
	suffix string
	mult   int64
}{
	{"KIB", 1 << 10},
	{"MIB", 1 << 20},
	{"GIB", 1 << 30},
	{"TIB", 1 << 40},
	{"KB", 1000},
	{"MB", 1000 * 1000},
	{"GB", 1000 * 1000 * 1000},
	{"TB", 1000 * 1000 * 1000 * 1000},
	{"B", 1},
	{"K", 1 << 10},
	{"M", 1 << 20},
	{"G", 1 << 30},
	{"T", 1 << 40},
}

// ParseSize reads a size with an optional unit suffix. It is strict about
// junk: "512MiBx" is an error, not 512MiB with a trailing letter.
func ParseSize(s string) (Size, error) {
	raw := strings.TrimSpace(s)
	if raw == "" {
		return 0, fmt.Errorf("invalid size %q: empty", s)
	}
	upper := strings.ToUpper(raw)
	for _, u := range sizeUnits {
		if !strings.HasSuffix(upper, u.suffix) {
			continue
		}
		num := strings.TrimSpace(upper[:len(upper)-len(u.suffix)])
		if num == "" {
			return 0, fmt.Errorf("invalid size %q: no number before the unit", s)
		}
		n, err := strconv.ParseFloat(num, 64)
		if err != nil {
			return 0, fmt.Errorf("invalid size %q", s)
		}
		return scaleSize(n, u.mult, s)
	}
	n, err := strconv.ParseFloat(raw, 64)
	if err != nil {
		return 0, fmt.Errorf("invalid size %q", s)
	}
	return scaleSize(n, 1, s)
}

// scaleSize applies the unit multiplier. A bare number has mult 1, which is
// why the multiplier is a parameter rather than a constant: it used to be
// forgotten on the suffixed path, so "512MiB" parsed as 512 *bytes* and every
// memory budget written in units was off by the unit's factor.
func scaleSize(n float64, mult int64, orig string) (Size, error) {
	if n < 0 {
		return 0, fmt.Errorf("invalid size %q: must not be negative", orig)
	}
	bytes := n * float64(mult)
	if bytes >= float64(math.MaxInt64) {
		return 0, fmt.Errorf("invalid size %q: out of range", orig)
	}
	return Size(int64(bytes)), nil
}

// UnmarshalYAML accepts both "512MiB" and a bare number.
func (s *Size) UnmarshalYAML(value *yaml.Node) error {
	var str string
	if err := value.Decode(&str); err != nil {
		return fmt.Errorf("invalid size: %w", err)
	}
	v, err := ParseSize(str)
	if err != nil {
		return err
	}
	*s = v
	return nil
}

// MarshalYAML writes the number back as a plain integer, so a config written
// by the API round-trips as something ParseSize still accepts.
func (s Size) MarshalYAML() (any, error) { return int64(s), nil }

// String renders the size in binary units, which is what an operator reading a
// log line wants ("512MiB" rather than "536870912").
func (s Size) String() string {
	switch {
	case s == 0:
		return "0"
	case s%(1<<30) == 0:
		return fmt.Sprintf("%dGiB", s/(1<<30))
	case s%(1<<20) == 0:
		return fmt.Sprintf("%dMiB", s/(1<<20))
	case s%(1<<10) == 0:
		return fmt.Sprintf("%dKiB", s/(1<<10))
	default:
		return fmt.Sprintf("%dB", int64(s))
	}
}

// Memory is the process-wide memory envelope.
//
// It exists because the only lever the Go runtime has before a cgroup OOM kill
// is the cgroup limit itself, and a cgroup kill is the worst possible outcome
// for a resolver: an outage, not a cache miss. GOMEMLIMIT turns the same limit
// into a soft one the GC works towards, and the cache is sized to yield before
// the GC has to start thrashing.
type Memory struct {
	// Limit is the soft heap ceiling handed to runtime/debug.SetMemoryLimit.
	// 0 means no ceiling.
	Limit Size `yaml:"limit"`
	// Reserve is subtracted from Limit before it is applied, covering what the
	// Go heap does not account for: goroutine stacks, runtime metadata, and —
	// importantly — the bbolt mmap for the cache snapshot and the query log,
	// which lives outside the Go heap entirely. A container limit with no
	// reserve is a limit the runtime will overshoot.
	Reserve Size `yaml:"reserve"`
	// CacheMaxBytes bounds the in-memory LRU by bytes as well as by entry
	// count. 0 means "derive from Limit", which is what an operator almost
	// always wants: the count budget alone does not bound the footprint,
	// because a single cached ANY or TXT answer can be tens of kilobytes.
	CacheMaxBytes Size `yaml:"cache_max_bytes"`
}

// DefaultReserve is the non-heap allowance applied when memory.reserve is not
// set. 25% is the usual rule of thumb for a process with a large mmapped file
// and a few thousand goroutines.
const DefaultReserve = Size(256 << 20)

// EffectiveLimit is the value to hand to debug.SetMemoryLimit, already reduced
// by the reserve. Zero means "no ceiling" (math.MaxInt64).
func (m Memory) EffectiveLimit() int64 {
	if m.Limit <= 0 {
		return math.MaxInt64
	}
	reserve := m.Reserve
	if reserve == 0 {
		reserve = DefaultReserve
	}
	limit := int64(m.Limit) - int64(reserve)
	if limit < 64<<20 {
		// A limit below the runtime's own floor just guarantees a death
		// spiral. Report the floor rather than the literal subtraction.
		return 64 << 20
	}
	return limit
}

// EffectiveCacheBytes is the byte budget for the LRU: the configured value, or
// 60% of the effective limit when unset. The remainder belongs to the GC, the
// query log's in-memory ring, the blocklist sets and the per-query working set.
func (m Memory) EffectiveCacheBytes() int64 {
	if m.CacheMaxBytes > 0 {
		return int64(m.CacheMaxBytes)
	}
	limit := m.EffectiveLimit()
	if limit == math.MaxInt64 {
		return 0 // unbounded
	}
	return limit * 3 / 5
}