Fix race condition that causes the test flaky. by shunping · Pull Request #36235 · apache/beam

shunping · 2025-09-22T20:10:41Z

The Precommit Go test has been flaky (https://github.com/apache/beam/actions/runs/17922206057/job/50959740198).

The error is related to TestServer_RunThenCancel and the error message is below.

    server_test.go:142: server.GetState() = RUNNING, want CANCELLED

After some investigation, I found it is related to a race condition between the main thread and the Server go routine.

The cancel call is made after the server starts to run.

beam/sdks/go/pkg/beam/runners/prism/internal/jobservices/server_test.go

Lines 114 to 124 in 08b0572

    
           runResp, err := undertest.Run(ctx, &jobpb.RunJobRequest{ 
        
           	PreparationId: resp.GetPreparationId(), 
        
           }) 
        
           if err != nil { 
        
           	t.Fatalf("server.Run() = %v, want nil", err) 
        
           } 
        
           if got := runResp.GetJobId(); got == "" { 
        
           	t.Fatalf("server.Run() = returned empty preparation ID, want non-empty") 
        
           } 
        
           cancelResp, err := undertest.Cancel(ctx, &jobpb.CancelJobRequest{

.

The race condition occurs in the following order:

The server go routine started to run in the call of undertest.Run() to the point of setting the state to RUNNING.
The main thread starts to trigger a cancel request by calling undertest.Cancel.

However, before job.CancelFn(ErrCancel) is called, the server go rountine finishes the check on Line 89 and return.

beam/sdks/go/pkg/beam/runners/prism/internal/jobservices/server_test.go

Lines 86 to 94 in 08b0572

    
           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 
        
           	} 
        
           })

.

Then when the main thread tries to get the server state at the end, the state will still be RUNNING.

In this PR, I added the logic to wait until the context is done before checking the cause.

Furthermore, after fixing the previous race condition, there is yet another: when undertest.Cancel is called, it is not guaranteed that the job reaches the RUNNING state. So additional logic is needed to wait for the RUNNING state before triggering the cancel.

github-actions · 2025-09-22T21:06:24Z

Assigning reviewers:

R: @jrmccluskey for label go.

Note: If you would like to opt out of this review, comment assign to next reviewer.

Available commands:

stop reviewer notifications - opt out of the automated review tooling
remind me after tests pass - tag the comment author after tests pass
waiting on author - shift the attention set back to the author (any comment or push by the author will return the attention set to the reviewers)

The PR bot will only process comments in the main thread (not review comments).

lostluck

Seems reasonable. Have you run this 100x with the race detector on and off to ensure it's deflaked?

Alternatively (or simultaneously), there's a gap we discussed offline about how the job.sendState function works. If we have it return the "decided" state, some of the other checks would be less racy as well.

lostluck · 2025-09-22T21:44:12Z

+	)
+	var jobIsRunning bool
+	for range maxRetries {
+		stateResp, err := undertest.GetState(ctx, &jobpb.GetJobStateRequest{JobId: runResp.GetJobId()})


I'd use the streaming version of the call, not the one-off . Then there's less dynamic waiting for things to happen for up to a second.

I have replaced this waiting part with a waitgroup, which is more concise and clear.

shunping · 2025-09-23T01:31:19Z

Seems reasonable. Have you run this 100x with the race detector on and off to ensure it's deflaked?

Yes I tried the fix 1000x with the race detector and do not see any failure.

Alternatively (or simultaneously), there's a gap we discussed offline about how the job.sendState function works. If we have it return the "decided" state, some of the other checks would be less racy as well.

If we change job.sendState to return a state after processing through a state machine, then a lot of exported apis in job.go, such as the following, would need changing to returning the state. Is that ok?

func (j *Job) Start() {
	j.sendState(jobpb.JobState_STARTING)
}

// Running indicates that the job is executing.
func (j *Job) Running() {
	j.sendState(jobpb.JobState_RUNNING)
}

// Done indicates that the job completed successfully.
func (j *Job) Done() {
	j.sendState(jobpb.JobState_DONE)
}

// Canceling indicates that the job is canceling.
func (j *Job) Canceling() {
	j.sendState(jobpb.JobState_CANCELLING)
}

lostluck · 2025-09-23T01:39:50Z

If we change job.sendState to return a state after processing through a state machine, then a lot of exported apis in job.go, such as the following, would need changing to returning the state. Is that ok?

func (j *Job) Start() {
	j.sendState(jobpb.JobState_STARTING)
}

// Running indicates that the job is executing.
func (j *Job) Running() {
	j.sendState(jobpb.JobState_RUNNING)
}

// Done indicates that the job completed successfully.
func (j *Job) Done() {
	j.sendState(jobpb.JobState_DONE)
}

// Canceling indicates that the job is canceling.
func (j *Job) Canceling() {
	j.sendState(jobpb.JobState_CANCELLING)
}

Mechanically, there's no problem with it, as these are internal packages, and we can change all the callers.

But there's not need to change all the callers right away, just for the cancellation call where it appears to matter, and we look up the state immediately. Most other calls are never inspecting the job status immediately afterwards, setting and forgetting as the job progresses through the pipeline execution.

It's OK not to change the other calls.

…ting waitgroup for clarity.

shunping · 2025-09-23T02:58:23Z

Mechanically, there's no problem with it, as these are internal packages, and we can change all the callers.

But there's not need to change all the callers right away, just for the cancellation call where it appears to matter, and we look up the state immediately. Most other calls are never inspecting the job status immediately afterwards, setting and forgetting as the job progresses through the pipeline execution.

It's OK not to change the other calls.

Let's use this PR to deflake a test first. I opened a new task for adding the state machine in sendState: 36241.

Fix race condition that causes the test flaky.

dbed708

github-actions Bot added go runners prism labels Sep 22, 2025

shunping marked this pull request as ready for review September 22, 2025 20:17

github-actions Bot added the Next Action: Reviewers label Sep 22, 2025

lostluck approved these changes Sep 22, 2025

View reviewed changes

Use another waitgroup to ensure a job is running. Also rename an exis…

04d79a0

…ting waitgroup for clarity.

shunping mentioned this pull request Sep 23, 2025

[Task]: Change sendState to handle state transition in a state machine #36241

Closed

17 tasks

shunping merged commit a32f2a3 into apache:master Sep 23, 2025
10 checks passed

shunping deleted the fix-test-server-run-then-cancel branch September 23, 2025 03:19

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Fix race condition that causes the test flaky. - #36235

Fix race condition that causes the test flaky.#36235
shunping merged 2 commits into
apache:masterfrom
shunping:fix-test-server-run-then-cancel

shunping commented Sep 22, 2025 •

edited

Loading

Uh oh!

github-actions Bot commented Sep 22, 2025

Uh oh!

lostluck left a comment

Uh oh!

lostluck Sep 22, 2025

Uh oh!

shunping Sep 23, 2025

Uh oh!

shunping commented Sep 23, 2025 •

edited

Loading

Uh oh!

lostluck commented Sep 23, 2025

Uh oh!

shunping commented Sep 23, 2025 •

edited

Loading

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

	runResp, err := undertest.Run(ctx, &jobpb.RunJobRequest{
	PreparationId: resp.GetPreparationId(),
	})
	if err != nil {
	t.Fatalf("server.Run() = %v, want nil", err)
	}
	if got := runResp.GetJobId(); got == "" {
	t.Fatalf("server.Run() = returned empty preparation ID, want non-empty")
	}

	cancelResp, err := undertest.Cancel(ctx, &jobpb.CancelJobRequest{

	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
	}
	})

Uh oh!

Conversation

shunping commented Sep 22, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

github-actions Bot commented Sep 22, 2025

Uh oh!

lostluck left a comment

Choose a reason for hiding this comment

Uh oh!

lostluck Sep 22, 2025

Choose a reason for hiding this comment

Uh oh!

shunping Sep 23, 2025

Choose a reason for hiding this comment

Uh oh!

shunping commented Sep 23, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

lostluck commented Sep 23, 2025

Uh oh!

shunping commented Sep 23, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

shunping commented Sep 22, 2025 •

edited

Loading

shunping commented Sep 23, 2025 •

edited

Loading

shunping commented Sep 23, 2025 •

edited

Loading