flaky test nad explicit dispatch fuction

Author: FarhanAliRaza

reflex/compiler/plugins.py

@@ -6,7 +6,7 @@
 from collections.abc import AsyncGenerator, Sequence
 from contextvars import ContextVar, Token
 from types import TracebackType
-from typing import Any, ClassVar, Protocol, TypeAlias
+from typing import Any, ClassVar, Literal, Protocol, TypeAlias, cast, overload
 
 from reflex_core.components.component import BaseComponent, Component, StatefulComponent
 from reflex_core.utils.imports import ParsedImportDict, collapse_imports, merge_imports
@@ -62,18 +62,67 @@ class CompilerHooks:
 
     plugins: tuple[CompilerPlugin, ...] = ()
 
+    @overload
+    async def _dispatch(
+        self,
+        hook_name: str,
+        *args: Any,
+        stop_on_result: Literal[False] = False,
+        **kwargs: Any,
+    ) -> list[Any]: ...
+
+    @overload
+    async def _dispatch(
+        self,
+        hook_name: str,
+        *args: Any,
+        stop_on_result: Literal[True],
+        **kwargs: Any,
+    ) -> Any | None: ...
+
+    async def _dispatch(
+        self,
+        hook_name: str,
+        *args: Any,
+        stop_on_result: bool = False,
+        **kwargs: Any,
+    ) -> list[Any] | Any | None:
+        """Dispatch a coroutine hook across all plugins in registration order.
+
+        Args:
+            hook_name: The plugin hook attribute to invoke.
+            *args: Positional arguments forwarded to the hook.
+            stop_on_result: Whether to return immediately on the first non-None
+                result instead of collecting all results.
+            **kwargs: Keyword arguments forwarded to the hook.
+
+        Returns:
+            When ``stop_on_result`` is false, a list of hook return values in
+            registration order. Otherwise, the first non-None result, or
+            ``None`` if every plugin returned ``None``.
+        """
+        results: list[Any] = []
+        for plugin in self.plugins:
+            result = await getattr(plugin, hook_name)(*args, **kwargs)
+            if stop_on_result and result is not None:
+                return result
+            results.append(result)
+        return None if stop_on_result else results
+
     async def eval_page(
         self,
         page_fn: Any,
         /,
         **kwargs: Any,
     ) -> PageContext | None:
         """Return the first page context produced by the plugin chain."""
-        for plugin in self.plugins:
-            result = await plugin.eval_page(page_fn, **kwargs)
-            if result is not None:
-                return result
-        return None
+        result = await self._dispatch(
+            "eval_page",
+            page_fn,
+            stop_on_result=True,
+            **kwargs,
+        )
+        return cast(PageContext | None, result)
 
     async def compile_page(
         self,
@@ -82,8 +131,7 @@ async def compile_page(
         **kwargs: Any,
     ) -> None:
         """Run all ``compile_page`` hooks in plugin order."""
-        for plugin in self.plugins:
-            await plugin.compile_page(page_ctx, **kwargs)
+        await self._dispatch("compile_page", page_ctx, **kwargs)
 
     async def compile_component(
         self,

tests/units/istate/manager/test_redis.py

@@ -306,9 +306,16 @@ async def test_oplock_contention_queue(
     state_manager_2._oplock_enabled = True
 
     modify_started = asyncio.Event()
-    modify_2_started = asyncio.Event()
+    contenders_started = asyncio.Event()
     modify_1_continue = asyncio.Event()
     modify_2_continue = asyncio.Event()
+    contender_started_count = 0
+
+    def mark_contender_started() -> None:
+        nonlocal contender_started_count
+        contender_started_count += 1
+        if contender_started_count == 2:
+            contenders_started.set()
 
     async def modify_1():
         async with state_manager_redis.modify_state(
@@ -321,7 +328,7 @@ async def modify_1():
 
     async def modify_2():
         await modify_started.wait()
-        modify_2_started.set()
+        mark_contender_started()
         async with state_manager_2.modify_state(
             _substate_key(token, root_state),
         ) as new_state:
@@ -331,7 +338,7 @@ async def modify_2():
 
     async def modify_3():
         await modify_started.wait()
-        modify_2_started.set()
+        mark_contender_started()
         async with state_manager_2.modify_state(
             _substate_key(token, root_state),
         ) as new_state:
@@ -343,7 +350,7 @@ async def modify_3():
     task_2 = asyncio.create_task(modify_2())
     task_3 = asyncio.create_task(modify_3())
 
-    await modify_2_started.wait()
+    await contenders_started.wait()
 
     # Let modify 1 complete
     modify_1_continue.set()
@@ -407,9 +414,16 @@ async def test_oplock_contention_no_lease(
     state_manager_3._oplock_enabled = True
 
     modify_started = asyncio.Event()
-    modify_2_started = asyncio.Event()
+    contenders_started = asyncio.Event()
     modify_1_continue = asyncio.Event()
     modify_2_continue = asyncio.Event()
+    contender_started_count = 0
+
+    def mark_contender_started() -> None:
+        nonlocal contender_started_count
+        contender_started_count += 1
+        if contender_started_count == 2:
+            contenders_started.set()
 
     async def modify_1():
         async with state_manager_redis.modify_state(
@@ -422,7 +436,7 @@ async def modify_1():
 
     async def modify_2():
         await modify_started.wait()
-        modify_2_started.set()
+        mark_contender_started()
         async with state_manager_2.modify_state(
             _substate_key(token, root_state),
         ) as new_state:
@@ -432,7 +446,7 @@ async def modify_2():
 
     async def modify_3():
         await modify_started.wait()
-        modify_2_started.set()
+        mark_contender_started()
         async with state_manager_3.modify_state(
             _substate_key(token, root_state),
         ) as new_state:
@@ -444,7 +458,7 @@ async def modify_3():
     task_2 = asyncio.create_task(modify_2())
     task_3 = asyncio.create_task(modify_3())
 
-    await modify_2_started.wait()
+    await contenders_started.wait()
 
     # Let modify 1 complete
     modify_1_continue.set()