From 604ff1bc0cad567110c16e4a73ae390fc1335aa3 Mon Sep 17 00:00:00 2001 From: David Date: Sun, 19 Jul 2026 17:36:33 +0100 Subject: [PATCH] Invalidate cached ContinuationIterable when its continuation fails GlobalCachingStateHandler caches partially materialized bag state (ContinuationIterable) across bundles under the user state cache token. The iterable's lazy remainder holds the continuation token minted by the read that created it. If that token becomes permanently invalid (for example the runner revoked it together with the work item that created it), the cached iterable fails on every subsequent read of that state and the cache never recovers: each retry is served the same poisoned object and replays the dead token instead of re-fetching. On Dataflow this manifests as an unbounded 'work token no longer valid' retry storm that stalls the stage. Drop the cache entry when iterating the continuation fails so the next read loads fresh state. The invalidation is identity-guarded via a new StateCache.invalidate_if_value so a stale iterable failing again later cannot evict a newer value cached under the same key. Co-Authored-By: Claude Fable 5 --- .../apache_beam/runners/worker/sdk_worker.py | 37 ++++-- .../runners/worker/sdk_worker_test.py | 108 ++++++++++++++++++ .../apache_beam/runners/worker/statecache.py | 16 +++ .../runners/worker/statecache_test.py | 36 ++++++ 4 files changed, 190 insertions(+), 7 deletions(-) diff --git a/sdks/python/apache_beam/runners/worker/sdk_worker.py b/sdks/python/apache_beam/runners/worker/sdk_worker.py index db25a40a405c..d348683a80c7 100644 --- a/sdks/python/apache_beam/runners/worker/sdk_worker.py +++ b/sdks/python/apache_beam/runners/worker/sdk_worker.py @@ -1235,7 +1235,7 @@ def blocking_get( cache_state_key = self._convert_to_cache_key(state_key) return self._state_cache.get( (cache_state_key, cache_token), - lambda key: self._partially_cached_iterable(state_key, coder)) + lambda key: self._partially_cached_iterable(state_key, coder, key)) def extend( self, @@ -1362,7 +1362,8 @@ def _get_cache_token(self, state_key): def _partially_cached_iterable( self, state_key, # type: beam_fn_api_pb2.StateKey - coder # type: coder_impl.CoderImpl + coder, # type: coder_impl.CoderImpl + cache_key # type: Any ): # type: (...) -> Iterable[Any] @@ -1380,20 +1381,42 @@ def _partially_cached_iterable( return self.ContinuationIterable( head, functools.partial( - self._lazy_iterator, state_key, coder, continuation_token)) + self._lazy_iterator, state_key, coder, continuation_token), + on_failure=functools.partial( + self._state_cache.invalidate_if_value, cache_key)) class ContinuationIterable(Generic[T], CacheAware): - def __init__(self, head, continue_iterator_fn): - # type: (Iterable[T], Callable[[], Iterable[T]]) -> None + def __init__( + self, + head, # type: Iterable[T] + continue_iterator_fn, # type: Callable[[], Iterable[T]] + on_failure=None # type: Optional[Callable[[Any], None]] + ): + # type: (...) -> None self.head = head self.continue_iterator_fn = continue_iterator_fn + self.on_failure = on_failure def __iter__(self): # type: () -> Iterator[T] for item in self.head: yield item - for item in self.continue_iterator_fn(): - yield item + try: + for item in self.continue_iterator_fn(): + yield item + except Exception: + # The continuation token bound into this iterable may be permanently + # invalid; drop this instance from the cache so the next read + # fetches fresh state instead of replaying a dead token forever. + if self.on_failure is not None: + try: + self.on_failure(self) + except Exception: + _LOGGER.warning( + 'Failed to invalidate cached state after continuation ' + 'failure.', + exc_info=True) + raise def get_referents_for_cache(self): # type: () -> List[Any] diff --git a/sdks/python/apache_beam/runners/worker/sdk_worker_test.py b/sdks/python/apache_beam/runners/worker/sdk_worker_test.py index bea313a4d2fd..346e08ad34e3 100644 --- a/sdks/python/apache_beam/runners/worker/sdk_worker_test.py +++ b/sdks/python/apache_beam/runners/worker/sdk_worker_test.py @@ -558,6 +558,31 @@ def clear(self, *args): def process_instruction_id(self, bundle_id): yield + class RevokableStateHandler(UnderlyingStateHandler): + """State handler whose continuation tokens can be permanently revoked, + as happens when a runner loses the work item that created them. + """ + def __init__(self): + super().__init__() + self._token_epoch = 0 + self.initial_read_count = 0 + + def revoke_outstanding_tokens(self): + self._token_epoch += 1 + + def get_raw(self, _state_key, continuation_token=None): + if continuation_token: + epoch, token = continuation_token.split(':') + if int(epoch) != self._token_epoch: + raise RuntimeError('continuation token no longer valid') + continuation_token = token + else: + self.initial_read_count += 1 + data, next_token = super().get_raw(_state_key, continuation_token) + if next_token is not None: + next_token = '%d:%s' % (self._token_epoch, next_token) + return data, next_token + def test_append_clear_with_preexisting_state(self): state = beam_fn_api_pb2.StateKey( bag_user_state=beam_fn_api_pb2.StateKey.BagUserState( @@ -649,6 +674,89 @@ def clear(): self.assertEqual(get_type(), list) self.assertEqual(get(), [i for i in range(1000)]) + def _make_revokable_handler(self): + underlying_state_handler = self.RevokableStateHandler() + state_cache = statecache.StateCache(100 << 20) + handler = GlobalCachingStateHandler(state_cache, underlying_state_handler) + + coder = VarIntCoder() + + state = beam_fn_api_pb2.StateKey( + bag_user_state=beam_fn_api_pb2.StateKey.BagUserState( + user_state_id='state1')) + + cache_token = beam_fn_api_pb2.ProcessBundleRequest.CacheToken( + token=b'state_token1', + user_state=beam_fn_api_pb2.ProcessBundleRequest.CacheToken.UserState()) + + underlying_state_handler.set_continuations(True) + underlying_state_handler.set_values([45, 46, 47], coder) + return underlying_state_handler, handler, coder, state, cache_token + + def test_failed_continuation_invalidates_cached_iterable(self): + (underlying_state_handler, handler, coder, state, + cache_token) = self._make_revokable_handler() + + with handler.process_instruction_id('bundle1', [cache_token]): + iterable = handler.blocking_get(state, coder.get_impl()) + self.assertIsInstance( + iterable, GlobalCachingStateHandler.ContinuationIterable) + underlying_state_handler.revoke_outstanding_tokens() + with self.assertRaises(RuntimeError): + list(iterable) + + with handler.process_instruction_id('bundle2', [cache_token]): + self.assertEqual([45, 46, 47], + list(handler.blocking_get(state, coder.get_impl()))) + + def test_failed_continuation_invalidates_after_partial_yield(self): + (underlying_state_handler, handler, coder, state, + cache_token) = self._make_revokable_handler() + + with handler.process_instruction_id('bundle1', [cache_token]): + iterator = iter(handler.blocking_get(state, coder.get_impl())) + self.assertEqual(45, next(iterator)) + self.assertEqual(46, next(iterator)) + underlying_state_handler.revoke_outstanding_tokens() + with self.assertRaises(RuntimeError): + next(iterator) + + with handler.process_instruction_id('bundle2', [cache_token]): + self.assertEqual([45, 46, 47], + list(handler.blocking_get(state, coder.get_impl()))) + + def test_stale_failed_iterable_does_not_evict_replacement(self): + (underlying_state_handler, handler, coder, state, + cache_token) = self._make_revokable_handler() + + with handler.process_instruction_id('bundle1', [cache_token]): + stale = handler.blocking_get(state, coder.get_impl()) + underlying_state_handler.revoke_outstanding_tokens() + with self.assertRaises(RuntimeError): + list(stale) + + with handler.process_instruction_id('bundle2', [cache_token]): + replacement = handler.blocking_get(state, coder.get_impl()) + self.assertEqual([45, 46, 47], list(replacement)) + # Failing the stale iterable again must not evict the replacement. + with self.assertRaises(RuntimeError): + list(stale) + self.assertIs(replacement, handler.blocking_get(state, coder.get_impl())) + self.assertEqual(2, underlying_state_handler.initial_read_count) + + def test_invalidation_failure_does_not_mask_continuation_error(self): + def failing_continuation(): + raise RuntimeError('continuation token no longer valid') + + def failing_invalidation(_value): + raise ValueError('invalidation failed') + + iterable = GlobalCachingStateHandler.ContinuationIterable( + [45, 46], failing_continuation, on_failure=failing_invalidation) + # The continuation error must propagate despite the failing callback. + with self.assertRaises(RuntimeError): + list(iterable) + class ShortIdCacheTest(unittest.TestCase): def testShortIdAssignment(self): diff --git a/sdks/python/apache_beam/runners/worker/statecache.py b/sdks/python/apache_beam/runners/worker/statecache.py index 33e1c55deeff..7313635efa96 100644 --- a/sdks/python/apache_beam/runners/worker/statecache.py +++ b/sdks/python/apache_beam/runners/worker/statecache.py @@ -338,6 +338,22 @@ def invalidate(self, key: Any) -> None: if weighted_value is not None: self._current_weight -= weighted_value.weight() + def invalidate_if_value(self, key: Any, expected_value: Any) -> None: + """Removes the entry for key only if it still holds expected_value. + + Lets callers that discover a previously returned value is unusable drop + it without evicting a newer value cached under the same key. + """ + assert self.is_cache_enabled() + with self._lock: + weighted_value = self._cache.get(key, None) + if weighted_value is None or _safe_isinstance(weighted_value, + _LoadingValue): + return + if weighted_value.value() is not expected_value: + return + self.invalidate(key) + def invalidate_all(self) -> None: with self._lock: self._cache.clear() diff --git a/sdks/python/apache_beam/runners/worker/statecache_test.py b/sdks/python/apache_beam/runners/worker/statecache_test.py index 7e02ff6a416a..ffb1faf93c15 100644 --- a/sdks/python/apache_beam/runners/worker/statecache_test.py +++ b/sdks/python/apache_beam/runners/worker/statecache_test.py @@ -158,6 +158,42 @@ def test_max_size(self): 'used/max 2/2 MB, hit 100.00%, lookups 0, ' 'avg load time 0 ns, loads 0, evictions 1')) + def test_invalidate_if_value(self): + cache = StateCache(5 << 20) + value = ['value'] + cache.put("key", WeightedValue(value, 1 << 20)) + # Equal but not identical: must not be removed. + cache.invalidate_if_value("key", ['value']) + self.assertEqual(cache.peek("key"), value) + cache.invalidate_if_value("key2", value) + cache.invalidate_if_value("key", value) + self.assertEqual(cache.peek("key"), None) + self.assertEqual(cache.size(), 0) + + def test_invalidate_if_value_ignores_in_flight_load(self): + cache = StateCache(5 << 20) + value = 'value' + load_started = threading.Event() + finish_load = threading.Event() + + def load(_key): + load_started.set() + finish_load.wait() + return value + + result = [] + loader = threading.Thread( + target=lambda: result.append(cache.get("key", load))) + loader.start() + load_started.wait() + # Must neither block on nor remove the in-flight load. + cache.invalidate_if_value("key", value) + self.assertEqual(cache.size(), 1) + finish_load.set() + loader.join() + self.assertEqual(result, [value]) + self.assertEqual(cache.peek("key"), value) + def test_invalidate_all(self): cache = StateCache(5 << 20) cache.put("key", WeightedValue("value", 1 << 20))