This is an automated email from the ASF dual-hosted git repository.
jrmccluskey 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 0b407cd78ce Fix lost state updates in Set state (#39175)
0b407cd78ce is described below
commit 0b407cd78ce967f2345c88d6b51233a9eeabcc2c
Author: Shunping Huang <[email protected]>
AuthorDate: Mon Jul 6 10:25:29 2026 -0400
Fix lost state updates in Set state (#39175)
* Add a test to reproduce the problem.
* Fix race condition in SetState compaction by awaiting outstanding state
requests
* Reformat
* Simply logic
* Fix lints
* Do not force to use prism for the new test.
---
.../apache_beam/runners/worker/bundle_processor.py | 37 ++++++++----
.../apache_beam/transforms/userstate_test.py | 69 ++++++++++++++++++++++
2 files changed, 95 insertions(+), 11 deletions(-)
diff --git a/sdks/python/apache_beam/runners/worker/bundle_processor.py
b/sdks/python/apache_beam/runners/worker/bundle_processor.py
index 1a852aa19d9..851efc81221 100644
--- a/sdks/python/apache_beam/runners/worker/bundle_processor.py
+++ b/sdks/python/apache_beam/runners/worker/bundle_processor.py
@@ -641,6 +641,8 @@ class SynchronousSetRuntimeState(userstate.SetRuntimeState):
self._value_coder = value_coder
self._cleared = False
self._added_elements: set[Any] = set()
+ # Track outstanding async state requests to await them at commit time.
+ self._futures = []
def _compact_data(self, rewrite=True):
accumulator = set(
@@ -650,9 +652,11 @@ class
SynchronousSetRuntimeState(userstate.SetRuntimeState):
self._added_elements))
if rewrite and accumulator:
- self._state_handler.clear(self._state_key)
- self._state_handler.extend(
- self._state_key, self._value_coder.get_impl(), accumulator)
+ # Compaction writes are asynchronous; queue them so they are not lost.
+ self._futures.append(self._state_handler.clear(self._state_key))
+ self._futures.append(
+ self._state_handler.extend(
+ self._state_key, self._value_coder.get_impl(), accumulator))
# Since everthing is already committed so we can safely reinitialize
# added_elements here.
@@ -666,7 +670,7 @@ class SynchronousSetRuntimeState(userstate.SetRuntimeState):
def add(self, value: Any) -> None:
if self._cleared:
# This is a good time explicitly clear.
- self._state_handler.clear(self._state_key)
+ self._futures.append(self._state_handler.clear(self._state_key))
self._cleared = False
self._added_elements.add(value)
@@ -678,15 +682,26 @@ class
SynchronousSetRuntimeState(userstate.SetRuntimeState):
self._added_elements = set()
def commit(self) -> None:
- to_await = None
if self._cleared:
- to_await = self._state_handler.clear(self._state_key)
+ self._futures.append(self._state_handler.clear(self._state_key))
+ self._cleared = False
if self._added_elements:
- to_await = self._state_handler.extend(
- self._state_key, self._value_coder.get_impl(), self._added_elements)
- if to_await:
- # To commit, we need to wait on the last state request future to
complete.
- to_await.get()
+ self._futures.append(
+ self._state_handler.extend(
+ self._state_key,
+ self._value_coder.get_impl(),
+ self._added_elements))
+ self._added_elements = set()
+
+ # Block on all outstanding async state requests to ensure data is
committed.
+ # We must swap and clear self._futures before awaiting them. Awaiting a
future
+ # yields control, during which new futures could be appended to
self._futures.
+ all_futures = self._futures
+ self._futures = []
+
+ for f in all_futures:
+ if f:
+ f.get()
class RangeSet:
diff --git a/sdks/python/apache_beam/transforms/userstate_test.py
b/sdks/python/apache_beam/transforms/userstate_test.py
index 6a8efd1a536..45dba5c9e9e 100644
--- a/sdks/python/apache_beam/transforms/userstate_test.py
+++ b/sdks/python/apache_beam/transforms/userstate_test.py
@@ -18,6 +18,9 @@
"""Unit tests for the Beam State and Timer API interfaces."""
# pytype: skip-file
+import queue
+import threading
+import time
import unittest
from typing import Any
@@ -34,6 +37,8 @@ from apache_beam.portability import common_urns
from apache_beam.portability.api import beam_runner_api_pb2
from apache_beam.runners import pipeline_context
from apache_beam.runners.common import DoFnSignature
+from apache_beam.runners.worker.sdk_worker import GrpcStateHandler
+from apache_beam.runners.worker.sdk_worker import _Future
from apache_beam.testing.test_pipeline import TestPipeline
from apache_beam.testing.test_stream import TestStream
from apache_beam.testing.util import assert_that
@@ -719,6 +724,70 @@ class StatefulDoFnOnDirectRunnerTest(unittest.TestCase):
actual_values = (values | beam.ParDo(SetStatefulDoFn()))
assert_that(actual_values, equal_to([1, 3, 6, 10, 10]))
+ # Mock random to always return 1.0 to force compaction on every add.
+ @mock.patch(
+ 'apache_beam.runners.worker.bundle_processor.random.random', lambda: 1.0)
+ def test_stateful_set_state_compaction_race_portably(self):
+ old_request = GrpcStateHandler._request
+ request_queue = queue.Queue()
+
+ def worker():
+ while True:
+ handler, request, future, instruction_id = request_queue.get()
+ time.sleep(0.1) # Simulate latency for each request sequentially.
+ handler._context.process_instruction_id = instruction_id
+ underlying_future = old_request(handler, request)
+ underlying_future.wait()
+ future.set(underlying_future.get())
+ request_queue.task_done()
+
+ t = threading.Thread(target=worker, daemon=True)
+ t.start()
+
+ def delayed_request(self, request):
+ if request.HasField('append') or request.HasField('clear'):
+ future = _Future()
+ instruction_id = getattr(self._context, 'process_instruction_id', None)
+ request_queue.put((self, request, future, instruction_id))
+ return future
+ else:
+ return old_request(self, request)
+
+ GrpcStateHandler._request = delayed_request
+
+ class SetStatefulDoFn(beam.DoFn):
+
+ SET_STATE = SetStateSpec('buffer', VarIntCoder())
+
+ def process(self, element, set_state=beam.DoFn.StateParam(SET_STATE)):
+ _, value = element
+ aggregated_value = 0
+ set_state.add(value)
+ for saved_value in set_state.read():
+ aggregated_value += saved_value
+ yield aggregated_value
+
+ try:
+ options = PipelineOptions([
+ '--max_cache_memory_usage_mb=100',
+ '--environment_type=LOOPBACK',
+ ])
+ with TestPipeline(options=options) as p:
+ test_stream = (
+ TestStream(
+ coder=beam.coders.TupleCoder((
+ beam.coders.StrUtf8Coder(), beam.coders.VarIntCoder()
+ ))).advance_watermark_to(10).add_elements([
+ ('key', 1)
+ ]).advance_watermark_to(20).add_elements([
+ ('key', 2)
+ ]).advance_watermark_to(30))
+ actual_values = (p | test_stream | beam.ParDo(SetStatefulDoFn()))
+ assert_that(actual_values, equal_to([1, 3]))
+
+ finally:
+ GrpcStateHandler._request = old_request
+
def test_stateful_set_state_clean_portably(self):
class SetStateClearingStatefulDoFn(beam.DoFn):
