fix prism runner race condition (#35024)

tosolveit · web-flow · commit 697e054b8e8d · 2025-05-30T10:08:06.000-07:00
diff --git a/sdks/go/pkg/beam/runners/prism/internal/jobservices/artifact.go b/sdks/go/pkg/beam/runners/prism/internal/jobservices/artifact.go
@@ -34,7 +34,9 @@ func (s *Server) ReverseArtifactRetrievalService(stream jobpb.ArtifactStagingSer
 	if err != nil {
 		return err
 	}
+	s.mu.RLock()
 	job := s.jobs[in.GetStagingToken()]
+	s.mu.RUnlock()
 
 	envs := job.Pipeline.GetComponents().GetEnvironments()
 	for _, env := range envs {
diff --git a/sdks/go/pkg/beam/runners/prism/internal/jobservices/server.go b/sdks/go/pkg/beam/runners/prism/internal/jobservices/server.go
@@ -43,7 +43,7 @@ type Server struct {
 	server *grpc.Server
 
 	// Job Management
-	mu    sync.Mutex
+	mu    sync.RWMutex
 	index uint32 // Use with atomics.
 	jobs  map[string]*Job
 

Original file line number	Diff line number	Diff line change
`@@ -34,7 +34,9 @@ func (s *Server) ReverseArtifactRetrievalService(stream jobpb.ArtifactStagingSer`
`34`	`34`	`if err != nil {`
`35`	`35`	`return err`
`36`	`36`	`}`
	`37`	`+ s.mu.RLock()`
`37`	`38`	`job := s.jobs[in.GetStagingToken()]`
	`39`	`+ s.mu.RUnlock()`
`38`	`40`
`39`	`41`	`envs := job.Pipeline.GetComponents().GetEnvironments()`
`40`	`42`	`for _, env := range envs {`