diff --git a/dapr/actor/runtime/_state_provider.py b/dapr/actor/runtime/_state_provider.py index eeb1e4995..86752d150 100644 --- a/dapr/actor/runtime/_state_provider.py +++ b/dapr/actor/runtime/_state_provider.py @@ -51,6 +51,10 @@ async def try_load_state( result = self._state_serializer.deserialize(raw_state_value, state_type) return (True, result) + def round_trip_state_value(self, value: Any) -> Any: + """Returns value as try_load_state would decode it after a save, without a store read.""" + return self._state_serializer.deserialize(self._state_serializer.serialize(value), object) + async def contains_state(self, actor_type: str, actor_id: str, state_name: str) -> bool: raw_state_value = await self._state_client.get_state(actor_type, actor_id, state_name) return (raw_state_value is not None) and len(raw_state_value) > 0 diff --git a/dapr/actor/runtime/state_manager.py b/dapr/actor/runtime/state_manager.py index c2882debb..185eacbed 100644 --- a/dapr/actor/runtime/state_manager.py +++ b/dapr/actor/runtime/state_manager.py @@ -60,6 +60,21 @@ def ttl_in_seconds(self, new_ttl_in_seconds: Optional[int]) -> None: self._ttl_in_seconds = new_ttl_in_seconds +class _NotFoundStateMetadata(StateMetadata[None]): + # A cached miss: the key is known to be absent from the state store. It reports the + # public 'none' change kind, so it is never saved, and is told apart by its type. + def __init__(self) -> None: + super().__init__(None, StateChangeKind.none) + + +def _is_not_found(state_metadata: StateMetadata) -> bool: + return isinstance(state_metadata, _NotFoundStateMetadata) + + +def _is_absent(state_metadata: StateMetadata) -> bool: + return _is_not_found(state_metadata) or state_metadata.change_kind == StateChangeKind.remove + + class ActorStateManager(Generic[T]): def __init__(self, actor: 'Actor'): self._actor = actor @@ -77,6 +92,9 @@ async def try_add_state(self, state_name: str, value: T) -> bool: state_change_tracker = self._get_contextual_state_tracker() if state_name in state_change_tracker: state_metadata = state_change_tracker[state_name] + if _is_not_found(state_metadata): + state_change_tracker[state_name] = StateMetadata(value, StateChangeKind.add) + return True if state_metadata.change_kind == StateChangeKind.remove: state_change_tracker[state_name] = StateMetadata(value, StateChangeKind.update) return True @@ -102,7 +120,7 @@ async def try_get_state(self, state_name: str) -> Tuple[bool, Optional[T]]: state_change_tracker = self._get_contextual_state_tracker() if state_name in state_change_tracker: state_metadata = state_change_tracker[state_name] - if state_metadata.change_kind == StateChangeKind.remove: + if _is_absent(state_metadata): return False, None return True, state_metadata.value has_value, val = await self._actor.runtime_ctx.state_provider.try_load_state( @@ -110,6 +128,10 @@ async def try_get_state(self, state_name: str) -> Tuple[bool, Optional[T]]: ) if has_value: state_change_tracker[state_name] = StateMetadata(val, StateChangeKind.none) + elif state_change_tracker is self._default_state_change_tracker: + # Only the default tracker is refreshed by reentrant saves, so a miss cached + # in a reentrant tracker could hide a nested call's write and overwrite it. + state_change_tracker[state_name] = _NotFoundStateMetadata() return has_value, val async def set_state(self, state_name: str, value: T) -> None: @@ -122,6 +144,11 @@ async def set_state_ttl(self, state_name: str, value: T, ttl_in_seconds: Optiona state_change_tracker = self._get_contextual_state_tracker() if state_name in state_change_tracker: state_metadata = state_change_tracker[state_name] + if _is_not_found(state_metadata): + state_change_tracker[state_name] = StateMetadata( + value, StateChangeKind.add, ttl_in_seconds + ) + return state_metadata.value = value state_metadata.ttl_in_seconds = ttl_in_seconds @@ -133,17 +160,11 @@ async def set_state_ttl(self, state_name: str, value: T, ttl_in_seconds: Optiona state_change_tracker[state_name] = state_metadata return - existed = await self._actor.runtime_ctx.state_provider.contains_state( - self._type_name, self._actor.id.id, state_name + # Add and update are both saved as an upsert, so skip the store read. Update also + # makes a later remove send a delete, since the key may already be in the store. + state_change_tracker[state_name] = StateMetadata( + value, StateChangeKind.update, ttl_in_seconds ) - if existed: - state_change_tracker[state_name] = StateMetadata( - value, StateChangeKind.update, ttl_in_seconds - ) - else: - state_change_tracker[state_name] = StateMetadata( - value, StateChangeKind.add, ttl_in_seconds - ) async def remove_state(self, state_name: str) -> None: if not await self.try_remove_state(state_name): @@ -153,7 +174,7 @@ async def try_remove_state(self, state_name: str) -> bool: state_change_tracker = self._get_contextual_state_tracker() if state_name in state_change_tracker: state_metadata = state_change_tracker[state_name] - if state_metadata.change_kind == StateChangeKind.remove: + if _is_absent(state_metadata): return False elif state_metadata.change_kind == StateChangeKind.add: state_change_tracker.pop(state_name, None) @@ -172,8 +193,7 @@ async def try_remove_state(self, state_name: str) -> bool: async def contains_state(self, state_name: str) -> bool: state_change_tracker = self._get_contextual_state_tracker() if state_name in state_change_tracker: - state_metadata = state_change_tracker[state_name] - return state_metadata.change_kind != StateChangeKind.remove + return not _is_absent(state_change_tracker[state_name]) return await self._actor.runtime_ctx.state_provider.contains_state( self._type_name, self._actor.id.id, state_name ) @@ -200,6 +220,9 @@ async def add_or_update_state( state_change_tracker = self._get_contextual_state_tracker() if state_name in state_change_tracker: state_metadata = state_change_tracker[state_name] + if _is_not_found(state_metadata): + state_change_tracker[state_name] = StateMetadata(value, StateChangeKind.add) + return value if state_metadata.change_kind == StateChangeKind.remove: state_change_tracker[state_name] = StateMetadata(value, StateChangeKind.update) return value @@ -225,11 +248,7 @@ async def get_state_names(self) -> List[str]: # TODO: Get all state names from Dapr once implemented. def append_names_sync(): state_change_tracker = self._get_contextual_state_tracker() - return [ - key - for key, value in state_change_tracker.items() - if value.change_kind != StateChangeKind.remove - ] + return [key for key, value in state_change_tracker.items() if not _is_absent(value)] default_loop = asyncio.get_running_loop() return await default_loop.run_in_executor(None, append_names_sync) @@ -267,6 +286,32 @@ async def save_state(self) -> None: ) for state_name in states_to_remove: state_change_tracker.pop(state_name, None) + if state_change_tracker is not self._default_state_change_tracker: + self._refresh_default_tracker(state_changes) + + def _refresh_default_tracker(self, state_changes: List[ActorStateChange]) -> None: + # Writes made through a reentrancy-scoped tracker are invisible to the default + # tracker, which activation, reminders and timers read from. Refresh its clean + # copies of the written keys, cached misses included, in the shape a fresh read + # would return, and drop removed keys. Entries with pending changes are left alone. + state_provider = self._actor.runtime_ctx.state_provider + for change in state_changes: + metadata = self._default_state_change_tracker.get(change.state_name) + if metadata is None or metadata.change_kind != StateChangeKind.none: + continue + # A None value is not written to the store, so let the next read reload it. + if change.change_kind == StateChangeKind.remove or change.value is None: + self._default_state_change_tracker.pop(change.state_name) + continue + try: + value = state_provider.round_trip_state_value(change.value) + except Exception: + # The save has already committed; fall back to reloading on the next read. + self._default_state_change_tracker.pop(change.state_name) + continue + self._default_state_change_tracker[change.state_name] = StateMetadata( + value, StateChangeKind.none, change.ttl_in_seconds + ) def is_state_marked_for_remove(self, state_name: str) -> bool: state_change_tracker = self._get_contextual_state_tracker() diff --git a/tests/actor/test_state_manager.py b/tests/actor/test_state_manager.py index 5ed4e24e6..03cf4a8b3 100644 --- a/tests/actor/test_state_manager.py +++ b/tests/actor/test_state_manager.py @@ -13,13 +13,16 @@ limitations under the License. """ +import asyncio import base64 +import json import unittest from unittest import mock from dapr.actor.id import ActorId from dapr.actor.runtime._type_information import ActorTypeInformation from dapr.actor.runtime.context import ActorRuntimeContext +from dapr.actor.runtime.reentrancy_context import reentrancy_ctx from dapr.actor.runtime.state_change import StateChangeKind from dapr.actor.runtime.state_manager import ActorStateManager, StateMetadata from dapr.serializers import DefaultJSONSerializer @@ -28,6 +31,10 @@ from tests.actor.utils import _async_mock, _run +def _operations(data): + return [(op['operation'], op['request']['key']) for op in json.loads(data)] + + class ActorStateManagerTests(unittest.TestCase): def setUp(self): # Create mock client @@ -118,6 +125,322 @@ def test_get_state_for_removed_value(self): self.assertFalse(has_value) self.assertIsNone(val) + def _run_reentrant(self, state_manager, coro_fn): + # Runs coro_fn inside a reentrancy-scoped call, then saves its tracker. + token = reentrancy_ctx.set('reentrancy-id') + try: + state_manager.set_state_context('ctx1') + try: + _run(coro_fn()) + _run(state_manager.save_state()) + finally: + state_manager.set_state_context(None) + finally: + reentrancy_ctx.reset(token) + + @mock.patch( + 'tests.actor.fake_client.FakeDaprActorClient.get_state', + new=_async_mock(return_value=b'"value1"'), + ) + @mock.patch( + 'tests.actor.fake_client.FakeDaprActorClient.save_state_transactionally', new=_async_mock() + ) + def test_reentrant_update_refreshes_default_tracker(self): + state_manager = ActorStateManager(self._fake_actor) + _run(state_manager.try_get_state('state1')) + + self._run_reentrant(state_manager, lambda: state_manager.set_state_ttl('state1', 'v2', 60)) + + # The default read is served from the refreshed entry, not the state store. + calls = self._fake_client.get_state.mock.call_count + has_value, val = _run(state_manager.try_get_state('state1')) + self.assertTrue(has_value) + self.assertEqual('v2', val) + self.assertEqual(calls, self._fake_client.get_state.mock.call_count) + state = state_manager._default_state_change_tracker['state1'] + self.assertEqual(StateChangeKind.none, state.change_kind) + self.assertEqual(60, state.ttl_in_seconds) + + @mock.patch( + 'tests.actor.fake_client.FakeDaprActorClient.get_state', + new=_async_mock(return_value=b'"value1"'), + ) + @mock.patch( + 'tests.actor.fake_client.FakeDaprActorClient.save_state_transactionally', new=_async_mock() + ) + def test_reentrant_remove_evicts_default_tracker(self): + state_manager = ActorStateManager(self._fake_actor) + _run(state_manager.try_get_state('state1')) + + self._run_reentrant(state_manager, lambda: state_manager.remove_state('state1')) + + self.assertNotIn('state1', state_manager._default_state_change_tracker) + self._fake_client.get_state.mock.return_value = None + has_value, val = _run(state_manager.try_get_state('state1')) + self.assertFalse(has_value) + self.assertIsNone(val) + + @mock.patch( + 'tests.actor.fake_client.FakeDaprActorClient.get_state', + new=_async_mock(return_value=b'"value1"'), + ) + @mock.patch( + 'tests.actor.fake_client.FakeDaprActorClient.save_state_transactionally', new=_async_mock() + ) + def test_reentrant_save_keeps_dirty_default_entry(self): + state_manager = ActorStateManager(self._fake_actor) + _run(state_manager.set_state('state1', 'pending')) + + self._run_reentrant(state_manager, lambda: state_manager.set_state('state1', 'v2')) + + state = state_manager._default_state_change_tracker['state1'] + self.assertEqual('pending', state.value) + self.assertEqual(StateChangeKind.update, state.change_kind) + + @mock.patch( + 'tests.actor.fake_client.FakeDaprActorClient.get_state', + new=_async_mock(return_value=b'[1, 2]'), + ) + @mock.patch( + 'tests.actor.fake_client.FakeDaprActorClient.save_state_transactionally', new=_async_mock() + ) + def test_reentrant_refresh_matches_fresh_read(self): + state_manager = ActorStateManager(self._fake_actor) + _run(state_manager.try_get_state('state1')) + + self._run_reentrant(state_manager, lambda: state_manager.set_state('state1', (3, 4))) + + _, refreshed = _run(state_manager.try_get_state('state1')) + self._fake_client.get_state.mock.return_value = self._serializer.serialize((3, 4)) + _, fresh = _run(ActorStateManager(self._fake_actor).try_get_state('state1')) + self.assertEqual([3, 4], fresh) + self.assertEqual(fresh, refreshed) + + @mock.patch( + 'tests.actor.fake_client.FakeDaprActorClient.get_state', + new=_async_mock(return_value=b'"value1"'), + ) + @mock.patch( + 'tests.actor.fake_client.FakeDaprActorClient.save_state_transactionally', new=_async_mock() + ) + def test_reentrant_save_of_none_evicts_default_tracker(self): + state_manager = ActorStateManager(self._fake_actor) + _run(state_manager.try_get_state('state1')) + + self._run_reentrant(state_manager, lambda: state_manager.set_state('state1', None)) + + self.assertNotIn('state1', state_manager._default_state_change_tracker) + + @mock.patch( + 'tests.actor.fake_client.FakeDaprActorClient.get_state', + new=_async_mock(return_value=b'"value1"'), + ) + @mock.patch( + 'tests.actor.fake_client.FakeDaprActorClient.save_state_transactionally', new=_async_mock() + ) + def test_reentrant_refresh_failure_evicts_default_tracker(self): + state_manager = ActorStateManager(self._fake_actor) + _run(state_manager.try_get_state('state1')) + _run(state_manager.try_get_state('state2')) + + async def set_both(): + await state_manager.set_state('state1', 'v2') + await state_manager.set_state('state2', 'v2') + + # The save must not raise after the write has committed, and every key is handled. + with mock.patch.object( + self._runtime_ctx.state_provider, + 'round_trip_state_value', + side_effect=ValueError('cannot decode'), + ): + self._run_reentrant(state_manager, set_both) + + self.assertNotIn('state1', state_manager._default_state_change_tracker) + self.assertNotIn('state2', state_manager._default_state_change_tracker) + + def _cache_miss(self, state_manager, state_name='state1'): + # Reads an absent key so the tracker caches the miss; returns the store read count. + has_value, _ = _run(state_manager.try_get_state(state_name)) + self.assertFalse(has_value) + return self._fake_client.get_state.mock.call_count + + @mock.patch('tests.actor.fake_client.FakeDaprActorClient.get_state', new=_async_mock()) + def test_missing_state_is_cached_as_absent(self): + state_manager = ActorStateManager(self._fake_actor) + calls = self._cache_miss(state_manager) + + self.assertEqual((False, None), _run(state_manager.try_get_state('state1'))) + with self.assertRaises(KeyError): + _run(state_manager.get_state('state1')) + self.assertFalse(_run(state_manager.contains_state('state1'))) + self.assertFalse(state_manager.is_state_marked_for_remove('state1')) + self.assertEqual([], _run(state_manager.get_state_names())) + self.assertEqual(calls, self._fake_client.get_state.mock.call_count) + + @mock.patch('tests.actor.fake_client.FakeDaprActorClient.get_state', new=_async_mock()) + def test_try_add_state_after_cached_miss(self): + state_manager = ActorStateManager(self._fake_actor) + calls = self._cache_miss(state_manager) + + self.assertTrue(_run(state_manager.try_add_state('state1', 'value1'))) + self.assertEqual(calls, self._fake_client.get_state.mock.call_count) + state = state_manager._default_state_change_tracker['state1'] + self.assertEqual(StateChangeKind.add, state.change_kind) + self.assertEqual('value1', state.value) + self.assertEqual((True, 'value1'), _run(state_manager.try_get_state('state1'))) + + @mock.patch('tests.actor.fake_client.FakeDaprActorClient.get_state', new=_async_mock()) + def test_set_state_after_cached_miss(self): + state_manager = ActorStateManager(self._fake_actor) + self._cache_miss(state_manager, 'state1') + calls = self._cache_miss(state_manager, 'state2') + + _run(state_manager.set_state('state1', 'value1')) + _run(state_manager.set_state_ttl('state2', 'value2', 60)) + self.assertEqual(calls, self._fake_client.get_state.mock.call_count) + tracker = state_manager._default_state_change_tracker + self.assertEqual(StateChangeKind.add, tracker['state1'].change_kind) + self.assertEqual('value1', tracker['state1'].value) + self.assertIsNone(tracker['state1'].ttl_in_seconds) + self.assertEqual(StateChangeKind.add, tracker['state2'].change_kind) + self.assertEqual('value2', tracker['state2'].value) + self.assertEqual(60, tracker['state2'].ttl_in_seconds) + + @mock.patch('tests.actor.fake_client.FakeDaprActorClient.get_state', new=_async_mock()) + def test_remove_state_after_cached_miss(self): + state_manager = ActorStateManager(self._fake_actor) + calls = self._cache_miss(state_manager) + + self.assertFalse(_run(state_manager.try_remove_state('state1'))) + with self.assertRaises(KeyError): + _run(state_manager.remove_state('state1')) + self.assertEqual(calls, self._fake_client.get_state.mock.call_count) + self.assertFalse(state_manager.is_state_marked_for_remove('state1')) + + @mock.patch('tests.actor.fake_client.FakeDaprActorClient.get_state', new=_async_mock()) + def test_get_or_add_state_after_cached_miss(self): + state_manager = ActorStateManager(self._fake_actor) + calls = self._cache_miss(state_manager) + + self.assertEqual('value1', _run(state_manager.get_or_add_state('state1', 'value1'))) + self.assertEqual(calls, self._fake_client.get_state.mock.call_count) + state = state_manager._default_state_change_tracker['state1'] + self.assertEqual(StateChangeKind.add, state.change_kind) + self.assertEqual('value1', state.value) + + @mock.patch('tests.actor.fake_client.FakeDaprActorClient.get_state', new=_async_mock()) + def test_add_or_update_state_after_cached_miss(self): + state_manager = ActorStateManager(self._fake_actor) + calls = self._cache_miss(state_manager) + update_value = mock.MagicMock(return_value='updated') + + val = _run(state_manager.add_or_update_state('state1', 'value1', update_value)) + self.assertEqual('value1', val) + update_value.assert_not_called() + self.assertEqual(calls, self._fake_client.get_state.mock.call_count) + state = state_manager._default_state_change_tracker['state1'] + self.assertEqual(StateChangeKind.add, state.change_kind) + self.assertEqual('value1', state.value) + + @mock.patch('tests.actor.fake_client.FakeDaprActorClient.get_state', new=_async_mock()) + @mock.patch( + 'tests.actor.fake_client.FakeDaprActorClient.save_state_transactionally', new=_async_mock() + ) + def test_save_state_skips_cached_miss(self): + state_manager = ActorStateManager(self._fake_actor) + self._cache_miss(state_manager, 'state1') + save = self._fake_client.save_state_transactionally.mock + save.reset_mock() + + _run(state_manager.save_state()) + save.assert_not_called() + + self._cache_miss(state_manager, 'state2') + _run(state_manager.set_state('state2', 'value2')) + _run(state_manager.save_state()) + save.assert_called_once() + self.assertEqual( + b'[{"operation":"upsert","request":{"key":"state2","value":"value2"}}]', + save.call_args.args[2], + ) + self.assertEqual((False, None), _run(state_manager.try_get_state('state1'))) + + @mock.patch('tests.actor.fake_client.FakeDaprActorClient.get_state', new=_async_mock()) + @mock.patch( + 'tests.actor.fake_client.FakeDaprActorClient.save_state_transactionally', new=_async_mock() + ) + def test_reentrant_write_replaces_cached_miss_in_default_tracker(self): + state_manager = ActorStateManager(self._fake_actor) + self._cache_miss(state_manager) + + # A reentrant call reads the absent key in its own tracker, then creates it. + async def create(): + self.assertFalse(await state_manager.contains_state('state1')) + await state_manager.set_state('state1', 'value2') + + self._run_reentrant(state_manager, create) + + # The reminder-facing default tracker must not keep reporting the key as absent. + calls = self._fake_client.get_state.mock.call_count + self.assertEqual((True, 'value2'), _run(state_manager.try_get_state('state1'))) + self.assertTrue(_run(state_manager.contains_state('state1'))) + self.assertEqual(calls, self._fake_client.get_state.mock.call_count) + + @mock.patch('tests.actor.fake_client.FakeDaprActorClient.get_state', new=_async_mock()) + @mock.patch( + 'tests.actor.fake_client.FakeDaprActorClient.save_state_transactionally', new=_async_mock() + ) + def test_reentrant_remove_evicts_cached_miss_from_default_tracker(self): + state_manager = ActorStateManager(self._fake_actor) + self._cache_miss(state_manager) + + # The key is created elsewhere, then removed by a reentrant call. + self._fake_client.get_state.mock.return_value = b'"value1"' + self._run_reentrant(state_manager, lambda: state_manager.remove_state('state1')) + + self.assertNotIn('state1', state_manager._default_state_change_tracker) + + @mock.patch('tests.actor.fake_client.FakeDaprActorClient.get_state', new=_async_mock()) + @mock.patch( + 'tests.actor.fake_client.FakeDaprActorClient.save_state_transactionally', new=_async_mock() + ) + def test_nested_reentrant_write_is_seen_by_outer_call(self): + state_manager = ActorStateManager(self._fake_actor) + get_state = self._fake_client.get_state.mock + save = self._fake_client.save_state_transactionally.mock + save.reset_mock() + + async def inner(): + state_manager.set_state_context('inner') + await state_manager.set_state('state1', 'from-inner') + await state_manager.save_state() + get_state.return_value = b'"from-inner"' + + async def outer(): + state_manager.set_state_context('outer') + self.assertFalse(await state_manager.contains_state('state1')) + self.assertEqual((False, None), await state_manager.try_get_state('state1')) + # A nested reentrant call (A -> B -> A) creates the key in its own tracker. + await asyncio.create_task(inner()) + self.assertTrue(await state_manager.contains_state('state1')) + self.assertEqual( + 'from-inner', await state_manager.get_or_add_state('state1', 'default') + ) + await state_manager.save_state() + + token = reentrancy_ctx.set('reentrancy-id') + try: + _run(outer()) + finally: + state_manager.set_state_context(None) + reentrancy_ctx.reset(token) + + save.assert_called_once() + self.assertEqual( + b'[{"operation":"upsert","request":{"key":"state1","value":"from-inner"}}]', + save.call_args.args[2], + ) + @mock.patch('tests.actor.fake_client.FakeDaprActorClient.get_state', new=_async_mock()) def test_set_state_for_new_state(self): state_manager = ActorStateManager(self._fake_actor) @@ -125,7 +448,7 @@ def test_set_state_for_new_state(self): _run(state_manager.set_state('state1', 'value1')) state = state_change_tracker['state1'] - self.assertEqual(StateChangeKind.add, state.change_kind) + self.assertEqual(StateChangeKind.update, state.change_kind) self.assertEqual('value1', state.value) self.assertEqual(None, state.ttl_in_seconds) @@ -136,12 +459,12 @@ def test_set_state_for_existing_state_only_in_mem(self): _run(state_manager.set_state('state1', 'value1')) state = state_change_tracker['state1'] - self.assertEqual(StateChangeKind.add, state.change_kind) + self.assertEqual(StateChangeKind.update, state.change_kind) self.assertEqual('value1', state.value) _run(state_manager.set_state('state1', 'value2')) state = state_change_tracker['state1'] - self.assertEqual(StateChangeKind.add, state.change_kind) + self.assertEqual(StateChangeKind.update, state.change_kind) self.assertEqual('value2', state.value) self.assertEqual(None, state.ttl_in_seconds) @@ -165,7 +488,7 @@ def test_set_state_ttl_for_new_state(self): _run(state_manager.set_state_ttl('state1', 'value1', 3600)) state = state_change_tracker['state1'] - self.assertEqual(StateChangeKind.add, state.change_kind) + self.assertEqual(StateChangeKind.update, state.change_kind) self.assertEqual('value1', state.value) self.assertEqual(3600, state.ttl_in_seconds) @@ -176,13 +499,13 @@ def test_set_state_ttl_for_existing_state_only_in_mem(self): _run(state_manager.set_state_ttl('state1', 'value1', 3600)) state = state_change_tracker['state1'] - self.assertEqual(StateChangeKind.add, state.change_kind) + self.assertEqual(StateChangeKind.update, state.change_kind) self.assertEqual('value1', state.value) self.assertEqual(3600, state.ttl_in_seconds) _run(state_manager.set_state_ttl('state1', 'value2', 7200)) state = state_change_tracker['state1'] - self.assertEqual(StateChangeKind.add, state.change_kind) + self.assertEqual(StateChangeKind.update, state.change_kind) self.assertEqual('value2', state.value) self.assertEqual(7200, state.ttl_in_seconds) @@ -200,6 +523,113 @@ def test_set_state_ttl_for_existing_state(self): self.assertEqual('value2', state.value) self.assertEqual(3600, state.ttl_in_seconds) + @mock.patch('tests.actor.fake_client.FakeDaprActorClient.get_state', new=_async_mock()) + @mock.patch( + 'tests.actor.fake_client.FakeDaprActorClient.save_state_transactionally', new=_async_mock() + ) + def test_set_state_does_not_read_state_store(self): + state_manager = ActorStateManager(self._fake_actor) + get_state = self._fake_client.get_state.mock + save = self._fake_client.save_state_transactionally.mock + _run(state_manager.set_state('state1', 'value1')) + _run(state_manager.set_state_ttl('state2', 'value2', 60)) + + self.assertEqual('value1', _run(state_manager.get_state('state1'))) + self.assertTrue(_run(state_manager.contains_state('state2'))) + self.assertFalse(_run(state_manager.try_add_state('state1', 'other'))) + _run(state_manager.save_state()) + + get_state.assert_not_called() + self.assertEqual( + b'[{"operation":"upsert","request":{"key":"state1","value":"value1"}},' + b'{"operation":"upsert","request":{"key":"state2","value":"value2",' + b'"metadata":{"ttlInSeconds":"60"}}}]', + save.call_args.args[2], + ) + + @mock.patch( + 'tests.actor.fake_client.FakeDaprActorClient.get_state', + new=_async_mock(return_value=b'"stored"'), + ) + @mock.patch( + 'tests.actor.fake_client.FakeDaprActorClient.save_state_transactionally', new=_async_mock() + ) + def test_set_then_remove_in_one_turn_sends_delete(self): + # The key may already be in the store, so removing it must still send a delete. + state_manager = ActorStateManager(self._fake_actor) + save = self._fake_client.save_state_transactionally.mock + _run(state_manager.set_state('state1', 'value1')) + _run(state_manager.set_state_ttl('state2', 'value2', 60)) + + self.assertTrue(_run(state_manager.try_remove_state('state1'))) + _run(state_manager.remove_state('state2')) + self.assertFalse(_run(state_manager.contains_state('state1'))) + _run(state_manager.save_state()) + + self._fake_client.get_state.mock.assert_not_called() + self.assertEqual( + [('delete', 'state1'), ('delete', 'state2')], _operations(save.call_args.args[2]) + ) + self.assertEqual({}, state_manager._default_state_change_tracker) + + @mock.patch('tests.actor.fake_client.FakeDaprActorClient.get_state', new=_async_mock()) + @mock.patch( + 'tests.actor.fake_client.FakeDaprActorClient.save_state_transactionally', new=_async_mock() + ) + def test_add_then_remove_in_one_turn_sends_nothing(self): + # try_add_state and a cached miss both know the key is absent, so nothing is sent. + state_manager = ActorStateManager(self._fake_actor) + save = self._fake_client.save_state_transactionally.mock + self.assertTrue(_run(state_manager.try_add_state('state1', 'value1'))) + self._cache_miss(state_manager, 'state2') + _run(state_manager.set_state('state2', 'value2')) + + self.assertTrue(_run(state_manager.try_remove_state('state1'))) + self.assertTrue(_run(state_manager.try_remove_state('state2'))) + _run(state_manager.save_state()) + + save.assert_not_called() + + @mock.patch( + 'tests.actor.fake_client.FakeDaprActorClient.get_state', + new=_async_mock(return_value=b'"stored"'), + ) + @mock.patch( + 'tests.actor.fake_client.FakeDaprActorClient.save_state_transactionally', new=_async_mock() + ) + def test_reentrant_set_then_remove_sends_delete_and_evicts_default(self): + state_manager = ActorStateManager(self._fake_actor) + save = self._fake_client.save_state_transactionally.mock + _run(state_manager.try_get_state('state1')) + calls = self._fake_client.get_state.mock.call_count + + async def set_then_remove(): + await state_manager.set_state('state1', 'value1') + await state_manager.remove_state('state1') + + self._run_reentrant(state_manager, set_then_remove) + + self.assertEqual(calls, self._fake_client.get_state.mock.call_count) + self.assertEqual([('delete', 'state1')], _operations(save.call_args.args[2])) + self.assertNotIn('state1', state_manager._default_state_change_tracker) + + @mock.patch( + 'tests.actor.fake_client.FakeDaprActorClient.get_state', + new=_async_mock(return_value=b'"stored"'), + ) + def test_try_add_and_try_remove_still_check_state_store(self): + state_manager = ActorStateManager(self._fake_actor) + get_state = self._fake_client.get_state.mock + self.assertFalse(_run(state_manager.try_add_state('state1', 'value1'))) + self.assertTrue(_run(state_manager.try_remove_state('state2'))) + self.assertEqual(2, get_state.call_count) + + get_state.return_value = None + state_manager = ActorStateManager(self._fake_actor) + self.assertFalse(_run(state_manager.try_remove_state('state1'))) + self.assertTrue(_run(state_manager.try_add_state('state2', 'value2'))) + self.assertEqual(4, get_state.call_count) + @mock.patch('tests.actor.fake_client.FakeDaprActorClient.get_state', new=_async_mock()) def test_set_state_ttl_lt_0_for_new_state(self): state_manager = ActorStateManager(self._fake_actor) @@ -214,13 +644,13 @@ def test_set_state_ttl_lt_0_for_existing_state_only_in_mem(self): _run(state_manager.set_state_ttl('state1', 'value1', 3600)) state = state_change_tracker['state1'] - self.assertEqual(StateChangeKind.add, state.change_kind) + self.assertEqual(StateChangeKind.update, state.change_kind) self.assertEqual('value1', state.value) self.assertEqual(3600, state.ttl_in_seconds) _run(state_manager.set_state_ttl('state1', 'value2', -3600)) state = state_change_tracker['state1'] - self.assertEqual(StateChangeKind.add, state.change_kind) + self.assertEqual(StateChangeKind.update, state.change_kind) self.assertEqual('value1', state.value) self.assertEqual(3600, state.ttl_in_seconds) @@ -395,16 +825,7 @@ def test_get_state_names(self): _run(state_manager.set_state('state3', 'value3')) names = _run(state_manager.get_state_names()) self.assertEqual(['state1', 'state2', 'state3'], names) - - self._fake_client.get_state.mock.assert_any_call( - self._test_type_info._name, self._test_actor_id.id, 'state1' - ) - self._fake_client.get_state.mock.assert_any_call( - self._test_type_info._name, self._test_actor_id.id, 'state2' - ) - self._fake_client.get_state.mock.assert_any_call( - self._test_type_info._name, self._test_actor_id.id, 'state3' - ) + self._fake_client.get_state.mock.assert_not_called() def test_get_state_names_excludes_only_removed(self): """States that still exist (add/update/none) must be listed; only a