From dbed708b83553664061add45628eabd8b164047e Mon Sep 17 00:00:00 2001
From: Shunping Huang <shunping@google.com>
Date: Mon, 22 Sep 2025 15:43:06 -0400
Subject: [PATCH 1/2] Fix race condition that causes the test flaky.

---
 .../prism/internal/jobservices/server_test.go | 43 +++++++++++++++++--
 1 file changed, 39 insertions(+), 4 deletions(-)

diff --git a/sdks/go/pkg/beam/runners/prism/internal/jobservices/server_test.go b/sdks/go/pkg/beam/runners/prism/internal/jobservices/server_test.go
index fb72048d478c..4deaab9e1381 100644
--- a/sdks/go/pkg/beam/runners/prism/internal/jobservices/server_test.go
+++ b/sdks/go/pkg/beam/runners/prism/internal/jobservices/server_test.go
@@ -20,6 +20,7 @@ import (
 	"errors"
 	"sync"
 	"testing"
+	"time"
 
 	jobpb "github.com/apache/beam/sdks/v2/go/pkg/beam/model/jobmanagement_v1"
 	pipepb "github.com/apache/beam/sdks/v2/go/pkg/beam/model/pipeline_v1"
@@ -86,10 +87,20 @@ func TestServer_RunThenCancel(t *testing.T) {
 	undertest := NewServer(0, func(j *Job) {
 		defer called.Done()
 		j.state.Store(jobpb.JobState_RUNNING)
-		if errors.Is(context.Cause(j.RootCtx), ErrCancel) {
-			j.SendMsg("pipeline canceled " + j.String())
-			j.Canceled()
-			return
+		for {
+			select {
+			case <-j.RootCtx.Done():
+				// The context was canceled. The goroutine "woke up."
+				// We check the reason for the cancellation.
+				if errors.Is(context.Cause(j.RootCtx), ErrCancel) {
+					j.SendMsg("pipeline canceled " + j.String())
+					j.Canceled()
+				}
+				return
+
+			case <-time.After(1 * time.Second):
+				// Just wait a little bit to receive the cancel signal
+			}
 		}
 	})
 	ctx := context.Background()
@@ -121,6 +132,30 @@ func TestServer_RunThenCancel(t *testing.T) {
 		t.Fatalf("server.Run() = returned empty preparation ID, want non-empty")
 	}
 
+	// Wait for the job to be in the RUNNING state before we cancel it.
+	const (
+		maxRetries = 10
+		retrySleep = 100 * time.Millisecond
+	)
+	var jobIsRunning bool
+	for range maxRetries {
+		stateResp, err := undertest.GetState(ctx, &jobpb.GetJobStateRequest{JobId: runResp.GetJobId()})
+		if err != nil {
+			t.Fatalf("server.GetState() during poll = %v, want nil", err)
+		}
+
+		if stateResp.State == jobpb.JobState_RUNNING {
+			jobIsRunning = true
+			break // Success! Job is running.
+		}
+		// Wait a bit before polling again
+		time.Sleep(retrySleep)
+	}
+
+	if !jobIsRunning {
+		t.Fatalf("Job did not enter RUNNING state after %v", maxRetries*retrySleep)
+	}
+
 	cancelResp, err := undertest.Cancel(ctx, &jobpb.CancelJobRequest{
 		JobId: runResp.GetJobId(),
 	})

From 04d79a020fdcbd6bba936ce3614af9536fb85305 Mon Sep 17 00:00:00 2001
From: Shunping Huang <shunping@google.com>
Date: Mon, 22 Sep 2025 22:05:19 -0400
Subject: [PATCH 2/2] Use another waitgroup to ensure a job is running. Also
 rename an existing waitgroup for clarity.

---
 .../prism/internal/jobservices/server_test.go | 39 ++++++-------------
 1 file changed, 11 insertions(+), 28 deletions(-)

diff --git a/sdks/go/pkg/beam/runners/prism/internal/jobservices/server_test.go b/sdks/go/pkg/beam/runners/prism/internal/jobservices/server_test.go
index 4deaab9e1381..80b38507539b 100644
--- a/sdks/go/pkg/beam/runners/prism/internal/jobservices/server_test.go
+++ b/sdks/go/pkg/beam/runners/prism/internal/jobservices/server_test.go
@@ -82,11 +82,14 @@ func TestServer_JobLifecycle(t *testing.T) {
 
 // Validates that invoking Cancel cancels a running job.
 func TestServer_RunThenCancel(t *testing.T) {
-	var called sync.WaitGroup
-	called.Add(1)
+	var canceled sync.WaitGroup
+	var running sync.WaitGroup
+	canceled.Add(1)
+	running.Add(1)
 	undertest := NewServer(0, func(j *Job) {
-		defer called.Done()
-		j.state.Store(jobpb.JobState_RUNNING)
+		defer canceled.Done()
+		j.Running()
+		running.Done()
 		for {
 			select {
 			case <-j.RootCtx.Done():
@@ -132,29 +135,8 @@ func TestServer_RunThenCancel(t *testing.T) {
 		t.Fatalf("server.Run() = returned empty preparation ID, want non-empty")
 	}
 
-	// Wait for the job to be in the RUNNING state before we cancel it.
-	const (
-		maxRetries = 10
-		retrySleep = 100 * time.Millisecond
-	)
-	var jobIsRunning bool
-	for range maxRetries {
-		stateResp, err := undertest.GetState(ctx, &jobpb.GetJobStateRequest{JobId: runResp.GetJobId()})
-		if err != nil {
-			t.Fatalf("server.GetState() during poll = %v, want nil", err)
-		}
-
-		if stateResp.State == jobpb.JobState_RUNNING {
-			jobIsRunning = true
-			break // Success! Job is running.
-		}
-		// Wait a bit before polling again
-		time.Sleep(retrySleep)
-	}
-
-	if !jobIsRunning {
-		t.Fatalf("Job did not enter RUNNING state after %v", maxRetries*retrySleep)
-	}
+	// wait until the job is running (i.e. j.Running() is called)
+	running.Wait()
 
 	cancelResp, err := undertest.Cancel(ctx, &jobpb.CancelJobRequest{
 		JobId: runResp.GetJobId(),
@@ -167,7 +149,8 @@ func TestServer_RunThenCancel(t *testing.T) {
 		t.Fatalf("server.Canceling() = %v, want %v", cancelResp.State, jobpb.JobState_CANCELLING)
 	}
 
-	called.Wait()
+	// wait until the job is canceled (i.e. j.Canceled() is called)
+	canceled.Wait()
 
 	stateResp, err := undertest.GetState(ctx, &jobpb.GetJobStateRequest{JobId: runResp.GetJobId()})
 	if err != nil {