Events: Horizontal scaling support in EventStream.subscribe

ayon_server/addons/library.py

@@ -29,7 +29,7 @@ def __init__(self) -> None:
             logging.error(f"Addons directory does not exist: {addons_dir}")
             return None
 
-        for addon_name in os.listdir(addons_dir):
+        for addon_name in sorted(os.listdir(addons_dir)):
             # ignore hidden directories (such as .git)
             if addon_name.startswith("."):
                 continue

ayon_server/api/messaging.py

@@ -13,6 +13,7 @@
 from ayon_server.background.background_worker import BackgroundWorker
 from ayon_server.config import ayonconfig
 from ayon_server.entities import UserEntity
+from ayon_server.events.eventstream import EventStream, HandlerType
 from ayon_server.lib.redis import Redis
 from ayon_server.utils import get_nickname, json_dumps, json_loads, obscure
 
@@ -22,6 +23,28 @@
 ]
 
 
+async def _handle_subscribers_task(event_id: str, handlers: list[HandlerType]) -> None:
+    event = await EventStream.get(event_id)
+    for handler in handlers:
+        try:
+            await handler(event)
+        except Exception:
+            logging.debug(f"Error in global event handler: {handler}")
+
+
+async def handle_subscribers(message: dict[str, Any]) -> None:
+    event_id = message.get("id", None)
+    store = message.get("store", None)
+    topic = message.get("topic", None)
+    if not (event_id and store):
+        return
+
+    handlers = EventStream.global_hooks.get(topic, {}).values()
+    if not handlers:
+        return
+    asyncio.create_task(_handle_subscribers_task(event_id, list(handlers)))
+
+
 class Client:
     id: str
     sock: WebSocket
@@ -143,6 +166,8 @@ async def run(self) -> None:
                 else:
                     message = json_loads(raw_message["data"])
 
+                await handle_subscribers(message)
+
                 # TODO: much much smarter logic here
                 for _client_id, client in self.clients.items():
                     project_name = message.get("project", None)

ayon_server/events/eventstream.py

@@ -15,25 +15,44 @@
 
 class EventStream:
     model: Type[EventModel] = EventModel
-    hooks: dict[str, dict[str, HandlerType]] = {}
+    local_hooks: dict[str, dict[str, HandlerType]] = {}
+    global_hooks: dict[str, dict[str, HandlerType]] = {}
 
     @classmethod
-    def subscribe(cls, topic: str, handler: HandlerType) -> str:
+    def subscribe(
+        cls, topic: str, handler: HandlerType, all_nodes: bool = False
+    ) -> str:
         token = create_id()
-        if topic not in cls.hooks:
-            cls.hooks[topic] = {}
-        cls.hooks[topic][token] = handler
-        return token
+        if all_nodes:
+            if topic not in cls.global_hooks:
+                cls.global_hooks[topic] = {}
+            cls.global_hooks[topic][token] = handler
+            return token
+        else:
+            if topic not in cls.local_hooks:
+                cls.local_hooks[topic] = {}
+            cls.local_hooks[topic][token] = handler
+            return token
 
     @classmethod
     def unsubscribe(cls, token: str) -> None:
+        # Local hooks
+        topics_to_remove = []
+        for topic in cls.local_hooks:
+            cls.local_hooks[topic].pop(token, None)
+            if not cls.local_hooks[topic]:
+                topics_to_remove.append(topic)
+        for topic in topics_to_remove:
+            cls.local_hooks.pop(topic)
+
+        # Global hooks
         topics_to_remove = []
-        for topic in cls.hooks:
-            cls.hooks[topic].pop(token, None)
-            if not cls.hooks[topic]:
+        for topic in cls.global_hooks:
+            cls.global_hooks[topic].pop(token, None)
+            if not cls.global_hooks[topic]:
                 topics_to_remove.append(topic)
         for topic in topics_to_remove:
-            cls.hooks.pop(topic)
+            cls.global_hooks.pop(topic)
 
     @classmethod
     async def dispatch(
@@ -192,7 +211,7 @@ async def dispatch(
             )
         )
 
-        handlers = cls.hooks.get(event.topic, {}).values()
+        handlers = cls.local_hooks.get(event.topic, {}).values()
         for handler in handlers:
             try:
                 await handler(event)