kill egress when usage too high (#574)

* kill egress when usage too high

* increase kill threshold if under max utilization

* only kill if using more than allowed

pkg/config/service.go

@@ -97,7 +97,7 @@ func NewServiceConfig(confString string) (*ServiceConfig, error) {
 	if conf.TrackCpuCost <= 0 {
 		conf.TrackCpuCost = trackCpuCost
 	}
-	if conf.MaxCpuUtilization <= 0 {
+	if conf.MaxCpuUtilization <= 0 || conf.MaxCpuUtilization > 1 {
 		conf.MaxCpuUtilization = maxCpuUtilization
 	}
 

pkg/service/process.go

@@ -17,6 +17,7 @@ package service
 import (
 	"context"
 	"os/exec"
+	"syscall"
 
 	"github.com/frostbyte73/core"
 	dto "github.com/prometheus/client_model/go"
@@ -76,7 +77,14 @@ func (p *Process) Gather() ([]*dto.MetricFamily, error) {
 
 	// Parse the result to match the Gatherer interface
 	return deserializeMetrics(p.info.EgressId, metricsResponse.Metrics)
+}
 
+func (p *Process) kill() {
+	if !p.closed.IsBroken() {
+		if err := p.cmd.Process.Signal(syscall.SIGINT); err != nil {
+			logger.Errorw("failed to kill Process", err, "egressID", p.req.EgressId)
+		}
+	}
 }
 
 func applyDefaultLabel(egressID string, families map[string]*dto.MetricFamily) {

pkg/service/service.go

@@ -23,7 +23,6 @@ import (
 	"os"
 	"path"
 	"sync"
-	"syscall"
 	"time"
 
 	"github.com/frostbyte73/core"
@@ -88,6 +87,7 @@ func NewService(conf *config.ServiceConfig, ioClient rpc.IOInfoClient) (*Service
 	if err := s.Start(s.conf,
 		s.promIsIdle,
 		s.promCanAcceptRequest,
+		s.killProcess,
 	); err != nil {
 		return nil, err
 	}
@@ -179,11 +179,16 @@ func (s *Service) KillAll() {
 	defer s.mu.RUnlock()
 
 	for _, h := range s.activeHandlers {
-		if !h.closed.IsBroken() {
-			if err := h.cmd.Process.Signal(syscall.SIGINT); err != nil {
-				logger.Errorw("failed to kill Process", err, "egressID", h.req.EgressId)
-			}
-		}
+		h.kill()
+	}
+}
+
+func (s *Service) killProcess(egressID string) {
+	s.mu.RLock()
+	defer s.mu.RUnlock()
+
+	if h, ok := s.activeHandlers[egressID]; ok {
+		h.kill()
 	}
 }
 

pkg/stats/monitor.go

@@ -41,27 +41,40 @@ type Monitor struct {
 	cpuStats *utils.CPUStats
 	requests atomic.Int32
 
-	mu        sync.Mutex
-	pending   map[string]*processStats
-	procStats map[int]*processStats
+	mu            sync.Mutex
+	highCPUCount  int
+	killThreshold float64
+	killProcess   func(string)
+	pending       map[string]*processStats
+	procStats     map[int]*processStats
 }
 
 type processStats struct {
 	egressID string
 
 	pendingUsage float64
 	lastUsage    float64
+	allowedUsage float64
 
 	totalCPU   float64
 	cpuCounter int
 	maxCPU     float64
 }
 
-const cpuHoldDuration = time.Second * 30
+const (
+	cpuHoldDuration      = time.Second * 30
+	defaultKillThreshold = 0.95
+)
 
 func NewMonitor(conf *config.ServiceConfig) *Monitor {
+	killThreshold := defaultKillThreshold
+	if killThreshold <= conf.CPUCostConfig.MaxCpuUtilization {
+		killThreshold = (1 + conf.CPUCostConfig.MaxCpuUtilization) / 2
+	}
+
 	return &Monitor{
 		cpuCostConfig: conf.CPUCostConfig,
+		killThreshold: killThreshold,
 		pending:       make(map[string]*processStats),
 		procStats:     make(map[int]*processStats),
 	}
@@ -71,7 +84,10 @@ func (m *Monitor) Start(
 	conf *config.ServiceConfig,
 	isIdle func() float64,
 	canAcceptRequest func() float64,
+	killProcess func(string),
 ) error {
+	m.killProcess = killProcess
+
 	procStats, err := utils.NewProcCPUStats(m.updateEgressStats)
 	if err != nil {
 		return err
@@ -183,11 +199,14 @@ func (m *Monitor) UpdatePID(egressID string, pid int) {
 }
 
 func (m *Monitor) updateEgressStats(idle float64, usage map[int]float64) {
-	m.promCPULoad.Set(1 - idle/m.cpuStats.NumCPU())
+	load := 1 - idle/m.cpuStats.NumCPU()
+	m.promCPULoad.Set(load)
 
 	m.mu.Lock()
 	defer m.mu.Unlock()
 
+	maxUsage := 0.0
+	var maxEgress string
 	for pid, cpuUsage := range usage {
 		if m.procStats[pid] == nil {
 			m.procStats[pid] = &processStats{}
@@ -203,8 +222,30 @@ func (m *Monitor) updateEgressStats(idle float64, usage map[int]float64) {
 
 		if procStats.egressID != "" {
 			m.promProcCPULoad.With(prometheus.Labels{"egress_id": procStats.egressID}).Set(cpuUsage)
+			if cpuUsage > procStats.allowedUsage && cpuUsage > maxUsage {
+				maxUsage = cpuUsage
+				maxEgress = procStats.egressID
+			}
 		}
 	}
+
+	if load > m.killThreshold {
+		logger.Warnw("high cpu usage", nil,
+			"load", load,
+			"requests", m.requests.Load(),
+		)
+
+		if m.requests.Load() > 1 {
+			if m.highCPUCount < 3 {
+				m.highCPUCount++
+				return
+			} else {
+				m.killProcess(maxEgress)
+			}
+		}
+	}
+
+	m.highCPUCount = 0
 }
 
 func (m *Monitor) CloseEgressStats(egressID string) (float64, float64) {
@@ -308,6 +349,7 @@ func (m *Monitor) AcceptRequest(req *rpc.StartEgressRequest) error {
 	ps := &processStats{
 		egressID:     req.EgressId,
 		pendingUsage: cpuHold,
+		allowedUsage: cpuHold,
 	}
 	time.AfterFunc(cpuHoldDuration, func() { ps.pendingUsage = 0 })
 	m.pending[req.EgressId] = ps