Set b.BundleErr correctly when dynamic split happens.

Author: shunping

sdks/go/pkg/beam/runners/prism/internal/execute_test.go

@@ -541,6 +541,27 @@ func TestFailureHang(t *testing.T) {
 	}
 }
 
+func TestFailure_SplitError(t *testing.T) {
+	initRunner(t)
+
+	p, s := beam.NewPipelineWithRoot()
+	configs := beam.Create(s, SourceConfig{NumElements: 100, InitialSplits: 1})
+	col := beam.ParDo(s, &slowFailSDF{}, configs)
+	beam.ParDo(s, &int64Check{Name: "sdf_fail", Want: []int{}}, col)
+
+	// Set a short timeout so the test fails quickly if the pipeline hangs due to split error handling bug
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	defer cancel()
+
+	_, err := executeWithT(ctx, t, p)
+	if err == nil {
+		t.Fatalf("expected pipeline failure, but got a success")
+	}
+	if want := "intentional split error from tracker"; !strings.Contains(err.Error(), want) {
+		t.Fatalf("expected pipeline failure with %q, but was %v", want, err)
+	}
+}
+
 func TestRunner_Passert(t *testing.T) {
 	initRunner(t)
 	tests := []struct {

sdks/go/pkg/beam/runners/prism/internal/stage.go

@@ -240,7 +240,7 @@ progress:
 				sr, err := b.Split(ctx, wk, 0.5 /* fraction of remainder */, nil /* allowed splits */)
 				if err != nil {
 					slog.Warn("SDK Error from split, aborting splits and failing bundle", "bundle", rb, "error", err.Error())
-					if b.BundleErr != nil {
+					if b.BundleErr == nil {
 						b.BundleErr = err
 					}
 					return b.BundleErr

sdks/go/pkg/beam/runners/prism/internal/testdofns_test.go

@@ -64,6 +64,7 @@ func init() {
 	register.Function2x1(combineIntSum)
 
 	register.DoFn3x1[*sdf.LockRTracker, SourceConfig, func(int64), error]((*intRangeFn)(nil))
+	register.DoFn4x1[context.Context, *sdf.LockRTracker, SourceConfig, func(int64), error]((*slowFailSDF)(nil))
 	register.Emitter1[int64]()
 	register.Emitter2[int64, int64]()
 }
@@ -404,3 +405,37 @@ func (fn *selfCheckpointingDoFn) ProcessElement(rt *sdf.LockRTracker, _ []byte,
 		}
 	}
 }
+
+type errorSplitTracker struct {
+	*offsetrange.Tracker
+}
+
+func (t *errorSplitTracker) TrySplit(fraction float64) (any, any, error) {
+	return nil, nil, fmt.Errorf("intentional split error from tracker")
+}
+
+type slowFailSDF struct{}
+
+func (fn *slowFailSDF) CreateInitialRestriction(config SourceConfig) offsetrange.Restriction {
+	return offsetrange.Restriction{Start: 0, End: config.NumElements}
+}
+
+func (fn *slowFailSDF) SplitRestriction(config SourceConfig, rest offsetrange.Restriction) []offsetrange.Restriction {
+	return rest.EvenSplits(config.InitialSplits)
+}
+
+func (fn *slowFailSDF) RestrictionSize(_ SourceConfig, rest offsetrange.Restriction) float64 {
+	return rest.Size()
+}
+
+func (fn *slowFailSDF) CreateTracker(rest offsetrange.Restriction) *sdf.LockRTracker {
+	return sdf.NewLockRTracker(&errorSplitTracker{Tracker: offsetrange.NewTracker(rest)})
+}
+
+func (fn *slowFailSDF) ProcessElement(ctx context.Context, rt *sdf.LockRTracker, config SourceConfig, emit func(int64)) error {
+	fmt.Println("DEBUG: slowFailSDF.ProcessElement invoked")
+	for i := rt.GetRestriction().(offsetrange.Restriction).Start; rt.TryClaim(i); i++ {
+		<-ctx.Done()
+	}
+	return nil
+}