core/restream/resources/resources.go

package resources

import (
	"context"
	"fmt"
	"sync"
	"time"

	"github.com/datarhei/core/v16/log"
	"github.com/datarhei/core/v16/psutil"
)

type resources struct {
	ncpu      float64
	maxCPU    float64
	maxMemory uint64

	consumerCPU    float64
	consumerMemory uint64

	limit      chan bool
	isLimiting bool

	cancelObserver context.CancelFunc

	lock      sync.Mutex
	startOnce sync.Once
	stopOnce  sync.Once

	logger log.Logger
}

type Resources interface {
	Start()
	Stop()

	Limit() <-chan bool

	Add(cpu float64, memory uint64) bool
	Remove(cpu float64, memory uint64)
}

type Config struct {
	MaxCPU    float64
	MaxMemory float64
	Logger    log.Logger
}

func New(config Config) (Resources, error) {
	r := &resources{
		maxCPU: config.MaxCPU,
		logger: config.Logger,
	}

	vmstat, err := psutil.VirtualMemory()
	if err != nil {
		return nil, fmt.Errorf("unable to determine available memory: %w", err)
	}

	ncpu, err := psutil.CPUCounts(true)
	if err != nil {
		return nil, fmt.Errorf("unable to determine number of logical CPUs: %w", err)
	}

	r.ncpu = ncpu

	r.maxMemory = uint64(float64(vmstat.Total) * config.MaxMemory / 100)

	if r.logger == nil {
		r.logger = log.New("")
	}

	r.logger = r.logger.WithFields(log.Fields{
		"max_cpu":    r.maxCPU,
		"max_memory": r.maxMemory,
	})

	r.logger.Debug().Log("Created")

	r.stopOnce.Do(func() {})

	return r, nil
}

func (r *resources) Start() {
	r.startOnce.Do(func() {
		r.limit = make(chan bool, 10)

		ctx, cancel := context.WithCancel(context.Background())
		r.cancelObserver = cancel

		go r.observe(ctx, time.Second)

		r.stopOnce = sync.Once{}

		r.logger.Debug().Log("Started")
	})
}

func (r *resources) Stop() {
	r.stopOnce.Do(func() {
		r.cancelObserver()

		r.startOnce = sync.Once{}

		r.logger.Debug().Log("Stopped")
	})
}

func (r *resources) Limit() <-chan bool {
	return r.limit
}

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 := psutil.CPUPercent()
			if err != nil {
				r.logger.Warn().WithError(err).Log("Failed to determine CPU load")
				continue
			}

			cpuload := (cpustat.User + cpustat.System + cpustat.Other) * r.ncpu

			vmstat, err := psutil.VirtualMemory()
			if err != nil {
				continue
			}

			r.logger.Debug().WithFields(log.Fields{
				"cur_cpu":    cpuload,
				"cur_memory": vmstat.Used,
			}).Log("Observation")

			limit := false

			if !r.isLimiting {
				if cpuload > r.maxCPU {
					r.logger.Debug().WithField("cpu", cpuload).Log("CPU limit reached")
					limit = true
				}

				if vmstat.Used > r.maxMemory {
					r.logger.Debug().WithField("memory", vmstat.Used).Log("Memory limit reached")
					limit = true
				}
			} else {
				limit = true
				if cpuload <= r.maxCPU*0.8 {
					r.logger.Debug().WithField("cpu", cpuload).Log("CPU limit released")
					limit = false
				}

				if vmstat.Used <= uint64(float64(r.maxMemory)*0.8) {
					r.logger.Debug().WithField("memory", vmstat.Used).Log("Memory limit reached")
					limit = false
				}
			}

			r.lock.Lock()
			if r.isLimiting != limit {
				r.logger.Debug().WithField("enabled", limit).Log("Limiting")
				r.isLimiting = limit
				select {
				case r.limit <- limit:
				default:
				}
			}
			r.lock.Unlock()
		}
	}
}

func (r *resources) Add(cpu float64, memory uint64) bool {
	r.lock.Lock()
	defer r.lock.Unlock()

	logger := r.logger.WithFields(log.Fields{
		"cpu":    cpu,
		"memory": memory,
	})

	logger.Debug().WithFields(log.Fields{
		"used_cpu":    r.consumerCPU,
		"used_memory": r.consumerMemory,
	}).Log("Request for acquiring resources")

	if r.isLimiting {
		logger.Debug().Log("Rejected, currently limiting")
		return false
	}

	if cpu <= 0 || memory == 0 {
		logger.Debug().Log("Rejected, invalid values")
		return false
	}

	if r.consumerCPU+cpu > r.maxCPU {
		logger.Debug().Log("Rejected, CPU limit exceeded")
		return false
	}

	if r.consumerMemory+memory > r.maxMemory {
		logger.Debug().Log("Rejected, memory limit exceeded")
		return false
	}

	r.consumerCPU += cpu
	r.consumerMemory += memory

	logger.Debug().WithFields(log.Fields{
		"used_cpu":    r.consumerCPU,
		"used_memory": r.consumerMemory,
	}).Log("Acquiring approved")

	return true
}

func (r *resources) Remove(cpu float64, memory uint64) {
	r.lock.Lock()
	defer r.lock.Unlock()

	logger := r.logger.WithFields(log.Fields{
		"cpu":    cpu,
		"memory": memory,
	})

	logger.Debug().WithFields(log.Fields{
		"used_cpu":    r.consumerCPU,
		"used_memory": r.consumerMemory,
	}).Log("Request for releasing resources")

	r.consumerCPU -= cpu
	r.consumerMemory -= memory

	if r.consumerCPU < 0 {
		logger.Warn().WithField("used_cpu", r.consumerCPU).Log("Used CPU resources below 0")
		r.consumerCPU = 0
	}

	logger.Debug().WithFields(log.Fields{
		"used_cpu":    r.consumerCPU,
		"used_memory": r.consumerMemory,
	}).Log("Releasing approved")
}