[Prism] Schedule consumers of a self checkpointing source

Fixes #39446.

An unbounded source SDF that returns a process continuation residual has
that residual re-queued as a pending element carrying the input element's
event time, MinTimestamp for an Impulse rooted source, so the stage's
watermark never advances. updateWatermarks returns no refreshes when the
output watermark does not move, and PersistBundle marked only the
producing stage as changed, so a consumer that was just handed pending
elements was never surfaced to the scheduler. Its elements accumulated
forever. checkForQuiescence cannot catch this, because the source stays
schedulable, so the job live locks instead of failing fast.

PersistBundle now records the consumers that accepted data, and
updateWatermarks hands them to the scheduler on the path where the output
watermark does not advance. Recording is gated on any bundle having
returned a residual, so a pipeline that never self checkpoints keeps its
previous bundle scheduling.

bundleReady additionally lets a stateful stage with no side inputs run on
pending data alone, under the same gate, since
statefulStageKind.buildEventTimeBundle takes data at any watermark and
gates only timers. Its stillSchedulable now requires buildable work, and
a key that supplies nothing no longer consumes the OneKeyPerBundle slot,
is not marked in progress, and does not hold the bundle's minimum
timestamp.

A consumer that reads a side input still waits on the watermark, since
side input readiness is derived from it.

Author: Eliaaazzz

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

@@ -235,6 +235,10 @@ type ElementManager struct {
 	livePending     atomic.Int64   // An accessible live pending count. DEBUG USE ONLY
 	pendingElements sync.WaitGroup // pendingElements counts all unprocessed elements in a job. Jobs with no pending elements terminate successfully.
 
+	// Latched once any bundle returns a residual, after which a stage watermark
+	// can be pinned indefinitely and can't be relied on to schedule consumers.
+	sawResidual atomic.Bool
+
 	processTimeEvents *stageRefreshQueue // Manages sequence of stage updates when interfacing with processing time. Callers must hold refreshCond.L lock.
 	testStreamHandler *testStreamHandler // Optional test stream handler when a test stream is in the pipeline.
 }
@@ -841,6 +845,17 @@ func reElementResiduals(residuals []Residual, inputInfo PColInfo, rb RunBundle)
 // input elements, and the committed output elements.
 func (em *ElementManager) PersistBundle(rb RunBundle, col2Coders map[string]PColInfo, d TentativeData, inputInfo PColInfo, residuals Residuals) {
 	stage := em.stages[rb.StageID]
+	// Consumers that received data from this bundle, recorded so they can still
+	// be scheduled when this stage's output watermark is held back. Only needed
+	// once something self checkpoints, so pipelines that never do keep their
+	// previous bundle scheduling exactly.
+	if len(residuals.Data) > 0 {
+		em.sawResidual.Store(true)
+	}
+	var changedConsumers set[string]
+	if em.sawResidual.Load() {
+		changedConsumers = set[string]{}
+	}
 	var seq int
 	for output, data := range d.Raw {
 		info := col2Coders[output]
@@ -910,6 +925,9 @@ func (em *ElementManager) PersistBundle(rb RunBundle, col2Coders map[string]PCol
 				count = consumer.AddPending(em, newPending)
 			}
 			em.addPending(count)
+			if changedConsumers != nil && count > 0 {
+				changedConsumers.insert(sID)
+			}
 		}
 		for _, link := range sideConsumers {
 			consumer := em.stages[link.Global]
@@ -942,6 +960,12 @@ func (em *ElementManager) PersistBundle(rb RunBundle, col2Coders map[string]PCol
 		// even if a panic occurs during `em.addPending`. This prevents potential deadlocks
 		// if the waitgroup unexpectedly drops below zero due to a runner bug.
 		defer stage.mu.Unlock()
+		if len(changedConsumers) > 0 {
+			if stage.consumersWithNewData == nil {
+				stage.consumersWithNewData = set[string]{}
+			}
+			stage.consumersWithNewData.merge(changedConsumers)
+		}
 		completed := stage.inprogress[rb.BundleID]
 		em.addPending(-len(completed.es))
 		delete(stage.inprogress, rb.BundleID)
@@ -1101,6 +1125,9 @@ func (em *ElementManager) ReturnResiduals(rb RunBundle, firstRsIndex int, inputI
 	stage := em.stages[rb.StageID]
 
 	stage.splitBundle(rb, firstRsIndex, em)
+	if len(residuals.Data) > 0 {
+		em.sawResidual.Store(true)
+	}
 	unprocessedElements := reElementResiduals(residuals.Data, inputInfo, rb)
 	if len(unprocessedElements) > 0 {
 		slog.Debug("ReturnResiduals: unprocessed elements", "bundle", rb, "count", len(unprocessedElements))
@@ -1211,6 +1238,9 @@ type stageState struct {
 	estimatedOutput    mtime.Time // Estimated watermark output from DoFns
 	previousInput      mtime.Time // input watermark before the latest watermark refresh
 
+	// Consumers handed data by a bundle, pending delivery to the scheduler.
+	consumersWithNewData set[string]
+
 	pending    elementHeap                          // pending input elements for this stage that are to be processesd
 	inprogress map[string]elements                  // inprogress elements by active bundles, keyed by bundle
 	sideInputs map[LinkID]map[typex.Window][][]byte // side input data for this stage, from {tid, inputID} -> window
@@ -1822,12 +1852,6 @@ keysPerBundle:
 		if ss.inprogressKeys.present(k) {
 			continue
 		}
-		newKeys.insert(k)
-		// Track the min-timestamp for later watermark handling.
-		if dnt.elements[0].timestamp < minTs {
-			minTs = dnt.elements[0].timestamp
-		}
-
 		dataInBundle := false
 
 		var toProcessForKey []element
@@ -1874,22 +1898,51 @@ keysPerBundle:
 				break
 			}
 		}
+		if len(toProcessForKey) > 0 {
+			newKeys.insert(k)
+			// Track the min-timestamp for later watermark handling. Elements pop
+			// in timestamp order, so the first selected one is the earliest.
+			if ts := toProcessForKey[0].timestamp; ts < minTs {
+				minTs = ts
+			}
+		}
 		toProcess = append(toProcess, toProcessForKey...)
 
 		if dnt.elements.Len() == 0 {
 			delete(ss.pendingByKeys, k)
 		}
-		if OneKeyPerBundle {
+		// A key that yielded nothing, such as one headed by a timer above the
+		// watermark, must not consume the single key slot, or the bundle is empty
+		// and the stage keeps rescheduling on it.
+		if OneKeyPerBundle && len(toProcessForKey) > 0 {
 			break keysPerBundle
 		}
 	}
 
-	// If we're out of data, and timers were not cleared then the watermark is accurate.
-	stillSchedulable := !(len(ss.pendingByKeys) == 0 && !timerCleared)
+	// Reschedule only when a later bundle could build something, or a cleared
+	// timer may have held back the minimum pending timestamp.
+	stillSchedulable := timerCleared || ss.hasBuildableDataLocked(watermark)
 
 	return toProcess, minTs, newKeys, holdsInBundle, nil, stillSchedulable, 0
 }
 
+// hasBuildableDataLocked reports whether a key that isn't in progress heads its
+// heap with data, or with a timer the watermark has reached. Callers hold ss.mu.
+func (ss *stageState) hasBuildableDataLocked(watermark mtime.Time) bool {
+	for k, dnt := range ss.pendingByKeys {
+		if ss.inprogressKeys.present(k) {
+			continue
+		}
+		if dnt.elements.Len() == 0 {
+			continue
+		}
+		if e := dnt.elements[0]; e.IsData() || e.timestamp <= watermark {
+			return true
+		}
+	}
+	return false
+}
+
 // buildEventTimeBundle for aggregation stages, processes all elements that are within the watermark for completed windows.
 func (*aggregateStageKind) buildEventTimeBundle(ss *stageState, watermark mtime.Time) (toProcess elementHeap, _ mtime.Time, _ set[string], _ map[mtime.Time]int, panesInBundle []bundlePane, schedulable bool, pendingAdjustment int) {
 	minTs := mtime.MaxTimestamp
@@ -2286,9 +2339,13 @@ func (ss *stageState) updateWatermarks(em *ElementManager) set[string] {
 	if minWatermarkHold < newOut {
 		newOut = minWatermarkHold
 	}
-	// If the newOut is smaller, then don't change downstream watermarks.
+	// If the newOut is smaller, then don't change downstream watermarks. Any
+	// consumer that received data still needs scheduling, since an unadvancing
+	// watermark is otherwise the only thing that would surface it.
 	if newOut <= ss.output {
-		return nil
+		refreshes := ss.consumersWithNewData
+		ss.consumersWithNewData = nil
+		return refreshes
 	}
 
 	// If bigger, advance the output watermark
@@ -2328,6 +2385,7 @@ func (ss *stageState) updateWatermarks(em *ElementManager) set[string] {
 
 	// Update this stage's output watermark, and then propagate that to downstream stages
 	refreshes := set[string]{}
+	ss.consumersWithNewData = nil
 	ss.output = newOut
 	for _, outputCol := range ss.outputIDs {
 		consumers := em.consumers[outputCol]
@@ -2435,13 +2493,20 @@ func (ss *stageState) bundleReady(em *ElementManager, emNow mtime.Time) (mtime.T
 	previousInputW := ss.previousInput
 
 	_, isOrdinaryStage := ss.kind.(*ordinaryStageKind)
-	if isOrdinaryStage && len(ss.sides) == 0 {
+	_, isStatefulStage := ss.kind.(*statefulStageKind)
+	switch {
+	case isOrdinaryStage && len(ss.sides) == 0:
 		// For ordinary stage with no side inputs, we use whether there are pending elements to determine
 		// whether a bundle is ready or not.
 		if len(ss.pending) == 0 {
 			return mtime.MinTimestamp, false, ptimeEventsReady, injectedReady
 		}
-	} else if inputW == upstreamW && previousInputW == inputW {
+	case isStatefulStage && len(ss.sides) == 0 && em.sawResidual.Load() && ss.hasBuildableDataLocked(upstreamW):
+		// A stateful stage processes pending data at whatever the current
+		// watermark is, so data alone makes it ready once something is self
+		// checkpointing. Side input readiness comes from the watermark, so
+		// stages that read one keep waiting.
+	case inputW == upstreamW && previousInputW == inputW:
 		// Otherwise, use the progression of watermark to determine the bundle readiness.
 		slog.Debug("bundleReady: unchanged upstream watermark",
 			slog.String("stage", ss.ID),