Skip to content

Commit

Permalink
PubSubEndpoint: Notify local subscribers through broadcaster
Browse files Browse the repository at this point in the history 
To have better ordering safety of messages (local subscribers won't be notified much before remote ones)
  • Loading branch information
@roekatz
roekatz committed Jun 20, 2023
1 parent ea16b34 commit a2e77b5
Show file tree
Hide file tree
Showing 2 changed files with 88 additions and 54 deletions.
101 changes: 65 additions & 36 deletions fastapi_websocket_pubsub/event_broadcaster.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,7 @@
from fastapi_websocket_rpc.utils import gen_uid


logger = get_logger('EventBroadcaster')
logger = get_logger("EventBroadcaster")


# Cross service broadcast consts
Expand All @@ -35,9 +35,14 @@ class EventBroadcasterContextManager:
Friend-like class of EventBroadcaster (accessing "protected" members )
"""

def __init__(self, event_broadcaster: "EventBroadcaster", listen: bool = True, share: bool = True) -> None:
def __init__(
self,
event_broadcaster: "EventBroadcaster",
listen: bool = True,
share: bool = True,
) -> None:
"""
Provide a context manager for an EventBroadcaster, managing if it listens to events coming from the broadcaster
Provide a context manager for an EventBroadcaster, managing if it listens to events coming from the broadcaster
and if it subscribes to the internal notifier to share its events with the broadcaster
Args:
Expand All @@ -56,7 +61,9 @@ async def __aenter__(self):
self._event_broadcaster._listen_count += 1
if self._event_broadcaster._listen_count == 1:
# We have our first listener start the read-task for it (And all those who'd follow)
logger.info("Listening for incoming events from broadcast channel (first listener started)")
logger.info(
"Listening for incoming events from broadcast channel (first listener started)"
)
# Start task listening on incoming broadcasts
self._event_broadcaster.start_reader_task()

Expand All @@ -66,11 +73,15 @@ async def __aenter__(self):
# We have our first publisher
# Init the broadcast used for sharing (reading has its own)
self._event_broadcaster._acquire_sharing_broadcast_channel()
logger.debug("Subscribing to ALL_TOPICS, and sharing messages with broadcast channel")
logger.debug(
"Subscribing to ALL_TOPICS, and sharing messages with broadcast channel"
)
# Subscribe to internal events form our own event notifier and broadcast them
await self._event_broadcaster._subscribe_to_all_topics()
else:
logger.debug(f"Did not subscribe to ALL_TOPICS: share count == {self._event_broadcaster._share_count}")
logger.debug(
f"Did not subscribe to ALL_TOPICS: share count == {self._event_broadcaster._share_count}"
)
return self

async def __aexit__(self, exc_type, exc, tb):
Expand All @@ -87,7 +98,7 @@ async def __aexit__(self, exc_type, exc, tb):
self._event_broadcaster._subscription_task = None

if self._share:
self._event_broadcaster._share_count -= 1
self._event_broadcaster._share_count -= 1
# if this was last sharer - we can stop subscribing to internal events - we aren't sharing anymore
if self._event_broadcaster._share_count == 0:
# Unsubscribe from internal events
Expand All @@ -110,8 +121,14 @@ class EventBroadcaster:
<Your Code>
"""

def __init__(self, broadcast_url: str, notifier: EventNotifier, channel="EventNotifier",
broadcast_type=None, is_publish_only=False) -> None:
def __init__(
self,
broadcast_url: str,
notifier: EventNotifier,
channel="EventNotifier",
broadcast_type=None,
is_publish_only=False,
) -> None:
"""
Args:
Expand All @@ -138,26 +155,30 @@ def __init__(self, broadcast_url: str, notifier: EventNotifier, channel="EventNo
self._publish_lock = None
# used to track creation / removal of resources needed per type (reader task->listen, and subscription to internal events->share)
self._listen_count: int = 0
self._share_count: int = 0
# If we opt to manage the context directly (i.e. call async with on the event broadcaster itself)
self._share_count: int = 0
# If we opt to manage the context directly (i.e. call async with on the event broadcaster itself)
self._context_manager = None


async def __broadcast_notifications__(self, subscription: Subscription, data):
async def __broadcast_notifications__(self, topics: TopicList, data):
"""
Share incoming internal notifications with the entire broadcast channel
Args:
subscription (Subscription): the subscription that got triggered
data: the event data
"""
logger.info("Broadcasting incoming event: {}".format({'topic': subscription.topic, 'notifier_id': self._id}))
note = BroadcastNotification(notifier_id=self._id, topics=[
subscription.topic], data=data)
logger.info(
"Broadcasting incoming event: {}".format(
{"topics": topics, "notifier_id": self._id}
)
)
note = BroadcastNotification(notifier_id=self._id, topics=topics, data=data)
# Publish event to broadcast
async with self._publish_lock:
async with self._sharing_broadcast_channel:
await self._sharing_broadcast_channel.publish(self._channel, note.json())
await self._sharing_broadcast_channel.publish(
self._channel, note.json()
)

def _acquire_sharing_broadcast_channel(self):
"""
Expand All @@ -167,9 +188,9 @@ def _acquire_sharing_broadcast_channel(self):
self._sharing_broadcast_channel = self._broadcast_type(self._broadcast_url)

async def _subscribe_to_all_topics(self):
return await self._notifier.subscribe(self._id,
ALL_TOPICS,
self.__broadcast_notifications__)
return await self._notifier.subscribe(
self._id, ALL_TOPICS, self.__broadcast_notifications__
)

async def _unsubscribe_from_topics(self):
return await self._notifier.unsubscribe(self._id)
Expand All @@ -183,16 +204,16 @@ def get_context(self, listen=True, share=True):
share (bool, optional): Should we share events with the broadcast channel. Defaults to True.
Returns:
EventBroadcasterContextManager: the context
EventBroadcasterContextManager: the context
"""
return EventBroadcasterContextManager(self, listen=listen, share=share)

def get_listening_context(self):
return EventBroadcasterContextManager(self, listen=True, share=False)

def get_sharing_context(self):
return EventBroadcasterContextManager(self, listen=False, share=True)

async def __aenter__(self):
"""
Convince caller (also backward compaltability)
Expand All @@ -201,7 +222,6 @@ async def __aenter__(self):
self._context_manager = self.get_context(listen=not self._is_publish_only)
return await self._context_manager.__aenter__()


async def __aexit__(self, exc_type, exc, tb):
await self._context_manager.__aexit__(exc_type, exc, tb)

Expand All @@ -215,14 +235,15 @@ def start_reader_task(self):
# Make sure a task wasn't started already
if self._subscription_task is not None:
# we already started a task for this worker process
logger.debug("No need for listen task, already started broadcast listen task for this notifier")
logger.debug(
"No need for listen task, already started broadcast listen task for this notifier"
)
return
# Trigger the task
logger.debug("Spawning broadcast listen task")
self._subscription_task = asyncio.create_task(
self.__read_notifications__())
self._subscription_task = asyncio.create_task(self.__read_notifications__())
return self._subscription_task

def get_reader_task(self):
return self._subscription_task

Expand All @@ -235,15 +256,23 @@ async def __read_notifications__(self):
listening_broadcast_channel = self._broadcast_type(self._broadcast_url)
async with listening_broadcast_channel:
# Subscribe to our channel
async with listening_broadcast_channel.subscribe(channel=self._channel) as subscriber:
async with listening_broadcast_channel.subscribe(
channel=self._channel
) as subscriber:
async for event in subscriber:
try:
notification = BroadcastNotification.parse_raw(
event.message)
# Avoid re-publishing our own broadcasts
if notification.notifier_id != self._id:
logger.info("Handling incoming broadcast event: {}".format({'topics': notification.topics, 'src': notification.notifier_id}))
# Notify subscribers of message received from broadcast
await self._notifier.notify(notification.topics, notification.data, notifier_id=self._id)
notification = BroadcastNotification.parse_raw(event.message)
logger.info(
"Handling incoming broadcast event: {}".format(
{
"topics": notification.topics,
"src": notification.notifier_id,
}
)
)
# Notify subscribers of message received from broadcast
await self._notifier.notify(
notification.topics, notification.data, notifier_id=self._id
)
except:
logger.exception("Failed handling incoming broadcast")
41 changes: 23 additions & 18 deletions fastapi_websocket_pubsub/pub_sub_server.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -92,6 +92,26 @@ async def subscribe(
) -> List[Subscription]:
return await self.notifier.subscribe(self._subscriber_id, topics, callback)

async def _publish_through_broadcaster(
self, topics: Union[TopicList, Topic], data=None
):
if self._broadcaster_sharing_context is None:
logger.debug(f"Getting new broadcaster sharing context")
self._broadcaster_sharing_context = self.broadcaster.get_context(
listen=True, share=True
)
logger.debug(f"Acquiring broadcaster sharing context")
try:
async with self._broadcaster_sharing_context:
# We don't notify notifier here, broadcaster listens to its own back-channel and notifies notifier,
# Thus all subscribers are notified in the same order (local and remote)
await self.broadcaster.__broadcast_notifications__(topics, data)
except Exception:
# Could check if the exception has to do with disconnection, but just in case better to restart sharing context anyway
logger.warning(f"Exception in publish, resetting sharing context")
self._broadcaster_sharing_context = None
raise

async def publish(self, topics: Union[TopicList, Topic], data=None):
"""
Publish events to subscribres of given topics currently connected to the endpoint
Expand All @@ -100,27 +120,12 @@ async def publish(self, topics: Union[TopicList, Topic], data=None):
topics (Union[TopicList, Topic]): topics to publish to relevant subscribers
data (Any, optional): Event data to be passed to each subscriber. Defaults to None.
"""
# if we have a broadcaster make sure we share with it (no matter where this call comes from)
# sharing here means - the broadcaster listens in to the notifier as well
logger.debug(f"Publishing message to topics: {topics}")
if self.broadcaster is not None:
if self._broadcaster_sharing_context is None:
logger.debug(f"Getting new broadcaster sharing context")
self._broadcaster_sharing_context = self.broadcaster.get_context(
listen=True, share=True
)
logger.debug(f"Acquiring broadcaster sharing context")
try:
async with self._broadcaster_sharing_context:
await self.notifier.notify(topics, data, notifier_id=self._id)
except Exception:
# Could check if the exception has to do with disconnection, but just in case better to restart sharing context anyway
logger.warning(f"Exception in publish, resetting sharing context")
self._broadcaster_sharing_context = None
raise

# otherwise just notify
# Use broadcaster if available, notifier would be notified when broadcaster reads its own notification
await self._publish_through_broadcaster(topics, data)
else:
# otherwise just notify
await self.notifier.notify(topics, data, notifier_id=self._id)

# canonical name (backward compatability)
Expand Down

0 comments on commit a2e77b5

Please sign in to comment.