From 4e9d046acb36570228e5c7b185328add9257c13c Mon Sep 17 00:00:00 2001
From: David Colburn <xero73@gmail.com>
Date: Thu, 24 Oct 2024 12:10:59 -0400
Subject: [PATCH] use estimated cpu from request if available

---
 go.mod               |  2 +-
 go.sum               |  4 ++--
 pkg/stats/monitor.go | 55 ++++++++++++++++++++++----------------------
 3 files changed, 31 insertions(+), 30 deletions(-)

diff --git a/go.mod b/go.mod
index 54fc2336..84ee2630 100644
--- a/go.mod
+++ b/go.mod
@@ -23,7 +23,7 @@ require (
 	github.com/gorilla/websocket v1.5.3
 	github.com/livekit/livekit-server v1.7.3-0.20241017190429-44a74fc06ae7
 	github.com/livekit/mageutil v0.0.0-20230125210925-54e8a70427c1
-	github.com/livekit/protocol v1.26.1-0.20241017190602-ef6fc8f9c752
+	github.com/livekit/protocol v1.27.1-0.20241022061022-caa595ed3292
 	github.com/livekit/psrpc v0.6.1-0.20240924010758-9f0a4268a3b9
 	github.com/livekit/server-sdk-go/v2 v2.2.2-0.20241015094126-b8538ae5d67b
 	github.com/pion/rtp v1.8.9
diff --git a/go.sum b/go.sum
index 87f2c7ed..fed35c3b 100644
--- a/go.sum
+++ b/go.sum
@@ -218,8 +218,8 @@ github.com/livekit/mageutil v0.0.0-20230125210925-54e8a70427c1 h1:jm09419p0lqTkD
 github.com/livekit/mageutil v0.0.0-20230125210925-54e8a70427c1/go.mod h1:Rs3MhFwutWhGwmY1VQsygw28z5bWcnEYmS1OG9OxjOQ=
 github.com/livekit/mediatransportutil v0.0.0-20240730083616-559fa5ece598 h1:yLlkHk2feSLHstD9n4VKg7YEBR4rLODTI4WE8gNBEnQ=
 github.com/livekit/mediatransportutil v0.0.0-20240730083616-559fa5ece598/go.mod h1:jwKUCmObuiEDH0iiuJHaGMXwRs3RjrB4G6qqgkr/5oE=
-github.com/livekit/protocol v1.26.1-0.20241017190602-ef6fc8f9c752 h1:EgULMfdFSW/3ZZckhiF+CwDApYTD3SkqR3MYazKeE5w=
-github.com/livekit/protocol v1.26.1-0.20241017190602-ef6fc8f9c752/go.mod h1:nxRzmQBKSYK64gqr7ABWwt78hvrgiO2wYuCojRYb7Gs=
+github.com/livekit/protocol v1.27.1-0.20241022061022-caa595ed3292 h1:wVzOLGSjJpCsdKHKKpPxYhXW/JL90l0XYFQbeINSdP4=
+github.com/livekit/protocol v1.27.1-0.20241022061022-caa595ed3292/go.mod h1:nxRzmQBKSYK64gqr7ABWwt78hvrgiO2wYuCojRYb7Gs=
 github.com/livekit/psrpc v0.6.1-0.20240924010758-9f0a4268a3b9 h1:33oBjGpVD9tYkDXQU42tnHl8eCX9G6PVUToBVuCUyOs=
 github.com/livekit/psrpc v0.6.1-0.20240924010758-9f0a4268a3b9/go.mod h1:CQUBSPfYYAaevg1TNCc6/aYsa8DJH4jSRFdCeSZk5u0=
 github.com/livekit/server-sdk-go/v2 v2.2.2-0.20241015094126-b8538ae5d67b h1:R1GpKwVbSYsG08k5sIkNCukvnrkOE18R8IO1YeujR8o=
diff --git a/pkg/stats/monitor.go b/pkg/stats/monitor.go
index aa451063..38755bf9 100644
--- a/pkg/stats/monitor.go
+++ b/pkg/stats/monitor.go
@@ -163,36 +163,37 @@ func (m *Monitor) canAcceptRequestLocked(req *rpc.StartEgressRequest) ([]interfa
 		"activeWeb", m.webRequests.Load(),
 	}
 
-	var accept bool
-	var required float64
-	switch r := req.Request.(type) {
-	case *rpc.StartEgressRequest_RoomComposite:
-		if m.webRequests.Load() >= m.cpuCostConfig.MaxConcurrentWeb {
-			return fields, false
-		}
-		if r.RoomComposite.AudioOnly {
-			required = m.cpuCostConfig.AudioRoomCompositeCpuCost
-		} else {
-			required = m.cpuCostConfig.RoomCompositeCpuCost
-		}
-	case *rpc.StartEgressRequest_Web:
-		if m.webRequests.Load() >= m.cpuCostConfig.MaxConcurrentWeb {
-			return fields, false
-		}
-		if r.Web.AudioOnly {
-			required = m.cpuCostConfig.AudioWebCpuCost
-		} else {
-			required = m.cpuCostConfig.WebCpuCost
+	required := req.EstimatedCpu
+	if required == 0 {
+		switch r := req.Request.(type) {
+		case *rpc.StartEgressRequest_RoomComposite:
+			if m.webRequests.Load() >= m.cpuCostConfig.MaxConcurrentWeb {
+				return fields, false
+			}
+			if r.RoomComposite.AudioOnly {
+				required = m.cpuCostConfig.AudioRoomCompositeCpuCost
+			} else {
+				required = m.cpuCostConfig.RoomCompositeCpuCost
+			}
+		case *rpc.StartEgressRequest_Web:
+			if m.webRequests.Load() >= m.cpuCostConfig.MaxConcurrentWeb {
+				return fields, false
+			}
+			if r.Web.AudioOnly {
+				required = m.cpuCostConfig.AudioWebCpuCost
+			} else {
+				required = m.cpuCostConfig.WebCpuCost
+			}
+		case *rpc.StartEgressRequest_Participant:
+			required = m.cpuCostConfig.ParticipantCpuCost
+		case *rpc.StartEgressRequest_TrackComposite:
+			required = m.cpuCostConfig.TrackCompositeCpuCost
+		case *rpc.StartEgressRequest_Track:
+			required = m.cpuCostConfig.TrackCpuCost
 		}
-	case *rpc.StartEgressRequest_Participant:
-		required = m.cpuCostConfig.ParticipantCpuCost
-	case *rpc.StartEgressRequest_TrackComposite:
-		required = m.cpuCostConfig.TrackCompositeCpuCost
-	case *rpc.StartEgressRequest_Track:
-		required = m.cpuCostConfig.TrackCpuCost
 	}
-	accept = available >= required
 
+	accept := available >= required
 	fields = append(fields,
 		"required", required,
 		"canAccept", accept,