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

@@ -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]

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):

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()

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))