ttipo/comrelay
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Merge commit '6bd7fac875d9e9009915053d8a590abb372c5679' into feature/livekit-upgrade-1.13.1

Browse Source
This commit is contained in:
ttipo
2026-06-25 23:54:33 +09:00
parent 7f6baff425 6bd7fac875
commit 4ad762881c
291 changed files with 37631 additions and 15381 deletions
Download Patch File Download Diff File Expand all files Collapse all files
livekit-server/pkg/agent/agent_test.go
+44 -7
View File
@@ -12,6 +12,7 @@ import (
"github.com/livekit/livekit-server/pkg/agent"
"github.com/livekit/livekit-server/pkg/agent/testutils"
"github.com/livekit/protocol/auth"
"github.com/livekit/protocol/livekit"
"github.com/livekit/protocol/rpc"
"github.com/livekit/protocol/utils/guid"
@@ -20,6 +21,7 @@ import (
)
func TestAgent(t *testing.T) {
testAgentName := "test_agent"
t.Run("dispatched jobs are assigned to a worker", func(t *testing.T) {
bus := psrpc.NewLocalMessageBus()
@@ -28,7 +30,7 @@ func TestAgent(t *testing.T) {
t.Cleanup(server.Close)
worker := server.SimulateAgentWorker()
worker.Register("test", livekit.JobType_JT_ROOM)
worker.Register(testAgentName, livekit.JobType_JT_ROOM)
jobAssignments := worker.JobAssignments.Observe()
job := &livekit.Job{
@@ -36,14 +38,19 @@ func TestAgent(t *testing.T) {
DispatchId: guid.New(guid.AgentDispatchPrefix),
Type: livekit.JobType_JT_ROOM,
Room: &livekit.Room{},
AgentName: "test",
AgentName: testAgentName,
}
_, err := client.JobRequest(context.Background(), "test", agent.RoomAgentTopic, job)
_, err := client.JobRequest(context.Background(), testAgentName, agent.RoomAgentTopic, job)
require.NoError(t, err)
select {
case a := <-jobAssignments.Events():
require.EqualValues(t, job.Id, a.Job.Id)
v, err := auth.ParseAPIToken(a.Token)
require.NoError(t, err)
_, claims, err := v.Verify(server.TestAPISecret)
require.NoError(t, err)
require.Equal(t, testAgentName, claims.Attributes[agent.AgentNameAttributeKey])
case <-time.After(time.Second):
require.Fail(t, "job assignment timeout")
}
@@ -108,7 +115,7 @@ func TestAgentLoadBalancing(t *testing.T) {
t.Cleanup(server.Close)
agents := make([]*testutils.AgentWorker, totalWorkers)
for i := 0; i < totalWorkers; i++ {
for i := range totalWorkers {
agents[i] = server.SimulateAgentWorker(
testutils.WithLabel(fmt.Sprintf("agent-%d", i)),
testutils.WithJobLoad(testutils.NewStableJobLoad(0.01)),
@@ -128,7 +135,7 @@ func TestAgentLoadBalancing(t *testing.T) {
}
// check that jobs are distributed normally
for i := 0; i < totalWorkers; i++ {
for i := range totalWorkers {
label := fmt.Sprintf("agent-%d", i)
require.GreaterOrEqual(t, jobCount[label], 0)
require.Less(t, jobCount[label], 35) // three std deviations from the mean is 32
@@ -147,7 +154,7 @@ func TestAgentLoadBalancing(t *testing.T) {
t.Cleanup(server.Close)
agents := make([]*testutils.AgentWorker, totalWorkers)
for i := 0; i < totalWorkers; i++ {
for i := range totalWorkers {
label := fmt.Sprintf("agent-%d", i)
if i%2 == 0 {
// make sure we have some workers that can accept jobs
@@ -169,7 +176,7 @@ func TestAgentLoadBalancing(t *testing.T) {
jobCount[w.Label] = len(w.Jobs())
}
for i := 0; i < totalWorkers; i++ {
for i := range totalWorkers {
label := fmt.Sprintf("agent-%d", i)
if i%2 == 0 {
@@ -181,3 +188,33 @@ func TestAgentLoadBalancing(t *testing.T) {
}
})
}
func TestConnectionClosedOnDispatchError(t *testing.T) {
t.Run("connection closed when unknown message type received", func(t *testing.T) {
bus := psrpc.NewLocalMessageBus()
server := testutils.NewTestServer(bus)
t.Cleanup(server.Close)
// register agent
worker := server.SimulateAgentWorker()
worker.Register("test_agent", livekit.JobType_JT_ROOM)
responses := worker.RegisterWorkerResponses.Observe()
select {
case <-responses.Events():
// registered
case <-time.After(time.Second):
require.Fail(t, "registration timeout")
}
responses.Stop()
// send invalid message (nil Message field triggers ErrUnknownWorkerSignal)
worker.SendMessage(&livekit.WorkerMessage{Message: nil})
select {
case <-worker.Closed():
// connection closed
case <-time.After(time.Second):
require.Fail(t, "connection should have been closed after dispatch error")
}
})
}
livekit-server/pkg/agent/client.go
+25 -15
View File
@@ -62,11 +62,13 @@ type JobRequest struct {
// only set for participant jobs
Participant *livekit.ParticipantInfo
Metadata string
AgentName string
AgentName string
Deployment string
}
type agentClient struct {
client rpc.AgentInternalClient
config Config
mu sync.RWMutex
@@ -87,7 +89,7 @@ type agentClient struct {
subDone chan struct{}
}
func NewAgentClient(bus psrpc.MessageBus) (Client, error) {
func NewAgentClient(bus psrpc.MessageBus, config Config) (Client, error) {
client, err := rpc.NewAgentInternalClient(bus)
if err != nil {
return nil, err
@@ -95,6 +97,7 @@ func NewAgentClient(bus psrpc.MessageBus) (Client, error) {
c := &agentClient{
client: client,
config: config,
workers: workerpool.New(50),
subDone: make(chan struct{}),
}
@@ -152,21 +155,23 @@ func (c *agentClient) LaunchJob(ctx context.Context, desc *JobRequest) *serverut
}
dispatcher.ForEach(func(curNs string) {
topic := GetAgentTopic(desc.AgentName, curNs)
topic := GetAgentTopic(desc.AgentName, curNs, desc.Deployment)
wg.Add(1)
c.workers.Submit(func() {
defer wg.Done()
// The cached agent parameters do not provide the exact combination of available job type/agent name/namespace, so some of the JobRequest RPC may not trigger any worker
job := &livekit.Job{
Id: utils.NewGuid(utils.AgentJobPrefix),
DispatchId: desc.DispatchId,
Type: desc.JobType,
Room: desc.Room,
Participant: desc.Participant,
Namespace: curNs,
AgentName: desc.AgentName,
Metadata: desc.Metadata,
Id: utils.NewGuid(utils.AgentJobPrefix),
DispatchId: desc.DispatchId,
Type: desc.JobType,
Room: desc.Room,
Participant: desc.Participant,
Namespace: curNs,
AgentName: desc.AgentName,
Metadata: desc.Metadata,
EnableRecording: c.config.EnableUserDataRecording,
Deployment: desc.Deployment,
}
resp, err := c.client.JobRequest(context.Background(), topic, jobTypeTopic, job)
if err != nil {
@@ -323,14 +328,19 @@ func (c *agentClient) Stop() error {
return nil
}
func GetAgentTopic(agentName, namespace string) string {
func GetAgentTopic(agentName, namespace, deployment string) string {
var topic string
if agentName == "" {
// Backward compatibility
return namespace
topic = namespace
} else if namespace == "" {
// Forward compatibility once the namespace field is removed from the worker SDK
return agentName
topic = agentName
} else {
return fmt.Sprintf("%s_%s", agentName, namespace)
topic = fmt.Sprintf("%s_%s", agentName, namespace)
}
if deployment != "" {
topic += "_" + deployment
}
return topic
}
livekit-server/pkg/agent/config.go
+8
View File
@@ -0,0 +1,8 @@
package agent
const DefaultTargetLoad = 0.7
type Config struct {
EnableUserDataRecording bool `yaml:"enable_user_data_recording"`
TargetLoad float32 `yaml:"target_load,omitempty"`
}
livekit-server/pkg/agent/testutils/server.go
+31 -14
View File
@@ -4,14 +4,15 @@ import (
"context"
"errors"
"io"
"maps"
"math"
"math/rand/v2"
"slices"
"sync"
"time"
"github.com/frostbyte73/core"
"github.com/gammazero/deque"
"golang.org/x/exp/maps"
"github.com/livekit/livekit-server/pkg/agent"
"github.com/livekit/livekit-server/pkg/config"
@@ -27,18 +28,25 @@ import (
)
type AgentService interface {
HandleConnection(context.Context, agent.SignalConn, agent.WorkerProtocolVersion)
HandleConnection(context.Context, agent.SignalConn, agent.WorkerRegistration)
DrainConnections(time.Duration)
}
type TestServer struct {
AgentService
TestAPIKey string
TestAPISecret string
}
func NewTestServer(bus psrpc.MessageBus) *TestServer {
localNode, _ := routing.NewLocalNode(nil)
return NewTestServerWithService(must.Get(service.NewAgentService(
&config.Config{Region: "test"},
&config.Config{
Region: "test",
Agents: agent.Config{
TargetLoad: agent.DefaultTargetLoad,
},
},
localNode,
bus,
auth.NewSimpleKeyProvider("test", "verysecretsecret"),
@@ -46,7 +54,7 @@ func NewTestServer(bus psrpc.MessageBus) *TestServer {
}
func NewTestServerWithService(s AgentService) *TestServer {
return &TestServer{s}
return &TestServer{s, "test", "verysecretsecret"}
}
type SimulatedWorkerOptions struct {
@@ -126,7 +134,7 @@ func (h *TestServer) SimulateAgentWorker(opts ...SimulatedWorkerOption) *AgentWo
}
ctx := service.WithAPIKey(o.Context, &auth.ClaimGrants{}, "test")
go h.HandleConnection(ctx, w, agent.CurrentProtocol)
go h.HandleConnection(ctx, w, agent.MakeWorkerRegistration())
return w
}
@@ -204,6 +212,15 @@ func (w *AgentWorker) Close() error {
return nil
}
func (w *AgentWorker) CloseWithReason(_ string) error {
return w.Close()
}
// Closed returns a channel that is closed when the connection is closed by the server
func (w *AgentWorker) Closed() <-chan struct{} {
return w.fuse.Watch()
}
func (w *AgentWorker) SetReadDeadline(t time.Time) error {
w.mu.Lock()
defer w.mu.Unlock()
@@ -315,7 +332,7 @@ func (w *AgentWorker) handlePong(m *livekit.WorkerPong) {
w.WorkerPongs.Emit(m)
}
func (w *AgentWorker) sendMessage(m *livekit.WorkerMessage) {
func (w *AgentWorker) SendMessage(m *livekit.WorkerMessage) {
select {
case <-w.fuse.Watch():
case w.workerMessages <- m:
@@ -323,43 +340,43 @@ func (w *AgentWorker) sendMessage(m *livekit.WorkerMessage) {
}
func (w *AgentWorker) SendRegister(m *livekit.RegisterWorkerRequest) {
w.sendMessage(&livekit.WorkerMessage{Message: &livekit.WorkerMessage_Register{
w.SendMessage(&livekit.WorkerMessage{Message: &livekit.WorkerMessage_Register{
Register: m,
}})
}
func (w *AgentWorker) SendAvailability(m *livekit.AvailabilityResponse) {
w.sendMessage(&livekit.WorkerMessage{Message: &livekit.WorkerMessage_Availability{
w.SendMessage(&livekit.WorkerMessage{Message: &livekit.WorkerMessage_Availability{
Availability: m,
}})
}
func (w *AgentWorker) SendUpdateWorker(m *livekit.UpdateWorkerStatus) {
w.sendMessage(&livekit.WorkerMessage{Message: &livekit.WorkerMessage_UpdateWorker{
w.SendMessage(&livekit.WorkerMessage{Message: &livekit.WorkerMessage_UpdateWorker{
UpdateWorker: m,
}})
}
func (w *AgentWorker) SendUpdateJob(m *livekit.UpdateJobStatus) {
w.sendMessage(&livekit.WorkerMessage{Message: &livekit.WorkerMessage_UpdateJob{
w.SendMessage(&livekit.WorkerMessage{Message: &livekit.WorkerMessage_UpdateJob{
UpdateJob: m,
}})
}
func (w *AgentWorker) SendPing(m *livekit.WorkerPing) {
w.sendMessage(&livekit.WorkerMessage{Message: &livekit.WorkerMessage_Ping{
w.SendMessage(&livekit.WorkerMessage{Message: &livekit.WorkerMessage_Ping{
Ping: m,
}})
}
func (w *AgentWorker) SendSimulateJob(m *livekit.SimulateJobRequest) {
w.sendMessage(&livekit.WorkerMessage{Message: &livekit.WorkerMessage_SimulateJob{
w.SendMessage(&livekit.WorkerMessage{Message: &livekit.WorkerMessage_SimulateJob{
SimulateJob: m,
}})
}
func (w *AgentWorker) SendMigrateJob(m *livekit.MigrateJobRequest) {
w.sendMessage(&livekit.WorkerMessage{Message: &livekit.WorkerMessage_MigrateJob{
w.SendMessage(&livekit.WorkerMessage{Message: &livekit.WorkerMessage_MigrateJob{
MigrateJob: m,
}})
}
@@ -411,7 +428,7 @@ func (w *AgentWorker) SimulateRoomJob(roomName string) {
func (w *AgentWorker) Jobs() []*AgentJob {
w.mu.Lock()
defer w.mu.Unlock()
return maps.Values(w.jobs)
return slices.Collect(maps.Values(w.jobs))
}
type stableJobLoad struct {
livekit-server/pkg/agent/worker.go
+72 -34
View File
@@ -18,12 +18,14 @@ import (
"context"
"errors"
"fmt"
"maps"
"sync"
"time"
"go.uber.org/multierr"
"google.golang.org/protobuf/proto"
pagent "github.com/livekit/protocol/agent"
protoagent "github.com/livekit/protocol/agent"
"github.com/livekit/protocol/livekit"
"github.com/livekit/protocol/logger"
"github.com/livekit/protocol/rpc"
@@ -43,6 +45,8 @@ var (
ErrDuplicateJobAssignment = errors.New("duplicate job assignment")
)
const AgentNameAttributeKey = "lk.agent_name"
type WorkerProtocolVersion int
const CurrentProtocol = 1
@@ -57,12 +61,15 @@ type SignalConn interface {
ReadWorkerMessage() (*livekit.WorkerMessage, int, error)
SetReadDeadline(time.Time) error
Close() error
CloseWithReason(reason string) error
}
func JobStatusIsEnded(s livekit.JobStatus) bool {
return s == livekit.JobStatus_JS_SUCCESS || s == livekit.JobStatus_JS_FAILED
}
type AssignmentHook func(next func(*livekit.JobAssignment) error) func(*livekit.JobAssignment) error
type WorkerSignalHandler interface {
HandleRegister(*livekit.RegisterWorkerRequest) error
HandleAvailability(*livekit.AvailabilityResponse) error
@@ -141,10 +148,20 @@ type WorkerRegistration struct {
Protocol WorkerProtocolVersion
ID string
Version string
AgentID string
AgentName string
Namespace string
JobType livekit.JobType
Permissions *livekit.ParticipantPermission
ClientIP string
Deployment string
}
func MakeWorkerRegistration() WorkerRegistration {
return WorkerRegistration{
ID: guid.New(guid.AgentWorkerPrefix),
Protocol: CurrentProtocol,
}
}
var _ WorkerSignalHandler = (*WorkerRegisterer)(nil)
@@ -152,18 +169,17 @@ var _ WorkerSignalHandler = (*WorkerRegisterer)(nil)
type WorkerRegisterer struct {
WorkerPingHandler
serverInfo *livekit.ServerInfo
protocol WorkerProtocolVersion
deadline time.Time
registration WorkerRegistration
registered bool
}
func NewWorkerRegisterer(conn SignalConn, serverInfo *livekit.ServerInfo, protocol WorkerProtocolVersion) *WorkerRegisterer {
func NewWorkerRegisterer(conn SignalConn, serverInfo *livekit.ServerInfo, base WorkerRegistration) *WorkerRegisterer {
return &WorkerRegisterer{
WorkerPingHandler: WorkerPingHandler{conn: conn},
serverInfo: serverInfo,
protocol: protocol,
registration: base,
deadline: time.Now().Add(RegisterTimeout),
}
}
@@ -185,6 +201,10 @@ func (h *WorkerRegisterer) HandleRegister(req *livekit.RegisterWorkerRequest) er
return ErrUnknownJobType
}
if err := protoagent.ValidateDeployment(req.GetDeployment()); err != nil {
return err
}
permissions := req.AllowedPermissions
if permissions == nil {
permissions = &livekit.ParticipantPermission{
@@ -195,15 +215,12 @@ func (h *WorkerRegisterer) HandleRegister(req *livekit.RegisterWorkerRequest) er
}
}
h.registration = WorkerRegistration{
Protocol: h.protocol,
ID: guid.New(guid.AgentWorkerPrefix),
Version: req.Version,
AgentName: req.AgentName,
Namespace: req.GetNamespace(),
JobType: req.GetType(),
Permissions: permissions,
}
h.registration.Version = req.Version
h.registration.AgentName = req.AgentName
h.registration.Namespace = req.GetNamespace()
h.registration.JobType = req.GetType()
h.registration.Permissions = permissions
h.registration.Deployment = req.GetDeployment()
h.registered = true
_, err := h.conn.WriteServerMessage(&livekit.ServerMessage{
@@ -256,6 +273,9 @@ func NewWorker(
logger: logger.WithValues(
"workerID", registration.ID,
"agentName", registration.AgentName,
"deployment", registration.Deployment,
"agentID", registration.AgentID,
"version", registration.Version,
"jobType", registration.JobType.String(),
),
@@ -268,10 +288,8 @@ func NewWorker(
}
}
func (w *Worker) sendRequest(req *livekit.ServerMessage) {
if _, err := w.conn.WriteServerMessage(req); err != nil {
w.logger.Warnw("error writing to websocket", err)
}
func (w *Worker) APIKey() string {
return w.apiKey
}
func (w *Worker) Status() livekit.WorkerStatus {
@@ -293,11 +311,7 @@ func (w *Worker) Logger() logger.Logger {
func (w *Worker) RunningJobs() map[livekit.JobID]*livekit.Job {
w.mu.Lock()
defer w.mu.Unlock()
jobs := make(map[livekit.JobID]*livekit.Job, len(w.runningJobs))
for k, v := range w.runningJobs {
jobs[k] = v
}
return jobs
return maps.Clone(w.runningJobs)
}
func (w *Worker) RunningJobCount() int {
@@ -316,7 +330,7 @@ func (w *Worker) GetJobState(jobID livekit.JobID) (*livekit.JobState, error) {
return utils.CloneProto(j.State), nil
}
func (w *Worker) AssignJob(ctx context.Context, job *livekit.Job) (*livekit.JobState, error) {
func (w *Worker) AssignJob(ctx context.Context, job *livekit.Job, hook AssignmentHook) (*livekit.JobState, error) {
availCh := make(chan *livekit.AvailabilityResponse, 1)
job = utils.CloneProto(job)
jobID := livekit.JobID(job.Id)
@@ -340,13 +354,17 @@ func (w *Worker) AssignJob(ctx context.Context, job *livekit.Job) (*livekit.JobS
job.State = &livekit.JobState{}
}
now := time.Now()
job.State.WorkerId = w.ID
job.State.AgentId = w.AgentID
job.State.UpdatedAt = now.UnixNano()
job.State.StartedAt = now.UnixNano()
job.State.Status = livekit.JobStatus_JS_RUNNING
w.sendRequest(&livekit.ServerMessage{Message: &livekit.ServerMessage_Availability{
if _, err := w.conn.WriteServerMessage(&livekit.ServerMessage{Message: &livekit.ServerMessage_Availability{
Availability: &livekit.AvailabilityRequest{Job: job},
}})
}}); err != nil {
return nil, err
}
timeout := time.NewTimer(AssignJobTimeout)
defer timeout.Stop()
@@ -354,20 +372,31 @@ func (w *Worker) AssignJob(ctx context.Context, job *livekit.Job) (*livekit.JobS
// See handleAvailability for the response
select {
case res := <-availCh:
if res.Terminate {
job.State.EndedAt = now.UnixNano()
job.State.Status = livekit.JobStatus_JS_SUCCESS
return job.State, nil
}
if !res.Available {
return nil, ErrWorkerNotAvailable
}
job.State.ParticipantIdentity = res.ParticipantIdentity
attributes := res.ParticipantAttributes
if attributes == nil {
attributes = make(map[string]string)
}
attributes[AgentNameAttributeKey] = w.AgentName
token, err := pagent.BuildAgentToken(
token, err := protoagent.BuildAgentToken(
w.apiKey,
w.apiSecret,
job.Room.Name,
res.ParticipantIdentity,
res.ParticipantName,
res.ParticipantMetadata,
res.ParticipantAttributes,
attributes,
w.Permissions,
)
if err != nil {
@@ -375,10 +404,18 @@ func (w *Worker) AssignJob(ctx context.Context, job *livekit.Job) (*livekit.JobS
return nil, err
}
// In OSS, Url is nil, and the used API Key is the same as the one used to connect the worker
w.sendRequest(&livekit.ServerMessage{Message: &livekit.ServerMessage_Assignment{
Assignment: &livekit.JobAssignment{Job: job, Url: nil, Token: token},
}})
send := func(a *livekit.JobAssignment) error {
_, err := w.conn.WriteServerMessage(&livekit.ServerMessage{
Message: &livekit.ServerMessage_Assignment{Assignment: a},
})
return err
}
if hook != nil {
send = hook(send)
}
if err := send(&livekit.JobAssignment{Job: job, Token: token}); err != nil {
return nil, err
}
state := utils.CloneProto(job.State)
@@ -407,7 +444,7 @@ func (w *Worker) TerminateJob(jobID livekit.JobID, reason rpc.JobTerminateReason
return nil, ErrJobNotFound
}
w.sendRequest(&livekit.ServerMessage{Message: &livekit.ServerMessage_Termination{
_, writeErr := w.conn.WriteServerMessage(&livekit.ServerMessage{Message: &livekit.ServerMessage_Termination{
Termination: &livekit.JobTermination{
JobId: string(jobID),
},
@@ -420,11 +457,12 @@ func (w *Worker) TerminateJob(jobID livekit.JobID, reason rpc.JobTerminateReason
errorStr = "agent worker left the room"
}
return w.UpdateJobStatus(&livekit.UpdateJobStatus{
state, updateErr := w.UpdateJobStatus(&livekit.UpdateJobStatus{
JobId: string(jobID),
Status: status,
Error: errorStr,
})
return state, multierr.Combine(writeErr, updateErr)
}
func (w *Worker) UpdateMetadata(metadata string) {
@@ -530,7 +568,7 @@ func (w *Worker) HandleSimulateJob(simulate *livekit.SimulateJobRequest) error {
}
go func() {
_, err := w.AssignJob(w.ctx, job)
_, err := w.AssignJob(w.ctx, job, nil)
if err != nil {
w.logger.Errorw("unable to simulate job", err, "jobID", job.Id)
}