mi
66 lines
2.2 KiB
Python
66 lines
2.2 KiB
Python
import asyncio
|
|
import json
|
|
import uuid
|
|
from collections.abc import AsyncIterator
|
|
from contextlib import suppress
|
|
from typing import Any
|
|
|
|
from redis.asyncio import Redis
|
|
|
|
CHANNEL_PREFIX = "han:rt:dialog:"
|
|
|
|
|
|
class LocalFanout:
|
|
def __init__(self) -> None:
|
|
self._queues: set[asyncio.Queue[dict[str, Any]]] = set()
|
|
|
|
async def publish(self, event: dict[str, Any]) -> None:
|
|
for queue in tuple(self._queues):
|
|
with suppress(asyncio.QueueFull):
|
|
queue.put_nowait(event)
|
|
|
|
async def subscribe(self) -> AsyncIterator[dict[str, Any]]:
|
|
queue: asyncio.Queue[dict[str, Any]] = asyncio.Queue(maxsize=256)
|
|
self._queues.add(queue)
|
|
try:
|
|
while True:
|
|
yield await queue.get()
|
|
finally:
|
|
self._queues.discard(queue)
|
|
|
|
|
|
class RealtimeFanout:
|
|
def __init__(self, redis: Redis, local: LocalFanout | None = None) -> None:
|
|
self.redis = redis
|
|
self.local = local or LocalFanout()
|
|
|
|
async def publish(self, event: dict[str, Any]) -> None:
|
|
event = {"event_id": str(uuid.uuid4()), **event}
|
|
channel = CHANNEL_PREFIX + str(event["dialog_id"])
|
|
try:
|
|
await self.redis.publish(channel, json.dumps(event, default=str, separators=(",", ":")))
|
|
except Exception:
|
|
await self.local.publish(event)
|
|
|
|
async def events(self, dialog_ids: set[uuid.UUID]) -> AsyncIterator[dict[str, Any]]:
|
|
channels = [CHANNEL_PREFIX + str(dialog_id) for dialog_id in dialog_ids]
|
|
pubsub = self.redis.pubsub()
|
|
try:
|
|
await pubsub.subscribe(*channels)
|
|
except Exception:
|
|
await pubsub.aclose()
|
|
async for event in self.local.subscribe():
|
|
if uuid.UUID(str(event["dialog_id"])) in dialog_ids:
|
|
yield event
|
|
return
|
|
try:
|
|
while True:
|
|
message = await pubsub.get_message(ignore_subscribe_messages=True, timeout=1)
|
|
if message:
|
|
yield json.loads(message["data"])
|
|
else:
|
|
await asyncio.sleep(0)
|
|
finally:
|
|
await pubsub.unsubscribe(*channels)
|
|
await pubsub.aclose()
|