Invalidate cached ContinuationIterable when its continuation fails

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