finish initial valkey

Author: benedikt-bartscher

reflex/state.py

@@ -40,6 +40,7 @@
 )
 
 from glide import (
+    OK,
     ConditionalChange,
     ExpirySet,
     ExpiryType,
@@ -77,7 +78,6 @@
     BaseModelV1 = BaseModelV2
 
 import wrapt
-from redis.exceptions import ResponseError
 
 import reflex.istate.dynamic
 from reflex import constants
@@ -100,6 +100,7 @@
     ImmutableStateError,
     InvalidStateManagerMode,
     LockExpiredError,
+    RedisConfigError,
     ReflexRuntimeError,
     SetUndefinedStateVarError,
     StateSchemaMismatchError,
@@ -3217,29 +3218,33 @@ class StateManagerRedis(StateManager):
     }
 
     # This lock is used to ensure we only subscribe to keyspace events once per token and worker
-    # _pubsub_locks: Dict[bytes, asyncio.Lock] = pydantic.PrivateAttr({})
-
-    _pubsub_clients: Dict[bytes, GlideClient] = pydantic.PrivateAttr({})
+    _pubsub_locks: Dict[bytes, asyncio.Lock] = pydantic.PrivateAttr({})
 
     async def get_redis(self) -> GlideClient:
         """Get the redis client.
 
         Returns:
             The redis client.
+
+        Raises:
+            RedisConfigError: If the redis client could not be configured.
         """
         if self.redis is not None:
             return self.redis
         redis = await prerequisites.get_redis()
         assert redis is not None
-        try:
-            _ = await redis.config_set(
-                {"notify-keyspace-events": self._redis_notify_keyspace_events},
+        config_result = await redis.config_set(
+            {"notify-keyspace-events": self._redis_notify_keyspace_events},
+        )
+        # Some redis servers only allow out-of-band configuration, so ignore errors here.
+        if (
+            config_result != OK
+            and not environment.REFLEX_IGNORE_REDIS_CONFIG_ERROR.get()
+        ):
+            raise RedisConfigError(
+                f"Failed to set notify-keyspace-events: {config_result}"
             )
-        # TODO: adjust exception for glide
-        except ResponseError:
-            # Some redis servers only allow out-of-band configuration, so ignore errors here.
-            if not environment.REFLEX_IGNORE_REDIS_CONFIG_ERROR.get():
-                raise
+
         self.redis = redis
         return redis
 
@@ -3407,6 +3412,7 @@ async def set_state(
         """
         # Check that we're holding the lock.
         redis = await self.get_redis()
+
         if lock_id is not None and await redis.get(self._lock_key(token)) != lock_id:
             raise LockExpiredError(
                 f"Lock expired for token {token} while processing. Consider increasing "
@@ -3440,9 +3446,15 @@ async def set_state(
                 _ = await redis.set(
                     _substate_key(client_token, state),
                     pickle_state,
-                    expiry=self.expiry,
-                    # ex=self.token_expiration,
+                    expiry=ExpirySet(
+                        expiry_type=ExpiryType.MILLSEC,
+                        value=self.token_expiration,
+                    ),
                 )
+                # if str(res) != OK:
+                #     raise RuntimeError(
+                #         f"Failed to set state for token {token}. {res} {OK}"
+                #     )
 
         # Wait for substates to be persisted.
         for t in tasks:
@@ -3478,18 +3490,6 @@ def _lock_key(token: str) -> bytes:
         client_token = _split_substate_key(token)[0]
         return f"{client_token}_lock".encode()
 
-    @property
-    def expiry(self) -> ExpirySet:
-        """Get the expiry set for the token.
-
-        Returns:
-            The expiry set for the token.
-        """
-        return ExpirySet(
-            expiry_type=ExpiryType.SEC,
-            value=self.token_expiration,
-        )
-
     async def _try_get_lock(self, lock_key: bytes, lock_id: bytes) -> bool | None:
         """Try to get a redis lock for a token.
 
@@ -3504,10 +3504,13 @@ async def _try_get_lock(self, lock_key: bytes, lock_id: bytes) -> bool | None:
         response = await redis.set(
             lock_key,
             lock_id,
-            expiry=self.expiry,
+            expiry=ExpirySet(
+                expiry_type=ExpiryType.MILLSEC,
+                value=self.lock_expiration,
+            ),
             conditional_set=ConditionalChange.ONLY_IF_DOES_NOT_EXIST,
         )
-        return bool(response)
+        return str(response) == OK
 
     async def get_pubsub(self, lock_key: bytes) -> GlideClient:
         """Get the pubsub client for a lock key channel.
@@ -3519,12 +3522,10 @@ async def get_pubsub(self, lock_key: bytes) -> GlideClient:
             The pubsub client.
         """
         lock_key_channel = f"__keyspace@0__:{lock_key.decode()}"
-        if lock_key_channel in self._pubsub_clients:
-            return self._pubsub_clients[lock_key_channel]
         pubsub_config = GlideClientConfiguration.PubSubSubscriptions(
             channels_and_patterns={
                 GlideClientConfiguration.PubSubChannelModes.Pattern: {lock_key_channel},
-                GlideClientConfiguration.PubSubChannelModes.Exact: {lock_key_channel},
+                # GlideClientConfiguration.PubSubChannelModes.Exact: {lock_key_channel},
             },
             callback=None,
             context=None,
@@ -3534,7 +3535,6 @@ async def get_pubsub(self, lock_key: bytes) -> GlideClient:
         )
         assert config is not None
         pubsub = await GlideClient.create(config)
-        self._pubsub_clients[lock_key] = pubsub
         return pubsub
 
     async def _wait_lock(self, lock_key: bytes, lock_id: bytes) -> None:
@@ -3545,58 +3545,48 @@ async def _wait_lock(self, lock_key: bytes, lock_id: bytes) -> None:
         Args:
             lock_key: The redis key for the lock.
             lock_id: The ID of the lock.
-
-        Raises:
-            ResponseError: when the keyspace config cannot be set.
         """
         state_is_locked = False
         # Enable keyspace notifications for the lock key, so we know when it is available.
         redis = await self.get_redis()
-        pubsub = await self.get_pubsub(lock_key)
-        # async with self.redis.pubsub() as pubsub:
-        # await pubsub.psubscribe(lock_key_channel)
-        # await pubsub.get_pubsub_message()
-        count = 0
-        while not state_is_locked:
-            count += 1
-            if count > 10000:
-                raise Exception("Could not obtain lock")
-            # wait for the lock to be released
-            print("waiting for lock to be released")
-            while True:
-                if not await redis.exists([lock_key]):
-                    # if not pubsub.try_get_pubsub_message():
-                    break  # key was removed, try to get the lock again
-                message = await pubsub.get_pubsub_message(
-                    # ignore_subscribe_messages=True,
-                    # timeout=self.lock_expiration / 1000.0,
-                )
-                # if message.pattern is None:
-                #     # raise Exception("Pattern is None")
-                #     continue
-                # raise Exception(message)
-                if message.message in self._redis_keyspace_lock_release_events:
-                    break
-            state_is_locked = await self._try_get_lock(lock_key, lock_id)
+        if lock_key not in self._pubsub_locks:
+            self._pubsub_locks[lock_key] = asyncio.Lock()
+        async with self._pubsub_locks[lock_key]:
+            pubsub = await self.get_pubsub(lock_key)
+            while not state_is_locked:
+                # wait for the lock to be released
+                while True:
+                    # check if we missed lock release events
+                    if await redis.exists([lock_key]) == 0:
+                        break  # key was removed, try to get the lock again
+
+                    try:
+                        # TODO: alternative to ignore_subscribe_messages?
+                        message = await asyncio.wait_for(
+                            pubsub.get_pubsub_message(),
+                            timeout=self.lock_expiration / 1000.0,
+                        )
+                    except asyncio.TimeoutError:
+                        continue
+                    if message.message in self._redis_keyspace_lock_release_events:
+                        break
+                state_is_locked = await self._try_get_lock(lock_key, lock_id)
 
     @override
-    async def disconnect(self, token: str):
+    async def disconnect(self, token: str) -> None:
         """Disconnect the token from the redis client.
 
         Args:
             token: The token to disconnect.
         """
         lock_key = self._lock_key(token)
-        # if lock := self._pubsub_locks.get(lock_key):
-        #     if lock.locked():
-        #         lock.release()
-        #     del self._pubsub_locks[lock_key]
-        if client := self._pubsub_clients.get(self._lock_key(token)):
-            await client.close()
-            del self._pubsub_clients[lock_key]
+        if lock := self._pubsub_locks.get(lock_key):
+            if lock.locked():
+                lock.release()
+            del self._pubsub_locks[lock_key]
 
     @contextlib.asynccontextmanager
-    async def _lock(self, token: str):
+    async def _lock(self, token: str) -> AsyncIterator[bytes]:
         """Obtain a redis lock for a token.
 
         Args:
@@ -3626,8 +3616,10 @@ async def _lock(self, token: str):
                 # only delete our lock
                 redis = await self.get_redis()
                 _ = await redis.delete([lock_key])
+                # if not res:
+                #     raise RuntimeError(f"Failed to release lock for token {token}")
 
-    async def close(self):
+    async def close(self) -> None:
         """Explicitly close the redis connection and connection_pool.
 
         It is necessary in testing scenarios to close between asyncio test cases
@@ -3636,14 +3628,9 @@ async def close(self):
 
         Note: Connections will be automatically reopened when needed.
         """
-        # await self.redis.aclose(close_connection_pool=True)
-        # TODO: is this needed with glide?
-        redis = await self.get_redis()
-        await redis.close()
-
-        for pubsub in self._pubsub_clients.values():
-            await pubsub.close()
-        self._pubsub_clients = {}
+        if self.redis is not None:
+            await self.redis.close()
+            self.redis = None
 
 
 def get_state_manager() -> StateManager:

reflex/utils/exceptions.py

@@ -87,6 +87,10 @@ class LockExpiredError(ReflexError):
     """Raised when the state lock expires while an event is being processed."""
 
 
+class RedisConfigError(ReflexError):
+    """Raised when the Redis configuration is not applied correctly."""
+
+
 class MatchTypeError(ReflexError, TypeError):
     """Raised when the return types of match cases are different."""
 

tests/integration/test_background_task.py

@@ -44,7 +44,7 @@ async def handle_event_yield_only(self):
 
         @rx.event(background=True)
         async def fast_yielding(self):
-            for _ in range(100000):
+            for _ in range(1000):
                 yield State.increment()
 
         @rx.event
@@ -409,4 +409,4 @@ def test_fast_yielding(
     assert background_task._poll_for(lambda: counter.text == "0", timeout=5)
 
     fast_yielding_button.click()
-    assert background_task._poll_for(lambda: counter.text == "100000", timeout=1200)
+    assert background_task._poll_for(lambda: counter.text == "1000", timeout=20)

tests/units/test_state.py

@@ -1928,6 +1928,14 @@ async def test_state_proxy(grandchild_state: GrandchildState, mock_app: rx.App):
         # Cannot access substates
         sp.substates[""]
 
+    assert (
+        sp.router.session.client_token == grandchild_state.router.session.client_token
+    )
+    assert (
+        sp.__wrapped__.router.session.client_token
+        == grandchild_state.router.session.client_token
+    )
+    assert sp.router.session.client_token is not None
     async with sp:
         assert sp._self_actx is not None
         assert sp._self_mutable  # proxy is mutable inside context