[Prism] Fix SDF splits and checkpoints in ElementManager.ReturnResiduals

Fixes a bug where splittable DoFN (SDF) splits and checkpoints would completely lose
their deferred residual restrictions, causing the pipeline to terminate prematurely
and downstream assertions to fail due to missing elements (e.g., test_register_finalizations).

Root Cause:
- The previous patch in v2 fixed a double-counting bug of livePending by removing the
  decoding and rescheduling of split residuals (stage.AddPending and em.addPending)
  from ReturnResiduals(), assuming they were already placed back by stage.splitBundle().
- This holds true for normal non-SDF channel splits, where stage.splitBundle already
  puts the unprocessed original elements back.
- However, when a splittable DoFn (SDF) checkpoints itself, the active element splits on
  its restriction rather than simple unprocessed channel elements. In this case, the original
  remaining elements (res) in splitBundle() has a length of 0, but the SDK worker
  returns a new restriction in the residuals.Data (e.g. unprocessedElements length 1).
- Because the previous patch completely removed rescheduling from ReturnResiduals, this new
  residual restriction was completely lost and never added back to the pending queue.

Solution:
- Inside ReturnResiduals(), we dynamically calculate the original remaining elements in
  the bundle: originalRemainingCount := len(completed.es) - firstRsIndex.
- We compare the total returned residuals (unprocessedElements) against originalRemainingCount.
- If len(unprocessedElements) > originalRemainingCount, the difference represents the new SDF
  residual restrictions. We selectively add ONLY these new residuals back to the stage pending
  heap and safely increment em.livePending by this difference.
- This elegantly preserves the fix for normal channel splits (preventing double-counting), while
  ensuring SDF checkpoint residuals are correctly scheduled.
- Also includes detailed slog.Info logging during execution to track livePending state changes
  accurately in addPending, ReturnResiduals, and splitBundle.

Author: shunping

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

@@ -240,8 +240,10 @@ type ElementManager struct {
 }
 
 func (em *ElementManager) addPending(v int) {
+	prev := em.livePending.Load()
 	em.livePending.Add(int64(v))
 	em.pendingElements.Add(v)
+	slog.Info("em.addPending", "delta", v, "prev", prev, "current", em.livePending.Load())
 }
 
 // LinkID represents a fully qualified input or output.
@@ -1091,16 +1093,25 @@ func (em *ElementManager) FailBundle(rb RunBundle) {
 	em.markChangedAndClearBundle(rb.StageID, rb.BundleID, nil)
 }
 
-// ReturnResiduals is called after a successful split, so the remaining work
-// can be re-assigned to a new bundle.
 func (em *ElementManager) ReturnResiduals(rb RunBundle, firstRsIndex int, inputInfo PColInfo, residuals Residuals) {
 	stage := em.stages[rb.StageID]
 
+	slog.Info("ElementManager.ReturnResiduals start", "bundle", rb, "firstRsIndex", firstRsIndex)
+
+	stage.mu.Lock()
+	completed := stage.inprogress[rb.BundleID]
+	originalRemainingCount := len(completed.es) - firstRsIndex
+	stage.mu.Unlock()
+
 	stage.splitBundle(rb, firstRsIndex, em)
 	unprocessedElements := reElementResiduals(residuals.Data, inputInfo, rb)
-	if len(unprocessedElements) > 0 {
-		slog.Debug("ReturnResiduals: unprocessed elements", "bundle", rb, "count", len(unprocessedElements))
+	if len(unprocessedElements) > originalRemainingCount {
+		newResiduals := unprocessedElements[originalRemainingCount:]
+		slog.Info("ReturnResiduals: new residuals added back", "bundle", rb, "count", len(newResiduals))
+		count := stage.AddPending(em, newResiduals)
+		em.addPending(count)
 	}
+	slog.Info("ElementManager.ReturnResiduals end", "bundle", rb, "unprocessedCount", len(unprocessedElements), "livePending", em.livePending.Load())
 	em.markStagesAsChanged(singleSet(rb.StageID))
 }
 
@@ -2185,7 +2196,7 @@ func (ss *stageState) splitBundle(rb RunBundle, firstResidual int, em *ElementMa
 	defer ss.mu.Unlock()
 
 	es := ss.inprogress[rb.BundleID]
-	slog.Debug("split elements", "bundle", rb, "elem count", len(es.es), "res", firstResidual)
+	slog.Info("splitBundle start", "bundle", rb, "elem count", len(es.es), "firstResidual", firstResidual, "livePending", em.livePending.Load())
 
 	prim := es.es[:firstResidual]
 	res := es.es[firstResidual:]
@@ -2205,6 +2216,7 @@ func (ss *stageState) splitBundle(rb RunBundle, firstResidual int, em *ElementMa
 	// we don't need to increment pending count in em, since it is already pending
 	ss.kind.addPending(ss, em, res)
 	ss.inprogress[rb.BundleID] = es
+	slog.Info("splitBundle completed", "bundle", rb, "primaryCount", len(prim), "residualCount", len(res), "livePending", em.livePending.Load())
 }
 
 // minimumPendingTimestamp returns the minimum pending timestamp from all pending elements,

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

@@ -686,53 +686,71 @@ func TestElementManager_OnWindowExpiration(t *testing.T) {
 }
 
 func TestElementManager_ReturnResidualsPendingCount(t *testing.T) {
-	em := NewElementManager(Config{})
-	em.AddStage("impulse", nil, []string{"input"}, nil)
-	em.AddStage("dofn", []string{"input"}, nil, nil)
-	em.Impulse("impulse")
-
-	stage := em.stages["dofn"]
-	info := PColInfo{
-		GlobalID: "generic_info",
-		WDec:     exec.MakeWindowDecoder(coder.NewGlobalWindow()),
-		WEnc:     exec.MakeWindowEncoder(coder.NewGlobalWindow()),
-		EDec: func(r io.Reader) []byte {
-			b, _ := io.ReadAll(r)
-			return b
+	tests := []struct {
+		name             string
+		firstRsIndex     int
+		wantFinalPending int64
+	}{
+		{
+			name:             "ChannelSplit",
+			firstRsIndex:     0,
+			wantFinalPending: 1,
+		},
+		{
+			name:             "SDFCheckpoint",
+			firstRsIndex:     1,
+			wantFinalPending: 2, // Incremented by 1 because the active portion (index 0) is still in progress and will be completed/decremented in PersistBundle.
 		},
 	}
 
-	// Initial state should have 1 pending element from impulse
-	if got, want := em.livePending.Load(), int64(1); got != want {
-		t.Fatalf("initial livePending = %v, want %v", got, want)
-	}
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			em := NewElementManager(Config{})
+			em.AddStage("impulse", nil, []string{"input"}, nil)
+			em.AddStage("dofn", []string{"input"}, nil, nil)
+			em.Impulse("impulse")
+
+			stage := em.stages["dofn"]
+			info := PColInfo{
+				GlobalID: "generic_info",
+				WDec:     exec.MakeWindowDecoder(coder.NewGlobalWindow()),
+				WEnc:     exec.MakeWindowEncoder(coder.NewGlobalWindow()),
+				EDec: func(r io.Reader) []byte {
+					b, _ := io.ReadAll(r)
+					return b
+				},
+			}
 
-	// Start a bundle
-	bundID, ok, _, _ := stage.startEventTimeBundle(mtime.MaxTimestamp, func() string { return "inst0" })
-	if !ok {
-		t.Fatalf("failed to start bundle")
-	}
+			// Initial state should have 1 pending element from impulse
+			if got, want := em.livePending.Load(), int64(1); got != want {
+				t.Fatalf("initial livePending = %v, want %v", got, want)
+			}
 
-	// Waitgroup/livePending shouldn't change on starting a bundle (it's still pending)
-	if got, want := em.livePending.Load(), int64(1); got != want {
-		t.Fatalf("livePending after startEventTimeBundle = %v, want %v", got, want)
-	}
+			// Start a bundle
+			bundID, ok, _, _ := stage.startEventTimeBundle(mtime.MaxTimestamp, func() string { return "inst0" })
+			if !ok {
+				t.Fatalf("failed to start bundle")
+			}
 
-	// Prepare residuals
-	residBytes := []byte{127, 223, 59, 100, 90, 28, 172, 9, 0, 0, 0, 1, 15, 3, 65, 66, 67} // windowed value header + ABC
-	residuals := Residuals{
-		Data: []Residual{{Element: residBytes}},
-	}
+			// Waitgroup/livePending shouldn't change on starting a bundle (it's still pending)
+			if got, want := em.livePending.Load(), int64(1); got != want {
+				t.Fatalf("livePending after startEventTimeBundle = %v, want %v", got, want)
+			}
+
+			// Prepare residuals
+			residBytes := []byte{127, 223, 59, 100, 90, 28, 172, 9, 0, 0, 0, 1, 15, 3, 65, 66, 67} // windowed value header + ABC
+			residuals := Residuals{
+				Data: []Residual{{Element: residBytes}},
+			}
+
+			rb := RunBundle{StageID: "dofn", BundleID: bundID}
 
-	rb := RunBundle{StageID: "dofn", BundleID: bundID}
-	
-	// Return residuals (Simulates splitting)
-	em.ReturnResiduals(rb, 0, info, residuals)
+			// Return residuals (Simulates splitting)
+			em.ReturnResiduals(rb, test.firstRsIndex, info, residuals)
 
-	// Since we split the bundle (0 primary completed, 1 residual returned), 
-	// the element remains pending. The pending count MUST still be exactly 1!
-	if got, want := em.livePending.Load(), int64(1); got != want {
-		t.Errorf("BUG DETECTED: livePending after ReturnResiduals = %v, want 1! (Elements counted twice)", got)
+			if got, want := em.livePending.Load(), test.wantFinalPending; got != want {
+				t.Errorf("livePending after ReturnResiduals = %v, want %v", got, want)
+			}
+		})
 	}
 }
-