Skip to content

[Prism] Handle pending adjustment for processing time bundle correctly #36384

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
shunping merged 1 commit into from
Oct 3, 2025
Merged

[Prism] Handle pending adjustment for processing time bundle correctly #36384

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 1 addition & 2 deletions runners/prism/java/build.gradle
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -86,9 +86,8 @@ def sickbayTests = [
'org.apache.beam.sdk.metrics.MetricsTest$CommittedMetricTests.testCommittedStringSetMetrics',
'org.apache.beam.sdk.metrics.MetricsTest$CommittedMetricTests.testCommittedGaugeMetrics',

// negative WaitGroup counter when failing bundle
// Instead of 42, Prism got 84, which suggests two early panes of 42 are fired.

@lostluck lostluck Oct 3, 2025

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is what's happening in the "testAfterProcessingTimeContinuationTriggerUsingState" after this change?

@shunping shunping Oct 3, 2025

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes. I think the previous negative WaitGroup counter masked the true error in this test.

The interesting thing is that it has a sister test at

beam/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/GroupByKeyTest.java

Line 206 in 1e97363

public void testAfterProcessingTimeContinuationTriggerEarly() throws Exception {
, which uses TestStream rather than stateful DoFn to generate the upstream data. That test passes.

'org.apache.beam.sdk.transforms.GroupByKeyTest$BasicTests.testAfterProcessingTimeContinuationTriggerUsingState',
'org.apache.beam.sdk.testing.TestStreamTest.testEarlyPanesOfWindow',

// A regression introduced when we use number of pending elements rather than watermark to determine
// the bundle readiness of a stateless stage.
Expand Down
44 changes: 26 additions & 18 deletions sdks/go/pkg/beam/runners/prism/internal/engine/elementmanager.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -477,12 +477,15 @@ func (em *ElementManager) Bundles(ctx context.Context, upstreamCancelFn context.
}
}
if ptimeEventsReady {
bundleID, ok, reschedule := ss.startProcessingTimeBundle(em, emNow, nextBundID)
bundleID, ok, reschedule, pendingAdjustment := ss.startProcessingTimeBundle(em, emNow, nextBundID)
// Handle the reschedule even when there's no bundle.
if reschedule {
em.changedStages.insert(stageID)
}
if ok {
if pendingAdjustment > 0 {
em.addPending(pendingAdjustment)
}
rb := RunBundle{StageID: stageID, BundleID: bundleID, Watermark: watermark}

em.inprogressBundles.insert(rb.BundleID)
Expand Down Expand Up @@ -1218,7 +1221,7 @@ type stageKind interface {
holdsInBundle map[mtime.Time]int, panesInBundle []bundlePane, schedulable bool, pendingAdjustment int)
// buildProcessingTimeBundle handles building processing-time bundles for the stage per it's kind.
buildProcessingTimeBundle(ss *stageState, em *ElementManager, emNow mtime.Time) (toProcess elementHeap, minTs mtime.Time, newKeys set[string],
holdsInBundle map[mtime.Time]int, panesInBundle []bundlePane, schedulable bool)
holdsInBundle map[mtime.Time]int, panesInBundle []bundlePane, schedulable bool, pendingAdjustment int)
// getPaneOrDefault based on the stage state, element metadata, and bundle id.
getPaneOrDefault(ss *stageState, defaultPane typex.PaneInfo, w typex.Window, keyBytes []byte, bundID string) typex.PaneInfo
}
Expand Down Expand Up @@ -1983,24 +1986,24 @@ keysPerBundle:
return toProcess, minTs, newKeys, holdsInBundle, panesInBundle, stillSchedulable, accumulatingPendingAdjustment
}

func (ss *stageState) startProcessingTimeBundle(em *ElementManager, emNow mtime.Time, genBundID func() string) (string, bool, bool) {
func (ss *stageState) startProcessingTimeBundle(em *ElementManager, emNow mtime.Time, genBundID func() string) (string, bool, bool, int) {
ss.mu.Lock()
defer ss.mu.Unlock()

toProcess, minTs, newKeys, holdsInBundle, panesInBundle, stillSchedulable := ss.kind.buildProcessingTimeBundle(ss, em, emNow)
toProcess, minTs, newKeys, holdsInBundle, panesInBundle, stillSchedulable, accumulatingPendingAdjustment := ss.kind.buildProcessingTimeBundle(ss, em, emNow)

if len(toProcess) == 0 {
// If we have nothing
return "", false, stillSchedulable
return "", false, stillSchedulable, accumulatingPendingAdjustment
}
bundID := ss.makeInProgressBundle(genBundID, toProcess, minTs, newKeys, holdsInBundle, panesInBundle)
slog.Debug("started a processing time bundle", "stageID", ss.ID, "bundleID", bundID, "size", len(toProcess), "emNow", emNow)
return bundID, true, stillSchedulable
return bundID, true, stillSchedulable, accumulatingPendingAdjustment
}

// handleProcessingTimeTimer contains the common code for handling processing-time timers for aggregation stages and stateful stages.
func handleProcessingTimeTimer(ss *stageState, em *ElementManager, emNow mtime.Time,
processTimerFn func(e element, toProcess []element, holdsInBundle map[mtime.Time]int, panesInBundle []bundlePane) ([]element, []bundlePane)) (elementHeap, mtime.Time, set[string], map[mtime.Time]int, []bundlePane, bool) {
processTimerFn func(e element, toProcess []element, holdsInBundle map[mtime.Time]int, panesInBundle []bundlePane) ([]element, []bundlePane, int)) (elementHeap, mtime.Time, set[string], map[mtime.Time]int, []bundlePane, bool, int) {
// TODO: Determine if it's possible and a good idea to treat all EventTime processing as a MinTime
// Special Case for ProcessingTime handling.
// Eg. Always queue EventTime elements at minTime.
Expand All @@ -2010,6 +2013,8 @@ func handleProcessingTimeTimer(ss *stageState, em *ElementManager, emNow mtime.T

var toProcess []element
var panesInBundle []bundlePane
var pendingAdjustment int
accumulatingPendingAdjustment := 0
minTs := mtime.MaxTimestamp
holdsInBundle := map[mtime.Time]int{}

Expand Down Expand Up @@ -2044,7 +2049,8 @@ func handleProcessingTimeTimer(ss *stageState, em *ElementManager, emNow mtime.T
minTs = e.timestamp
}

toProcess, panesInBundle = processTimerFn(e, toProcess, holdsInBundle, panesInBundle)
toProcess, panesInBundle, pendingAdjustment = processTimerFn(e, toProcess, holdsInBundle, panesInBundle)
accumulatingPendingAdjustment += pendingAdjustment
}

nextTime = ss.processingTimeTimers.Peek()
Expand All @@ -2065,24 +2071,26 @@ func handleProcessingTimeTimer(ss *stageState, em *ElementManager, emNow mtime.T
// Add a refresh if there are still processing time events to process.
stillSchedulable := (nextTime < emNow && nextTime != mtime.MaxTimestamp || len(notYet) > 0)

return toProcess, minTs, newKeys, holdsInBundle, panesInBundle, stillSchedulable
return toProcess, minTs, newKeys, holdsInBundle, panesInBundle, stillSchedulable, accumulatingPendingAdjustment
}

// buildProcessingTimeBundle for stateful stages prepares bundles for processing-time timers
func (*statefulStageKind) buildProcessingTimeBundle(ss *stageState, em *ElementManager, emNow mtime.Time) (elementHeap, mtime.Time, set[string], map[mtime.Time]int, []bundlePane, bool) {
return handleProcessingTimeTimer(ss, em, emNow, func(e element, toProcess []element, holdsInBundle map[mtime.Time]int, panesInBundle []bundlePane) ([]element, []bundlePane) {
func (*statefulStageKind) buildProcessingTimeBundle(ss *stageState, em *ElementManager, emNow mtime.Time) (elementHeap, mtime.Time, set[string], map[mtime.Time]int, []bundlePane, bool, int) {
return handleProcessingTimeTimer(ss, em, emNow, func(e element, toProcess []element, holdsInBundle map[mtime.Time]int, panesInBundle []bundlePane) ([]element, []bundlePane, int) {
holdsInBundle[e.holdTimestamp]++
// We're going to process this timer!
toProcess = append(toProcess, e)
return toProcess, nil
return toProcess, nil, 0
})
}

// buildProcessingTimeBundle for aggregation stages prepares bundles for after-processing-time triggers
func (*aggregateStageKind) buildProcessingTimeBundle(ss *stageState, em *ElementManager, emNow mtime.Time) (elementHeap, mtime.Time, set[string], map[mtime.Time]int, []bundlePane, bool) {
return handleProcessingTimeTimer(ss, em, emNow, func(e element, toProcess []element, holdsInBundle map[mtime.Time]int, panesInBundle []bundlePane) ([]element, []bundlePane) {
func (*aggregateStageKind) buildProcessingTimeBundle(ss *stageState, em *ElementManager, emNow mtime.Time) (elementHeap, mtime.Time, set[string], map[mtime.Time]int, []bundlePane, bool, int) {
return handleProcessingTimeTimer(ss, em, emNow, func(e element, toProcess []element, holdsInBundle map[mtime.Time]int, panesInBundle []bundlePane) ([]element, []bundlePane, int) {
// Different from `buildProcessingTimeBundle` for stateful stage,
// triggers don't hold back the watermark, so no holds are in the triggered bundle.
var pendingAdjustment int
var elems []element
state := ss.state[LinkID{}][e.window][string(e.keyBytes)]
endOfWindowReached := e.window.MaxTimestamp() < ss.input
ready := ss.strat.IsTriggerReady(triggerInput{
Expand All @@ -2095,22 +2103,22 @@ func (*aggregateStageKind) buildProcessingTimeBundle(ss *stageState, em *Element
state.Pane = computeNextTriggeredPane(state.Pane, endOfWindowReached)

// We're going to process this trigger!
elems, _ := ss.buildTriggeredBundle(em, string(e.keyBytes), e.window)

@shunping shunping Oct 3, 2025

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The pendingAdjustment from buildTriggerBundle was discarded previously, which leads to the issue when handling processing time bundle in accumulating mode.

elems, pendingAdjustment = ss.buildTriggeredBundle(em, string(e.keyBytes), e.window)
toProcess = append(toProcess, elems...)

ss.state[LinkID{}][e.window][string(e.keyBytes)] = state

panesInBundle = append(panesInBundle, bundlePane{})
}

return toProcess, panesInBundle
return toProcess, panesInBundle, pendingAdjustment
})
}

// buildProcessingTimeBundle for stateless stages is not supposed to be called currently
func (*ordinaryStageKind) buildProcessingTimeBundle(ss *stageState, em *ElementManager, emNow mtime.Time) (elementHeap, mtime.Time, set[string], map[mtime.Time]int, []bundlePane, bool) {
func (*ordinaryStageKind) buildProcessingTimeBundle(ss *stageState, em *ElementManager, emNow mtime.Time) (elementHeap, mtime.Time, set[string], map[mtime.Time]int, []bundlePane, bool, int) {
slog.Error("ordinary stages can't have processing time elements")
return nil, mtime.MinTimestamp, nil, nil, nil, false
return nil, mtime.MinTimestamp, nil, nil, nil, false, 0
}

// makeInProgressBundle is common code to store a set of elements as a bundle in progress.
Expand Down
Loading