[Prism] Terminate Job with CancelFn instead of panic (#31599)

* Refactor elementmanager Bundles with upstream CancelCauseFunc

* Fix minor edits

sdks/go/pkg/beam/runners/prism/internal/engine/elementmanager.go

@@ -34,6 +34,7 @@ import (
  "github.com/apache/beam/sdks/v2/go/pkg/beam/core/graph/window"
  "github.com/apache/beam/sdks/v2/go/pkg/beam/core/runtime/exec"
  "github.com/apache/beam/sdks/v2/go/pkg/beam/core/typex"
+ "github.com/apache/beam/sdks/v2/go/pkg/beam/internal/errors"
  "golang.org/x/exp/maps"
  "golang.org/x/exp/slog"
 )
@@ -290,7 +291,7 @@ func (rb RunBundle) LogValue() slog.Value {
 // Bundles is the core execution loop. It produces a sequences of bundles able to be executed.
 // The returned channel is closed when the context is canceled, or there are no pending elements
 // remaining.
-func (em *ElementManager) Bundles(ctx context.Context, nextBundID func() string) <-chan RunBundle {
+func (em *ElementManager) Bundles(ctx context.Context, upstreamCancelFn context.CancelCauseFunc, nextBundID func() string) <-chan RunBundle {
  runStageCh := make(chan RunBundle)
  ctx, cancelFn := context.WithCancelCause(ctx)
  go func() {
@@ -384,7 +385,9 @@ func (em *ElementManager) Bundles(ctx context.Context, nextBundID func() string)
  }
  }
  }
- em.checkForQuiescence(advanced)
+ if err := em.checkForQuiescence(advanced); err != nil {
+ upstreamCancelFn(err)
+ }
  }
  }()
  return runStageCh
@@ -400,11 +403,11 @@ func (em *ElementManager) Bundles(ctx context.Context, nextBundID func() string)
 // executing off the next TestStream event.
 //
 // Must be called while holding em.refreshCond.L.
-func (em *ElementManager) checkForQuiescence(advanced set[string]) {
+func (em *ElementManager) checkForQuiescence(advanced set[string]) error {
  defer em.refreshCond.L.Unlock()
  if len(em.inprogressBundles) > 0 {
  // If there are bundles in progress, then there may be watermark refreshes when they terminate.
- return
+ return nil
  }
  if len(em.watermarkRefreshes) > 0 {
  // If there are watermarks to refresh, we aren't yet stuck.
@@ -414,12 +417,12 @@ func (em *ElementManager) checkForQuiescence(advanced set[string]) {
  slog.Int("refreshCount", len(em.watermarkRefreshes)),
  slog.Int64("pendingElementCount", v),
  )
- return
+ return nil
  }
  if em.testStreamHandler == nil && len(em.processTimeEvents.events) > 0 {
  // If there's no test stream involved, and processing time events exist, then
  // it's only a matter of time.
- return
+ return nil
  }
  // The job has quiesced!
 
@@ -433,20 +436,20 @@ func (em *ElementManager) checkForQuiescence(advanced set[string]) {
  // Note: it's a prism bug if test stream never causes a refresh to occur for a given event.
  // It's not correct to move to the next event if no refreshes would occur.
  if len(em.watermarkRefreshes) > 0 {
- return
+ return nil
  } else if _, ok := nextEvent.(tsProcessingTimeEvent); ok {
  // It's impossible to fully control processing time SDK side handling for processing time
  // Runner side, so we specialize refresh handling here to avoid spuriously getting stuck.
  em.watermarkRefreshes.insert(em.testStreamHandler.ID)
- return
+ return nil
  }
  // If there are no refreshes, then there's no mechanism to make progress, so it's time to fast fail.
  }
 
  v := em.livePending.Load()
  if v == 0 {
  // Since there are no further pending elements, the job will be terminating successfully.
- return
+ return nil
  }
  // The job is officially stuck. Fail fast and produce debugging information.
  // Jobs must never get stuck so this indicates a bug in prism to be investigated.
@@ -469,7 +472,7 @@ func (em *ElementManager) checkForQuiescence(advanced set[string]) {
  upS := em.pcolParents[upPCol]
  stageState = append(stageState, fmt.Sprintln(id, "watermark in", inW, "out", outW, "upstream", upW, "from", upS, "pending", ss.pending, "byKey", ss.pendingByKeys, "inprogressKeys", ss.inprogressKeys, "byBundle", ss.inprogressKeysByBundle, "holds", ss.watermarkHolds.heap, "holdCounts", ss.watermarkHolds.counts, "holdsInBundle", ss.inprogressHoldsByBundle, "pttEvents", ss.processingTimeTimers.toFire))
  }
- panic(fmt.Sprintf("nothing in progress and no refreshes with non zero pending elements: %v\n%v", v, strings.Join(stageState, "")))
+ return errors.Errorf("nothing in progress and no refreshes with non zero pending elements: %v\n%v", v, strings.Join(stageState, ""))
 }
 
 // InputForBundle returns pre-allocated data for the given bundle, encoding the elements using

sdks/go/pkg/beam/runners/prism/internal/engine/elementmanager_test.go

@@ -316,6 +316,7 @@ func TestStageState_updateWatermarks(t *testing.T) {
 
 func TestElementManager(t *testing.T) {
  t.Run("impulse", func(t *testing.T) {
+ ctx, cancelFn := context.WithCancelCause(context.Background())
  em := NewElementManager(Config{})
  em.AddStage("impulse", nil, []string{"output"}, nil)
  em.AddStage("dofn", []string{"output"}, nil, nil)
@@ -327,7 +328,7 @@ func TestElementManager(t *testing.T) {
  }
 
  var i int
- ch := em.Bundles(context.Background(), func() string {
+ ch := em.Bundles(ctx, cancelFn, func() string {
  defer func() { i++ }()
  return fmt.Sprintf("%v", i)
  })
@@ -371,14 +372,15 @@ func TestElementManager(t *testing.T) {
  }
 
  t.Run("dofn", func(t *testing.T) {
+ ctx, cancelFn := context.WithCancelCause(context.Background())
  em := NewElementManager(Config{})
  em.AddStage("impulse", nil, []string{"input"}, nil)
  em.AddStage("dofn1", []string{"input"}, []string{"output"}, nil)
  em.AddStage("dofn2", []string{"output"}, nil, nil)
  em.Impulse("impulse")
 
  var i int
- ch := em.Bundles(context.Background(), func() string {
+ ch := em.Bundles(ctx, cancelFn, func() string {
  defer func() { i++ }()
  t.Log("generating bundle", i)
  return fmt.Sprintf("%v", i)
@@ -422,14 +424,15 @@ func TestElementManager(t *testing.T) {
  })
 
  t.Run("side", func(t *testing.T) {
+ ctx, cancelFn := context.WithCancelCause(context.Background())
  em := NewElementManager(Config{})
  em.AddStage("impulse", nil, []string{"input"}, nil)
  em.AddStage("dofn1", []string{"input"}, []string{"output"}, nil)
  em.AddStage("dofn2", []string{"input"}, nil, []LinkID{{Transform: "dofn2", Global: "output", Local: "local"}})
  em.Impulse("impulse")
 
  var i int
- ch := em.Bundles(context.Background(), func() string {
+ ch := em.Bundles(ctx, cancelFn, func() string {
  defer func() { i++ }()
  t.Log("generating bundle", i)
  return fmt.Sprintf("%v", i)
@@ -473,13 +476,14 @@ func TestElementManager(t *testing.T) {
  }
  })
  t.Run("residual", func(t *testing.T) {
+ ctx, cancelFn := context.WithCancelCause(context.Background())
  em := NewElementManager(Config{})
  em.AddStage("impulse", nil, []string{"input"}, nil)
  em.AddStage("dofn", []string{"input"}, nil, nil)
  em.Impulse("impulse")
 
  var i int
- ch := em.Bundles(context.Background(), func() string {
+ ch := em.Bundles(ctx, cancelFn, func() string {
  defer func() { i++ }()
  t.Log("generating bundle", i)
  return fmt.Sprintf("%v", i)

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

@@ -335,7 +335,7 @@ func executePipeline(ctx context.Context, wks map[string]*worker.W, j *jobservic
  eg.SetLimit(8)
 
  var instID uint64
- bundles := em.Bundles(egctx, func() string {
+ bundles := em.Bundles(egctx, j.CancelFn, func() string {
  return fmt.Sprintf("inst%03d", atomic.AddUint64(&instID, 1))
  })
  for {