From c7dd2b66a1737dcb8e595e41b7e2162dfdde1782 Mon Sep 17 00:00:00 2001 From: Pablo Rodriguez Defino Date: Thu, 15 Oct 2026 18:07:41 -0700 Subject: [PATCH] [Prism] Keep a stateful key in progress while a split primary holds it When the SDK accepts a channel split of a bundle of a stateful stage, stageState.splitBundle removed the key of every residual element from inprogressKeys and inprogressKeysByBundle, even when elements of the primary still carry that key. The residual of that key was then scheduled at once in a second bundle that ran concurrently with the primary. Prism gives each bundle a copy of the key's state at bundle start and replaces the key's state with the bundle's copy at commit, so the last bundle to commit wins. State updates, and the outputs that depend on them, are lost without any error. Serializing the processing of a key is the runner's job. The model defines state per key and window: ReadModifyWriteState "holds a single value per key and window. The value can be read, modified, and written back atomically" (beam_runner_api.proto). The stateful processing post says the runner "cannot parallelize processing of elements for a particular key+window", and the programming guide says "State is always fully scoped only to the current processing key". Splitting is valid SDK behaviour (beam_fn_api.proto), and the Java and Python harnesses accept channel splits on stateful stages. splitBundle now collects the keys of the primary elements and releases a residual's key only if no primary element has it. A held key is released when the primary bundle is persisted, as before. The handling of residual timer holds is unchanged. Reproduction with the Python SDK: all elements share one key, and a per-element sleep makes Prism split the bundle. class CountOnOneKey(beam.DoFn): COUNT = ReadModifyWriteStateSpec('count', VarIntCoder()) DONE = TimerSpec('done', TimeDomain.WATERMARK) def process(self, element, count=beam.DoFn.StateParam(COUNT), done=beam.DoFn.TimerParam(DONE)): time.sleep(0.25) count.write((count.read() or 0) + 1) done.set(GlobalWindow().max_timestamp()) @on_timer(DONE) def report(self, count=beam.DoFn.StateParam(COUNT)): yield count.read() options = PipelineOptions(['--runner=PrismRunner', '--prism_location=', '--environment_type=LOOPBACK']) with beam.Pipeline(options=options) as p: (p | beam.Create(range(40)) | beam.Map(lambda i: ('k', i)) | beam.ParDo(CountOnOneKey()) | beam.Map(print)) A correct runner prints 40. Prism built from master printed 20 in 10 of 10 runs; with this change it printed 40 in 10 of 10 runs. With --prism_log_level=debug, the same pipeline on master logs one split at element 20 of 40, and the residual bundle starts before the primary commits. TestElementManager_SplitStatefulBundle covers a split inside the only key, a split inside one of two keys, and a split between keys. The first two fail without this change. Not fixed here, follow-ups: - splitBundle slices the bundle's elements by the data channel index. A stateful bundle can hold data of one key and timers of another, and timers are not counted in the data channel, so the slice can misattribute elements. See #35771. - startTriggeredBundle, for aggregate stages, does not check inprogressKeys before it injects a bundle. --- .../prism/internal/engine/elementmanager.go | 15 +- .../engine/elementmanager_split_test.go | 194 ++++++++++++++++++ 2 files changed, 207 insertions(+), 2 deletions(-) create mode 100644 sdks/go/pkg/beam/runners/prism/internal/engine/elementmanager_split_test.go 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 80b4a75ac725..2ed73e2c6639 100644 --- a/sdks/go/pkg/beam/runners/prism/internal/engine/elementmanager.go +++ b/sdks/go/pkg/beam/runners/prism/internal/engine/elementmanager.go @@ -2318,9 +2318,20 @@ func (ss *stageState) splitBundle(rb RunBundle, firstResidual int, em *ElementMa es.es = prim + // A key stays assigned to this bundle while any primary element still + // carries it. Releasing it would let a residual of the same key run in a + // second bundle while the primary is still processing it, so both bundles + // would start from the same state and the last commit would drop the + // other's state changes. + primKeys := set[string]{} + for _, e := range prim { + primKeys.insert(string(e.keyBytes)) + } for _, e := range res { - delete(ss.inprogressKeysByBundle[rb.BundleID], string(e.keyBytes)) - delete(ss.inprogressKeys, string(e.keyBytes)) + if k := string(e.keyBytes); !primKeys.present(k) { + delete(ss.inprogressKeysByBundle[rb.BundleID], k) + delete(ss.inprogressKeys, k) + } if e.IsTimer() { slog.Warn("Unexpected split on a bundle with timers. See https://github.com/apache/beam/issues/35771 for information.") diff --git a/sdks/go/pkg/beam/runners/prism/internal/engine/elementmanager_split_test.go b/sdks/go/pkg/beam/runners/prism/internal/engine/elementmanager_split_test.go new file mode 100644 index 000000000000..d13552b0a363 --- /dev/null +++ b/sdks/go/pkg/beam/runners/prism/internal/engine/elementmanager_split_test.go @@ -0,0 +1,194 @@ +// Licensed to the Apache Software Foundation (ASF) under one or more +// contributor license agreements. See the NOTICE file distributed with +// this work for additional information regarding copyright ownership. +// The ASF licenses this file to You under the Apache License, Version 2.0 +// (the "License"); you may not use this file except in compliance with +// the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package engine + +import ( + "fmt" + "testing" + + "github.com/apache/beam/sdks/v2/go/pkg/beam/core/graph/mtime" + "github.com/apache/beam/sdks/v2/go/pkg/beam/core/graph/window" + "github.com/apache/beam/sdks/v2/go/pkg/beam/core/typex" +) + +// splitInsideKey returns a split point that keeps the first keep elements of +// key in the primary and returns the remaining elements of key, and every +// element after them, as the residual. +func splitInsideKey(key string, keep int) func(t *testing.T, bundleKeys []string) int { + return func(t *testing.T, bundleKeys []string) int { + t.Helper() + for i, k := range bundleKeys { + if k == key { + return i + keep + } + } + t.Fatalf("key %q not in bundle %v", key, bundleKeys) + return 0 + } +} + +// splitAfterFirstKey returns a split point that keeps exactly the elements of +// the first key of the bundle in the primary. +func splitAfterFirstKey(t *testing.T, bundleKeys []string) int { + t.Helper() + i := 0 + for i < len(bundleKeys) && bundleKeys[i] == bundleKeys[0] { + i++ + } + if i == len(bundleKeys) { + t.Fatalf("bundle %v has a single key, want at least two", bundleKeys) + } + return i +} + +// TestElementManager_SplitStatefulBundle checks that a channel split of a +// stateful stage's bundle keeps every key held by a primary element assigned +// to the primary bundle, so no other bundle processes that key concurrently. +// +// Prism gives each bundle a copy of the key's state at bundle start and +// replaces the key's state with the bundle's copy at commit. If two bundles +// process the same key at the same time, the last commit discards the state +// changes of the other. +func TestElementManager_SplitStatefulBundle(t *testing.T) { + tests := []struct { + name string + // keys has one entry per input element, all in the global window. + keys []string + // firstResidual picks the split point from the bundle's key order, + // which depends on map iteration order. + firstResidual func(t *testing.T, bundleKeys []string) int + }{ + { + name: "split inside the only key", + keys: []string{"K", "K", "K"}, + firstResidual: splitInsideKey("K", 1), + }, { + name: "split inside one of two keys", + keys: []string{"A", "B", "B"}, + firstResidual: splitInsideKey("B", 1), + }, { + name: "split between keys", + keys: []string{"A", "A", "B", "B"}, + firstResidual: splitAfterFirstKey, + }, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + em := NewElementManager(Config{}) + em.AddStage("impulse", nil, []string{"input"}, nil) + em.AddStage("stateful", []string{"input"}, nil, nil) + em.StageStateful("stateful", nil) + ss := em.stages["stateful"] + + var es []element + for i, k := range test.keys { + es = append(es, element{ + window: window.GlobalWindow{}, + timestamp: mtime.Time(i), + pane: typex.NoFiringPane(), + elmBytes: []byte(fmt.Sprintf("%s%d", k, i)), + keyBytes: []byte(k), + sequence: i, + }) + } + em.addPending(ss.AddPending(em, es)) + + var nextID int + genBundID := func() string { + defer func() { nextID++ }() + return fmt.Sprintf("inst%03d", nextID) + } + // startBundle builds the next event time bundle, if any, and + // returns it with the key of each of its elements. + startBundle := func() (RunBundle, []string, bool) { + bundID, ok, _, _ := ss.startEventTimeBundle(mtime.MaxTimestamp, genBundID) + if !ok { + return RunBundle{}, nil, false + } + var keys []string + for _, e := range ss.inprogress[bundID].es { + keys = append(keys, string(e.keyBytes)) + } + return RunBundle{StageID: ss.ID, BundleID: bundID, Watermark: mtime.MaxTimestamp}, keys, true + } + + primary, bundleKeys, ok := startBundle() + if !ok { + t.Fatal("startEventTimeBundle() built no bundle, want one") + } + if got, want := len(bundleKeys), len(test.keys); got != want { + t.Fatalf("bundle %v has %v elements, want %v", bundleKeys, got, want) + } + split := test.firstResidual(t, bundleKeys) + wantHeld := set[string]{} + for _, k := range bundleKeys[:split] { + wantHeld.insert(k) + } + + em.ReturnResiduals(primary, split, PColInfo{}, Residuals{}) + + checked := set[string]{} + for _, k := range bundleKeys { + if checked.present(k) { + continue + } + checked.insert(k) + want := wantHeld.present(k) + if got := ss.inprogressKeys.present(k); got != want { + t.Errorf("after split of %v at %v: key %v in inprogressKeys = %v, want %v", bundleKeys, split, k, got, want) + } + if got := ss.inprogressKeysByBundle[primary.BundleID].present(k); got != want { + t.Errorf("after split of %v at %v: key %v in inprogressKeysByBundle[%v] = %v, want %v", bundleKeys, split, k, primary.BundleID, got, want) + } + } + + // While the primary runs, only residuals of released keys may + // be scheduled. + processed := split + for { + rb, keys, ok := startBundle() + if !ok { + break + } + for _, k := range keys { + if wantHeld.present(k) { + t.Fatalf("bundle %v has key %v while primary bundle %v still processes it", rb.BundleID, k, primary.BundleID) + } + } + processed += len(keys) + em.PersistBundle(rb, nil, TentativeData{}, PColInfo{}, Residuals{}) + } + + // Once the primary commits, the residuals of its keys run. + em.PersistBundle(primary, nil, TentativeData{}, PColInfo{}, Residuals{}) + for { + rb, keys, ok := startBundle() + if !ok { + break + } + processed += len(keys) + em.PersistBundle(rb, nil, TentativeData{}, PColInfo{}, Residuals{}) + } + + if got, want := processed, len(test.keys); got != want { + t.Errorf("processed %v elements, want %v", got, want) + } + if got := em.livePending.Load(); got != 0 { + t.Errorf("live pending = %v after all bundles persisted, want 0", got) + } + }) + } +}