From 4faa5f02b769048cd740432aa48190b927152bec Mon Sep 17 00:00:00 2001 From: Jason Lernerman <43508353+u9g@users.noreply.github.com> Date: Wed, 21 Oct 2026 23:13:04 -0400 Subject: [PATCH] agent: route JobTerminate to the connection a job migrated to --- pkg/agent/agent_test.go | 81 +++++++++++++++++++++++++++++++++++ pkg/agent/testutils/server.go | 12 +++++- pkg/service/agenthandler.go | 18 +++++++- 3 files changed, 109 insertions(+), 2 deletions(-) diff --git a/pkg/agent/agent_test.go b/pkg/agent/agent_test.go index fab7656df..83fab87a5 100644 --- a/pkg/agent/agent_test.go +++ b/pkg/agent/agent_test.go @@ -279,6 +279,87 @@ func TestJobTerminateReleasesJob(t *testing.T) { }) } +// a reconnected worker reports its running jobs on its new connection, which must then receive their terminations. +func TestMigratedJobTerminate(t *testing.T) { + const agentName = "test_agent" + + // the server reads the ping only after handling the messages sent before it. + waitHandled := func(t *testing.T, w *testutils.AgentWorker) { + pongs := w.WorkerPongs.Observe() + defer pongs.Stop() + w.SendPing(&livekit.WorkerPing{Timestamp: time.Now().UnixMilli()}) + select { + case <-pongs.Events(): + case <-time.After(5 * time.Second): + require.Fail(t, "pong timeout") + } + } + + for _, tc := range []struct { + name string + closeBeforeMigrate bool + }{ + {"old connection closes before migration", true}, + {"old connection closes after migration", false}, + } { + t.Run(tc.name, func(t *testing.T) { + bus := psrpc.NewLocalMessageBus() + client := must.Get(rpc.NewAgentInternalClient(bus)) + server := testutils.NewTestServer(bus) + t.Cleanup(server.Close) + + old := server.SimulateAgentWorker() + old.Register(agentName, livekit.JobType_JT_ROOM) + waitHandled(t, old) + + job := &livekit.Job{ + Id: guid.New(guid.AgentJobPrefix), + DispatchId: guid.New(guid.AgentDispatchPrefix), + Type: livekit.JobType_JT_ROOM, + Room: &livekit.Room{}, + AgentName: agentName, + } + _, err := client.JobRequest(context.Background(), agentName, agent.RoomAgentTopic, job) + require.NoError(t, err) + + closeOld := func() { + require.NoError(t, old.Close()) + select { + case <-old.ConnectionDone(): + case <-time.After(5 * time.Second): + require.Fail(t, "old connection teardown timeout") + } + } + + if tc.closeBeforeMigrate { + closeOld() + } + reconnected := server.SimulateAgentWorker() + reconnected.Register(agentName, livekit.JobType_JT_ROOM) + reconnected.SendMigrateJob(&livekit.MigrateJobRequest{JobIds: []string{job.Id}}) + waitHandled(t, reconnected) + if !tc.closeBeforeMigrate { + closeOld() + } + + terminations := reconnected.JobTerminations.Observe() + defer terminations.Stop() + res, err := client.JobTerminate(context.Background(), job.Id, &rpc.JobTerminateRequest{ + JobId: job.Id, + Reason: rpc.JobTerminateReason_TERMINATION_REQUESTED, + }, psrpc.WithRequestTimeout(time.Second)) + require.NoError(t, err) + require.Equal(t, livekit.JobStatus_JS_SUCCESS, res.State.Status) + select { + case m := <-terminations.Events(): + require.Equal(t, job.Id, m.JobId) + case <-time.After(5 * time.Second): + require.Fail(t, "job termination timeout") + } + }) + } +} + func TestConnectionClosedOnDispatchError(t *testing.T) { t.Run("connection closed when unknown message type received", func(t *testing.T) { bus := psrpc.NewLocalMessageBus() diff --git a/pkg/agent/testutils/server.go b/pkg/agent/testutils/server.go index 30793d069..ed2ee9e4e 100644 --- a/pkg/agent/testutils/server.go +++ b/pkg/agent/testutils/server.go @@ -129,6 +129,7 @@ func (h *TestServer) SimulateAgentWorker(opts ...SimulatedWorkerOption) *AgentWo JobAssignments: events.NewObserverList[*livekit.JobAssignment](), JobTerminations: events.NewObserverList[*livekit.JobTermination](), WorkerPongs: events.NewObserverList[*livekit.WorkerPong](), + connectionDone: make(chan struct{}), } w.ctx, w.cancel = context.WithCancel(context.Background()) @@ -137,7 +138,10 @@ func (h *TestServer) SimulateAgentWorker(opts ...SimulatedWorkerOption) *AgentWo } ctx := service.WithAPIKey(o.Context, &auth.ClaimGrants{}, "test") - go h.HandleConnection(ctx, w, agent.MakeWorkerRegistration()) + go func() { + h.HandleConnection(ctx, w, agent.MakeWorkerRegistration()) + close(w.connectionDone) + }() return w } @@ -190,6 +194,7 @@ type AgentWorker struct { workerMessages chan *livekit.WorkerMessage serverMessages deque.Deque[*livekit.ServerMessage] jobs map[string]*AgentJob + connectionDone chan struct{} RegisterWorkerResponses *events.ObserverList[*livekit.RegisterWorkerResponse] AvailabilityRequests *events.ObserverList[*livekit.AvailabilityRequest] @@ -224,6 +229,11 @@ func (w *AgentWorker) Closed() <-chan struct{} { return w.fuse.Watch() } +// ConnectionDone returns a channel that is closed when the server has finished handling the connection +func (w *AgentWorker) ConnectionDone() <-chan struct{} { + return w.connectionDone +} + func (w *AgentWorker) SetReadDeadline(t time.Time) error { w.mu.Lock() defer w.mu.Unlock() diff --git a/pkg/service/agenthandler.go b/pkg/service/agenthandler.go index 81067896a..f8c7d0477 100644 --- a/pkg/service/agenthandler.go +++ b/pkg/service/agenthandler.go @@ -375,7 +375,10 @@ func (h *AgentHandler) deregisterWorker(w *agent.Worker) { jobs := w.RunningJobs() for jobID := range jobs { - h.deregisterJob(jobID) + // a job migrated to another connection is routed there. + if h.jobToWorker[jobID] == w { + h.deregisterJob(jobID) + } } } @@ -579,6 +582,19 @@ type agentHandlerWorker struct { *agent.Worker } +func (w *agentHandlerWorker) HandleMigrateJob(req *livekit.MigrateJobRequest) error { + if err := w.Worker.HandleMigrateJob(req); err != nil { + return err + } + + for _, id := range req.JobIds { + if err := w.h.registerJob(w.Worker, livekit.JobID(id)); err != nil { + w.Logger().Errorw("failed to register JobTerminate handler", err, "jobID", id) + } + } + return nil +} + func (w *agentHandlerWorker) HandleUpdateJob(update *livekit.UpdateJobStatus) error { if err := w.Worker.HandleUpdateJob(update); err != nil { return err