Merge remote-tracking branch 'origin/main' into benjamin/storage_metric_push

Author: biglittlebigben

.github/workflows/publish-egress.yaml

@@ -46,7 +46,7 @@ jobs:
             type=semver,pattern=v{{major}}.{{minor}}
 
       - name: Set up Go
-        uses: actions/setup-go@v4
+        uses: actions/setup-go@v5
         with:
           go-version: 1.21.4
 

pkg/service/process.go

@@ -35,9 +35,6 @@ type Process struct {
 	cmd              *exec.Cmd
 	ipcHandlerClient ipc.EgressHandlerClient
 	ready            chan struct{}
-	totalCPU         float64
-	cpuCounter       int
-	maxCPU           float64
 	closed           core.Fuse
 }
 
@@ -68,21 +65,6 @@ func NewProcess(
 	return p, nil
 }
 
-func (p *Process) updateCPU(cpu float64) {
-	p.totalCPU += cpu
-	p.cpuCounter++
-	if cpu > p.maxCPU {
-		p.maxCPU = cpu
-	}
-}
-
-func (p *Process) getUsageStats() (float64, float64) {
-	if p.cpuCounter == 0 {
-		return 0, 0
-	}
-	return p.totalCPU / float64(p.cpuCounter), p.maxCPU
-}
-
 // Gather implements the prometheus.Gatherer interface on server-side to allow aggregation of handler metrics
 func (p *Process) Gather() ([]*dto.MetricFamily, error) {
 	// Get the metrics from the handler via IPC

pkg/service/service.go

@@ -88,7 +88,6 @@ func NewService(conf *config.ServiceConfig, ioClient rpc.IOInfoClient) (*Service
 	if err := s.Start(s.conf,
 		s.promIsIdle,
 		s.promCanAcceptRequest,
-		s.promProcUpdate,
 	); err != nil {
 		return nil, err
 	}

pkg/service/service_prom.go

@@ -89,25 +89,6 @@ func (s *Service) promCanAcceptRequest() float64 {
 	return 0
 }
 
-func (s *Service) promProcUpdate(pUsage map[int]float64) map[string]float64 {
-	s.mu.RLock()
-	defer s.mu.RUnlock()
-
-	eUsage := make(map[string]float64)
-	for _, h := range s.activeHandlers {
-		if cmd := h.cmd; cmd != nil {
-			if process := cmd.Process; process != nil {
-				if usage, ok := pUsage[process.Pid]; ok {
-					eUsage[h.req.EgressId] = usage
-					h.updateCPU(usage)
-				}
-			}
-		}
-	}
-
-	return eUsage
-}
-
 func (s *Service) storeProcessEndedMetrics(egressID string, metrics string) error {
 	m, err := deserializeMetrics(egressID, metrics)
 	if err != nil {

pkg/service/service_rpc.go

@@ -154,6 +154,7 @@ func (s *Service) AddHandler(egressID string, p *Process) error {
 
 	select {
 	case <-p.ready:
+		s.UpdatePID(egressID, p.cmd.Process.Pid)
 		go func() {
 			err := p.cmd.Wait()
 			s.processEnded(p, err)
@@ -178,7 +179,7 @@ func (s *Service) processEnded(p *Process, err error) {
 		s.Stop(false)
 	}
 
-	avgCPU, maxCPU := p.getUsageStats()
+	avgCPU, maxCPU := s.CloseEgressStats(p.info.EgressId)
 	logger.Infow("egress stats", "egressID", p.req.EgressId, "avgCPU", avgCPU, "maxCPU", maxCPU)
 
 	s.EgressEnded(p.req)

pkg/stats/monitor.go

@@ -39,40 +39,62 @@ type Monitor struct {
 	requestGauge    *prometheus.GaugeVec
 
 	cpuStats *utils.CPUStats
+	requests atomic.Int32
 
-	pendingCPUs atomic.Float64
+	mu        sync.Mutex
+	pending   map[string]*processStats
+	procStats map[int]*processStats
+}
+
+type processStats struct {
+	egressID string
+
+	pendingUsage float64
+	lastUsage    float64
 
-	mu              sync.Mutex
-	requests        atomic.Int32
-	prevEgressUsage map[string]float64
+	totalCPU   float64
+	cpuCounter int
+	maxCPU     float64
 }
 
-const cpuHoldDuration = time.Second * 5
+const cpuHoldDuration = time.Second * 30
 
 func NewMonitor(conf *config.ServiceConfig) *Monitor {
 	return &Monitor{
 		cpuCostConfig: conf.CPUCostConfig,
+		pending:       make(map[string]*processStats),
+		procStats:     make(map[int]*processStats),
 	}
 }
 
 func (m *Monitor) Start(
 	conf *config.ServiceConfig,
 	isIdle func() float64,
 	canAcceptRequest func() float64,
-	procUpdate func(map[int]float64) map[string]float64,
 ) error {
 	procStats, err := utils.NewProcCPUStats(func(idle float64, usage map[int]float64) {
 		m.promCPULoad.Set(1 - idle/m.cpuStats.NumCPU())
-		egressUsage := procUpdate(usage)
-		for egressID, cpuUsage := range egressUsage {
-			m.promProcCPULoad.With(prometheus.Labels{"egress_id": egressID}).Set(cpuUsage)
-		}
-		for egressID := range m.prevEgressUsage {
-			if _, ok := egressUsage[egressID]; !ok {
-				m.promProcCPULoad.With(prometheus.Labels{"egress_id": egressID}).Set(0)
+
+		m.mu.Lock()
+		defer m.mu.Unlock()
+
+		for pid, cpuUsage := range usage {
+			if m.procStats[pid] == nil {
+				m.procStats[pid] = &processStats{}
+			}
+			procStats := m.procStats[pid]
+
+			procStats.lastUsage = cpuUsage
+			procStats.totalCPU += cpuUsage
+			procStats.cpuCounter++
+			if cpuUsage > procStats.maxCPU {
+				procStats.maxCPU = cpuUsage
+			}
+
+			if procStats.egressID != "" {
+				m.promProcCPULoad.With(prometheus.Labels{"egress_id": procStats.egressID}).Set(cpuUsage)
 			}
 		}
-		m.prevEgressUsage = egressUsage
 	})
 	if err != nil {
 		return err
@@ -168,33 +190,78 @@ func (m *Monitor) GetRequestCount() int {
 	return int(m.requests.Load())
 }
 
+func (m *Monitor) UpdatePID(egressID string, pid int) {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+
+	ps := m.pending[egressID]
+	delete(m.pending, egressID)
+
+	if existing := m.procStats[pid]; existing != nil {
+		ps.maxCPU = existing.maxCPU
+		ps.totalCPU = existing.totalCPU
+		ps.cpuCounter = existing.cpuCounter
+	}
+	m.procStats[pid] = ps
+}
+
+func (m *Monitor) CloseEgressStats(egressID string) (float64, float64) {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+
+	for pid, ps := range m.procStats {
+		if ps.egressID == egressID {
+			delete(m.procStats, pid)
+			return ps.totalCPU / float64(ps.cpuCounter), ps.maxCPU
+		}
+	}
+
+	return 0, 0
+}
+
 func (m *Monitor) CanAcceptRequest(req *rpc.StartEgressRequest) bool {
 	m.mu.Lock()
 	defer m.mu.Unlock()
 
-	return m.canAcceptRequest(req)
+	return m.canAcceptRequestLocked(req)
 }
 
-func (m *Monitor) canAcceptRequest(req *rpc.StartEgressRequest) bool {
+func (m *Monitor) canAcceptRequestLocked(req *rpc.StartEgressRequest) bool {
 	accept := false
 
 	total := m.cpuStats.NumCPU()
-	available := m.cpuStats.GetCPUIdle() - m.pendingCPUs.Load()
-
-	logger.Debugw("cpu check",
-		"total", total,
-		"available", available,
-		"active_requests", m.requests,
-	)
 
+	var available float64
 	if m.requests.Load() == 0 {
 		// if no requests, use total
 		available = total
 	} else {
+		var used float64
+		for _, ps := range m.pending {
+			if ps.pendingUsage > ps.lastUsage {
+				used += ps.pendingUsage
+			} else {
+				used += ps.lastUsage
+			}
+		}
+		for _, ps := range m.procStats {
+			if ps.pendingUsage > ps.lastUsage {
+				used += ps.pendingUsage
+			} else {
+				used += ps.lastUsage
+			}
+		}
+
 		// if already running requests, cap usage at MaxCpuUtilization
-		available -= (1 - m.cpuCostConfig.MaxCpuUtilization) * total
+		available = total - used - (total * (1 - m.cpuCostConfig.MaxCpuUtilization))
 	}
 
+	logger.Debugw("cpu check",
+		"total", total,
+		"available", available,
+		"active_requests", m.requests,
+	)
+
 	switch req.Request.(type) {
 	case *rpc.StartEgressRequest_RoomComposite:
 		accept = available >= m.cpuCostConfig.RoomCompositeCpuCost
@@ -215,10 +282,12 @@ func (m *Monitor) AcceptRequest(req *rpc.StartEgressRequest) error {
 	m.mu.Lock()
 	defer m.mu.Unlock()
 
-	if !m.canAcceptRequest(req) {
+	if !m.canAcceptRequestLocked(req) {
 		return errors.ErrResourceExhausted
 	}
 
+	m.requests.Inc()
+
 	var cpuHold float64
 	switch req.Request.(type) {
 	case *rpc.StartEgressRequest_RoomComposite:
@@ -233,9 +302,13 @@ func (m *Monitor) AcceptRequest(req *rpc.StartEgressRequest) error {
 		cpuHold = m.cpuCostConfig.TrackCpuCost
 	}
 
-	m.requests.Inc()
-	m.pendingCPUs.Add(cpuHold)
-	time.AfterFunc(cpuHoldDuration, func() { m.pendingCPUs.Sub(cpuHold) })
+	ps := &processStats{
+		egressID:     req.EgressId,
+		pendingUsage: cpuHold,
+	}
+	time.AfterFunc(cpuHoldDuration, func() { ps.pendingUsage = 0 })
+	m.pending[req.EgressId] = ps
+
 	return nil
 }