[Prism] Honor the resume delay of self-checkpointing SDF residuals

CHANGES.md

@@ -84,6 +84,7 @@
 
 * (Python) Fixed incorrect profiler options handling on portable runners ([#39613](https://github.com/apache/beam/issues/39613)).
 * (Java) KafkaIO dynamic reads no longer require the obsolete `beam_fn_api` experiment ([#29998](https://github.com/apache/beam/issues/29998)).
+* (Prism) Self-checkpointing splittable DoFns now resume after their requested delay instead of immediately, so polling SDFs no longer busy-spin ([#39848](https://github.com/apache/beam/issues/39848)).
 
 ## Security Fixes
 

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

@@ -458,6 +458,9 @@ func (em *ElementManager) Bundles(ctx context.Context, upstreamCancelFn context.
 			// Check each advanced stage, to see if it's able to execute based on the watermark.
 			for stageID := range advanced {
 				ss := em.stages[stageID]
+				if adj := ss.releaseDelayedResiduals(em, emNow); adj != 0 {
+					em.addPending(adj)
+				}
 				watermark, ready, ptimeEventsReady, injectedReady := ss.bundleReady(em, emNow)
 				if injectedReady {
 					ss.mu.Lock()
@@ -549,7 +552,7 @@ func (em *ElementManager) DumpStages() string {
 		if upS == "" {
 			upS = "IMPULSE  " // (extra spaces to allow print to align better.)
 		}
-		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, "bundlesToInject", ss.bundlesToInject))
+		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, "bundlesToInject", ss.bundlesToInject, "delayedResiduals", ss.delayedResiduals))
 
 		var outputConsumers, sideConsumers []string
 		for _, col := range ss.outputIDs {
@@ -794,6 +797,24 @@ type Residuals struct {
 
 // reElementResiduals extracts the windowed value header from residual bytes, and explodes them
 // back out to their windows.
+// partitionResiduals splits residuals into those that return to pending at
+// once and those to park, grouped by the processing time they become
+// schedulable per their SDK requested resume delay.
+func partitionResiduals(emNow mtime.Time, data []Residual) (immediate []Residual, delayed map[mtime.Time][]Residual) {
+	for _, r := range data {
+		if r.Delay <= 0 {
+			immediate = append(immediate, r)
+			continue
+		}
+		if delayed == nil {
+			delayed = map[mtime.Time][]Residual{}
+		}
+		fireAt := emNow.Add(r.Delay)
+		delayed[fireAt] = append(delayed[fireAt], r)
+	}
+	return immediate, delayed
+}
+
 func reElementResiduals(residuals []Residual, inputInfo PColInfo, rb RunBundle) []element {
 	var unprocessedElements []element
 	for _, residual := range residuals {
@@ -935,22 +956,35 @@ func (em *ElementManager) PersistBundle(rb RunBundle, col2Coders map[string]PCol
 		}
 	}
 
+	// Single processing time sample for this persist, for timer rebasing and
+	// residual delays alike. processTimeEvents requires the refresh lock.
+	em.refreshCond.L.Lock()
+	emNow := em.processingTimeNow()
+	em.refreshCond.L.Unlock()
+
 	// Triage timers into their time domains for scheduling.
 	// EventTime timers are handled with normal elements,
 	// ProcessingTime timers need to be scheduled into the processing time based queue.
-	newHolds, ptRefreshes := em.triageTimers(d, inputInfo, stage)
+	newHolds, ptRefreshes := em.triageTimers(d, inputInfo, stage, emNow)
 
-	// TODO(https://github.com/apache/beam/issues/39446)
-	// Return unprocessed to this stage's pending
-	// TODO sort out pending element watermark holds for process continuation residuals.
-	unprocessedElements := reElementResiduals(residuals.Data, inputInfo, rb)
-
-	// Add unprocessed back to the pending stack.
+	// A residual with an SDK requested resume delay is parked until that
+	// processing time arrives; the rest return to pending immediately.
+	immediate, delayed := partitionResiduals(emNow, residuals.Data)
+	unprocessedElements := reElementResiduals(immediate, inputInfo, rb)
 	if len(unprocessedElements) > 0 {
-		// TODO actually reschedule based on the residuals delay...
 		count := stage.AddPending(em, unprocessedElements)
 		em.addPending(count)
 	}
+	for fireAt, rs := range delayed {
+		elems := reElementResiduals(rs, inputInfo, rb)
+		if len(elems) == 0 {
+			continue
+		}
+		stage.parkResiduals(fireAt, elems)
+		em.addPending(len(elems))
+		// Schedules the release and, under a real-time clock, the wake-up.
+		ptRefreshes.insert(fireAt)
+	}
 	// Clear out the inprogress elements associated with the completed bundle.
 	// Must be done after adding the new pending elements to avoid an incorrect
 	// watermark advancement.
@@ -1035,7 +1069,7 @@ func (em *ElementManager) PersistBundle(rb RunBundle, col2Coders map[string]PCol
 }
 
 // triageTimers prepares received timers for eventual firing, as well as rebasing processing time timers as needed.
-func (em *ElementManager) triageTimers(d TentativeData, inputInfo PColInfo, stage *stageState) (map[mtime.Time]int, set[mtime.Time]) {
+func (em *ElementManager) triageTimers(d TentativeData, inputInfo PColInfo, stage *stageState, emNow mtime.Time) (map[mtime.Time]int, set[mtime.Time]) {
 	// Process each timer family in the order we received them, so we can filter to the last one.
 	// Since we're process each timer family individually, use a unique key for each userkey, tag, window.
 	// The last timer set for each combination is the next one we're keeping.
@@ -1044,10 +1078,6 @@ func (em *ElementManager) triageTimers(d TentativeData, inputInfo PColInfo, stag
 		tag string
 		win typex.Window
 	}
-	em.refreshCond.L.Lock()
-	emNow := em.processingTimeNow()
-	em.refreshCond.L.Unlock()
-
 	var pendingEventTimers []element
 	var pendingProcessingTimers []fireElement
 	stageRefreshTimes := set[mtime.Time]{}
@@ -1245,6 +1275,11 @@ type stageState struct {
 	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
 
+	// Residual elements parked until the processing time they become schedulable,
+	// per the SDK's requested resume delay. Parked elements pin the input
+	// watermark like pending elements until they are released.
+	delayedResiduals map[mtime.Time][]element
+
 	// Fields for stateful stages which need to be per key.
 	pendingByKeys          map[string]*dataAndTimers                             // pending input elements by Key, if stateful.
 	inprogressKeys         set[string]                                           // all keys that are assigned to bundles.
@@ -1373,6 +1408,38 @@ func (ss *stageState) AddPending(em *ElementManager, newPending []element) int {
 	return ss.kind.addPending(ss, em, newPending)
 }
 
+// parkResiduals defers residual elements until fireAt in processing time.
+// Parked elements pin the input watermark via minPendingTimestampLocked.
+// Callers must schedule a processing time refresh for the stage at fireAt.
+func (ss *stageState) parkResiduals(fireAt mtime.Time, elems []element) {
+	ss.mu.Lock()
+	defer ss.mu.Unlock()
+	if ss.delayedResiduals == nil {
+		ss.delayedResiduals = map[mtime.Time][]element{}
+	}
+	ss.delayedResiduals[fireAt] = append(ss.delayedResiduals[fireAt], elems...)
+}
+
+// releaseDelayedResiduals moves parked residuals that are schedulable at
+// emNow to pending, and returns the resulting pending count adjustment.
+// Callers must hold em.refreshCond.L.
+func (ss *stageState) releaseDelayedResiduals(em *ElementManager, emNow mtime.Time) int {
+	ss.mu.Lock()
+	var due []element
+	for t, elems := range ss.delayedResiduals {
+		if t <= emNow {
+			due = append(due, elems...)
+			delete(ss.delayedResiduals, t)
+		}
+	}
+	ss.mu.Unlock()
+	if len(due) == 0 {
+		return 0
+	}
+	// The parked elements were counted when parked; report only the delta.
+	return ss.AddPending(em, due) - len(due)
+}
+
 func (ss *stageState) injectTriggeredBundlesIfReady(em *ElementManager, window typex.Window, key string) int {
 	// Check on triggers for this key.
 	// Callers must hold em.refreshCond.L
@@ -2291,6 +2358,13 @@ func (ss *stageState) minPendingTimestampLocked() mtime.Time {
 	for _, es := range ss.inprogress {
 		minPending = mtime.Min(minPending, es.minTimestamp)
 	}
+	// Parked residuals are pending work that is not yet schedulable, and
+	// must pin the input watermark like any other pending element.
+	for _, elems := range ss.delayedResiduals {
+		for _, e := range elems {
+			minPending = mtime.Min(minPending, e.timestamp)
+		}
+	}
 	return minPending
 }
 

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

@@ -21,6 +21,7 @@ import (
 	"fmt"
 	"io"
 	"testing"
+	"time"
 
 	"github.com/apache/beam/sdks/v2/go/pkg/beam/core/graph/coder"
 	"github.com/apache/beam/sdks/v2/go/pkg/beam/core/graph/mtime"
@@ -297,3 +298,120 @@ func TestStatefulBuildEventTimeBundle_OneKeyPerBundle(t *testing.T) {
 		}
 	}
 }
+
+// delaySetup builds an impulse -> src pipeline and returns it with the first
+// source bundle already received, ready for the test to persist a residual.
+func delaySetup(t *testing.T, cfg Config) (*ElementManager, <-chan RunBundle, *stageState, RunBundle, context.CancelCauseFunc) {
+	t.Helper()
+	ctx, cancelFn := context.WithCancelCause(context.Background())
+	em := NewElementManager(cfg)
+	em.AddStage("impulse", nil, []string{"src_in"}, nil)
+	em.AddStage("src", []string{"src_in"}, nil, nil)
+	em.Impulse("impulse")
+	var i int
+	ch := em.Bundles(ctx, cancelFn, func() string {
+		defer func() { i++ }()
+		return fmt.Sprintf("%v", i)
+	})
+	rb, ok := <-ch
+	if !ok {
+		t.Fatal("bundles channel closed before the first source bundle")
+	}
+	if rb.StageID != "src" {
+		t.Fatalf("first bundle stage = %v, want src", rb.StageID)
+	}
+	return em, ch, em.stages["src"], rb, cancelFn
+}
+
+// TestPersistBundle_ResidualResumeDelay covers the SDK requested resume delay:
+// a delayed residual is parked with a watermark hold and released through the
+// processing time queue instead of returning to pending immediately.
+func TestPersistBundle_ResidualResumeDelay(t *testing.T) {
+	info := continuationInfo(t, false)
+	residual := func(delay time.Duration) Residuals {
+		return Residuals{
+			TransformID: "src",
+			InputID:     "i0",
+			Data:        []Residual{{Element: encodeElement(t, info, mtime.MinTimestamp), Delay: delay}},
+		}
+	}
+
+	t.Run("parks with hold and schedules", func(t *testing.T) {
+		em, _, src, rb, cancelFn := delaySetup(t, Config{EnableRTC: true})
+		defer cancelFn(nil)
+		const delay = 5 * time.Second
+		before := time.Now()
+		em.PersistBundle(rb, nil, TentativeData{}, info, residual(delay))
+
+		src.mu.Lock()
+		pendingLen := len(src.pending)
+		parked := 0
+		for _, elems := range src.delayedResiduals {
+			parked += len(elems)
+		}
+		minPending := src.minPendingTimestampLocked()
+		src.mu.Unlock()
+		if pendingLen != 0 {
+			t.Errorf("delayed residual returned to pending immediately: len(pending) = %v, want 0", pendingLen)
+		}
+		if parked != 1 {
+			t.Errorf("parked residuals = %v, want 1", parked)
+		}
+		if minPending != mtime.MinTimestamp {
+			t.Errorf("parked residual does not pin the input watermark: minPending = %v, want %v", minPending, mtime.MinTimestamp)
+		}
+		em.refreshCond.L.Lock()
+		fireAt, ok := em.processTimeEvents.Peek()
+		em.refreshCond.L.Unlock()
+		if !ok {
+			t.Fatal("no processing time event scheduled for the delayed residual")
+		}
+		if until := fireAt.ToTime().Sub(before); until <= 0 || until > delay {
+			t.Errorf("residual release scheduled %v after persisting, want within (0, %v]", until, delay)
+		}
+	})
+
+	t.Run("releases only after the delay", func(t *testing.T) {
+		em, ch, _, rb, cancelFn := delaySetup(t, Config{EnableRTC: true})
+		defer cancelFn(nil)
+		_ = em
+		const delay = 300 * time.Millisecond
+		start := time.Now()
+		em.PersistBundle(rb, nil, TentativeData{}, info, residual(delay))
+		select {
+		case rb2, ok := <-ch:
+			if !ok {
+				t.Fatal("bundles channel closed before the delayed residual fired")
+			}
+			if rb2.StageID != "src" {
+				t.Fatalf("resumed bundle stage = %v, want src", rb2.StageID)
+			}
+			if elapsed := time.Since(start); elapsed < 250*time.Millisecond {
+				t.Errorf("delayed residual fired after %v, want at least ~%v", elapsed, delay)
+			}
+		case <-time.After(10 * time.Second):
+			t.Fatal("delayed residual never fired; the real-time wake-up is missing")
+		}
+	})
+
+	t.Run("fast forwards without a real-time clock", func(t *testing.T) {
+		em, ch, _, rb, cancelFn := delaySetup(t, Config{})
+		defer cancelFn(nil)
+		start := time.Now()
+		em.PersistBundle(rb, nil, TentativeData{}, info, residual(5*time.Second))
+		select {
+		case rb2, ok := <-ch:
+			if !ok {
+				t.Fatal("bundles channel closed before the delayed residual fired")
+			}
+			if rb2.StageID != "src" {
+				t.Fatalf("resumed bundle stage = %v, want src", rb2.StageID)
+			}
+			if elapsed := time.Since(start); elapsed >= 4*time.Second {
+				t.Errorf("fast forward mode waited %v in real time for a synthetic delay", elapsed)
+			}
+		case <-time.After(10 * time.Second):
+			t.Fatal("delayed residual never fired in fast forward mode")
+		}
+	})
+}