[Prism] Schedule consumers of a self checkpointing source

.github/trigger_files/beam_PostCommit_Python_Xlang_Messaging_Direct.json

@@ -1,4 +1,4 @@
 {
   "comment": "Modify this file in a trivial way to cause this test suite to run",
-  "modification": 7
+  "modification": 8
 }

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