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 <noreply@anthropic.com>

Author: addenergyx

sdks/python/apache_beam/runners/worker/sdk_worker.py

@@ -1227,7 +1227,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,
@@ -1354,7 +1354,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]
 
@@ -1372,20 +1373,45 @@ 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 (e.g. the runner revoked it along with the work item that
+        # created it). If this instance stays cached it will fail every
+        # subsequent read of this state, so drop it from the cache (only if
+        # it is still the cached value) and let the next read fetch fresh
+        # state.
+        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]

sdks/python/apache_beam/runners/worker/sdk_worker_test.py

@@ -520,6 +520,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(
@@ -611,6 +636,85 @@ 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]):
+      # blocking_get caches a partially materialized iterable whose lazy
+      # remainder holds the continuation token minted for this read.
+      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)
+
+    # The failed iterable must not remain cached: a read in a later bundle
+    # must fetch fresh state instead of replaying the dead continuation
+    # token forever.
+    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()))
+      # The head page and the first continuation page succeed before the
+      # remaining tokens are revoked mid-iteration.
+      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 freshly cached
+      # 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)
+
 
 class ShortIdCacheTest(unittest.TestCase):
   def testShortIdAssignment(self):

sdks/python/apache_beam/runners/worker/statecache.py

@@ -338,6 +338,24 @@ 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.
+
+    Used by callers that discover a previously returned value is unusable
+    (e.g. it fails during iteration). The entry must only be removed if it
+    has not been replaced by a newer value in the meantime, which the stale
+    value's failure says nothing about.
+    """
+    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()

sdks/python/apache_beam/runners/worker/statecache_test.py

@@ -158,6 +158,20 @@ 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))
+    # A non-matching value must not cause removal.
+    cache.invalidate_if_value("key", 'other')
+    self.assertEqual(cache.peek("key"), value)
+    # A missing key is a no-op.
+    cache.invalidate_if_value("key2", value)
+    # The matching value is removed.
+    cache.invalidate_if_value("key", value)
+    self.assertEqual(cache.peek("key"), None)
+    self.assertEqual(cache.size(), 0)
+
   def test_invalidate_all(self):
     cache = StateCache(5 << 20)
     cache.put("key", WeightedValue("value", 1 << 20))