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