sing-box/service/oomkiller/timer.go

package oomkiller

import (
	"context"
	"runtime"
	runtimeDebug "runtime/debug"
	"runtime/metrics"
	"sync"
	"sync/atomic"
	"time"

	"github.com/sagernet/sing-box/adapter"
	"github.com/sagernet/sing-box/log"
	"github.com/sagernet/sing-box/option"
	"github.com/sagernet/sing-tun"
	"github.com/sagernet/sing/common/byteformats"
	E "github.com/sagernet/sing/common/exceptions"
	"github.com/sagernet/sing/common/memory"
)

const (
	defaultMinInterval               = 100 * time.Millisecond
	defaultMaxInterval               = 10 * time.Second
	defaultReleaseInterval           = time.Second
	defaultReleaseShrinkFloor        = 16
	defaultReleaseShrinkDivisor      = 10
	defaultMaxRateLookahead          = time.Second
	defaultSafetyMargin              = 5 * 1024 * 1024
	defaultAvailableTriggerMarginMin = 32 * 1024 * 1024
	defaultAvailableTriggerMarginMax = 128 * 1024 * 1024
)

type pressureState uint8

const (
	pressureStateNormal pressureState = iota
	pressureStateArmed
	pressureStateTriggered
)

type memorySample struct {
	usage          uint64
	available      uint64
	availableKnown bool
}

type pressureThresholds struct {
	trigger uint64
	armed   uint64
	resume  uint64
}

type timerConfig struct {
	memoryLimit     uint64
	safetyMargin    uint64
	hasSafetyMargin bool
	minInterval     time.Duration
	maxInterval     time.Duration
	policyMode      policyMode
	killerDisabled  bool
}

func buildTimerConfig(options option.OOMKillerServiceOptions, memoryLimit uint64, policyMode policyMode, killerDisabled bool) (timerConfig, error) {
	minInterval := defaultMinInterval
	if options.MinInterval != 0 {
		minInterval = time.Duration(options.MinInterval.Build())
		if minInterval <= 0 {
			return timerConfig{}, E.New("min_interval must be greater than 0")
		}
	}

	maxInterval := defaultMaxInterval
	if options.MaxInterval != 0 {
		maxInterval = time.Duration(options.MaxInterval.Build())
		if maxInterval <= 0 {
			return timerConfig{}, E.New("max_interval must be greater than 0")
		}
	}
	if maxInterval < minInterval {
		return timerConfig{}, E.New("max_interval must be greater than or equal to min_interval")
	}

	var (
		safetyMargin    uint64
		hasSafetyMargin bool
	)
	if options.SafetyMargin != nil && options.SafetyMargin.Value() > 0 {
		safetyMargin = options.SafetyMargin.Value()
		hasSafetyMargin = true
	} else if memoryLimit > 0 {
		safetyMargin = defaultSafetyMargin
		hasSafetyMargin = true
	}

	return timerConfig{
		memoryLimit:     memoryLimit,
		safetyMargin:    safetyMargin,
		hasSafetyMargin: hasSafetyMargin,
		minInterval:     minInterval,
		maxInterval:     maxInterval,
		policyMode:      policyMode,
		killerDisabled:  killerDisabled,
	}, nil
}

type timerState struct {
	state                   pressureState
	currentInterval         time.Duration
	forceMinInterval        bool
	pendingPressureBaseline bool
	pressureBaseline        memorySample
	pressureBaselineTime    time.Time
}

type adaptiveTimer struct {
	timerConfig
	logger          log.ContextLogger
	network         adapter.NetworkManager
	connections     adapter.ConnectionManager
	cacheFile       adapter.CacheFile
	recorder        *Recorder
	pressure        *atomic.Uint32
	limitThresholds pressureThresholds

	access          sync.Mutex
	timer           *time.Timer
	lastGoroutines  int
	lastConnections int
	lastGCCycles    uint64
	releaseGCCycles uint64
	lastRelease     time.Time
	timerState
}

func newAdaptiveTimer(logger log.ContextLogger, network adapter.NetworkManager, connections adapter.ConnectionManager, cacheFile adapter.CacheFile, recorder *Recorder, pressure *atomic.Uint32, config timerConfig) *adaptiveTimer {
	t := &adaptiveTimer{
		timerConfig: config,
		logger:      logger,
		network:     network,
		connections: connections,
		cacheFile:   cacheFile,
		recorder:    recorder,
		pressure:    pressure,
	}
	if config.policyMode == policyModeMemoryLimit || config.policyMode == policyModeNetworkExtension {
		t.limitThresholds = computeLimitThresholds(config.memoryLimit, config.safetyMargin)
	}
	return t
}

func (t *adaptiveTimer) start(carriedState *timerState) {
	t.access.Lock()
	defer t.access.Unlock()
	if t.timer != nil {
		return
	}
	if carriedState != nil {
		t.timerState = *carriedState
		t.timer = time.AfterFunc(t.minInterval, t.poll)
		return
	}
	t.startLocked()
}

func (t *adaptiveTimer) startLocked() {
	if t.timer != nil {
		return
	}
	t.state = pressureStateNormal
	t.forceMinInterval = false
	t.timer = time.AfterFunc(t.minInterval, t.poll)
}

func (t *adaptiveTimer) stop() timerState {
	t.access.Lock()
	defer t.access.Unlock()
	if t.timer != nil {
		t.timer.Stop()
		t.timer = nil
	}
	return t.timerState
}

func (t *adaptiveTimer) poll() {
	var triggered bool
	var rateTriggered bool
	sample := readMemorySample(t.policyMode)

	t.access.Lock()
	if t.timer == nil {
		t.access.Unlock()
		return
	}
	if t.pendingPressureBaseline {
		t.pressureBaseline = sample
		t.pressureBaselineTime = time.Now()
		t.pendingPressureBaseline = false
	}
	previousState := t.state
	t.state = t.nextState(sample)
	if t.state == pressureStateNormal {
		t.forceMinInterval = false
		if !t.pressureBaselineTime.IsZero() && time.Since(t.pressureBaselineTime) > t.maxInterval {
			t.pressureBaselineTime = time.Time{}
		}
	}
	interval := t.intervalForState()
	t.timer.Reset(interval)
	triggered = previousState != pressureStateTriggered && t.state == pressureStateTriggered
	if !triggered && !t.pressureBaselineTime.IsZero() && t.memoryLimit > 0 &&
		sample.usage > t.pressureBaseline.usage && sample.usage < t.memoryLimit {
		elapsed := time.Since(t.pressureBaselineTime)
		if elapsed >= t.minInterval/2 {
			growth := sample.usage - t.pressureBaseline.usage
			ratePerSecond := float64(growth) / elapsed.Seconds()
			headroom := t.memoryLimit - sample.usage
			secondsUntilLimit := float64(headroom) / ratePerSecond
			lookahead := min(2*interval, defaultMaxRateLookahead)
			if secondsUntilLimit < lookahead.Seconds() {
				triggered = true
				rateTriggered = true
				t.state = pressureStateTriggered
			}
		}
	}
	state := t.state
	t.access.Unlock()
	t.pressure.Store(uint32(state.memoryPressure()))
	var connections int
	if t.connections != nil {
		connections = t.connections.Count()
	}
	if t.recorder != nil {
		t.recorder.sample(sample, state, connections)
		if state != previousState {
			t.recorder.recordStateChange(state, sample)
		}
	}
	if !triggered {
		t.releaseIdleMemory(sample, connections)
		return
	}
	var reason string
	if rateTriggered {
		reason = resetReasonRate
		if t.killerDisabled {
			t.logger.Warn("memory growth rate critical (report only), usage: ", byteformats.FormatMemoryBytes(sample.usage), t.logDetails(sample))
		} else {
			t.logger.Error("memory growth rate critical, usage: ", byteformats.FormatMemoryBytes(sample.usage), t.logDetails(sample), ", resetting network")
			t.network.ReleaseMemory(context.Background())
		}
	} else {
		reason = resetReasonThreshold
		if t.killerDisabled {
			t.logger.Warn("memory threshold reached (report only), usage: ", byteformats.FormatMemoryBytes(sample.usage), t.logDetails(sample))
		} else {
			t.logger.Error("memory threshold reached, usage: ", byteformats.FormatMemoryBytes(sample.usage), t.logDetails(sample), ", resetting network")
			t.network.ReleaseMemory(context.Background())
		}
	}
	t.releaseMemory()
	if t.recorder != nil {
		after := readMemorySample(t.policyMode)
		t.recorder.recordReset(reason, sample, after, connections, t.killerDisabled)
		t.recorder.snapshot(SnapshotReasonReset, sample, t.belowTrigger(after), false)
	}
}

func (t *adaptiveTimer) releaseIdleMemory(sample memorySample, connections int) {
	goroutines := runtime.NumGoroutine()
	gcCycles := readGCCycles()
	now := time.Now()
	t.access.Lock()
	shrank := hasMeaningfulDrop(t.lastGoroutines, goroutines) || hasMeaningfulDrop(t.lastConnections, connections)
	idleGC := gcCycles == t.lastGCCycles && gcCycles != t.releaseGCCycles
	t.lastGoroutines = goroutines
	t.lastConnections = connections
	t.lastGCCycles = gcCycles
	if t.belowResume(sample) || now.Sub(t.lastRelease) < defaultReleaseInterval || (!shrank && !idleGC) {
		t.access.Unlock()
		return
	}
	t.lastRelease = now
	t.access.Unlock()
	runtimeDebug.FreeOSMemory()
	t.access.Lock()
	t.lastGCCycles = readGCCycles()
	t.releaseGCCycles = t.lastGCCycles
	t.access.Unlock()
	t.logger.Trace("released idle memory, usage: ", byteformats.FormatMemoryBytes(sample.usage), " -> ", byteformats.FormatMemoryBytes(memory.Total()))
}

func hasMeaningfulDrop(previous int, current int) bool {
	if current >= previous {
		return false
	}
	return previous-current >= max(defaultReleaseShrinkFloor, previous/defaultReleaseShrinkDivisor)
}

func readGCCycles() uint64 {
	samples := []metrics.Sample{{Name: "/gc/cycles/total:gc-cycles"}}
	metrics.Read(samples)
	if samples[0].Value.Kind() != metrics.KindUint64 {
		return 0
	}
	return samples[0].Value.Uint64()
}

func (t *adaptiveTimer) belowResume(sample memorySample) bool {
	switch t.policyMode {
	case policyModeMemoryLimit, policyModeNetworkExtension:
		return sample.usage < t.limitThresholds.resume
	case policyModeAvailable:
		return !sample.availableKnown || sample.available > t.availableThresholds(sample).resume
	default:
		return true
	}
}

func (t *adaptiveTimer) belowTrigger(sample memorySample) bool {
	switch t.policyMode {
	case policyModeMemoryLimit, policyModeNetworkExtension:
		return sample.usage < t.limitThresholds.trigger
	case policyModeAvailable:
		return !sample.availableKnown || sample.available > t.availableThresholds(sample).trigger
	default:
		return true
	}
}

func (t *adaptiveTimer) releaseMemory() {
	if t.cacheFile != nil {
		t.cacheFile.Flush()
	}
	t.pressure.Store(uint32(tun.MemoryPressureCritical))
	badCleanup()
	runtimeDebug.FreeOSMemory()
	t.access.Lock()
	t.releaseGCCycles = readGCCycles()
	t.access.Unlock()
}

func (s pressureState) memoryPressure() tun.MemoryPressure {
	switch s {
	case pressureStateArmed:
		return tun.MemoryPressureWarning
	case pressureStateTriggered:
		return tun.MemoryPressureCritical
	default:
		return tun.MemoryPressureNone
	}
}

func (t *adaptiveTimer) nextState(sample memorySample) pressureState {
	switch t.policyMode {
	case policyModeMemoryLimit, policyModeNetworkExtension:
		return nextPressureState(t.state,
			sample.usage >= t.limitThresholds.trigger,
			sample.usage >= t.limitThresholds.armed,
		)
	case policyModeAvailable:
		if !sample.availableKnown {
			return pressureStateNormal
		}
		thresholds := t.availableThresholds(sample)
		return nextPressureState(t.state,
			sample.available <= thresholds.trigger,
			sample.available <= thresholds.armed,
		)
	default:
		return pressureStateNormal
	}
}

func computeLimitThresholds(memoryLimit uint64, safetyMargin uint64) pressureThresholds {
	triggerMargin := min(safetyMargin, memoryLimit)
	armedMargin := min(triggerMargin*2, memoryLimit)
	resumeMargin := min(triggerMargin*4, memoryLimit)
	return pressureThresholds{
		trigger: memoryLimit - triggerMargin,
		armed:   memoryLimit - armedMargin,
		resume:  memoryLimit - resumeMargin,
	}
}

func RuntimeMemoryLimit(memoryLimit uint64) uint64 {
	return computeLimitThresholds(memoryLimit, defaultSafetyMargin).armed
}

func (t *adaptiveTimer) availableThresholds(sample memorySample) pressureThresholds {
	var triggerMargin uint64
	if t.hasSafetyMargin {
		triggerMargin = t.safetyMargin
	} else if sample.usage == 0 {
		triggerMargin = defaultAvailableTriggerMarginMin
	} else {
		triggerMargin = max(defaultAvailableTriggerMarginMin, min(sample.usage/4, defaultAvailableTriggerMarginMax))
	}
	return pressureThresholds{
		trigger: triggerMargin,
		armed:   triggerMargin * 2,
		resume:  triggerMargin * 4,
	}
}

func (t *adaptiveTimer) intervalForState() time.Duration {
	switch {
	case t.forceMinInterval || t.state != pressureStateNormal || !t.pressureBaselineTime.IsZero():
		t.currentInterval = t.minInterval
	default:
		if t.currentInterval == 0 {
			t.currentInterval = t.maxInterval
		} else {
			t.currentInterval = min(t.currentInterval*2, t.maxInterval)
		}
	}
	return t.currentInterval
}

func (t *adaptiveTimer) logDetails(sample memorySample) string {
	switch t.policyMode {
	case policyModeMemoryLimit, policyModeNetworkExtension:
		headroom := uint64(0)
		if sample.usage < t.memoryLimit {
			headroom = t.memoryLimit - sample.usage
		}
		return ", limit: " + byteformats.FormatMemoryBytes(t.memoryLimit) + ", headroom: " + byteformats.FormatMemoryBytes(headroom)
	case policyModeAvailable:
		if sample.availableKnown {
			return ", available: " + byteformats.FormatMemoryBytes(sample.available)
		}
	}
	return ""
}

func nextPressureState(current pressureState, shouldTrigger, shouldArm bool) pressureState {
	if current == pressureStateTriggered {
		if shouldArm {
			return pressureStateTriggered
		}
		return pressureStateNormal
	}
	if shouldTrigger {
		return pressureStateTriggered
	}
	if shouldArm {
		return pressureStateArmed
	}
	return pressureStateNormal
}

func readMemorySample(mode policyMode) memorySample {
	sample := memorySample{
		usage: memory.Total(),
	}
	if mode == policyModeAvailable || mode == policyModeNetworkExtension {
		sample.availableKnown = true
		sample.available = memory.Available()
	}
	return sample
}

func (s pressureState) String() string {
	switch s {
	case pressureStateArmed:
		return "armed"
	case pressureStateTriggered:
		return "triggered"
	default:
		return "normal"
	}
}