This is an automated email from the ASF dual-hosted git repository. dianfu pushed a commit to branch release-2.2 in repository https://gitbox.apache.org/repos/asf/flink.git
commit 1f920227136ba86d73d03fd4c25af4e3e6b90135 Author: wuyifan <[email protected]> AuthorDate: Mon Oct 26 15:40:16 2026 +0800 [FLINK-34625][python] Fix state TTL read cache disabling (#29576) Use per-state uncached handlers for TTL policies that require backend reads, preserving other states' shared caches and the existing write-cache policy. Merge pending MapState mutations in empty checks without flushing active iterators' caches. Query backend emptiness without TTL cleanup or timestamp refresh, excluding pending deleted keys and retaining TTL visibility. Restore the request's key context without record-triggered cleanup while preserving storage key prefixes across the supported backends. Cover mixed expiration, paginated iterator continuation, pending mutations, shared-cache isolation and backend key switching with regression tests. --- flink-python/pyflink/fn_execution/state_impl.py | 84 +++- .../pyflink/fn_execution/tests/test_state_impl.py | 414 +++++++++++++++++++- .../python/beam/state/BeamKeyedStateStore.java | 17 +- .../python/beam/state/BeamMapStateHandler.java | 44 +++ .../streaming/api/utils/PythonOperatorUtils.java | 12 +- .../python/beam/state/BeamMapStateHandlerTest.java | 425 +++++++++++++++++++++ .../runtime/state/AbstractKeyedStateBackend.java | 7 + .../flink/runtime/state/KeyedStateBackend.java | 9 + .../runtime/state/internal/InternalMapState.java | 24 +- .../state/metrics/MetricsTrackingMapState.java | 17 + .../flink/runtime/state/ttl/TtlMapState.java | 19 + .../flink/runtime/state/StateBackendTestBase.java | 31 ++ .../changelog/ChangelogKeyedStateBackend.java | 5 + .../restore/ChangelogMigrationRestoreTarget.java | 5 + .../forst/sync/ForStSyncKeyedStateBackend.java | 6 + .../state/rocksdb/RocksDBKeyedStateBackend.java | 6 + 16 files changed, 1097 insertions(+), 28 deletions(-) diff --git a/flink-python/pyflink/fn_execution/state_impl.py b/flink-python/pyflink/fn_execution/state_impl.py index 8ba850169db..ad9e6c279f8 100644 --- a/flink-python/pyflink/fn_execution/state_impl.py +++ b/flink-python/pyflink/fn_execution/state_impl.py @@ -21,6 +21,8 @@ from abc import ABC, abstractmethod from apache_beam.coders import coder_impl from apache_beam.portability.api import beam_fn_api_pb2 from apache_beam.runners.worker.bundle_processor import SynchronousBagRuntimeState +from apache_beam.runners.worker.sdk_worker import GlobalCachingStateHandler +from apache_beam.runners.worker.statecache import StateCache from apache_beam.transforms import userstate from enum import Enum from functools import partial @@ -120,10 +122,6 @@ class SynchronousKvRuntimeState(InternalKvState, ABC): elif ttl_config.get_update_type() == StateTtlConfig.UpdateType.OnReadAndWrite: self._cache_type = SynchronousKvRuntimeState.CacheType.ENABLE_WRITE_CACHE - if self._cache_type != SynchronousKvRuntimeState.CacheType.ENABLE_READ_WRITE_CACHE: - # disable read cache - self._remote_state_backend._state_handler._state_cache._cache._max_entries = 0 - @abstractmethod def get_internal_state(self): pass @@ -145,7 +143,7 @@ class SynchronousBagKvRuntimeState(SynchronousKvRuntimeState, ABC): def get_internal_state(self): if self._internal_state is None: self._internal_state = self._remote_state_backend._get_internal_bag_state( - self.name, self.namespace, self._value_coder, self._ttl_config) + self.name, self.namespace, self._value_coder, self._ttl_config, self._cache_type) return self._internal_state def _maybe_clear_write_cache(self): @@ -365,6 +363,7 @@ class CachingMapStateHandler(object): GET_FLAG = 0 ITERATE_FLAG = 1 CHECK_EMPTY_FLAG = 2 + CHECK_EMPTY_EXCLUDING_KEYS_FLAG = 3 # GET response flags EXIST_FLAG = 0 IS_NONE_FLAG = 1 @@ -555,6 +554,18 @@ class CachingMapStateHandler(object): return False return self._check_empty_raw(state_key) + def check_empty_excluding(self, state_key, excluded_keys, map_key_encoder): + output_stream = coder_impl.create_OutputStream() + output_stream.write_byte(self.CHECK_EMPTY_EXCLUDING_KEYS_FLAG) + output_stream.write_bigendian_int32(len(excluded_keys)) + for key in excluded_keys: + key_stream = coder_impl.create_OutputStream() + map_key_encoder(key, key_stream) + encoded_key = key_stream.get() + output_stream.write_bigendian_int32(len(encoded_key)) + output_stream.write(encoded_key) + return self._check_empty_raw(state_key, output_stream.get()) + def clear(self, state_key): self.clear_read_cache(state_key) return self._underlying.clear(state_key) @@ -577,10 +588,11 @@ class CachingMapStateHandler(object): def reset_cached_iterators_num(self): self._cached_iterator_num = 0 - def _check_empty_raw(self, state_key): - output_stream = coder_impl.create_OutputStream() - output_stream.write_byte(self.CHECK_EMPTY_FLAG) - continuation_token = output_stream.get() + def _check_empty_raw(self, state_key, continuation_token=None): + if continuation_token is None: + output_stream = coder_impl.create_OutputStream() + output_stream.write_byte(self.CHECK_EMPTY_FLAG) + continuation_token = output_stream.get() data, response_token = self._underlying.get_raw(state_key, continuation_token) if data[0] == self.IS_EMPTY_FLAG: return True @@ -736,8 +748,10 @@ class InternalSynchronousMapRuntimeState(object): state_key, map_key_coder, map_value_coder, - max_write_cache_entries): + max_write_cache_entries, + read_cache_enabled=True): self._map_state_handler = map_state_handler + self._read_cache_enabled = read_cache_enabled self._state_key = state_key self._map_key_coder = map_key_coder if isinstance(map_key_coder, FieldCoder): @@ -819,9 +833,22 @@ class InternalSynchronousMapRuntimeState(object): return exists def is_empty(self): - if self._is_empty is None: - if len(self._write_cache) > 0: - self.commit() + if self._is_empty: + return True + if self._write_cache: + if any(exists for exists, _ in self._write_cache.values()): + self._is_empty = False + elif self._cleared: + self._is_empty = True + else: + self._is_empty = self._map_state_handler.check_empty_excluding( + self._state_key, self._write_cache, self._map_key_encoder) + elif self._cleared: + self._is_empty = True + elif not self._read_cache_enabled: + self._is_empty = self._map_state_handler.check_empty_excluding( + self._state_key, (), self._map_key_encoder) + elif self._is_empty is None: self._is_empty = self._map_state_handler.check_empty(self._state_key) return self._is_empty @@ -964,6 +991,11 @@ class RemoteKeyedStateBackend(object): self._state_handler = state_handler self._map_state_handler = CachingMapStateHandler( state_handler, map_state_read_cache_size) + # TTL reads must reach the state backend without disabling other states' caches. + self._uncached_state_handler = GlobalCachingStateHandler( + StateCache(0), state_handler._underlying) + self._uncached_map_state_handler = CachingMapStateHandler( + self._uncached_state_handler, map_state_read_cache_size) self._key_coder_impl = key_coder.get_impl() self.namespace_coder = namespace_coder if namespace_coder: @@ -1056,7 +1088,7 @@ class RemoteKeyedStateBackend(object): self._all_states[name] = wrapped_state return wrapped_state - def _get_internal_bag_state(self, name, namespace, element_coder, ttl_config): + def _get_internal_bag_state(self, name, namespace, element_coder, ttl_config, cache_type): encoded_namespace = self._encode_namespace(namespace) cached_state = self._internal_state_cache.get( (name, self._encoded_current_key, encoded_namespace)) @@ -1069,7 +1101,8 @@ class RemoteKeyedStateBackend(object): if isinstance(element_coder, FieldCoder): element_coder = FlinkCoder(element_coder) state_spec = userstate.BagStateSpec(name, element_coder) - internal_state = self._create_bag_state(state_spec, encoded_namespace, ttl_config) + internal_state = self._create_bag_state( + state_spec, encoded_namespace, ttl_config, cache_type) return internal_state def _get_internal_map_state( @@ -1083,11 +1116,15 @@ class RemoteKeyedStateBackend(object): name, encoded_namespace, map_key_coder, map_value_coder, ttl_config, cache_type) return internal_map_state - def _create_bag_state(self, state_spec: userstate.StateSpec, encoded_namespace, ttl_config) \ + def _create_bag_state( + self, state_spec: userstate.StateSpec, encoded_namespace, ttl_config, cache_type) \ -> userstate.AccumulatingRuntimeState: if isinstance(state_spec, userstate.BagStateSpec): + state_handler = self._state_handler + if cache_type != SynchronousKvRuntimeState.CacheType.ENABLE_READ_WRITE_CACHE: + state_handler = self._uncached_state_handler bag_state = SynchronousBagRuntimeState( - self._state_handler, + state_handler, state_key=self.get_bag_state_key( state_spec.name, self._encoded_current_key, encoded_namespace, ttl_config), value_coder=state_spec.coder) @@ -1113,12 +1150,18 @@ class RemoteKeyedStateBackend(object): write_cache_size = 0 else: write_cache_size = self._map_state_write_cache_size + read_cache_enabled = \ + cache_type == SynchronousKvRuntimeState.CacheType.ENABLE_READ_WRITE_CACHE + map_state_handler = self._map_state_handler + if not read_cache_enabled: + map_state_handler = self._uncached_map_state_handler return InternalSynchronousMapRuntimeState( - self._map_state_handler, + map_state_handler, state_key, map_key_coder, map_value_coder, - write_cache_size) + write_cache_size, + read_cache_enabled) def _encode_namespace(self, namespace): if namespace is not None: @@ -1158,7 +1201,8 @@ class RemoteKeyedStateBackend(object): self.commit_internal_state(state._internal_state) def clear_cached_iterators(self): - if self._map_state_handler.get_cached_iterators_num() > 0: + if self._map_state_handler.get_cached_iterators_num() > 0 or \ + self._uncached_map_state_handler.get_cached_iterators_num() > 0: self._clear_iterator_mark.multimap_side_input.key = self._encoded_current_key self._map_state_handler.clear(self._clear_iterator_mark) diff --git a/flink-python/pyflink/fn_execution/tests/test_state_impl.py b/flink-python/pyflink/fn_execution/tests/test_state_impl.py index 30f42e73ff3..8522b9dde05 100644 --- a/flink-python/pyflink/fn_execution/tests/test_state_impl.py +++ b/flink-python/pyflink/fn_execution/tests/test_state_impl.py @@ -16,13 +16,421 @@ # limitations under the License. ################################################################################ +import contextlib import unittest +from unittest import mock -from apache_beam.coders import StrUtf8Coder +from apache_beam.coders import VarIntCoder, StrUtf8Coder from apache_beam.portability.api import beam_fn_api_pb2 +from apache_beam.runners.worker.sdk_worker import GlobalCachingStateHandler, StateHandler +from apache_beam.runners.worker.statecache import StateCache -from pyflink.fn_execution.state_impl import CachingMapStateHandler, \ - InternalSynchronousMapRuntimeState +from pyflink.common import Time, Types +from pyflink.datastream.state import StateTtlConfig, MapStateDescriptor +from pyflink.fn_execution.state_impl import RemoteKeyedStateBackend, RemoteOperatorStateBackend, \ + CachingMapStateHandler, InternalSynchronousMapRuntimeState + + +class StateTtlReadCacheTests(unittest.TestCase): + + @staticmethod + def _ttl_configs(): + for update_type in (StateTtlConfig.UpdateType.OnCreateAndWrite, + StateTtlConfig.UpdateType.OnReadAndWrite): + yield StateTtlConfig.new_builder(Time.seconds(5)) \ + .set_update_type(update_type) \ + .set_state_visibility(StateTtlConfig.StateVisibility.NeverReturnExpired) \ + .build() + yield StateTtlConfig.new_builder(Time.seconds(5)) \ + .update_ttl_on_read_and_write() \ + .set_state_visibility(StateTtlConfig.StateVisibility.ReturnExpiredIfNotCleanedUp) \ + .build() + + @staticmethod + def _create_backend(cache_size=1 << 20): + # Only remote I/O is mocked. Both Beam caches and Flink runtime states are real. + remote = mock.Mock(spec=StateHandler) + remote.process_instruction_id.return_value = contextlib.nullcontext() + remote.append_raw.return_value = None + remote.clear.return_value = None + handler = GlobalCachingStateHandler(StateCache(cache_size), remote) + backend = RemoteKeyedStateBackend( + handler, StrUtf8Coder(), StrUtf8Coder(), 16, 16, 16) + backend.set_current_key('key') + return backend, handler, remote + + @staticmethod + def _encode(value): + return VarIntCoder().get_impl().encode_nested(value) + + @classmethod + def _map_entry(cls, key, value): + return cls._encode(key) + b'\x01' + cls._encode(value) + + @classmethod + def _check_empty_excluding_request(cls, *keys): + encoded_keys = [cls._encode(key) for key in keys] + return b'\x03' + len(keys).to_bytes(4, 'big') + b''.join( + len(key).to_bytes(4, 'big') + key for key in encoded_keys) + + def test_value_ttl_bypasses_read_cache_without_invalidating_other_states(self): + for ttl in self._ttl_configs(): + with self.subTest(ttl=ttl): + backend, handler, remote = self._create_backend() + cache = handler._state_cache + remote.get_raw.side_effect = [ + (self._encode(42), b''), (self._encode(100), b''), (b'', b'')] + with handler.process_instruction_id('bundle', []): + cached = backend.get_value_state('cached', VarIntCoder()) + self.assertEqual(cached.value(), 42) + state = backend.get_value_state('ttl', VarIntCoder(), ttl) + self.assertEqual(state.value(), 100) + # An empty backend response represents expiration, without a timing dependency. + self.assertIsNone(state.value()) + self.assertEqual(cached.value(), 42) + self.assertEqual(remote.get_raw.call_count, 3) + self.assertIs(backend._state_handler, handler) + self.assertIs(handler._state_cache, cache) + self.assertTrue(cache.is_cache_enabled()) + self.assertIs(cached.get_internal_state()._state_handler, handler) + self.assertIsNot(state.get_internal_state()._state_handler, handler) + + def test_cache_permitting_states_keep_read_cache(self): + ttl = StateTtlConfig.new_builder(Time.seconds(5)) \ + .set_state_visibility(StateTtlConfig.StateVisibility.ReturnExpiredIfNotCleanedUp) \ + .build() + for config in (None, ttl): + with self.subTest(ttl=config): + backend, handler, remote = self._create_backend() + remote.get_raw.side_effect = [(self._encode(100), b'')] + with handler.process_instruction_id('bundle', []): + state = backend.get_value_state('value', VarIntCoder(), config) + self.assertEqual(state.value(), 100) + self.assertEqual(state.value(), 100) + remote.get_raw.assert_called_once() + self.assertIs(state.get_internal_state()._state_handler, handler) + + def test_ttl_write_cache_policy_is_preserved(self): + for ttl in self._ttl_configs(): + with self.subTest(ttl=ttl): + backend, handler, remote = self._create_backend() + with handler.process_instruction_id('bundle', []): + value = backend.get_value_state('value', VarIntCoder(), ttl) + map_state = backend.get_map_state('map', VarIntCoder(), VarIntCoder(), ttl) + value.update(100) + map_state.put(1, 100) + if ttl.get_state_visibility() == \ + StateTtlConfig.StateVisibility.NeverReturnExpired: + self.assertEqual(remote.append_raw.call_count, 2) + else: + remote.append_raw.assert_not_called() + backend.commit() + self.assertEqual(remote.append_raw.call_count, 2) + + def test_map_ttl_bypasses_read_cache_without_invalidating_other_states(self): + for ttl in self._ttl_configs(): + with self.subTest(ttl=ttl): + backend, handler, remote = self._create_backend() + exists = bytes([CachingMapStateHandler.EXIST_FLAG]) + missing = bytes([CachingMapStateHandler.NOT_EXIST_FLAG]) + remote.get_raw.side_effect = [ + (exists + self._encode(42), b''), + (exists + self._encode(100), b''), (missing, b'')] + with handler.process_instruction_id('bundle', []): + cached = backend.get_map_state('cached', VarIntCoder(), VarIntCoder()) + self.assertEqual(cached.get(1), 42) + state = backend.get_map_state('ttl', VarIntCoder(), VarIntCoder(), ttl) + self.assertEqual(state.get(1), 100) + self.assertIsNone(state.get(1)) + self.assertEqual(cached.get(1), 42) + self.assertEqual(remote.get_raw.call_count, 3) + self.assertIs(cached.get_internal_state()._map_state_handler._state_cache, + handler._state_cache) + self.assertFalse(state.get_internal_state()._map_state_handler + ._state_cache.is_cache_enabled()) + + def test_map_ttl_is_empty_after_committed_value_expires(self): + for ttl in self._ttl_configs(): + with self.subTest(ttl=ttl): + backend, handler, remote = self._create_backend() + remote.get_raw.side_effect = [ + (bytes([CachingMapStateHandler.NOT_EXIST_FLAG]), b''), + (bytes([CachingMapStateHandler.IS_EMPTY_FLAG]), b'')] + with handler.process_instruction_id('bundle', []): + state = backend.get_map_state('ttl', VarIntCoder(), VarIntCoder(), ttl) + state.put(1, 100) + backend.commit() + self.assertIsNone(state.get(1)) + self.assertTrue(state.is_empty()) + self.assertEqual(remote.get_raw.call_count, 2) + + def test_map_ttl_is_empty_refreshes_within_bundle(self): + for ttl in self._ttl_configs(): + with self.subTest(ttl=ttl): + backend, handler, remote = self._create_backend() + responses = [ + (bytes([CachingMapStateHandler.NOT_EMPTY_FLAG]), b''), + (bytes([CachingMapStateHandler.IS_EMPTY_FLAG]), b'')] + writes_are_uncached = ttl.get_state_visibility() == \ + StateTtlConfig.StateVisibility.NeverReturnExpired + if writes_are_uncached: + responses.append((bytes([CachingMapStateHandler.NOT_EMPTY_FLAG]), b'')) + remote.get_raw.side_effect = responses + with handler.process_instruction_id('bundle', []): + state = backend.get_map_state('ttl', VarIntCoder(), VarIntCoder(), ttl) + mod_count = state.get_internal_state()._mod_count + self.assertFalse(state.is_empty()) + self.assertTrue(state.is_empty()) + self.assertEqual(state.get_internal_state()._mod_count, mod_count) + self.assertTrue(state.is_empty()) + self.assertEqual(remote.get_raw.call_count, 2) + for call in remote.get_raw.call_args_list: + self.assertEqual(call.args[1], self._check_empty_excluding_request()) + state.put(1, 100) + self.assertFalse(state.is_empty()) + self.assertEqual(remote.get_raw.call_count, 3 if writes_are_uncached else 2) + + def test_map_is_empty_merges_pending_mutations(self): + ttl = list(self._ttl_configs())[-1] + for config in (None, ttl): + for operation in ('put', 'remove', 'clear_then_put', 'clear'): + with self.subTest(ttl=config, operation=operation): + backend, handler, remote = self._create_backend() + with handler.process_instruction_id('bundle', []): + state = backend.get_map_state('map', VarIntCoder(), VarIntCoder(), config) + if operation == 'remove': + state.put(1, 100) + backend.commit() + remote.reset_mock() + state.remove(1) + remote.get_raw.return_value = ( + bytes([CachingMapStateHandler.IS_EMPTY_FLAG]), b'') + elif operation == 'clear': + state.clear() + else: + if operation == 'clear_then_put': + state.clear() + state.put(1, 100) + self.assertEqual(state.is_empty(), operation in ('remove', 'clear')) + remote.append_raw.assert_not_called() + remote.clear.assert_not_called() + if operation == 'remove': + remote.get_raw.assert_called_once() + self.assertEqual(remote.get_raw.call_args.args[1], + self._check_empty_excluding_request(1)) + else: + remote.get_raw.assert_not_called() + + def test_map_ttl_is_empty_preserves_buffered_iterators(self): + ttl = list(self._ttl_configs())[-1] + expected = {'items': [(1, 100), (2, 200)], 'keys': [1, 2], 'values': [100, 200]} + for method in expected: + for cleared, remove in ((False, False), (False, True), (True, False)): + with self.subTest(method=method, cleared=cleared, remove=remove): + backend, handler, remote = self._create_backend() + remote.get_raw.side_effect = lambda _, request: ( + bytes([CachingMapStateHandler.IS_EMPTY_FLAG]) if request[0] == 3 else b'', + b'') + with handler.process_instruction_id('bundle', []): + state = backend.get_map_state('ttl', VarIntCoder(), VarIntCoder(), ttl) + if cleared: + state.clear() + state.put(1, 100) + state.put(2, 200) + iterator = getattr(state, method)() + self.assertEqual(next(iterator), expected[method][0]) + if remove: + iterator.remove() + self.assertFalse(state.is_empty()) + self.assertEqual(next(iterator), expected[method][1]) + if remove: + iterator.remove() + self.assertTrue(state.is_empty()) + self.assertEqual(list(iterator), []) + remote.append_raw.assert_not_called() + remote.clear.assert_not_called() + + def test_map_ttl_iterator_still_rejects_external_mutations(self): + ttl = list(self._ttl_configs())[-1] + for method in ('items', 'keys', 'values'): + for operation in ('put', 'remove', 'clear'): + with self.subTest(method=method, operation=operation): + backend, handler, remote = self._create_backend() + remote.get_raw.return_value = (b'', b'') + with handler.process_instruction_id('bundle', []): + state = backend.get_map_state('ttl', VarIntCoder(), VarIntCoder(), ttl) + state.put(1, 100) + state.put(2, 200) + iterator = getattr(state, method)() + next(iterator) + if operation == 'put': + state.put(3, 300) + elif operation == 'remove': + state.remove(2) + else: + state.clear() + with self.assertRaisesRegex(Exception, 'Concurrent modification detected'): + next(iterator) + + def test_map_ttl_is_empty_excludes_pending_deletions(self): + ttl = list(self._ttl_configs())[-1] + for backend_keys in ((1,), (1, 2)): + with self.subTest(backend_keys=backend_keys): + backend, handler, remote = self._create_backend() + response = CachingMapStateHandler.IS_EMPTY_FLAG if backend_keys == (1,) else \ + CachingMapStateHandler.NOT_EMPTY_FLAG + remote.get_raw.return_value = (bytes([response]), b'') + with handler.process_instruction_id('bundle', []): + state = backend.get_map_state('ttl', VarIntCoder(), VarIntCoder(), ttl) + state.remove(1) + self.assertEqual(state.is_empty(), backend_keys == (1,)) + remote.append_raw.assert_not_called() + remote.get_raw.assert_called_once() + self.assertEqual(remote.get_raw.call_args.args[1], + self._check_empty_excluding_request(1)) + + def test_map_ttl_is_empty_preserves_paged_remote_iterators(self): + ttl = list(self._ttl_configs())[-1] + expected = {'items': [(1, 100), (2, 200)], 'keys': [1, 2], 'values': [100, 200]} + for method in expected: + with self.subTest(method=method): + backend, handler, remote = self._create_backend() + encode_page = self._encode if method == 'keys' else \ + lambda key: self._map_entry(key, key * 100) + remote.get_raw.side_effect = [ + (encode_page(1), b'original-next'), + (bytes([CachingMapStateHandler.NOT_EMPTY_FLAG]), b''), + (encode_page(2), b'')] + with handler.process_instruction_id('bundle', []): + state = backend.get_map_state('ttl', VarIntCoder(), VarIntCoder(), ttl) + iterator = getattr(state, method)() + self.assertEqual(next(iterator), expected[method][0]) + iterator.remove() + self.assertFalse(state.is_empty()) + self.assertEqual(next(iterator), expected[method][1]) + self.assertEqual(list(iterator), []) + self.assertEqual( + backend._uncached_map_state_handler.get_cached_iterators_num(), 0) + self.assertEqual(remote.get_raw.call_count, 3) + self.assertEqual(remote.get_raw.call_args_list[1].args[1], + self._check_empty_excluding_request(1)) + self.assertIn(b'original-next', remote.get_raw.call_args_list[2].args[1]) + remote.append_raw.assert_not_called() + + def test_map_ttl_is_empty_does_not_start_another_ttl_iterator(self): + ttl = list(self._ttl_configs())[-1] + backend, handler, remote = self._create_backend() + remote.get_raw.return_value = (self._map_entry(1, 100), b'active') + with handler.process_instruction_id('bundle', []): + state = backend.get_map_state('ttl', VarIntCoder(), VarIntCoder(), ttl) + iterator = state.items() + self.assertEqual(next(iterator), (1, 100)) + iterator.remove() + internal = state.get_internal_state() + pending = dict(internal._write_cache) + mod_count = internal._mod_count + + def check_empty(_, request): + # Iterating mixed expired/live backend entries could invalidate the active token. + self.assertEqual(request, self._check_empty_excluding_request(1)) + return bytes([CachingMapStateHandler.NOT_EMPTY_FLAG]), b'' + + remote.get_raw.side_effect = check_empty + self.assertFalse(state.is_empty()) + self.assertEqual(internal._write_cache, pending) + self.assertEqual(internal._mod_count, mod_count) + self.assertEqual(backend._uncached_map_state_handler.get_cached_iterators_num(), 1) + remote.append_raw.assert_not_called() + remote.clear.assert_not_called() + remote.get_raw.side_effect = None + remote.get_raw.return_value = (self._map_entry(2, 200), b'') + self.assertEqual(list(iterator), [(2, 200)]) + self.assertEqual(backend._uncached_map_state_handler.get_cached_iterators_num(), 0) + + def test_cache_permitting_map_states_keep_is_empty_cache(self): + ttl = StateTtlConfig.new_builder(Time.seconds(5)) \ + .set_state_visibility(StateTtlConfig.StateVisibility.ReturnExpiredIfNotCleanedUp) \ + .build() + for config in (None, ttl): + with self.subTest(ttl=config): + backend, handler, remote = self._create_backend() + remote.get_raw.side_effect = [ + (bytes([CachingMapStateHandler.NOT_EMPTY_FLAG]), b'')] + with handler.process_instruction_id('bundle', []): + state = backend.get_map_state('map', VarIntCoder(), VarIntCoder(), config) + self.assertFalse(state.is_empty()) + self.assertFalse(state.is_empty()) + remote.get_raw.assert_called_once() + self.assertEqual(remote.get_raw.call_args.args[1], + bytes([CachingMapStateHandler.CHECK_EMPTY_FLAG])) + + def test_broadcast_map_state_keeps_default_cache_behavior(self): + _, handler, remote = self._create_backend() + backend = RemoteOperatorStateBackend(handler, 16, 16, 16) + with handler.process_instruction_id('bundle', []): + state = backend.get_broadcast_state( + MapStateDescriptor('broadcast', Types.INT(), Types.INT())) + state.put(1, 100) + backend.commit() + self.assertFalse(state._internal_map_state.is_empty()) + self.assertEqual(state.get(1), 100) + remote.get_raw.assert_not_called() + + def test_map_ttl_iteration_and_unfinished_iterator_cleanup(self): + ttl = next(self._ttl_configs()) + backend, handler, remote = self._create_backend() + remote.get_raw.side_effect = [ + (self._map_entry(1, 100), b''), (b'', b''), + (self._map_entry(1, 100), b'next-page'), + (self._map_entry(2, 200), b''), + (self._map_entry(1, 100), b'unfinished')] + with handler.process_instruction_id('bundle', []): + state = backend.get_map_state('ttl', VarIntCoder(), VarIntCoder(), ttl) + self.assertEqual(list(state.items()), [(1, 100)]) + self.assertEqual(list(state.items()), []) + self.assertEqual(list(state.items()), [(1, 100), (2, 200)]) + self.assertIn(b'next-page', remote.get_raw.call_args_list[3].args[1]) + self.assertEqual(next(state.items()), (1, 100)) + self.assertEqual(backend._map_state_handler.get_cached_iterators_num(), 0) + backend.clear_cached_iterators() + remote.clear.assert_called_once() + clear_key = remote.clear.call_args.args[0].multimap_side_input + self.assertEqual(clear_key.transform_id, 'clear_iterators') + self.assertEqual(clear_key.key, b'key') + + def test_ttl_handler_is_preserved_when_reusing_key_and_namespace(self): + ttl = next(self._ttl_configs()) + backend, handler, remote = self._create_backend() + remote.get_raw.side_effect = [ + (self._encode(100), b''), (self._encode(200), b''), (b'', b'')] + with handler.process_instruction_id('bundle', []): + state = backend.get_value_state('ttl', VarIntCoder(), ttl) + state.set_current_namespace('first') + self.assertEqual(state.value(), 100) + first_internal_state = state.get_internal_state() + backend.set_current_key('other') + state.set_current_namespace('second') + self.assertEqual(state.value(), 200) + backend.set_current_key('key') + state.set_current_namespace('first') + self.assertIs(state.get_internal_state(), first_internal_state) + self.assertIsNone(state.value()) + first_key, second_key, reused_key = [ + call.args[0] for call in remote.get_raw.call_args_list] + self.assertEqual(first_key, reused_key) + self.assertNotEqual(first_key, second_key) + + def test_repeated_state_registration_with_disabled_cache(self): + ttl = next(self._ttl_configs()) + backend, handler, remote = self._create_backend(cache_size=0) + remote.get_raw.side_effect = [(self._encode(100), b''), (b'', b'')] + with handler.process_instruction_id('bundle', []): + state = backend.get_value_state('ttl', VarIntCoder(), ttl) + self.assertEqual(state.value(), 100) + self.assertIs(backend.get_value_state('ttl', VarIntCoder(), ttl), state) + self.assertIsNone(state.value()) + self.assertEqual(remote.get_raw.call_count, 2) + self.assertFalse(handler._state_cache.is_cache_enabled()) class InternalSynchronousMapRuntimeStateTests(unittest.TestCase): diff --git a/flink-python/src/main/java/org/apache/flink/streaming/api/runners/python/beam/state/BeamKeyedStateStore.java b/flink-python/src/main/java/org/apache/flink/streaming/api/runners/python/beam/state/BeamKeyedStateStore.java index 83981df889b..b50242f3367 100644 --- a/flink-python/src/main/java/org/apache/flink/streaming/api/runners/python/beam/state/BeamKeyedStateStore.java +++ b/flink-python/src/main/java/org/apache/flink/streaming/api/runners/python/beam/state/BeamKeyedStateStore.java @@ -152,7 +152,12 @@ public class BeamKeyedStateStore implements BeamStateStore { byte[] keyBytes = mapUserState.getKey().toByteArray(); bais.setBuffer(keyBytes, 0, keyBytes.length); Object key = keySerializer.deserialize(baisWrapper); - setCurrentKey(key); + boolean readOnly = + request.hasGet() + && !request.getGet().getContinuationToken().isEmpty() + && request.getGet().getContinuationToken().byteAt(0) + == BeamMapStateHandler.CHECK_EMPTY_EXCLUDING_KEYS_FLAG; + setCurrentKey(key, readOnly); // Retrieve state descriptor byte[] data = Base64.getDecoder().decode(mapUserState.getSideInputId()); @@ -198,13 +203,19 @@ public class BeamKeyedStateStore implements BeamStateStore { } private void setCurrentKey(Object key) { + setCurrentKey(key, false); + } + + private void setCurrentKey(Object key, boolean forStateAccess) { if (keyedStateBackend.getKeySerializer() instanceof RowDataSerializer) { setCurrentKeyForStreaming( (KeyedStateBackend<? super BinaryRowData>) keyedStateBackend, ((RowDataSerializer) keyedStateBackend.getKeySerializer()) - .toBinaryRow((RowData) key)); + .toBinaryRow((RowData) key), + forStateAccess); } else { - setCurrentKeyForStreaming((KeyedStateBackend<Object>) keyedStateBackend, key); + setCurrentKeyForStreaming( + (KeyedStateBackend<Object>) keyedStateBackend, key, forStateAccess); } } } diff --git a/flink-python/src/main/java/org/apache/flink/streaming/api/runners/python/beam/state/BeamMapStateHandler.java b/flink-python/src/main/java/org/apache/flink/streaming/api/runners/python/beam/state/BeamMapStateHandler.java index 0b181f6bf41..8c8909b6b80 100644 --- a/flink-python/src/main/java/org/apache/flink/streaming/api/runners/python/beam/state/BeamMapStateHandler.java +++ b/flink-python/src/main/java/org/apache/flink/streaming/api/runners/python/beam/state/BeamMapStateHandler.java @@ -25,16 +25,21 @@ import org.apache.flink.core.memory.ByteArrayOutputStreamWithPos; import org.apache.flink.core.memory.DataInputViewStreamWrapper; import org.apache.flink.core.memory.DataOutputViewStreamWrapper; import org.apache.flink.python.PythonOptions; +import org.apache.flink.runtime.state.internal.InternalMapState; import org.apache.flink.streaming.api.utils.ByteArrayWrapper; import org.apache.beam.model.fnexecution.v1.BeamFnApi; import org.apache.beam.vendor.grpc.v1p60p1.com.google.protobuf.ByteString; import java.util.HashMap; +import java.util.HashSet; import java.util.Iterator; import java.util.Map; +import java.util.Set; import java.util.UUID; +import static org.apache.flink.util.Preconditions.checkArgument; + /** BeamMapStateHandler handles operations on a {@link MapState}. */ public class BeamMapStateHandler extends AbstractBeamStateHandler<MapState<ByteArrayWrapper, byte[]>> { @@ -46,6 +51,7 @@ public class BeamMapStateHandler private static final byte ITERATE_FLAG = 1; private static final byte CHECK_EMPTY_FLAG = 2; + static final byte CHECK_EMPTY_EXCLUDING_KEYS_FLAG = 3; /** Map state GET response flags. */ private static final byte EXIST_FLAG = 0; @@ -118,6 +124,8 @@ public class BeamMapStateHandler // [flag (1 byte)][serialized map key] // The continuation token structure of CHECK_EMPTY request is: // [flag (1 byte)] + // The continuation token structure of CHECK_EMPTY_EXCLUDING_KEYS request is: + // [flag (1 byte)][key count (int32)][key length (int32)][serialized key]... // The continuation token structure of ITERATE request is: // [flag (1 byte)][iterate type (1 byte)][iterator token length (int32)][iterator token] byte[] getRequest = request.getGet().getContinuationToken().toByteArray(); @@ -133,6 +141,22 @@ public class BeamMapStateHandler case CHECK_EMPTY_FLAG: response = handleCheckEmpty(mapState); break; + case CHECK_EMPTY_EXCLUDING_KEYS_FLAG: + bais.setBuffer(getRequest, 1, getRequest.length - 1); + final int keyCount = baisWrapper.readInt(); + checkArgument(keyCount >= 0, "The excluded key count must not be negative."); + final Set<ByteArrayWrapper> excludedKeys = new HashSet<>(); + for (int i = 0; i < keyCount; i++) { + final int keyLength = baisWrapper.readInt(); + checkArgument( + keyLength >= 0 && keyLength <= bais.available(), + "Invalid excluded key length."); + final byte[] key = new byte[keyLength]; + baisWrapper.readFully(key); + excludedKeys.add(new ByteArrayWrapper(key)); + } + response = handleCheckEmptyExcludingKeys(mapState, excludedKeys); + break; case ITERATE_FLAG: bais.setBuffer(getRequest, 1, getRequest.length - 1); IterateType iterateType = IterateType.fromOrd(baisWrapper.readByte()); @@ -245,6 +269,26 @@ public class BeamMapStateHandler } } + private BeamFnApi.StateGetResponse.Builder handleCheckEmptyExcludingKeys( + MapState<ByteArrayWrapper, byte[]> mapState, Set<ByteArrayWrapper> excludedKeys) + throws Exception { + if (mapState instanceof InternalMapState<?, ?, ?, ?>) { + return ((InternalMapState<?, ?, ?, ?>) mapState).isEmptyExcluding(excludedKeys) + ? IS_EMPTY_RESPONSE + : NOT_EMPTY_RESPONSE; + } + // Operator map state has no TTL decorator. + final Iterable<ByteArrayWrapper> keys = mapState.keys(); + if (keys != null) { + for (ByteArrayWrapper key : keys) { + if (!excludedKeys.contains(key)) { + return NOT_EMPTY_RESPONSE; + } + } + } + return IS_EMPTY_RESPONSE; + } + private BeamFnApi.StateGetResponse.Builder handleIterate( MapState<ByteArrayWrapper, byte[]> mapState, IterateType iterateType, diff --git a/flink-python/src/main/java/org/apache/flink/streaming/api/utils/PythonOperatorUtils.java b/flink-python/src/main/java/org/apache/flink/streaming/api/utils/PythonOperatorUtils.java index dc26f452ab5..7d657d2969f 100644 --- a/flink-python/src/main/java/org/apache/flink/streaming/api/utils/PythonOperatorUtils.java +++ b/flink-python/src/main/java/org/apache/flink/streaming/api/utils/PythonOperatorUtils.java @@ -31,8 +31,18 @@ public class PythonOperatorUtils { /** Set the current key for streaming operator. */ public static <K> void setCurrentKeyForStreaming( KeyedStateBackend<K> stateBackend, K currentKey) { + setCurrentKeyForStreaming(stateBackend, currentKey, false); + } + + /** Selects a state-request key without treating a read-only query as another input record. */ + public static <K> void setCurrentKeyForStreaming( + KeyedStateBackend<K> stateBackend, K currentKey, boolean forStateAccess) { if (!inBatchExecutionMode(stateBackend)) { - stateBackend.setCurrentKey(currentKey); + if (forStateAccess) { + stateBackend.setCurrentKeyForStateAccess(currentKey); + } else { + stateBackend.setCurrentKey(currentKey); + } } } diff --git a/flink-python/src/test/java/org/apache/flink/streaming/api/runners/python/beam/state/BeamMapStateHandlerTest.java b/flink-python/src/test/java/org/apache/flink/streaming/api/runners/python/beam/state/BeamMapStateHandlerTest.java new file mode 100644 index 00000000000..7e24de0ed93 --- /dev/null +++ b/flink-python/src/test/java/org/apache/flink/streaming/api/runners/python/beam/state/BeamMapStateHandlerTest.java @@ -0,0 +1,425 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.flink.streaming.api.runners.python.beam.state; + +import org.apache.flink.api.common.state.MapState; +import org.apache.flink.api.common.state.MapStateDescriptor; +import org.apache.flink.api.common.state.StateTtlConfig; +import org.apache.flink.api.common.typeinfo.Types; +import org.apache.flink.api.common.typeutils.base.IntSerializer; +import org.apache.flink.api.common.typeutils.base.array.BytePrimitiveArraySerializer; +import org.apache.flink.configuration.Configuration; +import org.apache.flink.configuration.StateLatencyTrackOptions; +import org.apache.flink.fnexecution.v1.FlinkFnApi.StateDescriptor; +import org.apache.flink.fnexecution.v1.FlinkFnApi.StateDescriptor.StateTTLConfig; +import org.apache.flink.fnexecution.v1.FlinkFnApi.StateDescriptor.StateTTLConfig.CleanupStrategies; +import org.apache.flink.python.PythonOptions; +import org.apache.flink.runtime.state.AbstractKeyedStateBackend; +import org.apache.flink.runtime.state.KeyGroupRangeAssignment; +import org.apache.flink.runtime.state.VoidNamespace; +import org.apache.flink.runtime.state.VoidNamespaceSerializer; +import org.apache.flink.runtime.state.hashmap.HashMapStateBackend; +import org.apache.flink.streaming.api.operators.StreamMap; +import org.apache.flink.streaming.api.utils.ByteArrayWrapper; +import org.apache.flink.streaming.api.utils.ByteArrayWrapperSerializer; +import org.apache.flink.streaming.util.KeyedOneInputStreamOperatorTestHarness; + +import org.apache.beam.model.fnexecution.v1.BeamFnApi; +import org.apache.beam.vendor.grpc.v1p60p1.com.google.protobuf.ByteString; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.CsvSource; +import org.junit.jupiter.params.provider.ValueSource; + +import java.io.ByteArrayInputStream; +import java.io.ByteArrayOutputStream; +import java.io.DataInputStream; +import java.io.DataOutputStream; +import java.time.Duration; +import java.util.ArrayList; +import java.util.Base64; +import java.util.Collections; +import java.util.List; +import java.util.concurrent.atomic.AtomicInteger; + +import static org.assertj.core.api.Assertions.assertThat; + +/** Tests the map-state protocol against a real heap backend and a controllable TTL clock. */ +class BeamMapStateHandlerTest { + + @ParameterizedTest + @CsvSource({ + "ITEMS, false, ReturnExpiredIfNotCleanedUp", + "KEYS, false, ReturnExpiredIfNotCleanedUp", + "VALUES, false, ReturnExpiredIfNotCleanedUp", + "ITEMS, true, ReturnExpiredIfNotCleanedUp", + "KEYS, true, ReturnExpiredIfNotCleanedUp", + "VALUES, true, ReturnExpiredIfNotCleanedUp", + "ITEMS, false, NeverReturnExpired", + "KEYS, false, NeverReturnExpired", + "VALUES, false, NeverReturnExpired", + "ITEMS, true, NeverReturnExpired", + "KEYS, true, NeverReturnExpired", + "VALUES, true, NeverReturnExpired" + }) + void testMixedExpiryDoesNotInvalidateActiveIterator( + IterateType type, boolean metrics, StateTtlConfig.StateVisibility visibility) + throws Exception { + try (Fixture fixture = new Fixture(metrics, visibility)) { + for (int key = 0; key < 64; key++) { + fixture.state.put(key(key), new byte[] {(byte) key}); + } + List<Integer> order = new ArrayList<>(); + for (ByteArrayWrapper key : fixture.state.keys()) { + order.add(Byte.toUnsignedInt(key.get(0))); + } + fixture.harness.setStateTtlProcessingTime(101); + for (int key : order.subList(0, 8)) { + fixture.state.put(key(key), new byte[] {(byte) key}); + } + + BeamFnApi.StateGetResponse page = iterate(fixture, type, ByteString.EMPTY); + List<Integer> firstPage = decodeKeys(page, type); + assertThat(firstPage).containsExactlyElementsOf(order.subList(0, 2)); + assertThat(page.getContinuationToken()).isNotEqualTo(ByteString.EMPTY); + assertThat(checkEmptyExcluding(fixture, Collections.singletonList(firstPage.get(0)))) + .isFalse(); + + List<Integer> visited = new ArrayList<>(firstPage); + while (!page.getContinuationToken().isEmpty()) { + page = iterate(fixture, type, page.getContinuationToken()); + visited.addAll(decodeKeys(page, type)); + } + // ReturnExpiredIfNotCleanedUp permits expired entries to be returned during cleanup. + List<Integer> expected = + visibility == StateTtlConfig.StateVisibility.ReturnExpiredIfNotCleanedUp + ? order + : order.subList(0, 8); + assertThat(visited).containsExactlyElementsOf(expected); + } + } + + @ParameterizedTest + @ValueSource(booleans = {false, true}) + void testExcludedKeysAndNullValues(boolean metrics) throws Exception { + try (Fixture fixture = + new Fixture(metrics, StateTtlConfig.StateVisibility.NeverReturnExpired)) { + fixture.state.put(key(1), null); + fixture.state.put(key(2), new byte[] {2}); + assertThat(checkEmptyExcluding(fixture, Collections.emptyList())).isFalse(); + assertThat(checkEmptyExcluding(fixture, Collections.singletonList(2))).isFalse(); + assertThat(checkEmptyExcluding(fixture, List.of(1, 2))).isTrue(); + fixture.state.clear(); + assertThat(checkEmptyExcluding(fixture, Collections.emptyList())).isTrue(); + } + } + + @ParameterizedTest + @ValueSource(booleans = {false, true}) + void testCheckingEmptyDoesNotRefreshTtl(boolean metrics) throws Exception { + try (Fixture fixture = + new Fixture(metrics, StateTtlConfig.StateVisibility.NeverReturnExpired)) { + fixture.state.put(key(1), new byte[] {1}); + fixture.harness.setStateTtlProcessingTime(50); + assertThat(checkEmptyExcluding(fixture, Collections.singletonList(2))).isFalse(); + fixture.harness.setStateTtlProcessingTime(101); + assertThat(fixture.state.get(key(1))).isNull(); + } + } + + @ParameterizedTest + @ValueSource(booleans = {false, true}) + void testEmptyExclusionsDoNotRunIncrementalCleanup(boolean metrics) throws Exception { + try (Fixture fixture = + new Fixture(metrics, StateTtlConfig.StateVisibility.NeverReturnExpired, true)) { + fixture.state.put(key(1), new byte[] {1}); + fixture.harness.setStateTtlProcessingTime(101); + assertThat(checkEmptyExcluding(fixture, Collections.emptyList())).isTrue(); + // A read-only query reports expiration while retaining the physical state entry. + assertThat(fixture.harness.numKeyedStateEntries()).isEqualTo(1); + } + } + + @ParameterizedTest + @ValueSource(booleans = {false, true}) + void testReadOnlyRpcDoesNotTriggerPerRecordCleanup(boolean metrics) throws Exception { + try (Fixture fixture = + new Fixture( + metrics, StateTtlConfig.StateVisibility.NeverReturnExpired, false, 16)) { + BeamStateRequestHandler requests = + BeamStateRequestHandler.of( + fixture.backend, null, IntSerializer.INSTANCE, null, fixture.config); + AtomicInteger selections = new AtomicInteger(); + fixture.backend.registerKeySelectionListener(ignored -> selections.incrementAndGet()); + BeamFnApi.StateKey firstKey = rpcStateKey(0); + BeamFnApi.StateKey secondKey = rpcStateKey(1); + appendThroughStore(requests, firstKey); + appendThroughStore(requests, secondKey); + BeamFnApi.StateGetResponse page = + getThroughStore( + requests, + firstKey, + iterateRequest(IterateType.ITEMS, ByteString.EMPTY)); + assertThat(page.getContinuationToken()).isNotEqualTo(ByteString.EMPTY); + int beforeQuery = selections.get(); + assertThat(beforeQuery).isEqualTo(3); + fixture.harness.setStateTtlProcessingTime(101); + + for (int key : new int[] {0, 1, 0}) { + BeamFnApi.StateGetResponse response = + getThroughStore( + requests, + rpcStateKey(key), + checkEmptyExcludingRequest(Collections.emptyList())); + assertThat(response.getData()).isEqualTo(ByteString.copyFrom(new byte[] {3})); + assertThat(selections.get()).isEqualTo(beforeQuery); + assertThat(fixture.backend.getCurrentKey()).isEqualTo(key); + assertThat(fixture.backend.getCurrentKeyGroupIndex()) + .isEqualTo(KeyGroupRangeAssignment.assignToKeyGroup(key, 16)); + assertThat(fixture.harness.numKeyedStateEntries()).isEqualTo(2); + } + + int ordinaryRequests = 0; + while (!page.getContinuationToken().isEmpty()) { + page = + getThroughStore( + requests, + firstKey, + iterateRequest(IterateType.ITEMS, page.getContinuationToken())); + ordinaryRequests++; + } + assertThat(ordinaryRequests).isPositive(); + assertThat(selections.get()).isEqualTo(beforeQuery + ordinaryRequests); + } + } + + private static BeamFnApi.StateKey rpcStateKey(int key) throws Exception { + CleanupStrategies cleanup = + CleanupStrategies.newBuilder() + .setIsCleanupInBackground(false) + .addStrategies( + CleanupStrategies.MapStrategiesEntry.newBuilder() + .setStrategy( + CleanupStrategies.Strategies.INCREMENTAL_CLEANUP) + .setIncrementalCleanupStrategy( + CleanupStrategies.IncrementalCleanupStrategy + .newBuilder() + .setCleanupSize(1) + .setRunCleanupForEveryRecord(true))) + .build(); + StateDescriptor descriptor = + StateDescriptor.newBuilder() + .setStateName("rpc-map") + .setStateTtlConfig( + StateTTLConfig.newBuilder() + .setTtl(100) + .setUpdateType(StateTTLConfig.UpdateType.OnReadAndWrite) + .setStateVisibility( + StateTTLConfig.StateVisibility.NeverReturnExpired) + .setCleanupStrategies(cleanup)) + .build(); + ByteArrayOutputStream keyBytes = new ByteArrayOutputStream(); + new DataOutputStream(keyBytes).writeInt(key); + return BeamFnApi.StateKey.newBuilder() + .setMultimapSideInput( + BeamFnApi.StateKey.MultimapSideInput.newBuilder() + .setKey(ByteString.copyFrom(keyBytes.toByteArray())) + .setSideInputId( + Base64.getEncoder() + .encodeToString(descriptor.toByteArray()))) + .build(); + } + + private static void appendThroughStore(BeamStateRequestHandler handler, BeamFnApi.StateKey key) + throws Exception { + ByteArrayOutputStream bytes = new ByteArrayOutputStream(); + DataOutputStream data = new DataOutputStream(bytes); + data.writeInt(4); + for (int mapKey = 0; mapKey < 4; mapKey++) { + data.writeByte(2); + data.writeInt(1); + data.writeByte(mapKey); + data.writeInt(1); + data.writeByte(mapKey); + } + handler.handle( + BeamFnApi.StateRequest.newBuilder() + .setId("append") + .setStateKey(key) + .setAppend( + BeamFnApi.StateAppendRequest.newBuilder() + .setData(ByteString.copyFrom(bytes.toByteArray()))) + .build()) + .toCompletableFuture() + .get(); + } + + private static BeamFnApi.StateGetResponse getThroughStore( + BeamStateRequestHandler handler, BeamFnApi.StateKey key, byte[] request) + throws Exception { + return handler.handle( + BeamFnApi.StateRequest.newBuilder() + .setId("get") + .setStateKey(key) + .setGet( + BeamFnApi.StateGetRequest.newBuilder() + .setContinuationToken(ByteString.copyFrom(request))) + .build()) + .toCompletableFuture() + .get() + .build() + .getGet(); + } + + private static boolean checkEmptyExcluding(Fixture fixture, List<Integer> excluded) + throws Exception { + BeamFnApi.StateGetResponse response = get(fixture, checkEmptyExcludingRequest(excluded)); + assertThat(response.getContinuationToken()).isEqualTo(ByteString.EMPTY); + assertThat(response.getData().size()).isEqualTo(1); + assertThat(response.getData().byteAt(0)).isIn((byte) 3, (byte) 4); + return response.getData().byteAt(0) == 3; + } + + private static byte[] checkEmptyExcludingRequest(List<Integer> excluded) throws Exception { + ByteArrayOutputStream bytes = new ByteArrayOutputStream(); + DataOutputStream request = new DataOutputStream(bytes); + request.writeByte(3); + request.writeInt(excluded.size()); + for (int key : excluded) { + request.writeInt(1); + request.writeByte(key); + } + return bytes.toByteArray(); + } + + private static BeamFnApi.StateGetResponse iterate( + Fixture fixture, IterateType type, ByteString token) throws Exception { + return get(fixture, iterateRequest(type, token)); + } + + private static byte[] iterateRequest(IterateType type, ByteString token) throws Exception { + ByteArrayOutputStream bytes = new ByteArrayOutputStream(); + DataOutputStream request = new DataOutputStream(bytes); + request.writeByte(1); + request.writeByte(type.getOrd()); + request.writeInt(token.size()); + token.writeTo(request); + return bytes.toByteArray(); + } + + private static BeamFnApi.StateGetResponse get(Fixture fixture, byte[] request) + throws Exception { + return fixture.handler + .handleGet( + BeamFnApi.StateRequest.newBuilder() + .setId("request") + .setGet( + BeamFnApi.StateGetRequest.newBuilder() + .setContinuationToken(ByteString.copyFrom(request))) + .build(), + fixture.state) + .build() + .getGet(); + } + + private static List<Integer> decodeKeys(BeamFnApi.StateGetResponse response, IterateType type) + throws Exception { + DataInputStream input = + new DataInputStream(new ByteArrayInputStream(response.getData().toByteArray())); + List<Integer> keys = new ArrayList<>(); + while (input.available() > 0) { + keys.add(input.readUnsignedByte()); + if (type != IterateType.KEYS && input.readBoolean()) { + input.readByte(); + } + } + return keys; + } + + private static ByteArrayWrapper key(int key) { + return new ByteArrayWrapper(new byte[] {(byte) key}); + } + + private static class Fixture implements AutoCloseable { + private final KeyedOneInputStreamOperatorTestHarness<Integer, Integer, Integer> harness; + private final MapState<ByteArrayWrapper, byte[]> state; + private final BeamMapStateHandler handler; + private final AbstractKeyedStateBackend<Integer> backend; + private final Configuration config; + + private Fixture(boolean metrics, StateTtlConfig.StateVisibility visibility) + throws Exception { + this(metrics, visibility, false); + } + + private Fixture( + boolean metrics, StateTtlConfig.StateVisibility visibility, boolean incremental) + throws Exception { + this(metrics, visibility, incremental, 1); + } + + private Fixture( + boolean metrics, + StateTtlConfig.StateVisibility visibility, + boolean incremental, + int maxParallelism) + throws Exception { + config = new Configuration(); + config.set(PythonOptions.MAP_STATE_ITERATE_RESPONSE_BATCH_SIZE, 2); + config.set(StateLatencyTrackOptions.LATENCY_TRACK_ENABLED, metrics); + config.set(StateLatencyTrackOptions.LATENCY_TRACK_SAMPLE_INTERVAL, 1); + StreamMap<Integer, Integer> operator = new StreamMap<>(value -> value); + harness = + new KeyedOneInputStreamOperatorTestHarness<>( + operator, value -> 0, Types.INT, maxParallelism, 1, 0); + harness.setStateBackend( + new HashMapStateBackend() + .configure(config, Thread.currentThread().getContextClassLoader())); + harness.open(); + harness.setStateTtlProcessingTime(0); + backend = (AbstractKeyedStateBackend<Integer>) operator.<Integer>getKeyedStateBackend(); + backend.setCurrentKey(0); + MapStateDescriptor<ByteArrayWrapper, byte[]> descriptor = + new MapStateDescriptor<>( + "map", + ByteArrayWrapperSerializer.INSTANCE, + BytePrimitiveArraySerializer.INSTANCE); + StateTtlConfig.Builder ttl = + StateTtlConfig.newBuilder(Duration.ofMillis(100)) + .setUpdateType(StateTtlConfig.UpdateType.OnReadAndWrite) + .setStateVisibility(visibility) + .disableCleanupInBackground(); + if (incremental) { + ttl.cleanupIncrementally(1, false); + } + descriptor.enableTimeToLive(ttl.build()); + state = + backend.getPartitionedState( + VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, descriptor); + if (metrics) { + assertThat(state.getClass().getSimpleName()).isEqualTo("MetricsTrackingMapState"); + } + handler = new BeamMapStateHandler(config); + } + + @Override + public void close() throws Exception { + harness.close(); + } + } +} diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/AbstractKeyedStateBackend.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/AbstractKeyedStateBackend.java index 1bec7b2218d..f34378a250a 100644 --- a/flink-runtime/src/main/java/org/apache/flink/runtime/state/AbstractKeyedStateBackend.java +++ b/flink-runtime/src/main/java/org/apache/flink/runtime/state/AbstractKeyedStateBackend.java @@ -262,6 +262,13 @@ public abstract class AbstractKeyedStateBackend<K> KeyGroupRangeAssignment.assignToKeyGroup(newKey, numberOfKeyGroups)); } + @Override + public void setCurrentKeyForStateAccess(K newKey) { + this.keyContext.setCurrentKey(newKey); + this.keyContext.setCurrentKeyGroupIndex( + KeyGroupRangeAssignment.assignToKeyGroup(newKey, numberOfKeyGroups)); + } + @Override public void setCurrentKeyAndKeyGroup(K newKey, int newKeyGroupIndex) { notifyKeySelected(newKey); diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/KeyedStateBackend.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/KeyedStateBackend.java index d773fdb4a0a..490e788b59d 100644 --- a/flink-runtime/src/main/java/org/apache/flink/runtime/state/KeyedStateBackend.java +++ b/flink-runtime/src/main/java/org/apache/flink/runtime/state/KeyedStateBackend.java @@ -43,6 +43,15 @@ public interface KeyedStateBackend<K> */ void setCurrentKey(K newKey); + /** + * Restores the key context for a state request rather than selecting a key for another input + * record. Backends with record-triggered maintenance should override this method to avoid + * running that maintenance for the request. + */ + default void setCurrentKeyForStateAccess(K newKey) { + setCurrentKey(newKey); + } + /** * @return Current key. */ diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/internal/InternalMapState.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/internal/InternalMapState.java index 92ae9c0c4c0..84572179c8a 100644 --- a/flink-runtime/src/main/java/org/apache/flink/runtime/state/internal/InternalMapState.java +++ b/flink-runtime/src/main/java/org/apache/flink/runtime/state/internal/InternalMapState.java @@ -21,6 +21,7 @@ package org.apache.flink.runtime.state.internal; import org.apache.flink.api.common.state.MapState; import java.util.Map; +import java.util.Set; /** * The peer to the {@link MapState} in the internal state type hierarchy. @@ -33,4 +34,25 @@ import java.util.Map; * @param <UV> Type of the value in the state */ public interface InternalMapState<K, N, UK, UV> - extends InternalKvState<K, N, Map<UK, UV>>, MapState<UK, UV> {} + extends InternalKvState<K, N, Map<UK, UV>>, MapState<UK, UV> { + + /** + * Returns whether this state has no visible entries other than the excluded keys, without + * updating TTL timestamps or cleaning up expired entries. + * + * <p>This allows clients with buffered deletions to check emptiness without applying those + * deletions or invalidating active iterators. TTL decorators must override this method instead + * of using their ordinary, potentially mutating key iteration. + */ + default boolean isEmptyExcluding(Set<?> excludedKeys) throws Exception { + final Iterable<UK> keys = keys(); + if (keys != null) { + for (UK key : keys) { + if (!excludedKeys.contains(key)) { + return false; + } + } + } + return true; + } +} diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/metrics/MetricsTrackingMapState.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/metrics/MetricsTrackingMapState.java index 872b12ea691..f693ffdd97e 100644 --- a/flink-runtime/src/main/java/org/apache/flink/runtime/state/metrics/MetricsTrackingMapState.java +++ b/flink-runtime/src/main/java/org/apache/flink/runtime/state/metrics/MetricsTrackingMapState.java @@ -29,6 +29,7 @@ import org.apache.flink.util.function.ThrowingConsumer; import java.io.IOException; import java.util.Iterator; import java.util.Map; +import java.util.Set; /** * This class wraps map state with latency tracking logic. @@ -261,6 +262,22 @@ class MetricsTrackingMapState<K, N, UK, UV> } } + @Override + public boolean isEmptyExcluding(Set<?> excludedKeys) throws Exception { + if (sizeTrackingStateMetric != null && sizeTrackingStateMetric.trackMetricsOnIsEmpty()) { + sizeTrackingStateMetric.updateMetrics( + MapStateMetrics.MAP_STATE_IS_EMPTY_KEY_SIZE, super.sizeOfKey()); + } + if (latencyTrackingStateMetric != null + && latencyTrackingStateMetric.trackMetricsOnIsEmpty()) { + return trackLatencyWithException( + () -> original.isEmptyExcluding(excludedKeys), + MapStateMetrics.MAP_STATE_IS_EMPTY_LATENCY); + } else { + return original.isEmptyExcluding(excludedKeys); + } + } + protected long sizeOfKeyAndUserKey(UK userKey) throws IOException { if (userKeySerializer == null || userKey == null) { diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/ttl/TtlMapState.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/ttl/TtlMapState.java index 2dc2df3a031..242fa43c23d 100644 --- a/flink-runtime/src/main/java/org/apache/flink/runtime/state/ttl/TtlMapState.java +++ b/flink-runtime/src/main/java/org/apache/flink/runtime/state/ttl/TtlMapState.java @@ -33,6 +33,7 @@ import java.util.HashMap; import java.util.Iterator; import java.util.Map; import java.util.NoSuchElementException; +import java.util.Set; import java.util.function.Function; /** @@ -132,6 +133,24 @@ class TtlMapState<K, N, UK, UV> return original.isEmpty(); } + @Override + public boolean isEmptyExcluding(Set<?> excludedKeys) throws Exception { + // Ordinary TTL reads can clean up entries and invalidate another active iterator. + final Iterable<Map.Entry<UK, TtlValue<UV>>> entries = original.entries(); + final long now = timeProvider.currentTimestamp(); + if (entries != null) { + for (Map.Entry<UK, TtlValue<UV>> entry : entries) { + final TtlValue<UV> value = entry.getValue(); + if (!excludedKeys.contains(entry.getKey()) + && value != null + && (returnExpired || !TtlUtils.expired(value, ttl, now))) { + return false; + } + } + } + return true; + } + @Nullable @Override public Map<UK, TtlValue<UV>> getUnexpiredOrNull(@Nonnull Map<UK, TtlValue<UV>> ttlValue) { diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/state/StateBackendTestBase.java b/flink-runtime/src/test/java/org/apache/flink/runtime/state/StateBackendTestBase.java index 92468cb6880..ba3515d1eb0 100644 --- a/flink-runtime/src/test/java/org/apache/flink/runtime/state/StateBackendTestBase.java +++ b/flink-runtime/src/test/java/org/apache/flink/runtime/state/StateBackendTestBase.java @@ -3588,6 +3588,37 @@ public abstract class StateBackendTestBase<B extends AbstractStateBackend> { } } + @TestTemplate + void testMapStateReadOnlyKeyContext() throws Exception { + CheckpointableKeyedStateBackend<Integer> backend = + createKeyedBackend(IntSerializer.INSTANCE); + try { + MapState<Integer, Long> state = + backend.getPartitionedState( + VoidNamespace.INSTANCE, + VoidNamespaceSerializer.INSTANCE, + new MapStateDescriptor<>("map", Integer.class, Long.class)); + backend.setCurrentKey(1); + state.put(1, 100L); + backend.setCurrentKey(2); + state.put(1, 200L); + + int[] selections = {0}; + backend.registerKeySelectionListener(ignored -> selections[0]++); + for (int key : new int[] {1, 3, 2, 1}) { + backend.setCurrentKeyForStateAccess(key); + assertThat(backend.getCurrentKey()).isEqualTo(key); + assertThat(state.get(1)).isEqualTo(key == 3 ? null : key * 100L); + } + assertThat(selections[0]).isZero(); + backend.setCurrentKey(1); + assertThat(selections[0]).isEqualTo(1); + } finally { + IOUtils.closeQuietly(backend); + backend.dispose(); + } + } + @TestTemplate void testMapStateIsEmpty() throws Exception { MapStateDescriptor<Integer, Long> kvId = diff --git a/flink-state-backends/flink-statebackend-changelog/src/main/java/org/apache/flink/state/changelog/ChangelogKeyedStateBackend.java b/flink-state-backends/flink-statebackend-changelog/src/main/java/org/apache/flink/state/changelog/ChangelogKeyedStateBackend.java index 4182023900b..469db21fa8d 100644 --- a/flink-state-backends/flink-statebackend-changelog/src/main/java/org/apache/flink/state/changelog/ChangelogKeyedStateBackend.java +++ b/flink-state-backends/flink-statebackend-changelog/src/main/java/org/apache/flink/state/changelog/ChangelogKeyedStateBackend.java @@ -294,6 +294,11 @@ public class ChangelogKeyedStateBackend<K> keyedStateBackend.setCurrentKey(newKey); } + @Override + public void setCurrentKeyForStateAccess(K newKey) { + keyedStateBackend.setCurrentKeyForStateAccess(newKey); + } + @Override public K getCurrentKey() { return keyedStateBackend.getCurrentKey(); diff --git a/flink-state-backends/flink-statebackend-changelog/src/main/java/org/apache/flink/state/changelog/restore/ChangelogMigrationRestoreTarget.java b/flink-state-backends/flink-statebackend-changelog/src/main/java/org/apache/flink/state/changelog/restore/ChangelogMigrationRestoreTarget.java index 8814e61a81a..19c0cdf6d2a 100644 --- a/flink-state-backends/flink-statebackend-changelog/src/main/java/org/apache/flink/state/changelog/restore/ChangelogMigrationRestoreTarget.java +++ b/flink-state-backends/flink-statebackend-changelog/src/main/java/org/apache/flink/state/changelog/restore/ChangelogMigrationRestoreTarget.java @@ -202,6 +202,11 @@ public class ChangelogMigrationRestoreTarget<K> implements ChangelogRestoreTarge keyedStateBackend.setCurrentKey(newKey); } + @Override + public void setCurrentKeyForStateAccess(K newKey) { + keyedStateBackend.setCurrentKeyForStateAccess(newKey); + } + @Override public void setCurrentKeyAndKeyGroup(K newKey, int newKeyGroupIndex) { keyedStateBackend.setCurrentKeyAndKeyGroup(newKey, newKeyGroupIndex); diff --git a/flink-state-backends/flink-statebackend-forst/src/main/java/org/apache/flink/state/forst/sync/ForStSyncKeyedStateBackend.java b/flink-state-backends/flink-statebackend-forst/src/main/java/org/apache/flink/state/forst/sync/ForStSyncKeyedStateBackend.java index 4f410ca4287..78b05b1c343 100644 --- a/flink-state-backends/flink-statebackend-forst/src/main/java/org/apache/flink/state/forst/sync/ForStSyncKeyedStateBackend.java +++ b/flink-state-backends/flink-statebackend-forst/src/main/java/org/apache/flink/state/forst/sync/ForStSyncKeyedStateBackend.java @@ -432,6 +432,12 @@ public class ForStSyncKeyedStateBackend<K> extends AbstractKeyedStateBackend<K> sharedRocksKeyBuilder.setKeyAndKeyGroup(getCurrentKey(), getCurrentKeyGroupIndex()); } + @Override + public void setCurrentKeyForStateAccess(K newKey) { + super.setCurrentKeyForStateAccess(newKey); + sharedRocksKeyBuilder.setKeyAndKeyGroup(getCurrentKey(), getCurrentKeyGroupIndex()); + } + @Override public void setCurrentKeyAndKeyGroup(K newKey, int newKeyGroupIndex) { super.setCurrentKeyAndKeyGroup(newKey, newKeyGroupIndex); diff --git a/flink-state-backends/flink-statebackend-rocksdb/src/main/java/org/apache/flink/state/rocksdb/RocksDBKeyedStateBackend.java b/flink-state-backends/flink-statebackend-rocksdb/src/main/java/org/apache/flink/state/rocksdb/RocksDBKeyedStateBackend.java index b77232cb067..eb6ea788204 100644 --- a/flink-state-backends/flink-statebackend-rocksdb/src/main/java/org/apache/flink/state/rocksdb/RocksDBKeyedStateBackend.java +++ b/flink-state-backends/flink-statebackend-rocksdb/src/main/java/org/apache/flink/state/rocksdb/RocksDBKeyedStateBackend.java @@ -523,6 +523,12 @@ public class RocksDBKeyedStateBackend<K> extends AbstractKeyedStateBackend<K> { sharedRocksKeyBuilder.setKeyAndKeyGroup(getCurrentKey(), getCurrentKeyGroupIndex()); } + @Override + public void setCurrentKeyForStateAccess(K newKey) { + super.setCurrentKeyForStateAccess(newKey); + sharedRocksKeyBuilder.setKeyAndKeyGroup(getCurrentKey(), getCurrentKeyGroupIndex()); + } + @Override public void setCurrentKeyAndKeyGroup(K newKey, int newKeyGroupIndex) { super.setCurrentKeyAndKeyGroup(newKey, newKeyGroupIndex);
