diff --git a/pkg/agent/agent_test.go b/pkg/agent/agent_test.go index fab7656df..eac0d8cb4 100644 --- a/pkg/agent/agent_test.go +++ b/pkg/agent/agent_test.go @@ -60,6 +60,41 @@ func TestAgent(t *testing.T) { require.Fail(t, "job assignment timeout") } }) + + t.Run("dispatch burst does not wait for the CheckEnabled timeout", func(t *testing.T) { + bus := psrpc.NewLocalMessageBus() + server := testutils.NewTestServer(bus) + t.Cleanup(server.Close) + + rpcClient := must.Get(rpc.NewAgentInternalClient(bus)) + registered := must.Get(rpcClient.SubscribeWorkerRegistered(context.Background(), agent.DefaultHandlerNamespace)) + defer registered.Close() + + worker := server.SimulateAgentWorker(testutils.WithJobLoad(testutils.NewStableJobLoad(0))) + worker.Register(testAgentName, livekit.JobType_JT_ROOM) + select { + case <-registered.Channel(): + case <-time.After(time.Second): + require.Fail(t, "registration timeout") + } + + client := must.Get(agent.NewAgentClient(bus, agent.Config{})) + t.Cleanup(func() { _ = client.Stop() }) + + const jobCount = 60 // more than the client's 50 pool workers + start := time.Now() + for range jobCount { + go client.LaunchJob(context.Background(), &agent.JobRequest{ + JobType: livekit.JobType_JT_ROOM, + Room: &livekit.Room{}, + AgentName: testAgentName, + }) + } + require.Eventually(t, func() bool { + return len(worker.Jobs()) == jobCount + }, 2*agent.CheckEnabledTimeout, 10*time.Millisecond) + require.Less(t, time.Since(start), 2*time.Second) + }) } func testBatchJobRequest(t require.TestingT, batchSize int, totalJobs int, client rpc.AgentInternalClient, workers []*testutils.AgentWorker) <-chan struct{} { diff --git a/pkg/agent/client.go b/pkg/agent/client.go index 383729f29..478d99e77 100644 --- a/pkg/agent/client.go +++ b/pkg/agent/client.go @@ -240,7 +240,7 @@ func (c *agentClient) getDispatcher(agName string, jobType livekit.JobType) *ser } done := make(chan *serverutils.IncrementalDispatcher[string], 1) - c.workers.Submit(func() { + go func() { agentNames.ForEach(func(ag string) { if ag == agName { select { @@ -253,7 +253,7 @@ func (c *agentClient) getDispatcher(agName string, jobType livekit.JobType) *ser case done <- nil: default: } - }) + }() return <-done }