package resources import ( "context" "fmt" "os" "sync" "time" "github.com/datarhei/core/v16/log" "github.com/datarhei/core/v16/psutil" ) type Info struct { Mem MemoryInfo CPU CPUInfo } type MemoryInfo struct { Total uint64 // bytes Available uint64 // bytes Used uint64 // bytes Limit uint64 // bytes Core uint64 // bytes Throttling bool Error error } type CPUInfo struct { NCPU float64 // number of cpus System float64 // percent 0-100 User float64 // percent 0-100 Idle float64 // percent 0-100 Other float64 // percent 0-100 Limit float64 // percent 0-100 Core float64 // percent 0-100 Throttling bool Error error } type resources struct { psutil psutil.Util ncpu float64 maxCPU float64 // percent 0-100*ncpu maxMemory uint64 // bytes isUnlimited bool isCPULimiting bool isMemoryLimiting bool self psutil.Process cancelObserver context.CancelFunc lock sync.RWMutex startOnce sync.Once stopOnce sync.Once logger log.Logger } type Resources interface { Start() Stop() // HasLimits returns whether any limits have been set. HasLimits() bool // Limits returns the CPU (percent 0-100) and memory (bytes) limits. Limits() (float64, uint64) // ShouldLimit returns whether cpu and/or memory is currently limited. ShouldLimit() (bool, bool) // Request checks whether the requested resources are available. Request(cpu float64, memory uint64) error // Info returns the current resource usage Info() Info } type Config struct { MaxCPU float64 // percent 0-100 MaxMemory float64 // percent 0-100 PSUtil psutil.Util Logger log.Logger } func New(config Config) (Resources, error) { isUnlimited := false if config.MaxCPU <= 0 && config.MaxMemory <= 0 { isUnlimited = true } if config.MaxCPU <= 0 { config.MaxCPU = 100 } if config.MaxMemory <= 0 { config.MaxMemory = 100 } if config.MaxCPU > 100 || config.MaxMemory > 100 { return nil, fmt.Errorf("both MaxCPU and MaxMemory must have a range of 0-100") } r := &resources{ maxCPU: config.MaxCPU, psutil: config.PSUtil, isUnlimited: isUnlimited, logger: config.Logger, } if r.logger == nil { r.logger = log.New("") } if r.psutil == nil { r.psutil = psutil.DefaultUtil } vmstat, err := r.psutil.VirtualMemory() if err != nil { return nil, fmt.Errorf("unable to determine available memory: %w", err) } ncpu, err := r.psutil.CPUCounts(true) if err != nil { return nil, fmt.Errorf("unable to determine number of logical CPUs: %w", err) } r.ncpu = ncpu r.maxCPU *= r.ncpu r.maxMemory = uint64(float64(vmstat.Total) * config.MaxMemory / 100) r.logger = r.logger.WithFields(log.Fields{ "ncpu": r.ncpu, "max_cpu": r.maxCPU, "max_memory": r.maxMemory, }) r.self, err = psutil.NewProcess(int32(os.Getpid()), false) if err != nil { return nil, fmt.Errorf("unable to create process observer for self: %w", err) } r.logger.Debug().Log("Created") r.stopOnce.Do(func() {}) return r, nil } func (r *resources) Start() { r.startOnce.Do(func() { ctx, cancel := context.WithCancel(context.Background()) r.cancelObserver = cancel go r.observe(ctx, time.Second) r.stopOnce = sync.Once{} r.logger.Info().Log("Started") }) } func (r *resources) Stop() { r.stopOnce.Do(func() { r.cancelObserver() r.self.Stop() r.startOnce = sync.Once{} r.logger.Info().Log("Stopped") }) } func (r *resources) observe(ctx context.Context, interval time.Duration) { ticker := time.NewTicker(interval) defer ticker.Stop() r.logger.Debug().Log("Observer started") for { select { case <-ctx.Done(): return case <-ticker.C: cpustat, err := r.psutil.CPUPercent() if err != nil { r.logger.Warn().WithError(err).Log("Failed to determine system CPU usage") continue } cpuload := (cpustat.User + cpustat.System + cpustat.Other) * r.ncpu vmstat, err := r.psutil.VirtualMemory() if err != nil { r.logger.Warn().WithError(err).Log("Failed to determine system memory usage") continue } r.logger.Debug().WithFields(log.Fields{ "cur_cpu": cpuload, "cur_memory": vmstat.Used, }).Log("Observation") doCPULimit := false if !r.isUnlimited { if !r.isCPULimiting { if cpuload >= r.maxCPU { r.logger.Debug().WithField("cpu", cpuload).Log("CPU limit reached") doCPULimit = true } } else { doCPULimit = true if cpuload < r.maxCPU { r.logger.Debug().WithField("cpu", cpuload).Log("CPU limit released") doCPULimit = false } } } doMemoryLimit := false if !r.isUnlimited { if !r.isMemoryLimiting { if vmstat.Used >= r.maxMemory { r.logger.Debug().WithField("memory", vmstat.Used).Log("Memory limit reached") doMemoryLimit = true } } else { doMemoryLimit = true if vmstat.Used < r.maxMemory { r.logger.Debug().WithField("memory", vmstat.Used).Log("Memory limit released") doMemoryLimit = false } } } r.lock.Lock() if r.isCPULimiting != doCPULimit { r.logger.Warn().WithFields(log.Fields{ "enabled": doCPULimit, }).Log("Limiting CPU") r.isCPULimiting = doCPULimit } if r.isMemoryLimiting != doMemoryLimit { r.logger.Warn().WithFields(log.Fields{ "enabled": doMemoryLimit, }).Log("Limiting memory") r.isMemoryLimiting = doMemoryLimit } r.lock.Unlock() } } } func (r *resources) HasLimits() bool { return !r.isUnlimited } func (r *resources) Limits() (float64, uint64) { return r.maxCPU / r.ncpu, r.maxMemory } func (r *resources) ShouldLimit() (bool, bool) { r.lock.RLock() defer r.lock.RUnlock() return r.isCPULimiting, r.isMemoryLimiting } func (r *resources) Request(cpu float64, memory uint64) error { r.lock.RLock() defer r.lock.RUnlock() logger := r.logger.WithFields(log.Fields{ "req_cpu": cpu, "req_memory": memory, }) logger.Debug().Log("Request for acquiring resources") if r.isCPULimiting || r.isMemoryLimiting { logger.Debug().Log("Rejected, currently limiting") return fmt.Errorf("resources are currenlty actively limited") } if cpu <= 0 || memory == 0 { logger.Debug().Log("Rejected, invalid values") return fmt.Errorf("the cpu and/or memory values are invalid: cpu=%f, memory=%d", cpu, memory) } cpustat, err := r.psutil.CPUPercent() if err != nil { r.logger.Warn().WithError(err).Log("Failed to determine system CPU usage") return fmt.Errorf("the system CPU usage couldn't be determined") } cpuload := (cpustat.User + cpustat.System + cpustat.Other) * r.ncpu vmstat, err := r.psutil.VirtualMemory() if err != nil { r.logger.Warn().WithError(err).Log("Failed to determine system memory usage") return fmt.Errorf("the system memory usage couldn't be determined") } if cpuload+cpu > r.maxCPU { logger.Debug().WithField("cur_cpu", cpuload).Log("Rejected, CPU limit exceeded") return fmt.Errorf("the CPU limit would be exceeded: %f + %f > %f", cpuload, cpu, r.maxCPU) } if vmstat.Used+memory > r.maxMemory { logger.Debug().WithField("cur_memory", vmstat.Used).Log("Rejected, memory limit exceeded") return fmt.Errorf("the memory limit would be exceeded: %d + %d > %d", vmstat.Used, memory, r.maxMemory) } logger.Debug().WithFields(log.Fields{ "cur_cpu": cpuload, "cur_memory": vmstat.Used, }).Log("Acquiring approved") return nil } func (r *resources) Info() Info { cpulimit, memlimit := r.Limits() cputhrottling, memthrottling := r.ShouldLimit() cpustat, cpuerr := r.psutil.CPUPercent() memstat, memerr := r.psutil.VirtualMemory() selfcpu, _ := r.self.CPUPercent() selfmem, _ := r.self.VirtualMemory() cpuinfo := CPUInfo{ NCPU: r.ncpu, System: cpustat.System, User: cpustat.User, Idle: cpustat.Idle, Other: cpustat.Other, Limit: cpulimit, Core: selfcpu.System + selfcpu.User + selfcpu.Other, Throttling: cputhrottling, Error: cpuerr, } meminfo := MemoryInfo{ Total: memstat.Total, Available: memstat.Available, Used: memstat.Used, Limit: memlimit, Core: selfmem, Throttling: memthrottling, Error: memerr, } i := Info{ CPU: cpuinfo, Mem: meminfo, } return i }