Add configurable affinity_mode (spread) for egress pod selection [port #1209]

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
 }
 

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
 }

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) {

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))
+		})
+	}
+}

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 {