This is an automated email from the ASF dual-hosted git repository.
Abacn pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git
The following commit(s) were added to refs/heads/master by this push:
new 73ea0b72c12 [Prism] Honor the resume delay of self-checkpointing SDF
residuals (#39849)
73ea0b72c12 is described below
commit 73ea0b72c128e4bb2dcf95c8fbe476b3baabf2df
Author: Elia Liu <[email protected]>
AuthorDate: Wed Aug 26 11:10:27 2026 +1000
[Prism] Honor the resume delay of self-checkpointing SDF residuals (#39849)
* [Prism] Honor the resume delay of self-checkpointing SDF residuals
* Un-sickbay tests now passing
---------
Co-authored-by: Eliaazzz <[email protected]>
Co-authored-by: Yi Hu <[email protected]>
---
CHANGES.md | 1 +
runners/prism/java/build.gradle | 6 --
.../prism/internal/engine/elementmanager.go | 102 +++++++++++++++---
.../engine/elementmanager_continuation_test.go | 118 +++++++++++++++++++++
4 files changed, 207 insertions(+), 20 deletions(-)
diff --git a/CHANGES.md b/CHANGES.md
index 3f097d96dab..48d053f2396 100644
--- a/CHANGES.md
+++ b/CHANGES.md
@@ -86,6 +86,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
diff --git a/runners/prism/java/build.gradle b/runners/prism/java/build.gradle
index 03049f29331..8f402425cd6 100644
--- a/runners/prism/java/build.gradle
+++ b/runners/prism/java/build.gradle
@@ -81,10 +81,6 @@ def sickbayTests = [
//
// There is not currently a category for excluding these _only_ in
committed mode
'org.apache.beam.sdk.metrics.MetricsTest$CommittedMetricTests.testAllCommittedMetrics',
-
'org.apache.beam.sdk.metrics.MetricsTest$CommittedMetricTests.testCommittedCounterMetrics',
-
'org.apache.beam.sdk.metrics.MetricsTest$CommittedMetricTests.testCommittedDistributionMetrics',
-
'org.apache.beam.sdk.metrics.MetricsTest$CommittedMetricTests.testCommittedStringSetMetrics',
-
'org.apache.beam.sdk.metrics.MetricsTest$CommittedMetricTests.testCommittedGaugeMetrics',
// Instead of 42, Prism got 84, which suggests two early panes of 42 are
fired.
'org.apache.beam.sdk.transforms.GroupByKeyTest$BasicTests.testAfterProcessingTimeContinuationTriggerUsingState',
@@ -112,8 +108,6 @@ def sickbayTests = [
'org.apache.beam.sdk.transforms.GroupIntoBatchesTest.testInStreamingMode',
'org.apache.beam.sdk.transforms.GroupIntoBatchesTest.testBufferingTimerInFixedWindow',
'org.apache.beam.sdk.transforms.GroupIntoBatchesTest.testBufferingTimerInGlobalWindow',
- // ShardedKey not yet implemented.
-
'org.apache.beam.sdk.transforms.GroupIntoBatchesTest.testWithShardedKeyInGlobalWindow',
// Some tests failed when using TestStream with keyed elements.
// https://github.com/apache/beam/issues/36984
diff --git a/sdks/go/pkg/beam/runners/prism/internal/engine/elementmanager.go
b/sdks/go/pkg/beam/runners/prism/internal/engine/elementmanager.go
index 7ebf174d9ec..80b4a75ac72 100644
--- a/sdks/go/pkg/beam/runners/prism/internal/engine/elementmanager.go
+++ b/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
}
diff --git
a/sdks/go/pkg/beam/runners/prism/internal/engine/elementmanager_continuation_test.go
b/sdks/go/pkg/beam/runners/prism/internal/engine/elementmanager_continuation_test.go
index 2c0eac52e5e..fd29c885f00 100644
---
a/sdks/go/pkg/beam/runners/prism/internal/engine/elementmanager_continuation_test.go
+++
b/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")
+ }
+ })
+}
