Use child context in watchAgents (#5888)

Signed-off-by: Kevin Su <[email protected]>
flyteorg · Oct 23, 2024 · 0472cbb · 0472cbb
1 parent 9d63d0a
commit 0472cbb
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/flyteplugins/go/tasks/plugins/webapi/agent/plugin.go b/flyteplugins/go/tasks/plugins/webapi/agent/plugin.go
@@ -337,8 +337,10 @@ func (p *Plugin) getAsyncAgentClient(ctx context.Context, agent *Deployment) (se
 
 func (p *Plugin) watchAgents(ctx context.Context, agentService *core.AgentService) {
 	go wait.Until(func() {
-		clientSet := getAgentClientSets(ctx)
-		agentRegistry := getAgentRegistry(ctx, clientSet)
+		childCtx, cancel := context.WithCancel(ctx)
+		defer cancel()
+		clientSet := getAgentClientSets(childCtx)
+		agentRegistry := getAgentRegistry(childCtx, clientSet)
 		p.setRegistry(agentRegistry)
 		agentService.SetSupportedTaskType(maps.Keys(agentRegistry))
 	}, p.cfg.PollInterval.Duration, ctx.Done())