shared event hub connection across sessions

Author: nemanjab17

source/ftrack_api/event/hub_proxy.py

@@ -0,0 +1,86 @@
+# :coding: utf-8
+# :copyright: Copyright (c) 2024 ftrack
+
+"""Session-scoped EventHub proxy for tracking subscriptions."""
+
+from __future__ import absolute_import
+
+from builtins import object
+
+
+class SessionEventHubProxy(object):
+    """Proxy for EventHub that tracks subscriptions per session.
+
+    This wrapper intercepts subscribe() and unsubscribe() calls to maintain
+    a session-local list of subscriber IDs. When the session is closed, only
+    the subscribers created by that session are unsubscribed, leaving other
+    sessions' subscribers intact on the shared EventHub.
+    """
+
+    def __init__(self, event_hub, session):
+        """Initialize the proxy.
+
+        Args:
+            event_hub: The underlying EventHub instance (may be shared).
+            session: The Session instance that owns this proxy.
+        """
+        self._event_hub = event_hub
+        self._session = session
+
+    def subscribe(self, subscription, callback, subscriber=None, priority=100):
+        """Subscribe to events and track the subscriber ID for this session.
+
+        All arguments are passed through to the underlying EventHub.subscribe().
+        The returned subscriber ID is tracked for cleanup when this session closes.
+
+        Args:
+            subscription (str): The subscription expression.
+            callback (callable): The callback to invoke when matching events occur.
+            subscriber (dict): Optional subscriber metadata.
+            priority (int): Optional priority (lower = earlier execution).
+
+        Returns:
+            str: The subscriber identifier.
+        """
+        subscriber_id = self._event_hub.subscribe(
+            subscription, callback, subscriber=subscriber, priority=priority
+        )
+
+        # Track this subscriber for session cleanup
+        if not hasattr(self._session, '_session_subscribers'):
+            self._session._session_subscribers = []
+        self._session._session_subscribers.append(subscriber_id)
+
+        return subscriber_id
+
+    def unsubscribe(self, subscriber_identifier):
+        """Unsubscribe and remove from session tracking.
+
+        Args:
+            subscriber_identifier (str): The subscriber identifier to unsubscribe.
+        """
+        self._event_hub.unsubscribe(subscriber_identifier)
+
+        # Remove from session tracking
+        if hasattr(self._session, '_session_subscribers'):
+            if subscriber_identifier in self._session._session_subscribers:
+                self._session._session_subscribers.remove(
+                    subscriber_identifier)
+
+    def __getattr__(self, name):
+        """Delegate all other method/attribute access to the underlying EventHub.
+
+        This allows the proxy to be used transparently as if it were the
+        actual EventHub instance.
+
+        Args:
+            name (str): The attribute name to access.
+
+        Returns:
+            The attribute from the underlying EventHub.
+        """
+        return getattr(self._event_hub, name)
+
+    def __repr__(self):
+        """Return string representation."""
+        return '<SessionEventHubProxy for {0}>'.format(self._event_hub)

source/ftrack_api/event/hub_registry.py

@@ -0,0 +1,190 @@
+# :coding: utf-8
+# :copyright: Copyright (c) 2024 ftrack
+
+"""EventHub singleton registry for connection sharing across sessions."""
+
+from __future__ import absolute_import
+
+import threading
+import weakref
+import logging
+
+
+class EventHubRegistry(object):
+    """Global registry for shared EventHub instances.
+
+    Manages EventHub instances to ensure that sessions with identical
+    credentials can share the same WebSocket connection to the server.
+
+    This prevents connection exhaustion when multiple Session instances
+    are created with the same credentials.
+    """
+
+    def __init__(self):
+        """Initialize the registry."""
+        self._hubs = {}  # (server_url, api_user, api_key) -> EventHub
+        self._hub_sessions = {}  # hub_key -> WeakSet of session weakrefs
+        self._hub_connect_locks = {}  # hub_key -> Lock (for first connection)
+        self._lock = threading.RLock()
+        self.logger = logging.getLogger(
+            __name__ + '.' + self.__class__.__name__)
+
+    def get_or_create(
+        self,
+        server_url,
+        api_user,
+        api_key,
+        headers=None,
+        cookies=None,
+        auto_connect=False,
+    ):
+        """Get existing EventHub or create new one for credentials.
+
+        Args:
+            server_url (str): The ftrack server URL.
+            api_user (str): The API user to authenticate as.
+            api_key (str): The API key to authenticate with.
+            headers (dict): Optional custom headers.
+            cookies (dict): Optional custom cookies.
+            auto_connect (bool): Whether to automatically connect the hub.
+
+        Returns:
+            EventHub: Shared EventHub instance for these credentials.
+
+        Note:
+            Headers and cookies are only used when creating a new hub.
+            If a hub already exists for these credentials, the provided
+            headers/cookies are ignored (the existing hub is returned).
+        """
+        # Create key from credentials
+        # Note: headers/cookies not included in key for simplicity
+        # This means sessions with same credentials share hub even if headers differ
+        key = (server_url, api_user, api_key)
+
+        with self._lock:
+            if key not in self._hubs:
+                # Import here to avoid circular dependency
+                import ftrack_api.event.hub
+
+                # Create new hub
+                hub = ftrack_api.event.hub.EventHub(
+                    server_url, api_user, api_key, headers=headers, cookies=cookies
+                )
+
+                self._hubs[key] = hub
+                self._hub_sessions[key] = weakref.WeakSet()
+                self._hub_connect_locks[key] = threading.Lock()
+
+                self.logger.debug(
+                    'Created new shared EventHub for {0}@{1}'.format(
+                        api_user, server_url
+                    )
+                )
+
+                # Connect if requested
+                if auto_connect:
+                    self._ensure_connected(key, hub)
+            else:
+                existing_hub = self._hubs[key]
+                self.logger.debug(
+                    'Reusing existing EventHub for {0}@{1} ({2} active sessions)'.format(
+                        api_user, server_url, len(self._hub_sessions[key])
+                    )
+                )
+
+                # If auto_connect requested and hub not connected, connect it
+                if auto_connect and not existing_hub.connected:
+                    self._ensure_connected(key, existing_hub)
+
+                hub = existing_hub
+
+            return hub
+
+    def _ensure_connected(self, hub_key, hub):
+        """Ensure the hub is connected (thread-safe).
+
+        Args:
+            hub_key: The hub registry key.
+            hub: The EventHub instance to connect.
+        """
+        # Use per-hub lock to ensure only one thread initiates connection
+        with self._hub_connect_locks[hub_key]:
+            if not hub.connected and not hub._connection_initialised:
+                hub.init_connection()
+                # Connect in background thread
+                connect_thread = threading.Thread(target=hub.connect)
+                connect_thread.daemon = True
+                connect_thread.start()
+
+    def register_session(self, hub_key, session_weakref):
+        """Register a session as using this hub.
+
+        Args:
+            hub_key: The hub registry key (server_url, api_user, api_key).
+            session_weakref: Weak reference to the session.
+        """
+        with self._lock:
+            if hub_key in self._hub_sessions:
+                self._hub_sessions[hub_key].add(session_weakref)
+
+    def on_session_deleted(self, hub_key):
+        """Callback when a session is garbage collected.
+
+        Args:
+            hub_key: The hub registry key.
+        """
+        with self._lock:
+            if hub_key not in self._hub_sessions:
+                return
+
+            # WeakSet automatically removes dead references
+            active_sessions = len(self._hub_sessions[hub_key])
+
+            if active_sessions == 0:
+                # No more sessions using this hub - disconnect and clean up
+                hub = self._hubs.pop(hub_key, None)
+                self._hub_sessions.pop(hub_key, None)
+                self._hub_connect_locks.pop(hub_key, None)
+
+                if hub:
+                    try:
+                        if hub.connected:
+                            # Disconnect without unsubscribing - subscribers already cleaned up
+                            hub.disconnect(unsubscribe=False)
+                            self.logger.debug(
+                                'Disconnected unused shared EventHub for {0}@{1}'.format(
+                                    hub_key[1], hub_key[0]
+                                )
+                            )
+                    except Exception as error:
+                        self.logger.debug(
+                            'Error disconnecting unused EventHub: {0}'.format(
+                                error)
+                        )
+
+    def get_session_count(self, hub_key):
+        """Get the number of active sessions for a hub.
+
+        Args:
+            hub_key: The hub registry key.
+
+        Returns:
+            int: Number of active sessions using this hub.
+        """
+        with self._lock:
+            if hub_key in self._hub_sessions:
+                return len(self._hub_sessions[hub_key])
+            return 0
+
+
+# Global singleton instance
+_GLOBAL_EVENT_HUB_REGISTRY = EventHubRegistry()
+
+
+def get_event_hub_registry():
+    """Get the global EventHub registry instance.
+
+    Returns:
+        EventHubRegistry: The global registry instance.
+    """
+    return _GLOBAL_EVENT_HUB_REGISTRY