Merge commit '0da97ebd2149ec2a0412a273b17270f540431d81' as 'livekit-server'
This commit is contained in:
@@ -0,0 +1,183 @@
|
||||
package agent_test
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
"go.uber.org/atomic"
|
||||
|
||||
"github.com/livekit/livekit-server/pkg/agent"
|
||||
"github.com/livekit/livekit-server/pkg/agent/testutils"
|
||||
"github.com/livekit/protocol/livekit"
|
||||
"github.com/livekit/protocol/rpc"
|
||||
"github.com/livekit/protocol/utils/guid"
|
||||
"github.com/livekit/protocol/utils/must"
|
||||
"github.com/livekit/psrpc"
|
||||
)
|
||||
|
||||
func TestAgent(t *testing.T) {
|
||||
t.Run("dispatched jobs are assigned to a worker", func(t *testing.T) {
|
||||
bus := psrpc.NewLocalMessageBus()
|
||||
|
||||
client := must.Get(rpc.NewAgentInternalClient(bus))
|
||||
server := testutils.NewTestServer(bus)
|
||||
t.Cleanup(server.Close)
|
||||
|
||||
worker := server.SimulateAgentWorker()
|
||||
worker.Register("test", livekit.JobType_JT_ROOM)
|
||||
jobAssignments := worker.JobAssignments.Observe()
|
||||
|
||||
job := &livekit.Job{
|
||||
Id: guid.New(guid.AgentJobPrefix),
|
||||
DispatchId: guid.New(guid.AgentDispatchPrefix),
|
||||
Type: livekit.JobType_JT_ROOM,
|
||||
Room: &livekit.Room{},
|
||||
AgentName: "test",
|
||||
}
|
||||
_, err := client.JobRequest(context.Background(), "test", agent.RoomAgentTopic, job)
|
||||
require.NoError(t, err)
|
||||
|
||||
select {
|
||||
case a := <-jobAssignments.Events():
|
||||
require.EqualValues(t, job.Id, a.Job.Id)
|
||||
case <-time.After(time.Second):
|
||||
require.Fail(t, "job assignment timeout")
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func testBatchJobRequest(t require.TestingT, batchSize int, totalJobs int, client rpc.AgentInternalClient, workers []*testutils.AgentWorker) <-chan struct{} {
|
||||
var assigned atomic.Uint32
|
||||
done := make(chan struct{})
|
||||
for _, w := range workers {
|
||||
assignments := w.JobAssignments.Observe()
|
||||
go func() {
|
||||
defer assignments.Stop()
|
||||
for {
|
||||
select {
|
||||
case <-done:
|
||||
case <-assignments.Events():
|
||||
if assigned.Inc() == uint32(totalJobs) {
|
||||
close(done)
|
||||
}
|
||||
}
|
||||
}
|
||||
}()
|
||||
}
|
||||
|
||||
// wait for agent registration
|
||||
time.Sleep(100 * time.Millisecond)
|
||||
|
||||
var wg sync.WaitGroup
|
||||
for i := 0; i < totalJobs; i += batchSize {
|
||||
wg.Add(1)
|
||||
go func(start int) {
|
||||
defer wg.Done()
|
||||
for j := start; j < start+batchSize && j < totalJobs; j++ {
|
||||
job := &livekit.Job{
|
||||
Id: guid.New(guid.AgentJobPrefix),
|
||||
DispatchId: guid.New(guid.AgentDispatchPrefix),
|
||||
Type: livekit.JobType_JT_ROOM,
|
||||
Room: &livekit.Room{},
|
||||
AgentName: "test",
|
||||
}
|
||||
_, err := client.JobRequest(context.Background(), "test", agent.RoomAgentTopic, job)
|
||||
require.NoError(t, err)
|
||||
}
|
||||
}(i)
|
||||
}
|
||||
wg.Wait()
|
||||
|
||||
return done
|
||||
}
|
||||
|
||||
func TestAgentLoadBalancing(t *testing.T) {
|
||||
t.Run("jobs are distributed normally with baseline worker load", func(t *testing.T) {
|
||||
totalWorkers := 5
|
||||
totalJobs := 100
|
||||
|
||||
bus := psrpc.NewLocalMessageBus()
|
||||
|
||||
client := must.Get(rpc.NewAgentInternalClient(bus))
|
||||
t.Cleanup(client.Close)
|
||||
server := testutils.NewTestServer(bus)
|
||||
t.Cleanup(server.Close)
|
||||
|
||||
agents := make([]*testutils.AgentWorker, totalWorkers)
|
||||
for i := 0; i < totalWorkers; i++ {
|
||||
agents[i] = server.SimulateAgentWorker(
|
||||
testutils.WithLabel(fmt.Sprintf("agent-%d", i)),
|
||||
testutils.WithJobLoad(testutils.NewStableJobLoad(0.01)),
|
||||
)
|
||||
agents[i].Register("test", livekit.JobType_JT_ROOM)
|
||||
}
|
||||
|
||||
select {
|
||||
case <-testBatchJobRequest(t, 10, totalJobs, client, agents):
|
||||
case <-time.After(time.Second):
|
||||
require.Fail(t, "job assignment timeout")
|
||||
}
|
||||
|
||||
jobCount := make(map[string]int)
|
||||
for _, w := range agents {
|
||||
jobCount[w.Label] = len(w.Jobs())
|
||||
}
|
||||
|
||||
// check that jobs are distributed normally
|
||||
for i := 0; i < totalWorkers; i++ {
|
||||
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
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("jobs are distributed with variable and overloaded worker load", func(t *testing.T) {
|
||||
totalWorkers := 4
|
||||
totalJobs := 15
|
||||
|
||||
bus := psrpc.NewLocalMessageBus()
|
||||
|
||||
client := must.Get(rpc.NewAgentInternalClient(bus))
|
||||
t.Cleanup(client.Close)
|
||||
server := testutils.NewTestServer(bus)
|
||||
t.Cleanup(server.Close)
|
||||
|
||||
agents := make([]*testutils.AgentWorker, totalWorkers)
|
||||
for i := 0; i < totalWorkers; i++ {
|
||||
label := fmt.Sprintf("agent-%d", i)
|
||||
if i%2 == 0 {
|
||||
// make sure we have some workers that can accept jobs
|
||||
agents[i] = server.SimulateAgentWorker(testutils.WithLabel(label))
|
||||
} else {
|
||||
agents[i] = server.SimulateAgentWorker(testutils.WithLabel(label), testutils.WithDefaultWorkerLoad(0.9))
|
||||
}
|
||||
agents[i].Register("test", livekit.JobType_JT_ROOM)
|
||||
}
|
||||
|
||||
select {
|
||||
case <-testBatchJobRequest(t, 1, totalJobs, client, agents):
|
||||
case <-time.After(time.Second):
|
||||
require.Fail(t, "job assignment timeout")
|
||||
}
|
||||
|
||||
jobCount := make(map[string]int)
|
||||
for _, w := range agents {
|
||||
jobCount[w.Label] = len(w.Jobs())
|
||||
}
|
||||
|
||||
for i := 0; i < totalWorkers; i++ {
|
||||
label := fmt.Sprintf("agent-%d", i)
|
||||
|
||||
if i%2 == 0 {
|
||||
require.GreaterOrEqual(t, jobCount[label], 2)
|
||||
} else {
|
||||
require.Equal(t, 0, jobCount[label])
|
||||
}
|
||||
require.GreaterOrEqual(t, jobCount[label], 0)
|
||||
}
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,336 @@
|
||||
// Copyright 2024 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 agent
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/gammazero/workerpool"
|
||||
"google.golang.org/protobuf/types/known/emptypb"
|
||||
|
||||
serverutils "github.com/livekit/livekit-server/pkg/utils"
|
||||
"github.com/livekit/protocol/livekit"
|
||||
"github.com/livekit/protocol/logger"
|
||||
"github.com/livekit/protocol/rpc"
|
||||
"github.com/livekit/protocol/utils"
|
||||
"github.com/livekit/psrpc"
|
||||
)
|
||||
|
||||
const (
|
||||
EnabledCacheTTL = 1 * time.Minute
|
||||
RoomAgentTopic = "room"
|
||||
PublisherAgentTopic = "publisher"
|
||||
ParticipantAgentTopic = "participant"
|
||||
DefaultHandlerNamespace = ""
|
||||
|
||||
CheckEnabledTimeout = 5 * time.Second
|
||||
)
|
||||
|
||||
var jobTypeTopics = map[livekit.JobType]string{
|
||||
livekit.JobType_JT_ROOM: RoomAgentTopic,
|
||||
livekit.JobType_JT_PUBLISHER: PublisherAgentTopic,
|
||||
livekit.JobType_JT_PARTICIPANT: ParticipantAgentTopic,
|
||||
}
|
||||
|
||||
type Client interface {
|
||||
// LaunchJob starts a room or participant job on an agent.
|
||||
// it will launch a job once for each worker in each namespace
|
||||
LaunchJob(ctx context.Context, desc *JobRequest) *serverutils.IncrementalDispatcher[*livekit.Job]
|
||||
TerminateJob(ctx context.Context, jobID string, reason rpc.JobTerminateReason) (*livekit.JobState, error)
|
||||
Stop() error
|
||||
}
|
||||
|
||||
type JobRequest struct {
|
||||
DispatchId string
|
||||
JobType livekit.JobType
|
||||
Room *livekit.Room
|
||||
// only set for participant jobs
|
||||
Participant *livekit.ParticipantInfo
|
||||
Metadata string
|
||||
AgentName string
|
||||
}
|
||||
|
||||
type agentClient struct {
|
||||
client rpc.AgentInternalClient
|
||||
|
||||
mu sync.RWMutex
|
||||
|
||||
// cache response to avoid constantly checking with controllers
|
||||
// cache is invalidated with AgentRegistered updates
|
||||
roomNamespaces *serverutils.IncrementalDispatcher[string] // deprecated
|
||||
publisherNamespaces *serverutils.IncrementalDispatcher[string] // deprecated
|
||||
participantNamespaces *serverutils.IncrementalDispatcher[string] // deprecated
|
||||
roomAgentNames *serverutils.IncrementalDispatcher[string]
|
||||
publisherAgentNames *serverutils.IncrementalDispatcher[string]
|
||||
participantAgentNames *serverutils.IncrementalDispatcher[string]
|
||||
|
||||
enabledExpiresAt time.Time
|
||||
|
||||
workers *workerpool.WorkerPool
|
||||
|
||||
invalidateSub psrpc.Subscription[*emptypb.Empty]
|
||||
subDone chan struct{}
|
||||
}
|
||||
|
||||
func NewAgentClient(bus psrpc.MessageBus) (Client, error) {
|
||||
client, err := rpc.NewAgentInternalClient(bus)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
c := &agentClient{
|
||||
client: client,
|
||||
workers: workerpool.New(50),
|
||||
subDone: make(chan struct{}),
|
||||
}
|
||||
|
||||
sub, err := c.client.SubscribeWorkerRegistered(context.Background(), DefaultHandlerNamespace)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
c.invalidateSub = sub
|
||||
|
||||
go func() {
|
||||
// invalidate cache
|
||||
for range sub.Channel() {
|
||||
c.mu.Lock()
|
||||
c.roomNamespaces = nil
|
||||
c.publisherNamespaces = nil
|
||||
c.participantNamespaces = nil
|
||||
c.roomAgentNames = nil
|
||||
c.publisherAgentNames = nil
|
||||
c.participantAgentNames = nil
|
||||
c.mu.Unlock()
|
||||
}
|
||||
|
||||
c.subDone <- struct{}{}
|
||||
}()
|
||||
|
||||
return c, nil
|
||||
}
|
||||
|
||||
func (c *agentClient) LaunchJob(ctx context.Context, desc *JobRequest) *serverutils.IncrementalDispatcher[*livekit.Job] {
|
||||
var wg sync.WaitGroup
|
||||
ret := serverutils.NewIncrementalDispatcher[*livekit.Job]()
|
||||
defer func() {
|
||||
c.workers.Submit(func() {
|
||||
wg.Wait()
|
||||
ret.Done()
|
||||
})
|
||||
}()
|
||||
|
||||
jobTypeTopic, ok := jobTypeTopics[desc.JobType]
|
||||
if !ok {
|
||||
return ret
|
||||
}
|
||||
|
||||
dispatcher := c.getDispatcher(desc.AgentName, desc.JobType)
|
||||
|
||||
if dispatcher == nil {
|
||||
logger.Infow("not dispatching agent job since no worker is available",
|
||||
"agentName", desc.AgentName,
|
||||
"jobType", desc.JobType,
|
||||
"room", desc.Room.Name,
|
||||
"roomID", desc.Room.Sid)
|
||||
return ret
|
||||
}
|
||||
|
||||
dispatcher.ForEach(func(curNs string) {
|
||||
topic := GetAgentTopic(desc.AgentName, curNs)
|
||||
|
||||
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,
|
||||
}
|
||||
resp, err := c.client.JobRequest(context.Background(), topic, jobTypeTopic, job)
|
||||
if err != nil {
|
||||
logger.Infow("failed to send job request", "error", err, "namespace", curNs, "jobType", desc.JobType, "agentName", desc.AgentName)
|
||||
return
|
||||
}
|
||||
job.State = resp.State
|
||||
ret.Add(job)
|
||||
})
|
||||
})
|
||||
|
||||
return ret
|
||||
}
|
||||
|
||||
func (c *agentClient) TerminateJob(ctx context.Context, jobID string, reason rpc.JobTerminateReason) (*livekit.JobState, error) {
|
||||
resp, err := c.client.JobTerminate(context.Background(), jobID, &rpc.JobTerminateRequest{
|
||||
JobId: jobID,
|
||||
Reason: reason,
|
||||
})
|
||||
if err != nil {
|
||||
logger.Infow("failed to send job request", "error", err, "jobID", jobID)
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return resp.State, nil
|
||||
}
|
||||
|
||||
func (c *agentClient) getDispatcher(agName string, jobType livekit.JobType) *serverutils.IncrementalDispatcher[string] {
|
||||
c.mu.Lock()
|
||||
|
||||
if time.Since(c.enabledExpiresAt) > EnabledCacheTTL || c.roomNamespaces == nil ||
|
||||
c.publisherNamespaces == nil || c.participantNamespaces == nil || c.roomAgentNames == nil || c.publisherAgentNames == nil || c.participantAgentNames == nil {
|
||||
c.enabledExpiresAt = time.Now()
|
||||
c.roomNamespaces = serverutils.NewIncrementalDispatcher[string]()
|
||||
c.publisherNamespaces = serverutils.NewIncrementalDispatcher[string]()
|
||||
c.participantNamespaces = serverutils.NewIncrementalDispatcher[string]()
|
||||
c.roomAgentNames = serverutils.NewIncrementalDispatcher[string]()
|
||||
c.publisherAgentNames = serverutils.NewIncrementalDispatcher[string]()
|
||||
c.participantAgentNames = serverutils.NewIncrementalDispatcher[string]()
|
||||
|
||||
go c.checkEnabled(c.roomNamespaces, c.publisherNamespaces, c.participantNamespaces, c.roomAgentNames, c.publisherAgentNames, c.participantAgentNames)
|
||||
}
|
||||
|
||||
var target *serverutils.IncrementalDispatcher[string]
|
||||
var agentNames *serverutils.IncrementalDispatcher[string]
|
||||
switch jobType {
|
||||
case livekit.JobType_JT_ROOM:
|
||||
target = c.roomNamespaces
|
||||
agentNames = c.roomAgentNames
|
||||
case livekit.JobType_JT_PUBLISHER:
|
||||
target = c.publisherNamespaces
|
||||
agentNames = c.publisherAgentNames
|
||||
case livekit.JobType_JT_PARTICIPANT:
|
||||
target = c.participantNamespaces
|
||||
agentNames = c.participantAgentNames
|
||||
}
|
||||
c.mu.Unlock()
|
||||
|
||||
if agName == "" {
|
||||
// if no agent name is given, we would need to dispatch backwards compatible mode
|
||||
// which means dispatching to each of the namespaces
|
||||
return target
|
||||
}
|
||||
|
||||
done := make(chan *serverutils.IncrementalDispatcher[string], 1)
|
||||
c.workers.Submit(func() {
|
||||
agentNames.ForEach(func(ag string) {
|
||||
if ag == agName {
|
||||
select {
|
||||
case done <- target:
|
||||
default:
|
||||
}
|
||||
}
|
||||
})
|
||||
select {
|
||||
case done <- nil:
|
||||
default:
|
||||
}
|
||||
})
|
||||
|
||||
return <-done
|
||||
}
|
||||
|
||||
func (c *agentClient) checkEnabled(roomNamespaces, publisherNamespaces, participantNamespaces, roomAgentNames, publisherAgentNames, participantAgentNames *serverutils.IncrementalDispatcher[string]) {
|
||||
defer roomNamespaces.Done()
|
||||
defer publisherNamespaces.Done()
|
||||
defer participantNamespaces.Done()
|
||||
defer roomAgentNames.Done()
|
||||
defer publisherAgentNames.Done()
|
||||
defer participantAgentNames.Done()
|
||||
|
||||
resChan, err := c.client.CheckEnabled(context.Background(), &rpc.CheckEnabledRequest{}, psrpc.WithRequestTimeout(CheckEnabledTimeout))
|
||||
if err != nil {
|
||||
logger.Errorw("failed to check enabled", err)
|
||||
return
|
||||
}
|
||||
|
||||
roomNSMap := make(map[string]bool)
|
||||
publisherNSMap := make(map[string]bool)
|
||||
participantNSMap := make(map[string]bool)
|
||||
roomAgMap := make(map[string]bool)
|
||||
publisherAgMap := make(map[string]bool)
|
||||
participantAgMap := make(map[string]bool)
|
||||
|
||||
for r := range resChan {
|
||||
if r.Result.GetRoomEnabled() {
|
||||
for _, ns := range r.Result.GetNamespaces() {
|
||||
if _, ok := roomNSMap[ns]; !ok {
|
||||
roomNamespaces.Add(ns)
|
||||
roomNSMap[ns] = true
|
||||
}
|
||||
}
|
||||
for _, ag := range r.Result.GetAgentNames() {
|
||||
if _, ok := roomAgMap[ag]; !ok {
|
||||
roomAgentNames.Add(ag)
|
||||
roomAgMap[ag] = true
|
||||
}
|
||||
}
|
||||
}
|
||||
if r.Result.GetPublisherEnabled() {
|
||||
for _, ns := range r.Result.GetNamespaces() {
|
||||
if _, ok := publisherNSMap[ns]; !ok {
|
||||
publisherNamespaces.Add(ns)
|
||||
publisherNSMap[ns] = true
|
||||
}
|
||||
}
|
||||
for _, ag := range r.Result.GetAgentNames() {
|
||||
if _, ok := publisherAgMap[ag]; !ok {
|
||||
publisherAgentNames.Add(ag)
|
||||
publisherAgMap[ag] = true
|
||||
}
|
||||
}
|
||||
}
|
||||
if r.Result.GetParticipantEnabled() {
|
||||
for _, ns := range r.Result.GetNamespaces() {
|
||||
if _, ok := participantNSMap[ns]; !ok {
|
||||
participantNamespaces.Add(ns)
|
||||
participantNSMap[ns] = true
|
||||
}
|
||||
}
|
||||
for _, ag := range r.Result.GetAgentNames() {
|
||||
if _, ok := participantAgMap[ag]; !ok {
|
||||
participantAgentNames.Add(ag)
|
||||
participantAgMap[ag] = true
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (c *agentClient) Stop() error {
|
||||
_ = c.invalidateSub.Close()
|
||||
<-c.subDone
|
||||
return nil
|
||||
}
|
||||
|
||||
func GetAgentTopic(agentName, namespace string) string {
|
||||
if agentName == "" {
|
||||
// Backward compatibility
|
||||
return namespace
|
||||
} else if namespace == "" {
|
||||
// Forward compatibility once the namespace field is removed from the worker SDK
|
||||
return agentName
|
||||
} else {
|
||||
return fmt.Sprintf("%s_%s", agentName, namespace)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,483 @@
|
||||
package testutils
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"io"
|
||||
"math"
|
||||
"math/rand/v2"
|
||||
"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"
|
||||
"github.com/livekit/livekit-server/pkg/routing"
|
||||
"github.com/livekit/livekit-server/pkg/service"
|
||||
"github.com/livekit/protocol/auth"
|
||||
"github.com/livekit/protocol/livekit"
|
||||
"github.com/livekit/protocol/utils/events"
|
||||
"github.com/livekit/protocol/utils/guid"
|
||||
"github.com/livekit/protocol/utils/must"
|
||||
"github.com/livekit/protocol/utils/options"
|
||||
"github.com/livekit/psrpc"
|
||||
)
|
||||
|
||||
type AgentService interface {
|
||||
HandleConnection(context.Context, agent.SignalConn, agent.WorkerProtocolVersion)
|
||||
DrainConnections(time.Duration)
|
||||
}
|
||||
|
||||
type TestServer struct {
|
||||
AgentService
|
||||
}
|
||||
|
||||
func NewTestServer(bus psrpc.MessageBus) *TestServer {
|
||||
localNode, _ := routing.NewLocalNode(nil)
|
||||
return NewTestServerWithService(must.Get(service.NewAgentService(
|
||||
&config.Config{Region: "test"},
|
||||
localNode,
|
||||
bus,
|
||||
auth.NewSimpleKeyProvider("test", "verysecretsecret"),
|
||||
)))
|
||||
}
|
||||
|
||||
func NewTestServerWithService(s AgentService) *TestServer {
|
||||
return &TestServer{s}
|
||||
}
|
||||
|
||||
type SimulatedWorkerOptions struct {
|
||||
Context context.Context
|
||||
Label string
|
||||
SupportResume bool
|
||||
DefaultJobLoad float32
|
||||
JobLoadThreshold float32
|
||||
DefaultWorkerLoad float32
|
||||
HandleAvailability func(AgentJobRequest)
|
||||
HandleAssignment func(*livekit.Job) JobLoad
|
||||
}
|
||||
|
||||
type SimulatedWorkerOption func(*SimulatedWorkerOptions)
|
||||
|
||||
func WithContext(ctx context.Context) SimulatedWorkerOption {
|
||||
return func(o *SimulatedWorkerOptions) {
|
||||
o.Context = ctx
|
||||
}
|
||||
}
|
||||
|
||||
func WithLabel(label string) SimulatedWorkerOption {
|
||||
return func(o *SimulatedWorkerOptions) {
|
||||
o.Label = label
|
||||
}
|
||||
}
|
||||
|
||||
func WithJobAvailabilityHandler(h func(AgentJobRequest)) SimulatedWorkerOption {
|
||||
return func(o *SimulatedWorkerOptions) {
|
||||
o.HandleAvailability = h
|
||||
}
|
||||
}
|
||||
|
||||
func WithJobAssignmentHandler(h func(*livekit.Job) JobLoad) SimulatedWorkerOption {
|
||||
return func(o *SimulatedWorkerOptions) {
|
||||
o.HandleAssignment = h
|
||||
}
|
||||
}
|
||||
|
||||
func WithJobLoad(l JobLoad) SimulatedWorkerOption {
|
||||
return WithJobAssignmentHandler(func(j *livekit.Job) JobLoad { return l })
|
||||
}
|
||||
|
||||
func WithDefaultWorkerLoad(load float32) SimulatedWorkerOption {
|
||||
return func(o *SimulatedWorkerOptions) {
|
||||
o.DefaultWorkerLoad = load
|
||||
}
|
||||
}
|
||||
|
||||
func (h *TestServer) SimulateAgentWorker(opts ...SimulatedWorkerOption) *AgentWorker {
|
||||
o := &SimulatedWorkerOptions{
|
||||
Context: context.Background(),
|
||||
Label: guid.New("TEST_AGENT_"),
|
||||
DefaultJobLoad: 0.1,
|
||||
JobLoadThreshold: 0.8,
|
||||
DefaultWorkerLoad: 0.0,
|
||||
HandleAvailability: func(r AgentJobRequest) { r.Accept() },
|
||||
HandleAssignment: func(j *livekit.Job) JobLoad { return nil },
|
||||
}
|
||||
options.Apply(o, opts)
|
||||
|
||||
w := &AgentWorker{
|
||||
workerMessages: make(chan *livekit.WorkerMessage, 1),
|
||||
jobs: map[string]*AgentJob{},
|
||||
SimulatedWorkerOptions: o,
|
||||
|
||||
RegisterWorkerResponses: events.NewObserverList[*livekit.RegisterWorkerResponse](),
|
||||
AvailabilityRequests: events.NewObserverList[*livekit.AvailabilityRequest](),
|
||||
JobAssignments: events.NewObserverList[*livekit.JobAssignment](),
|
||||
JobTerminations: events.NewObserverList[*livekit.JobTermination](),
|
||||
WorkerPongs: events.NewObserverList[*livekit.WorkerPong](),
|
||||
}
|
||||
w.ctx, w.cancel = context.WithCancel(context.Background())
|
||||
|
||||
if o.DefaultWorkerLoad > 0.0 {
|
||||
w.sendStatus()
|
||||
}
|
||||
|
||||
ctx := service.WithAPIKey(o.Context, &auth.ClaimGrants{}, "test")
|
||||
go h.HandleConnection(ctx, w, agent.CurrentProtocol)
|
||||
|
||||
return w
|
||||
}
|
||||
|
||||
func (h *TestServer) Close() {
|
||||
h.DrainConnections(1)
|
||||
}
|
||||
|
||||
var _ agent.SignalConn = (*AgentWorker)(nil)
|
||||
|
||||
type JobLoad interface {
|
||||
Load() float32
|
||||
}
|
||||
|
||||
type AgentJob struct {
|
||||
*livekit.Job
|
||||
JobLoad
|
||||
}
|
||||
|
||||
type AgentJobRequest struct {
|
||||
w *AgentWorker
|
||||
*livekit.AvailabilityRequest
|
||||
}
|
||||
|
||||
func (r AgentJobRequest) Accept() {
|
||||
identity := guid.New("PI_")
|
||||
r.w.SendAvailability(&livekit.AvailabilityResponse{
|
||||
JobId: r.Job.Id,
|
||||
Available: true,
|
||||
SupportsResume: r.w.SupportResume,
|
||||
ParticipantName: identity,
|
||||
ParticipantIdentity: identity,
|
||||
})
|
||||
}
|
||||
|
||||
func (r AgentJobRequest) Reject() {
|
||||
r.w.SendAvailability(&livekit.AvailabilityResponse{
|
||||
JobId: r.Job.Id,
|
||||
Available: false,
|
||||
})
|
||||
}
|
||||
|
||||
type AgentWorker struct {
|
||||
*SimulatedWorkerOptions
|
||||
|
||||
fuse core.Fuse
|
||||
mu sync.Mutex
|
||||
ctx context.Context
|
||||
cancel context.CancelFunc
|
||||
workerMessages chan *livekit.WorkerMessage
|
||||
serverMessages deque.Deque[*livekit.ServerMessage]
|
||||
jobs map[string]*AgentJob
|
||||
|
||||
RegisterWorkerResponses *events.ObserverList[*livekit.RegisterWorkerResponse]
|
||||
AvailabilityRequests *events.ObserverList[*livekit.AvailabilityRequest]
|
||||
JobAssignments *events.ObserverList[*livekit.JobAssignment]
|
||||
JobTerminations *events.ObserverList[*livekit.JobTermination]
|
||||
WorkerPongs *events.ObserverList[*livekit.WorkerPong]
|
||||
}
|
||||
|
||||
func (w *AgentWorker) statusWorker() {
|
||||
t := time.NewTicker(2 * time.Second)
|
||||
defer t.Stop()
|
||||
|
||||
for !w.fuse.IsBroken() {
|
||||
w.sendStatus()
|
||||
<-t.C
|
||||
}
|
||||
}
|
||||
|
||||
func (w *AgentWorker) Close() error {
|
||||
w.mu.Lock()
|
||||
defer w.mu.Unlock()
|
||||
w.fuse.Break()
|
||||
return nil
|
||||
}
|
||||
|
||||
func (w *AgentWorker) SetReadDeadline(t time.Time) error {
|
||||
w.mu.Lock()
|
||||
defer w.mu.Unlock()
|
||||
if !w.fuse.IsBroken() {
|
||||
cancel := w.cancel
|
||||
if t.IsZero() {
|
||||
w.ctx, w.cancel = context.WithCancel(context.Background())
|
||||
} else {
|
||||
w.ctx, w.cancel = context.WithDeadline(context.Background(), t)
|
||||
}
|
||||
cancel()
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (w *AgentWorker) ReadWorkerMessage() (*livekit.WorkerMessage, int, error) {
|
||||
for {
|
||||
w.mu.Lock()
|
||||
ctx := w.ctx
|
||||
w.mu.Unlock()
|
||||
|
||||
select {
|
||||
case <-w.fuse.Watch():
|
||||
return nil, 0, io.EOF
|
||||
case <-ctx.Done():
|
||||
if err := ctx.Err(); errors.Is(err, context.DeadlineExceeded) {
|
||||
return nil, 0, err
|
||||
}
|
||||
case m := <-w.workerMessages:
|
||||
return m, 0, nil
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (w *AgentWorker) WriteServerMessage(m *livekit.ServerMessage) (int, error) {
|
||||
w.mu.Lock()
|
||||
defer w.mu.Unlock()
|
||||
w.serverMessages.PushBack(m)
|
||||
if w.serverMessages.Len() == 1 {
|
||||
go w.handleServerMessages()
|
||||
}
|
||||
return 0, nil
|
||||
}
|
||||
|
||||
func (w *AgentWorker) handleServerMessages() {
|
||||
w.mu.Lock()
|
||||
for w.serverMessages.Len() != 0 {
|
||||
m := w.serverMessages.Front()
|
||||
w.mu.Unlock()
|
||||
|
||||
switch m := m.Message.(type) {
|
||||
case *livekit.ServerMessage_Register:
|
||||
w.handleRegister(m.Register)
|
||||
case *livekit.ServerMessage_Availability:
|
||||
w.handleAvailability(m.Availability)
|
||||
case *livekit.ServerMessage_Assignment:
|
||||
w.handleAssignment(m.Assignment)
|
||||
case *livekit.ServerMessage_Termination:
|
||||
w.handleTermination(m.Termination)
|
||||
case *livekit.ServerMessage_Pong:
|
||||
w.handlePong(m.Pong)
|
||||
}
|
||||
|
||||
w.mu.Lock()
|
||||
w.serverMessages.PopFront()
|
||||
}
|
||||
w.mu.Unlock()
|
||||
}
|
||||
|
||||
func (w *AgentWorker) handleRegister(m *livekit.RegisterWorkerResponse) {
|
||||
w.RegisterWorkerResponses.Emit(m)
|
||||
}
|
||||
|
||||
func (w *AgentWorker) handleAvailability(m *livekit.AvailabilityRequest) {
|
||||
w.AvailabilityRequests.Emit(m)
|
||||
if w.HandleAvailability != nil {
|
||||
w.HandleAvailability(AgentJobRequest{w, m})
|
||||
} else {
|
||||
AgentJobRequest{w, m}.Accept()
|
||||
}
|
||||
}
|
||||
|
||||
func (w *AgentWorker) handleAssignment(m *livekit.JobAssignment) {
|
||||
w.JobAssignments.Emit(m)
|
||||
|
||||
var load JobLoad
|
||||
if w.HandleAssignment != nil {
|
||||
load = w.HandleAssignment(m.Job)
|
||||
}
|
||||
|
||||
if load == nil {
|
||||
load = NewStableJobLoad(w.DefaultJobLoad)
|
||||
}
|
||||
|
||||
w.mu.Lock()
|
||||
defer w.mu.Unlock()
|
||||
w.jobs[m.Job.Id] = &AgentJob{m.Job, load}
|
||||
}
|
||||
|
||||
func (w *AgentWorker) handleTermination(m *livekit.JobTermination) {
|
||||
w.JobTerminations.Emit(m)
|
||||
|
||||
w.mu.Lock()
|
||||
defer w.mu.Unlock()
|
||||
delete(w.jobs, m.JobId)
|
||||
}
|
||||
|
||||
func (w *AgentWorker) handlePong(m *livekit.WorkerPong) {
|
||||
w.WorkerPongs.Emit(m)
|
||||
}
|
||||
|
||||
func (w *AgentWorker) sendMessage(m *livekit.WorkerMessage) {
|
||||
select {
|
||||
case <-w.fuse.Watch():
|
||||
case w.workerMessages <- m:
|
||||
}
|
||||
}
|
||||
|
||||
func (w *AgentWorker) SendRegister(m *livekit.RegisterWorkerRequest) {
|
||||
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{
|
||||
Availability: m,
|
||||
}})
|
||||
}
|
||||
|
||||
func (w *AgentWorker) SendUpdateWorker(m *livekit.UpdateWorkerStatus) {
|
||||
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{
|
||||
UpdateJob: m,
|
||||
}})
|
||||
}
|
||||
|
||||
func (w *AgentWorker) SendPing(m *livekit.WorkerPing) {
|
||||
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{
|
||||
SimulateJob: m,
|
||||
}})
|
||||
}
|
||||
|
||||
func (w *AgentWorker) SendMigrateJob(m *livekit.MigrateJobRequest) {
|
||||
w.sendMessage(&livekit.WorkerMessage{Message: &livekit.WorkerMessage_MigrateJob{
|
||||
MigrateJob: m,
|
||||
}})
|
||||
}
|
||||
|
||||
func (w *AgentWorker) sendStatus() {
|
||||
w.mu.Lock()
|
||||
var load float32
|
||||
jobCount := len(w.jobs)
|
||||
|
||||
if len(w.jobs) == 0 {
|
||||
load = w.DefaultWorkerLoad
|
||||
} else {
|
||||
for _, j := range w.jobs {
|
||||
load += j.Load()
|
||||
}
|
||||
}
|
||||
w.mu.Unlock()
|
||||
|
||||
status := livekit.WorkerStatus_WS_AVAILABLE
|
||||
if load > w.JobLoadThreshold {
|
||||
status = livekit.WorkerStatus_WS_FULL
|
||||
}
|
||||
|
||||
w.SendUpdateWorker(&livekit.UpdateWorkerStatus{
|
||||
Status: &status,
|
||||
Load: load,
|
||||
JobCount: uint32(jobCount),
|
||||
})
|
||||
}
|
||||
|
||||
func (w *AgentWorker) Register(agentName string, jobType livekit.JobType) {
|
||||
w.SendRegister(&livekit.RegisterWorkerRequest{
|
||||
Type: jobType,
|
||||
AgentName: agentName,
|
||||
})
|
||||
go w.statusWorker()
|
||||
}
|
||||
|
||||
func (w *AgentWorker) SimulateRoomJob(roomName string) {
|
||||
w.SendSimulateJob(&livekit.SimulateJobRequest{
|
||||
Type: livekit.JobType_JT_ROOM,
|
||||
Room: &livekit.Room{
|
||||
Sid: guid.New(guid.RoomPrefix),
|
||||
Name: roomName,
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
func (w *AgentWorker) Jobs() []*AgentJob {
|
||||
w.mu.Lock()
|
||||
defer w.mu.Unlock()
|
||||
return maps.Values(w.jobs)
|
||||
}
|
||||
|
||||
type stableJobLoad struct {
|
||||
load float32
|
||||
}
|
||||
|
||||
func NewStableJobLoad(load float32) JobLoad {
|
||||
return stableJobLoad{load}
|
||||
}
|
||||
|
||||
func (s stableJobLoad) Load() float32 {
|
||||
return s.load
|
||||
}
|
||||
|
||||
type periodicJobLoad struct {
|
||||
amplitude float64
|
||||
period time.Duration
|
||||
epoch time.Time
|
||||
}
|
||||
|
||||
func NewPeriodicJobLoad(max float32, period time.Duration) JobLoad {
|
||||
return periodicJobLoad{
|
||||
amplitude: float64(max / 2),
|
||||
period: period,
|
||||
epoch: time.Now().Add(-time.Duration(rand.Int64N(int64(period)))),
|
||||
}
|
||||
}
|
||||
|
||||
func (s periodicJobLoad) Load() float32 {
|
||||
a := math.Sin(time.Since(s.epoch).Seconds() / s.period.Seconds() * math.Pi * 2)
|
||||
return float32(s.amplitude + a*s.amplitude)
|
||||
}
|
||||
|
||||
type uniformRandomJobLoad struct {
|
||||
min, max float32
|
||||
rng func() float64
|
||||
}
|
||||
|
||||
func NewUniformRandomJobLoad(min, max float32) JobLoad {
|
||||
return uniformRandomJobLoad{min, max, rand.Float64}
|
||||
}
|
||||
|
||||
func NewUniformRandomJobLoadWithRNG(min, max float32, rng *rand.Rand) JobLoad {
|
||||
return uniformRandomJobLoad{min, max, rng.Float64}
|
||||
}
|
||||
|
||||
func (s uniformRandomJobLoad) Load() float32 {
|
||||
return rand.Float32()*(s.max-s.min) + s.min
|
||||
}
|
||||
|
||||
type normalRandomJobLoad struct {
|
||||
mean, stddev float64
|
||||
rng func() float64
|
||||
}
|
||||
|
||||
func NewNormalRandomJobLoad(mean, stddev float64) JobLoad {
|
||||
return normalRandomJobLoad{mean, stddev, rand.Float64}
|
||||
}
|
||||
|
||||
func NewNormalRandomJobLoadWithRNG(mean, stddev float64, rng *rand.Rand) JobLoad {
|
||||
return normalRandomJobLoad{mean, stddev, rng.Float64}
|
||||
}
|
||||
|
||||
func (s normalRandomJobLoad) Load() float32 {
|
||||
u := 1 - s.rng()
|
||||
v := s.rng()
|
||||
z := math.Sqrt(-2*math.Log(u)) * math.Cos(2*math.Pi*v)
|
||||
return float32(max(0, z*s.stddev+s.mean))
|
||||
}
|
||||
@@ -0,0 +1,559 @@
|
||||
// Copyright 2024 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 agent
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"google.golang.org/protobuf/proto"
|
||||
|
||||
pagent "github.com/livekit/protocol/agent"
|
||||
"github.com/livekit/protocol/livekit"
|
||||
"github.com/livekit/protocol/logger"
|
||||
"github.com/livekit/protocol/rpc"
|
||||
"github.com/livekit/protocol/utils"
|
||||
"github.com/livekit/protocol/utils/guid"
|
||||
"github.com/livekit/psrpc"
|
||||
)
|
||||
|
||||
var (
|
||||
ErrUnimplementedWrorkerSignal = errors.New("unimplemented worker signal")
|
||||
ErrUnknownWorkerSignal = errors.New("unknown worker signal")
|
||||
ErrUnknownJobType = errors.New("unknown job type")
|
||||
ErrJobNotFound = psrpc.NewErrorf(psrpc.NotFound, "no running job for given jobID")
|
||||
ErrWorkerClosed = errors.New("worker closed")
|
||||
ErrWorkerNotAvailable = errors.New("worker not available")
|
||||
ErrAvailabilityTimeout = errors.New("agent worker availability timeout")
|
||||
ErrDuplicateJobAssignment = errors.New("duplicate job assignment")
|
||||
)
|
||||
|
||||
type WorkerProtocolVersion int
|
||||
|
||||
const CurrentProtocol = 1
|
||||
|
||||
const (
|
||||
RegisterTimeout = 10 * time.Second
|
||||
AssignJobTimeout = 10 * time.Second
|
||||
)
|
||||
|
||||
type SignalConn interface {
|
||||
WriteServerMessage(msg *livekit.ServerMessage) (int, error)
|
||||
ReadWorkerMessage() (*livekit.WorkerMessage, int, error)
|
||||
SetReadDeadline(time.Time) error
|
||||
Close() error
|
||||
}
|
||||
|
||||
func JobStatusIsEnded(s livekit.JobStatus) bool {
|
||||
return s == livekit.JobStatus_JS_SUCCESS || s == livekit.JobStatus_JS_FAILED
|
||||
}
|
||||
|
||||
type WorkerSignalHandler interface {
|
||||
HandleRegister(*livekit.RegisterWorkerRequest) error
|
||||
HandleAvailability(*livekit.AvailabilityResponse) error
|
||||
HandleUpdateJob(*livekit.UpdateJobStatus) error
|
||||
HandleSimulateJob(*livekit.SimulateJobRequest) error
|
||||
HandlePing(*livekit.WorkerPing) error
|
||||
HandleUpdateWorker(*livekit.UpdateWorkerStatus) error
|
||||
HandleMigrateJob(*livekit.MigrateJobRequest) error
|
||||
}
|
||||
|
||||
func DispatchWorkerSignal(req *livekit.WorkerMessage, h WorkerSignalHandler) error {
|
||||
switch m := req.Message.(type) {
|
||||
case *livekit.WorkerMessage_Register:
|
||||
return h.HandleRegister(m.Register)
|
||||
case *livekit.WorkerMessage_Availability:
|
||||
return h.HandleAvailability(m.Availability)
|
||||
case *livekit.WorkerMessage_UpdateJob:
|
||||
return h.HandleUpdateJob(m.UpdateJob)
|
||||
case *livekit.WorkerMessage_SimulateJob:
|
||||
return h.HandleSimulateJob(m.SimulateJob)
|
||||
case *livekit.WorkerMessage_Ping:
|
||||
return h.HandlePing(m.Ping)
|
||||
case *livekit.WorkerMessage_UpdateWorker:
|
||||
return h.HandleUpdateWorker(m.UpdateWorker)
|
||||
case *livekit.WorkerMessage_MigrateJob:
|
||||
return h.HandleMigrateJob(m.MigrateJob)
|
||||
default:
|
||||
return ErrUnknownWorkerSignal
|
||||
}
|
||||
}
|
||||
|
||||
var _ WorkerSignalHandler = (*UnimplementedWorkerSignalHandler)(nil)
|
||||
|
||||
type UnimplementedWorkerSignalHandler struct{}
|
||||
|
||||
func (UnimplementedWorkerSignalHandler) HandleRegister(*livekit.RegisterWorkerRequest) error {
|
||||
return fmt.Errorf("%w: Register", ErrUnimplementedWrorkerSignal)
|
||||
}
|
||||
func (UnimplementedWorkerSignalHandler) HandleAvailability(*livekit.AvailabilityResponse) error {
|
||||
return fmt.Errorf("%w: Availability", ErrUnimplementedWrorkerSignal)
|
||||
}
|
||||
func (UnimplementedWorkerSignalHandler) HandleUpdateJob(*livekit.UpdateJobStatus) error {
|
||||
return fmt.Errorf("%w: UpdateJob", ErrUnimplementedWrorkerSignal)
|
||||
}
|
||||
func (UnimplementedWorkerSignalHandler) HandleSimulateJob(*livekit.SimulateJobRequest) error {
|
||||
return fmt.Errorf("%w: SimulateJob", ErrUnimplementedWrorkerSignal)
|
||||
}
|
||||
func (UnimplementedWorkerSignalHandler) HandlePing(*livekit.WorkerPing) error {
|
||||
return fmt.Errorf("%w: Ping", ErrUnimplementedWrorkerSignal)
|
||||
}
|
||||
func (UnimplementedWorkerSignalHandler) HandleUpdateWorker(*livekit.UpdateWorkerStatus) error {
|
||||
return fmt.Errorf("%w: UpdateWorker", ErrUnimplementedWrorkerSignal)
|
||||
}
|
||||
func (UnimplementedWorkerSignalHandler) HandleMigrateJob(*livekit.MigrateJobRequest) error {
|
||||
return fmt.Errorf("%w: MigrateJob", ErrUnimplementedWrorkerSignal)
|
||||
}
|
||||
|
||||
type WorkerPingHandler struct {
|
||||
UnimplementedWorkerSignalHandler
|
||||
conn SignalConn
|
||||
}
|
||||
|
||||
func (h WorkerPingHandler) HandlePing(ping *livekit.WorkerPing) error {
|
||||
_, err := h.conn.WriteServerMessage(&livekit.ServerMessage{
|
||||
Message: &livekit.ServerMessage_Pong{
|
||||
Pong: &livekit.WorkerPong{
|
||||
LastTimestamp: ping.Timestamp,
|
||||
Timestamp: time.Now().UnixMilli(),
|
||||
},
|
||||
},
|
||||
})
|
||||
return err
|
||||
}
|
||||
|
||||
type WorkerRegistration struct {
|
||||
Protocol WorkerProtocolVersion
|
||||
ID string
|
||||
Version string
|
||||
AgentName string
|
||||
Namespace string
|
||||
JobType livekit.JobType
|
||||
Permissions *livekit.ParticipantPermission
|
||||
}
|
||||
|
||||
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 {
|
||||
return &WorkerRegisterer{
|
||||
WorkerPingHandler: WorkerPingHandler{conn: conn},
|
||||
serverInfo: serverInfo,
|
||||
protocol: protocol,
|
||||
deadline: time.Now().Add(RegisterTimeout),
|
||||
}
|
||||
}
|
||||
|
||||
func (h *WorkerRegisterer) Deadline() time.Time {
|
||||
return h.deadline
|
||||
}
|
||||
|
||||
func (h *WorkerRegisterer) Registration() WorkerRegistration {
|
||||
return h.registration
|
||||
}
|
||||
|
||||
func (h *WorkerRegisterer) Registered() bool {
|
||||
return h.registered
|
||||
}
|
||||
|
||||
func (h *WorkerRegisterer) HandleRegister(req *livekit.RegisterWorkerRequest) error {
|
||||
if !livekit.IsJobType(req.GetType()) {
|
||||
return ErrUnknownJobType
|
||||
}
|
||||
|
||||
permissions := req.AllowedPermissions
|
||||
if permissions == nil {
|
||||
permissions = &livekit.ParticipantPermission{
|
||||
CanSubscribe: true,
|
||||
CanPublish: true,
|
||||
CanPublishData: true,
|
||||
CanUpdateMetadata: true,
|
||||
}
|
||||
}
|
||||
|
||||
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.registered = true
|
||||
|
||||
_, err := h.conn.WriteServerMessage(&livekit.ServerMessage{
|
||||
Message: &livekit.ServerMessage_Register{
|
||||
Register: &livekit.RegisterWorkerResponse{
|
||||
WorkerId: h.registration.ID,
|
||||
ServerInfo: h.serverInfo,
|
||||
},
|
||||
},
|
||||
})
|
||||
return err
|
||||
}
|
||||
|
||||
var _ WorkerSignalHandler = (*Worker)(nil)
|
||||
|
||||
type Worker struct {
|
||||
WorkerPingHandler
|
||||
WorkerRegistration
|
||||
|
||||
apiKey string
|
||||
apiSecret string
|
||||
logger logger.Logger
|
||||
|
||||
ctx context.Context
|
||||
cancel context.CancelFunc
|
||||
closed chan struct{}
|
||||
|
||||
mu sync.Mutex
|
||||
load float32
|
||||
status livekit.WorkerStatus
|
||||
|
||||
runningJobs map[livekit.JobID]*livekit.Job
|
||||
availability map[livekit.JobID]chan *livekit.AvailabilityResponse
|
||||
}
|
||||
|
||||
func NewWorker(
|
||||
registration WorkerRegistration,
|
||||
apiKey string,
|
||||
apiSecret string,
|
||||
conn SignalConn,
|
||||
logger logger.Logger,
|
||||
) *Worker {
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
|
||||
return &Worker{
|
||||
WorkerPingHandler: WorkerPingHandler{conn: conn},
|
||||
WorkerRegistration: registration,
|
||||
apiKey: apiKey,
|
||||
apiSecret: apiSecret,
|
||||
logger: logger.WithValues(
|
||||
"workerID", registration.ID,
|
||||
"agentName", registration.AgentName,
|
||||
"jobType", registration.JobType.String(),
|
||||
),
|
||||
|
||||
ctx: ctx,
|
||||
cancel: cancel,
|
||||
closed: make(chan struct{}),
|
||||
|
||||
runningJobs: make(map[livekit.JobID]*livekit.Job),
|
||||
availability: make(map[livekit.JobID]chan *livekit.AvailabilityResponse),
|
||||
}
|
||||
}
|
||||
|
||||
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) Status() livekit.WorkerStatus {
|
||||
w.mu.Lock()
|
||||
defer w.mu.Unlock()
|
||||
return w.status
|
||||
}
|
||||
|
||||
func (w *Worker) Load() float32 {
|
||||
w.mu.Lock()
|
||||
defer w.mu.Unlock()
|
||||
return w.load
|
||||
}
|
||||
|
||||
func (w *Worker) Logger() logger.Logger {
|
||||
return w.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
|
||||
}
|
||||
|
||||
func (w *Worker) RunningJobCount() int {
|
||||
w.mu.Lock()
|
||||
defer w.mu.Unlock()
|
||||
return len(w.runningJobs)
|
||||
}
|
||||
|
||||
func (w *Worker) GetJobState(jobID livekit.JobID) (*livekit.JobState, error) {
|
||||
w.mu.Lock()
|
||||
defer w.mu.Unlock()
|
||||
j, ok := w.runningJobs[jobID]
|
||||
if !ok {
|
||||
return nil, ErrJobNotFound
|
||||
}
|
||||
return utils.CloneProto(j.State), nil
|
||||
}
|
||||
|
||||
func (w *Worker) AssignJob(ctx context.Context, job *livekit.Job) (*livekit.JobState, error) {
|
||||
availCh := make(chan *livekit.AvailabilityResponse, 1)
|
||||
job = utils.CloneProto(job)
|
||||
jobID := livekit.JobID(job.Id)
|
||||
|
||||
w.mu.Lock()
|
||||
if _, ok := w.availability[jobID]; ok {
|
||||
w.mu.Unlock()
|
||||
return nil, ErrDuplicateJobAssignment
|
||||
}
|
||||
|
||||
w.availability[jobID] = availCh
|
||||
w.mu.Unlock()
|
||||
|
||||
defer func() {
|
||||
w.mu.Lock()
|
||||
delete(w.availability, jobID)
|
||||
w.mu.Unlock()
|
||||
}()
|
||||
|
||||
if job.State == nil {
|
||||
job.State = &livekit.JobState{}
|
||||
}
|
||||
now := time.Now()
|
||||
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{
|
||||
Availability: &livekit.AvailabilityRequest{Job: job},
|
||||
}})
|
||||
|
||||
timeout := time.NewTimer(AssignJobTimeout)
|
||||
defer timeout.Stop()
|
||||
|
||||
// See handleAvailability for the response
|
||||
select {
|
||||
case res := <-availCh:
|
||||
if !res.Available {
|
||||
return nil, ErrWorkerNotAvailable
|
||||
}
|
||||
|
||||
job.State.ParticipantIdentity = res.ParticipantIdentity
|
||||
|
||||
token, err := pagent.BuildAgentToken(
|
||||
w.apiKey,
|
||||
w.apiSecret,
|
||||
job.Room.Name,
|
||||
res.ParticipantIdentity,
|
||||
res.ParticipantName,
|
||||
res.ParticipantMetadata,
|
||||
res.ParticipantAttributes,
|
||||
w.Permissions,
|
||||
)
|
||||
if err != nil {
|
||||
w.logger.Errorw("failed to build agent token", err)
|
||||
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},
|
||||
}})
|
||||
|
||||
state := utils.CloneProto(job.State)
|
||||
|
||||
w.mu.Lock()
|
||||
w.runningJobs[jobID] = job
|
||||
w.mu.Unlock()
|
||||
|
||||
// TODO sweep jobs that are never started. We can't do this until all SDKs actually update the the JOB state
|
||||
|
||||
return state, nil
|
||||
case <-timeout.C:
|
||||
return nil, ErrAvailabilityTimeout
|
||||
case <-w.ctx.Done():
|
||||
return nil, ErrWorkerClosed
|
||||
case <-ctx.Done():
|
||||
return nil, ctx.Err()
|
||||
}
|
||||
}
|
||||
|
||||
func (w *Worker) TerminateJob(jobID livekit.JobID, reason rpc.JobTerminateReason) (*livekit.JobState, error) {
|
||||
w.mu.Lock()
|
||||
_, ok := w.runningJobs[jobID]
|
||||
w.mu.Unlock()
|
||||
|
||||
if !ok {
|
||||
return nil, ErrJobNotFound
|
||||
}
|
||||
|
||||
w.sendRequest(&livekit.ServerMessage{Message: &livekit.ServerMessage_Termination{
|
||||
Termination: &livekit.JobTermination{
|
||||
JobId: string(jobID),
|
||||
},
|
||||
}})
|
||||
|
||||
status := livekit.JobStatus_JS_SUCCESS
|
||||
errorStr := ""
|
||||
if reason == rpc.JobTerminateReason_AGENT_LEFT_ROOM {
|
||||
status = livekit.JobStatus_JS_FAILED
|
||||
errorStr = "agent worker left the room"
|
||||
}
|
||||
|
||||
return w.UpdateJobStatus(&livekit.UpdateJobStatus{
|
||||
JobId: string(jobID),
|
||||
Status: status,
|
||||
Error: errorStr,
|
||||
})
|
||||
}
|
||||
|
||||
func (w *Worker) UpdateMetadata(metadata string) {
|
||||
w.logger.Debugw("worker metadata updated", nil, "metadata", metadata)
|
||||
}
|
||||
|
||||
func (w *Worker) IsClosed() bool {
|
||||
select {
|
||||
case <-w.closed:
|
||||
return true
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
func (w *Worker) Close() {
|
||||
w.mu.Lock()
|
||||
if w.IsClosed() {
|
||||
w.mu.Unlock()
|
||||
return
|
||||
}
|
||||
|
||||
w.logger.Infow("closing worker", "workerID", w.ID, "jobType", w.JobType, "agentName", w.AgentName)
|
||||
|
||||
close(w.closed)
|
||||
w.cancel()
|
||||
_ = w.conn.Close()
|
||||
w.mu.Unlock()
|
||||
}
|
||||
|
||||
func (w *Worker) HandleAvailability(res *livekit.AvailabilityResponse) error {
|
||||
w.mu.Lock()
|
||||
defer w.mu.Unlock()
|
||||
|
||||
jobID := livekit.JobID(res.JobId)
|
||||
availCh, ok := w.availability[jobID]
|
||||
if !ok {
|
||||
w.logger.Warnw("received availability response for unknown job", nil, "jobID", jobID)
|
||||
return nil
|
||||
}
|
||||
|
||||
availCh <- res
|
||||
delete(w.availability, jobID)
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (w *Worker) HandleUpdateJob(update *livekit.UpdateJobStatus) error {
|
||||
_, err := w.UpdateJobStatus(update)
|
||||
if err != nil {
|
||||
// treating this as a debug message only
|
||||
// this can happen if the Room closes first, which would delete the agent dispatch
|
||||
// that would mark the job as successful. subsequent updates from the same worker
|
||||
// would not be able to find the same jobID.
|
||||
w.logger.Debugw("received job update for unknown job", "jobID", update.JobId)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (w *Worker) UpdateJobStatus(update *livekit.UpdateJobStatus) (*livekit.JobState, error) {
|
||||
w.mu.Lock()
|
||||
defer w.mu.Unlock()
|
||||
|
||||
jobID := livekit.JobID(update.JobId)
|
||||
job, ok := w.runningJobs[jobID]
|
||||
if !ok {
|
||||
return nil, psrpc.NewErrorf(psrpc.NotFound, "received job update for unknown job")
|
||||
}
|
||||
|
||||
now := time.Now()
|
||||
job.State.UpdatedAt = now.UnixNano()
|
||||
|
||||
if job.State.Status == livekit.JobStatus_JS_PENDING && update.Status != livekit.JobStatus_JS_PENDING {
|
||||
job.State.StartedAt = now.UnixNano()
|
||||
}
|
||||
|
||||
job.State.Status = update.Status
|
||||
job.State.Error = update.Error
|
||||
|
||||
if JobStatusIsEnded(update.Status) {
|
||||
job.State.EndedAt = now.UnixNano()
|
||||
delete(w.runningJobs, jobID)
|
||||
|
||||
w.logger.Infow("job ended", "jobID", update.JobId, "status", update.Status, "error", update.Error)
|
||||
}
|
||||
|
||||
return proto.Clone(job.State).(*livekit.JobState), nil
|
||||
}
|
||||
|
||||
func (w *Worker) HandleSimulateJob(simulate *livekit.SimulateJobRequest) error {
|
||||
jobType := livekit.JobType_JT_ROOM
|
||||
if simulate.Participant != nil {
|
||||
jobType = livekit.JobType_JT_PUBLISHER
|
||||
}
|
||||
|
||||
job := &livekit.Job{
|
||||
Id: guid.New(guid.AgentJobPrefix),
|
||||
Type: jobType,
|
||||
Room: simulate.Room,
|
||||
Participant: simulate.Participant,
|
||||
Namespace: w.Namespace,
|
||||
AgentName: w.AgentName,
|
||||
}
|
||||
|
||||
go func() {
|
||||
_, err := w.AssignJob(w.ctx, job)
|
||||
if err != nil {
|
||||
w.logger.Errorw("unable to simulate job", err, "jobID", job.Id)
|
||||
}
|
||||
}()
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (w *Worker) HandleUpdateWorker(update *livekit.UpdateWorkerStatus) error {
|
||||
w.mu.Lock()
|
||||
defer w.mu.Unlock()
|
||||
|
||||
if status := update.Status; status != nil && w.status != *status {
|
||||
w.status = *status
|
||||
w.Logger().Debugw("worker status changed", "status", w.status)
|
||||
}
|
||||
w.load = update.GetLoad()
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (w *Worker) HandleMigrateJob(req *livekit.MigrateJobRequest) error {
|
||||
// TODO(theomonnom): On OSS this is not implemented
|
||||
// We could maybe just move a specific job to another worker
|
||||
return nil
|
||||
}
|
||||
Reference in New Issue
Block a user