diff --git a/pkg/config/service.go b/pkg/config/service.go index 174c4ada..186db36d 100644 --- a/pkg/config/service.go +++ b/pkg/config/service.go @@ -68,6 +68,26 @@ type ServiceConfig struct { PrometheusPort int `yaml:"prometheus_port"` // prometheus handler port DebugHandlerPort int `yaml:"debug_handler_port"` // egress debug handler port + // AffinityMode controls job distribution across pods. + // "pack" — busy pods win (default; preserves upstream behaviour) + // "spread" — CPU-proportional; use for TrackEgress-only fleets + // "type_aware" — RoomComposite/Web prefer idle pods; Track/Participant spread by load + AffinityMode string `yaml:"affinity_mode"` + + // SoftRejectFloor is the affinity score returned when CanAcceptRequest is + // false but the pod has not yet reached MaxActiveRequests. This prevents + // the pod from going completely silent during transient CPU spikes (e.g. + // Chrome cold-start), keeping it in the dispatcher's selection pool as a + // last resort. Set to 0 (default) to disable and preserve upstream -1 behaviour. + // Recommended production value: 0.01 + SoftRejectFloor float32 `yaml:"soft_reject_floor"` + + // MaxActiveRequests is the hard capacity limit for this pod. When the pod + // has this many active recording jobs, it returns -1 (hard reject) even if + // SoftRejectFloor is set. Set to 0 to disable this guard. + // Recommended value: 16 (= interviews_per_pod × 4 tracks per interview) + MaxActiveRequests int32 `yaml:"max_active_requests"` + *CPUCostConfig `yaml:"cpu_cost"` // CPU costs for the different egress types } diff --git a/pkg/server/server.go b/pkg/server/server.go index a6cb24b9..8bf95fd1 100644 --- a/pkg/server/server.go +++ b/pkg/server/server.go @@ -57,6 +57,7 @@ type Server struct { ioClient info.SessionReporter activeRequests atomic.Int32 + pendingClaims atomic.Int32 // claimed-but-not-yet-accepted requests (spread/type_aware modes) terminating core.Fuse shutdown core.Fuse } diff --git a/pkg/server/server_rpc.go b/pkg/server/server_rpc.go index 3910dde2..aa9829b5 100644 --- a/pkg/server/server_rpc.go +++ b/pkg/server/server_rpc.go @@ -16,6 +16,7 @@ package server import ( "context" + "math/rand" "net/http" "os" "os/exec" @@ -44,8 +45,25 @@ var ( tracer = otel.Tracer("github.com/livekit/egress/pkg/server") ) +// consumePendingClaim decrements pendingClaims by 1, floored at 0. +// Called both from StartEgress (claim accepted) and the 2s self-decay timer +// set in StartEgressAffinity. The CompareAndSwap loop ensures exactly one +// decrement fires per increment even when both callers race. +func (s *Server) consumePendingClaim() { + for { + n := s.pendingClaims.Load() + if n <= 0 { + return + } + if s.pendingClaims.CompareAndSwap(n, n-1) { + return + } + } +} + func (s *Server) StartEgress(ctx context.Context, req *rpc.StartEgressRequest) (*livekit.EgressInfo, error) { s.activeRequests.Inc() + s.consumePendingClaim() // hand slot from pending to m.requests ctx, span := tracer.Start(ctx, "Service.StartEgress") defer span.End() @@ -206,19 +224,85 @@ func (s *Server) processEnded(req *rpc.StartEgressRequest, info *livekit.EgressI } func (s *Server) StartEgressAffinity(_ context.Context, req *rpc.StartEgressRequest) float32 { - if s.IsDisabled() || !s.monitor.CanAcceptRequest(req) { - // cannot accept + if s.IsDisabled() { + return -1 // pod is shutting down — always hard reject + } + + if !s.monitor.CanAcceptRequest(req) { + return s.softRejectScore() + } + + switch s.conf.AffinityMode { + + case "spread": + // Increment pendingClaims so subsequent affinity calls on this pod see a + // lower score before m.requests.Inc fires in StartEgress. The 2s self-decay + // guards against claims that are never accepted (lost RPCs, rejected requests). + // consumePendingClaim uses a CAS loop so exactly one of StartEgress or the + // timer decrements the counter — no double-decrement. + s.pendingClaims.Inc() + time.AfterFunc(2*time.Second, s.consumePendingClaim) + pending := s.pendingClaims.Load() + + // An idle pod returns ≥ 1.0 to trigger psrpc's ShortCircuitTimeout (500ms + // fast path). Without this, AvailableCPUFractionWithPending returns ~0.6 for + // an idle pod (2.4/4.0), which never crosses MaximumAffinity=1.0, so every + // dispatch waits the full AffinityTimeout even when pods are completely free. + // Jitter is added (not subtracted) so the score stays ≥ 1.0. + if s.activeRequests.Load() == 0 && pending <= 1 { + return 1.0 + rand.Float32()*0.001 + } + + return s.monitor.AvailableCPUFractionWithPending(pending) - rand.Float32()*0.001 + + case "type_aware": + if isHeavyEgressRequest(req) { + // RoomComposite/Web need ~4 CPUs. Strongly prefer idle pods. + if s.activeRequests.Load() == 0 { + return 1.0 - rand.Float32()*0.001 + } + return 0.5 + } + // Light requests use the same pending-counter + jitter logic as spread. + s.pendingClaims.Inc() + time.AfterFunc(2*time.Second, s.consumePendingClaim) + pending := s.pendingClaims.Load() + return s.monitor.AvailableCPUFractionWithPending(pending) - rand.Float32()*0.001 + + default: // "pack" or empty — upstream behaviour unchanged + if s.activeRequests.Load() == 0 { + return 0.5 + } + return 1 + } +} + +// softRejectScore returns the score to use when CanAcceptRequest is false. +// If SoftRejectFloor is non-zero and the pod is below its hard capacity limit, +// returns the floor so the pod still participates in dispatcher selection as a +// last resort. Otherwise returns -1 (silent / hard reject). +func (s *Server) softRejectScore() float32 { + floor := s.conf.SoftRejectFloor + + if floor <= 0 { + return -1 + } + + maxActive := s.conf.MaxActiveRequests + if maxActive > 0 && s.activeRequests.Load() >= maxActive { return -1 } - if s.activeRequests.Load() == 0 { - // group multiple track and track composite requests. - // if this instance is idle and another is already handling some, the request will go to that server. - // this avoids having many instances with one track request each, taking availability from room composite. - return 0.5 + return floor +} + +func isHeavyEgressRequest(req *rpc.StartEgressRequest) bool { + switch req.Request.(type) { + case *rpc.StartEgressRequest_RoomComposite, + *rpc.StartEgressRequest_Web: + return true } - // already handling a request and has available cpu - return 1 + return false } func (s *Server) ListActiveEgress(ctx context.Context, _ *rpc.ListActiveEgressRequest) (*rpc.ListActiveEgressResponse, error) { diff --git a/pkg/server/server_rpc_test.go b/pkg/server/server_rpc_test.go new file mode 100644 index 00000000..9847e889 --- /dev/null +++ b/pkg/server/server_rpc_test.go @@ -0,0 +1,142 @@ +// Copyright 2026 LiveKit, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package server + +import ( + "sync" + "testing" + + "github.com/stretchr/testify/require" + + "github.com/livekit/protocol/livekit" + "github.com/livekit/protocol/rpc" + + "github.com/livekit/egress/pkg/config" +) + +// TestConsumePendingClaim_FloorAtZero verifies the CAS loop never decrements below zero. +func TestConsumePendingClaim_FloorAtZero(t *testing.T) { + s := &Server{} + + // Counter starts at 0; extra consume calls must be no-ops. + s.consumePendingClaim() + s.consumePendingClaim() + require.Equal(t, int32(0), s.pendingClaims.Load()) + + // One increment; one consume; counter back to 0. + s.pendingClaims.Inc() + require.Equal(t, int32(1), s.pendingClaims.Load()) + s.consumePendingClaim() + require.Equal(t, int32(0), s.pendingClaims.Load()) + + // Second consume after counter is already 0 is a no-op. + s.consumePendingClaim() + require.Equal(t, int32(0), s.pendingClaims.Load()) +} + +// TestConsumePendingClaim_NoDoubleDecrement verifies that concurrent consumptions +// of the same claim (StartEgress + self-decay timer racing) each fire exactly once +// and together decrement the counter by exactly N, not 2N. +func TestConsumePendingClaim_NoDoubleDecrement(t *testing.T) { + const claims = 50 + s := &Server{} + + for i := 0; i < claims; i++ { + s.pendingClaims.Inc() + } + require.Equal(t, int32(claims), s.pendingClaims.Load()) + + // Simulate StartEgress and self-decay racing concurrently for all claims. + // Both sides try to decrement; together they must not go below zero. + var wg sync.WaitGroup + for i := 0; i < claims*2; i++ { // 2× concurrency, one genuine + one spurious per claim + wg.Add(1) + go func() { + defer wg.Done() + s.consumePendingClaim() + }() + } + wg.Wait() + + require.Equal(t, int32(0), s.pendingClaims.Load(), "counter must be exactly 0, not negative") +} + +// --- Tests for softRejectScore (TASK-03) --- + +func TestSoftRejectFloorDisabled(t *testing.T) { + // SoftRejectFloor=0 → feature disabled → always return -1 + s := &Server{conf: &config.ServiceConfig{SoftRejectFloor: 0}} + require.Equal(t, float32(-1), s.softRejectScore()) +} + +func TestSoftRejectFloorReturnedWhenBelowMax(t *testing.T) { + // Floor set, activeRequests < MaxActiveRequests → return floor + s := &Server{conf: &config.ServiceConfig{SoftRejectFloor: 0.01, MaxActiveRequests: 16}} + s.activeRequests.Store(8) + require.Equal(t, float32(0.01), s.softRejectScore()) +} + +func TestSoftRejectFloorHardRejectWhenAtMax(t *testing.T) { + // Floor set, activeRequests >= MaxActiveRequests → -1 (genuinely full) + s := &Server{conf: &config.ServiceConfig{SoftRejectFloor: 0.01, MaxActiveRequests: 16}} + s.activeRequests.Store(16) + require.Equal(t, float32(-1), s.softRejectScore()) +} + +func TestSoftRejectFloorNoGuardWhenMaxIsZero(t *testing.T) { + // MaxActiveRequests=0 → guard disabled → return floor regardless of load + s := &Server{conf: &config.ServiceConfig{SoftRejectFloor: 0.01, MaxActiveRequests: 0}} + s.activeRequests.Store(20) + require.Equal(t, float32(0.01), s.softRejectScore()) +} + + +func TestIsHeavyEgressRequest(t *testing.T) { + for _, tc := range []struct { + name string + req *rpc.StartEgressRequest + expected bool + }{ + { + name: "RoomComposite is heavy", + req: &rpc.StartEgressRequest{Request: &rpc.StartEgressRequest_RoomComposite{RoomComposite: &livekit.RoomCompositeEgressRequest{}}}, + expected: true, + }, + { + name: "Web is heavy", + req: &rpc.StartEgressRequest{Request: &rpc.StartEgressRequest_Web{Web: &livekit.WebEgressRequest{}}}, + expected: true, + }, + { + name: "Track is not heavy", + req: &rpc.StartEgressRequest{Request: &rpc.StartEgressRequest_Track{Track: &livekit.TrackEgressRequest{}}}, + expected: false, + }, + { + name: "TrackComposite is not heavy", + req: &rpc.StartEgressRequest{Request: &rpc.StartEgressRequest_TrackComposite{TrackComposite: &livekit.TrackCompositeEgressRequest{}}}, + expected: false, + }, + { + name: "Participant is not heavy", + req: &rpc.StartEgressRequest{Request: &rpc.StartEgressRequest_Participant{Participant: &livekit.ParticipantEgressRequest{}}}, + expected: false, + }, + } { + t.Run(tc.name, func(t *testing.T) { + require.Equal(t, tc.expected, isHeavyEgressRequest(tc.req)) + }) + } +} diff --git a/pkg/stats/monitor.go b/pkg/stats/monitor.go index c8d4af43..410f25a2 100644 --- a/pkg/stats/monitor.go +++ b/pkg/stats/monitor.go @@ -604,6 +604,38 @@ func (m *Monitor) GetAvailableCPU() float64 { return available } +// AvailableCPUFraction returns the fraction of CPU budget remaining (0.0–1.0). +// Returns 1.0 when the pod is idle. Used by StartEgressAffinity for spread/type_aware modes. +func (m *Monitor) AvailableCPUFraction() float32 { + m.mu.Lock() + defer m.mu.Unlock() + total, available, _, _ := m.getCPUUsageLocked() + if total == 0 { + return 1.0 + } + return float32(available / total) +} + +// AvailableCPUFractionWithPending is like AvailableCPUFraction but deducts +// pendingSlots extra track-cost units from the available budget. Used by +// StartEgressAffinity in spread/type_aware modes to account for claimed-but-not-yet-accepted +// requests during burst windows (where m.requests.Inc has not yet fired). +func (m *Monitor) AvailableCPUFractionWithPending(pendingSlots int32) float32 { + m.mu.Lock() + defer m.mu.Unlock() + total, available, _, _ := m.getCPUUsageLocked() + if total == 0 { + return 1.0 + } + if pendingSlots > 0 { + available -= float64(pendingSlots) * m.cpuCostConfig.TrackCpuCost + if available < 0 { + available = 0 + } + } + return float32(available / total) +} + func (m *Monitor) getCPUUsageLocked() (total, available, pending, used float64) { total = m.cpuStats.NumCPU() if m.requests.Load() == 0 {