From 8197bacca025208994b160591a3ac51faaba58eb Mon Sep 17 00:00:00 2001 From: Shunping Huang Date: Sun, 21 Sep 2025 15:28:27 -0400 Subject: [PATCH] Fix panic in teststream. --- .../pkg/beam/runners/prism/internal/engine/elementmanager.go | 2 +- sdks/go/pkg/beam/runners/prism/internal/engine/teststream.go | 3 +++ 2 files changed, 4 insertions(+), 1 deletion(-) 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 1e76d748809b..cba4774dd3f9 100644 --- a/sdks/go/pkg/beam/runners/prism/internal/engine/elementmanager.go +++ b/sdks/go/pkg/beam/runners/prism/internal/engine/elementmanager.go @@ -2056,7 +2056,7 @@ func (ss *stageState) String() string { return fmt.Sprintf("[%v] IN: %v OUT: %v UP: %q %v, kind: %v", ss.ID, ss.input, ss.output, pcol, up, ss.kind) } -// updateWatermarks performs the following operations: +// updateWatermarks performs the following operations and returns a possible set of stages to refresh next or nil. // // Watermark_In' = MAX(Watermark_In, MIN(U(TS_Pending), U(Watermark_InputPCollection))) // Watermark_Out' = MAX(Watermark_Out, MIN(Watermark_In', U(minWatermarkHold))) diff --git a/sdks/go/pkg/beam/runners/prism/internal/engine/teststream.go b/sdks/go/pkg/beam/runners/prism/internal/engine/teststream.go index 0af4e7dc41f0..bab9ff048889 100644 --- a/sdks/go/pkg/beam/runners/prism/internal/engine/teststream.go +++ b/sdks/go/pkg/beam/runners/prism/internal/engine/teststream.go @@ -253,6 +253,9 @@ func (ev tsFinalEvent) Execute(em *ElementManager) { em.testStreamHandler.UpdateHold(em, mtime.MaxTimestamp) ss := em.stages[ev.stageID] kickSet := ss.updateWatermarks(em) + if kickSet == nil { + kickSet = make(set[string]) + } kickSet.insert(ev.stageID) em.changedStages.merge(kickSet) }