Remember missing actor state keys instead of re-reading them

dapr/actor/runtime/state_manager.py

@@ -63,6 +63,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
@@ -80,6 +95,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
@@ -105,14 +123,18 @@ 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(
             self._type_name, self._actor.id.id, state_name
         )
         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:
@@ -125,6 +147,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
 
@@ -156,7 +183,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)
@@ -175,8 +202,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
         )
@@ -203,6 +229,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
@@ -228,11 +257,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)
@@ -276,8 +301,8 @@ async def save_state(self) -> None:
     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 in place, in the shape a fresh read would return,
-        # and drop removed keys. Entries with pending changes are left alone.
+        # copies of the written keys in place, 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)
@@ -295,10 +320,13 @@ def _refresh_default_tracker(self, state_changes: List[ActorStateChange]) -> Non
                 continue
             try:
                 value = state_provider.round_trip_state_value(change.value)
-            except Exception as exc:
+            except Exception:
                 # The save has already committed; fall back to reloading on the next read.
                 logger.warning(
-                    'could not refresh cached actor state %s: %s', change.state_name, exc
+                    'could not refresh cached actor state %s after reentrant save; '
+                    'it will be reloaded',
+                    change.state_name,
+                    exc_info=True,
                 )
                 self._default_state_change_tracker.pop(change.state_name)
                 continue

tests/actor/test_state_manager.py

@@ -13,6 +13,7 @@
 limitations under the License.
 """
 
+import asyncio
 import base64
 import unittest
 from unittest import mock
@@ -86,11 +87,15 @@ def test_get_state_for_no_state(self):
         self.assertFalse(has_value)
         self.assertIsNone(val)
 
-        # Test if the test value is empty string
-        self._fake_client.get_state.return_value = ''
-        has_value, val = _run(state_manager.try_get_state('state1'))
+        # A fresh manager, because the miss above is now cached in the default tracker.
+        get_state_mock = self._fake_client.get_state.mock
+        get_state_mock.return_value = ''
+        get_state_mock.reset_mock(return_value=False)
+        state_manager_empty = ActorStateManager(self._fake_actor)
+        has_value, val = _run(state_manager_empty.try_get_state('state1'))
         self.assertFalse(has_value)
         self.assertIsNone(val)
+        get_state_mock.assert_called_once()
 
     @mock.patch(
         'tests.actor.fake_client.FakeDaprActorClient.get_state',
@@ -269,6 +274,189 @@ async def 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)