diff --git a/pkg/agent/worker.go b/pkg/agent/worker.go index bdd5f5c40..7b7411803 100644 --- a/pkg/agent/worker.go +++ b/pkg/agent/worker.go @@ -662,8 +662,26 @@ func (w *Worker) Draining() bool { return w.draining } +// MigrateJobRequest carries only job IDs, so a migrated job is tracked without its room or dispatch details. 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 + w.mu.Lock() + defer w.mu.Unlock() + + now := time.Now().UnixNano() + for _, id := range req.JobIds { + jobID := livekit.JobID(id) + if _, ok := w.runningJobs[jobID]; ok { + continue + } + w.runningJobs[jobID] = &livekit.Job{ + Id: id, + State: &livekit.JobState{ + Status: livekit.JobStatus_JS_RUNNING, + UpdatedAt: now, + WorkerId: w.ID, + AgentId: w.AgentID, + }, + } + } return nil } diff --git a/pkg/agent/worker_test.go b/pkg/agent/worker_test.go index 5b9947892..6167f742d 100644 --- a/pkg/agent/worker_test.go +++ b/pkg/agent/worker_test.go @@ -12,6 +12,8 @@ import ( "github.com/livekit/livekit-server/pkg/agent" "github.com/livekit/livekit-server/pkg/agent/endpoint/wire" "github.com/livekit/protocol/livekit" + "github.com/livekit/protocol/logger" + "github.com/livekit/protocol/rpc" ) type nopSignalConn struct{} @@ -76,3 +78,33 @@ func TestHandleRegisterWithoutEndpointsIgnoresReservedNames(t *testing.T) { }), name) } } + +type recordingSignalConn struct { + nopSignalConn + sent []*livekit.ServerMessage +} + +func (c *recordingSignalConn) WriteServerMessage(m *livekit.ServerMessage) (int, error) { + c.sent = append(c.sent, m) + return 0, nil +} + +// jobs reported on a new connection can be terminated and updated there. +func TestHandleMigrateJob(t *testing.T) { + conn := &recordingSignalConn{} + w := agent.NewWorker(agent.WorkerRegistration{}, "key", "secret", conn, logger.GetLogger()) + + require.NoError(t, w.HandleMigrateJob(&livekit.MigrateJobRequest{JobIds: []string{"AJ_terminated", "AJ_updated"}})) + require.Equal(t, 2, w.RunningJobCount()) + + state, err := w.TerminateJob("AJ_terminated", rpc.JobTerminateReason_TERMINATION_REQUESTED) + require.NoError(t, err) + require.Equal(t, livekit.JobStatus_JS_SUCCESS, state.Status) + require.Len(t, conn.sent, 1) + require.Equal(t, "AJ_terminated", conn.sent[0].GetTermination().GetJobId()) + + state, err = w.UpdateJobStatus(&livekit.UpdateJobStatus{JobId: "AJ_updated", Status: livekit.JobStatus_JS_SUCCESS}) + require.NoError(t, err) + require.Equal(t, livekit.JobStatus_JS_SUCCESS, state.Status) + require.Zero(t, w.RunningJobCount()) +}