[Bug][Prism]: Prism gets stuck when trying to flatten 2 unbounded pcollections

[avery-nguyen-verily]

What happened?

What happened?

If you try running a pipeline that has multiple unbounded sources where 1 or more do not emit any data within the window and try to flatten the pcollections together Prism will not execute any subsequent ParDos until all sources emit something even after the window ends.

Expected Behaviour

Prism should finish processing the data once the window ends, even if not all the sources that are being flattened have emitted data.

Minimal Example

package main

import (
        "context"
        "log"
        "math"
        "time"

        "github.com/apache/beam/sdks/v2/go/pkg/beam"
        "github.com/apache/beam/sdks/v2/go/pkg/beam/core/graph/window"
        beamLog "github.com/apache/beam/sdks/v2/go/pkg/beam/log"
        "github.com/apache/beam/sdks/v2/go/pkg/beam/register"
        "github.com/apache/beam/sdks/v2/go/pkg/beam/transforms/periodic"
        "github.com/apache/beam/sdks/v2/go/pkg/beam/x/beamx"
)

func init() {
        register.Function3x0(count)
        register.Function3x0(initCount)
        register.Emitter1[int]()
}

func initCount(ctx context.Context, _ []byte, emit func(int)) {
        beamLog.Info(ctx, "Starting count")
        emit(1)
}

func count(ctx context.Context, currCount int, emit func(int)) {
        beamLog.Infof(ctx, "Current Count = %d", currCount)
        emit(currCount + 1)
}

func main() {
        beam.Init()

        ctx := context.Background()
        pipeline, scope := beam.NewPipelineWithRoot()

        duration := 5 * time.Second
        fixedWindow := window.NewFixedWindows(duration)
        unboundedSource1 := periodic.Impulse(scope, time.Now(), time.UnixMilli(math.MaxInt64), duration, false)
        unboundedSource2 := periodic.Impulse(scope, time.UnixMilli(math.MaxInt32), time.UnixMilli(math.MaxInt64), duration, false)
        trigger1 := beam.WindowInto(scope, fixedWindow, unboundedSource1)
        trigger2 := beam.WindowInto(scope, fixedWindow, unboundedSource2)

        s1 := scope.Scope("source")
        c1 := beam.ParDo(s1, initCount, trigger1)
        c2 := beam.ParDo(s1, initCount, trigger2)
        s := beam.Flatten(s1, c1, c2) // prism will never run any ParDo that depends on this
        s2 := scope.Scope("count")
        c3 := beam.ParDo(s2, count, s)
        c4 := beam.ParDo(s2, count, c3)
        beam.ParDo(s2, count, c4)

        if err := beamx.Run(ctx, pipeline); err != nil {
                log.Fatalf("Failed to execute job: %v", err)
        }
}

Output

Prism only ever runs initCount, none of the count ParDo functions are ever run because it is waiting for data from c2

2025/01/30 15:55:59 INFO log from SDK worker worker.ID=job-001[go-job-1-1738252559500852940]_go worker.endpoint=localhost:38569 sdk.transformID=e12 sdk.location=/home/user/workspace/beam-starter-go/main.go:25 sdk.time=2025-01-30T15:55:59.830Z sdk.msg="Starting count"
2025/01/30 15:55:59 INFO log from SDK worker worker.ID=job-001[go-job-1-1738252559500852940]_go worker.endpoint=localhost:38569 sdk.transformID=e12 sdk.location=/home/user/workspace/beam-starter-go/main.go:25 sdk.time=2025-01-30T15:55:59.830Z sdk.msg="Starting count"
2025/01/30 15:55:59 INFO log from SDK worker worker.ID=job-001[go-job-1-1738252559500852940]_go worker.endpoint=localhost:38569 sdk.transformID=e12 sdk.location=/home/user/workspace/beam-starter-go/main.go:25 sdk.time=2025-01-30T15:55:59.830Z sdk.msg="Starting count"
2025/01/30 15:55:59 INFO log from SDK worker worker.ID=job-001[go-job-1-1738252559500852940]_go worker.endpoint=localhost:38569 sdk.transformID=e12 sdk.location=/home/user/workspace/beam-starter-go/main.go:25 sdk.time=2025-01-30T15:55:59.830Z sdk.msg="Starting count"
2025/01/30 15:55:59 INFO log from SDK worker worker.ID=job-001[go-job-1-1738252559500852940]_go worker.endpoint=localhost:38569 sdk.transformID=e12 sdk.location=/home/user/workspace/beam-starter-go/main.go:25 sdk.time=2025-01-30T15:55:59.830Z sdk.msg="Starting count"

Issue Priority

Priority: 2 (default / most bugs should be filed as P2)

Issue Components

• Component: Python SDK
• Component: Java SDK
• Component: Go SDK
• Component: Typescript SDK
• Component: IO connector
• Component: Beam YAML
• Component: Beam examples
• Component: Beam playground
• Component: Beam katas
• Component: Website
• Component: Infrastructure
• Component: Spark Runner
• Component: Flink Runner
• Component: Samza Runner
• Component: Twister2 Runner
• Component: Hazelcast Jet Runner
• Component: Google Cloud Dataflow Runner

[shunping]

I am able to reproduce this with the code provided.

However, I can see the Pardo output after Flatten if I modify the start timestamp in unboundedSource2 as follows:

unboundedSource2 := periodic.Impulse(scope, time.Now(), time.UnixMilli(math.MaxInt64), duration, false)

[shunping]

I also tried to implement the same logic in python as follows, but prism does not get stuck when running this pipeline.

import apache_beam as beam
from apache_beam.options.pipeline_options import PipelineOptions
from apache_beam.transforms.periodicsequence import PeriodicImpulse
from apache_beam.transforms.window import FixedWindows
from apache_beam.utils.timestamp import Timestamp

options = PipelineOptions([
    "--job_server_timeout=600",
    "--environment_type=LOOPBACK",
    "--runner=PortableRunner",
    "--job_endpoint=localhost:8073",
])


class InitCount(beam.DoFn):

  def process(self, element):
    return [1]


class PlusOne(beam.DoFn):

  def process(self, element):
    print(element)
    return [element + 1]


INT32_MAX_TIMESTAMP = Timestamp(micros=2147483647 * 1000)

duration = 5
fixed_window = FixedWindows(size=duration)
with beam.Pipeline(options=options) as p:
  unboundedSource1 = p | "s1" >> PeriodicImpulse(fire_interval=duration)
  unboundedSource2 = p | "s2" >> PeriodicImpulse(
      start_timestamp=INT32_MAX_TIMESTAMP, fire_interval=duration)

  trigger1 = unboundedSource1 | "w1" >> beam.WindowInto(fixed_window)
  trigger2 = unboundedSource1 | "w2" >> beam.WindowInto(fixed_window)

  c1 = trigger1 | "i1" >> beam.ParDo(InitCount())
  c2 = trigger2 | "i2" >> beam.ParDo(InitCount())
  _ = ([c1, c2] | beam.Flatten() | "p1" >> beam.ParDo(PlusOne())
       | "p2" >> beam.ParDo(PlusOne())) | "p3" >> beam.ParDo(PlusOne())

Could that be a Go SDK problem? @lostluck

[shunping]

Going back to the go pipeline, I see the following related goroutine stacktraces from debugz

1 @ 0x102883ea8 0x102fbf89c 0x102fbf7e9 0x102fba364 0x102fb43c8 0x1028e8d24
#	0x102fbf89b	github.com/apache/beam/sdks/v2/go/pkg/beam/runners/prism/internal/engine.(*stageState).minPendingTimestampLocked+0x1ab	./Projects/beam-dev-go-prism/sdks/go/pkg/beam/runners/prism/internal/engine/elementmanager.go:1928
#	0x102fbf7e8	github.com/apache/beam/sdks/v2/go/pkg/beam/runners/prism/internal/engine.(*stageState).updateWatermarks+0xf8		./Projects/beam-dev-go-prism/sdks/go/pkg/beam/runners/prism/internal/engine/elementmanager.go:1948
#	0x102fba363	github.com/apache/beam/sdks/v2/go/pkg/beam/runners/prism/internal/engine.(*ElementManager).refreshWatermarks+0x133	./Projects/beam-dev-go-prism/sdks/go/pkg/beam/runners/prism/internal/engine/elementmanager.go:1073
#	0x102fb43c7	github.com/apache/beam/sdks/v2/go/pkg/beam/runners/prism/internal/engine.(*ElementManager).Bundles.func2+0x147		./Projects/beam-dev-go-prism/sdks/go/pkg/beam/runners/prism/internal/engine/elementmanager.go:407

1 @ 0x1028e0d08 0x1028c1094 0x10308ed70 0x10308d984 0x1028e8d24
#	0x10308ed6f	github.com/apache/beam/sdks/v2/go/pkg/beam/runners/prism/internal.executePipeline+0xf9f	 ./Projects/beam-dev-go-prism/sdks/go/pkg/beam/runners/prism/internal/execute.go:341
#	0x10308d983	github.com/apache/beam/sdks/v2/go/pkg/beam/runners/prism/internal.RunPipeline+0x393	./Projects/beam-dev-go-prism/sdks/go/pkg/beam/runners/prism/internal/execute.go:84

1 @ 0x1028e0d08 0x1028c1f04 0x1028c1ee5 0x1028e2288 0x1028f1ed0 0x1028f296c 0x1028f2919 0x102fba09c 0x102fb85fc 0x10309c978 0x103090950 0x103088404 0x1028e8d24
#	0x1028e2287	internal/sync.runtime_SemacquireMutex+0x27											./go/pkg/mod/golang.org/toolchain@v0.0.1-go1.24.1.darwin-arm64/src/runtime/sema.go:95
#	0x1028f1ecf	internal/sync.(*Mutex).lockSlow+0x16f												./go/pkg/mod/golang.org/toolchain@v0.0.1-go1.24.1.darwin-arm64/src/internal/sync/mutex.go:149
#	0x1028f296b	internal/sync.(*Mutex).Lock+0x6b												./go/pkg/mod/golang.org/toolchain@v0.0.1-go1.24.1.darwin-arm64/src/internal/sync/mutex.go:70
#	0x1028f2918	sync.(*Mutex).Lock+0x18														./go/pkg/mod/golang.org/toolchain@v0.0.1-go1.24.1.darwin-arm64/src/sync/mutex.go:46
#	0x102fba09b	github.com/apache/beam/sdks/v2/go/pkg/beam/runners/prism/internal/engine.(*ElementManager).markChangedAndClearBundle+0x4b	./Projects/beam-dev-go-prism/sdks/go/pkg/beam/runners/prism/internal/engine/elementmanager.go:1048
#	0x102fb85fb	github.com/apache/beam/sdks/v2/go/pkg/beam/runners/prism/internal/engine.(*ElementManager).PersistBundle+0xb2b			./Projects/beam-dev-go-prism/sdks/go/pkg/beam/runners/prism/internal/engine/elementmanager.go:933
#	0x10309c977	github.com/apache/beam/sdks/v2/go/pkg/beam/runners/prism/internal.(*stage).Execute+0x1a87					./Projects/beam-dev-go-prism/sdks/go/pkg/beam/runners/prism/internal/stage.go:333
#	0x10309094f	github.com/apache/beam/sdks/v2/go/pkg/beam/runners/prism/internal.executePipeline.func4+0xbf					./Projects/beam-dev-go-prism/sdks/go/pkg/beam/runners/prism/internal/execute.go:355
#	0x103088403	golang.org/x/sync/errgroup.(*Group).Go.func1+0x53										./go/pkg/mod/golang.org/x/sync@v0.13.0/errgroup/errgroup.go:79

1 @ 0x1028e0d08 0x1028c1f04 0x1028c1ee5 0x1028e23ac 0x1028f4460 0x102fb4ea4 0x1028e8d24
#	0x1028e23ab	sync.runtime_SemacquireWaitGroup+0x2b										./go/pkg/mod/golang.org/toolchain@v0.0.1-go1.24.1.darwin-arm64/src/runtime/sema.go:110
#	0x1028f445f	sync.(*WaitGroup).Wait+0x6f											./go/pkg/mod/golang.org/toolchain@v0.0.1-go1.24.1.darwin-arm64/src/sync/waitgroup.go:118
#	0x102fb4ea3	github.com/apache/beam/sdks/v2/go/pkg/beam/runners/prism/internal/engine.(*ElementManager).Bundles.func1+0x33	./Projects/beam-dev-go-prism/sdks/go/pkg/beam/runners/prism/internal/engine/elementmanager.go:350

[lostluck]

Offhand, that stack trace looks normal. One thread is updating the stage watermarks, (which requires holding the ElementManager condition lock, and the stage's personal lock) and PersistBundle is waiting on that lock to write to the a stage's state. A question here is which stages are being locked up (are they the same lock?).

Very interesting that Python has different behaviour here, and might be pointing to a Go SDK side bug, if it's not an issue with Prism.