check memory usage before accepting request (livekit#846)

tribe-health · Jan 20, 2025 · 1f29a53 · 1f29a53
1 parent 0540867
commit 1f29a53
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 1 deletion.
diff --git a/pkg/config/service.go b/pkg/config/service.go
@@ -59,6 +59,7 @@ type ServiceConfig struct {
 type CPUCostConfig struct {
 	MaxCpuUtilization         float64 `yaml:"max_cpu_utilization"` // maximum allowed CPU utilization when deciding to accept a request. Default to 80%
 	MaxMemory                 float64 `yaml:"max_memory"`          // maximum allowed memory usage in GB. 0 to disable
+	MemoryCost                float64 `yaml:"memory_cost"`         // minimum memory in GB
 	MaxConcurrentWeb          int32   `yaml:"max_concurrent_web"`  // maximum allowed chrome/x/pulse instances
 	RoomCompositeCpuCost      float64 `yaml:"room_composite_cpu_cost"`
 	AudioRoomCompositeCpuCost float64 `yaml:"audio_room_composite_cpu_cost"`

diff --git a/pkg/stats/monitor.go b/pkg/stats/monitor.go
@@ -62,6 +62,7 @@ type Monitor struct {
 	highCPUDuration int
 	pending         map[string]*processStats
 	procStats       map[int]*processStats
+	memoryUsage     float64
 }
 
 type processStats struct {
@@ -163,12 +164,19 @@ func (m *Monitor) canAcceptRequestLocked(req *rpc.StartEgressRequest) ([]interfa
 		"used", used,
 		"activeRequests", m.requests.Load(),
 		"activeWeb", m.webRequests.Load(),
+		"memory", m.memoryUsage,
+	}
+
+	if m.cpuCostConfig.MaxMemory > 0 && m.memoryUsage+m.cpuCostConfig.MemoryCost >= m.cpuCostConfig.MaxMemory {
+		fields = append(fields, "canAccept", false)
+		return fields, false
 	}
 
 	required := req.EstimatedCpu
 	switch r := req.Request.(type) {
 	case *rpc.StartEgressRequest_RoomComposite:
 		if m.webRequests.Load() >= m.cpuCostConfig.MaxConcurrentWeb {
+			fields = append(fields, "canAccept", false)
 			return fields, false
 		}
 		if required == 0 {
@@ -180,6 +188,7 @@ func (m *Monitor) canAcceptRequestLocked(req *rpc.StartEgressRequest) ([]interfa
 		}
 	case *rpc.StartEgressRequest_Web:
 		if m.webRequests.Load() >= m.cpuCostConfig.MaxConcurrentWeb {
+			fields = append(fields, "canAccept", false)
 			return fields, false
 		}
 		if required == 0 {
@@ -446,9 +455,10 @@ func (m *Monitor) updateEgressStats(stats *hwstats.ProcStats) {
 		}
 	}
 
+	m.memoryUsage = float64(totalMemory) / gb
 	if m.cpuCostConfig.MaxMemory > 0 && totalMemory > int(m.cpuCostConfig.MaxMemory*gb) {
 		logger.Warnw("high memory usage", nil,
-			"memory", float64(totalMemory)/gb,
+			"memory", m.memoryUsage,
 			"requests", m.requests.Load(),
 		)
 		m.svc.KillProcess(maxMemoryEgress, errors.ErrOOM(float64(maxMemory)/gb))