Move wk.Stop() to context cancel.

apache · Sep 18, 2023 · 6dc03eb · 6dc03eb
1 parent 3024ec2
commit 6dc03eb
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 3 deletions.
diff --git a/sdks/go/pkg/beam/runners/prism/internal/environments.go b/sdks/go/pkg/beam/runners/prism/internal/environments.go
@@ -99,6 +99,7 @@ func externalEnvironment(ctx context.Context, ep *pipepb.ExternalPayload, wk *wo
 	pool.StopWorker(context.Background(), &fnpb.StopWorkerRequest{
 		WorkerId: wk.ID,
 	})
+	wk.Stop()
 }
 
 func dockerEnvironment(ctx context.Context, logger *slog.Logger, dp *pipepb.DockerPayload, wk *worker.W, artifactEndpoint string) error {
@@ -170,6 +171,7 @@ func dockerEnvironment(ctx context.Context, logger *slog.Logger, dp *pipepb.Dock
 	// Start goroutine to wait on container state.
 	go func() {
 		defer cli.Close()
+		defer wk.Stop()
 
 		statusCh, errCh := cli.ContainerWait(ctx, containerID, container.WaitConditionNotRunning)
 		select {

diff --git a/sdks/go/pkg/beam/runners/prism/internal/execute.go b/sdks/go/pkg/beam/runners/prism/internal/execute.go
@@ -71,9 +71,6 @@ func RunPipeline(j *jobservices.Job) {
 	}
 	j.SendMsg("pipeline completed " + j.String())
 
-	// Stop the worker.
-	wk.Stop()
-
 	j.SendMsg("terminating " + j.String())
 	j.Done()
 }