diff --git a/.env.example b/.env.example index 8426e3a23e..d215147046 100644 --- a/.env.example +++ b/.env.example @@ -6,7 +6,7 @@ OPENAI_API_BASE_URL='' OPENAI_API_KEY='' # With WEBSOCKET_MANAGER=redis, room channels require PUBLISH, SUBSCRIBE and -# PSUBSCRIBE permissions plus channel ACLs &socketio and &socketio#*. +# UNSUBSCRIBE/PING permissions plus channel ACLs &socketio and &socketio#*. # Set false for shared-channel delivery. This must match on every instance; # changing it requires a full fleet restart, not a rolling restart. # WEBSOCKET_REDIS_ROOM_CHANNELS=true diff --git a/backend/open_webui/socket/redis_room_channels.py b/backend/open_webui/socket/redis_room_channels.py index 73ac602ae8..0ad8c9aa5f 100644 --- a/backend/open_webui/socket/redis_room_channels.py +++ b/backend/open_webui/socket/redis_room_channels.py @@ -11,7 +11,7 @@ class AsyncRedisRoomChannelManager(AsyncRedisManager): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) - self._local_room_channels = set() + self._subscription_queue = asyncio.Queue() # collision-free while namespaces contain no '#' (socket.io default '/'); rooms may contain '#' def _room_channel(self, namespace, room): @@ -20,12 +20,29 @@ class AsyncRedisRoomChannelManager(AsyncRedisManager): def basic_enter_room(self, sid, namespace, room, eio_sid=None): super().basic_enter_room(sid, namespace, room, eio_sid=eio_sid) if room is not None: - self._local_room_channels.add(self._room_channel(namespace, room)) + self._subscription_queue.put_nowait(asyncio.Event()) def basic_leave_room(self, sid, namespace, room): super().basic_leave_room(sid, namespace, room) if room is not None and room not in self.rooms.get(namespace, {}): - self._local_room_channels.discard(self._room_channel(namespace, room)) + self._subscription_queue.put_nowait(asyncio.Event()) + + async def _wait_for_rooms(self): + ready = asyncio.Event() + self._subscription_queue.put_nowait(ready) + await ready.wait() + + async def connect(self, eio_sid, namespace): + sid = await super().connect(eio_sid, namespace) + if sid is not None: + await self._wait_for_rooms() + return sid + + async def enter_room(self, sid, namespace, room, eio_sid=None): + local = self.is_connected(sid, namespace) + await super().enter_room(sid, namespace, room, eio_sid=eio_sid) + if local and room is not None: + await self._wait_for_rooms() async def _publish(self, data): if data.get('method') == 'emit' and isinstance(data.get('room'), str): @@ -57,23 +74,73 @@ class AsyncRedisRoomChannelManager(AsyncRedisManager): ) break + async def _sync_rooms(self, pubsub): + wanted = {self.channel.encode()} | { + self._room_channel(namespace, room) + for namespace, rooms in self.rooms.items() + for room in rooms + if room is not None + } + if added := wanted - (pubsub.channels.keys() - pubsub.pending_unsubscribe_channels): + await pubsub.subscribe(*added) + if removed := pubsub.channels.keys() - wanted - pubsub.pending_unsubscribe_channels: + await pubsub.unsubscribe(*removed) + + async def _listen_to_rooms(self, pubsub): + pubsub.ignore_subscribe_messages = False + self._subscription_queue.put_nowait(asyncio.Event()) + changed = asyncio.create_task(self._subscription_queue.get()) + message = asyncio.create_task(pubsub.get_message(timeout=None)) + waiters = [] + serial = 0 + try: + while True: + await asyncio.wait([message] if waiters else [message, changed], return_when=asyncio.FIRST_COMPLETED) + if not waiters and changed.done(): + waiters.append(changed.result()) + while not self._subscription_queue.empty(): + waiters.append(self._subscription_queue.get_nowait()) + changed = asyncio.create_task(self._subscription_queue.get()) + await self._sync_rooms(pubsub) + # PONG confirms Redis processed the subscriptions before a join returns. + serial += 1 + await pubsub.ping(str(serial).encode()) + if message.done(): + data = message.result() or {} + message = asyncio.create_task(pubsub.get_message(timeout=None)) + if data.get('type') == 'subscribe': + # redis-py can reconnect internally and restore a pending removal. + self._subscription_queue.put_nowait(asyncio.Event()) + elif data.get('type') == 'pong' and data['data'] == str(serial).encode(): + for ready in waiters: + ready.set() + waiters.clear() + else: + yield data + finally: + if changed.done() and not changed.cancelled(): + waiters.append(changed.result()) + for ready in waiters: + self._subscription_queue.put_nowait(ready) + message.cancel() + changed.cancel() + await asyncio.gather(message, changed, return_exceptions=True) + await pubsub.aclose() + async def _redis_listen_with_retries(self): _, error = self._get_redis_module_and_error() retry_sleep = 1 - subscribed = False while True: try: - if not subscribed: - self._redis_connect() - await self.pubsub.subscribe(self.channel) - await self.pubsub.psubscribe(f'{self.channel}#*') - retry_sleep = 1 - async for message in self.pubsub.listen(): + self._redis_connect() + await self.pubsub.subscribe(self.channel) + retry_sleep = 1 + async for message in self._listen_to_rooms(self.pubsub): yield message except error as exc: if isinstance(exc, NoPermissionError): self._get_logger().error( - 'Redis denied subscribing: %s. Check SUBSCRIBE/PSUBSCRIBE permissions and channel ACLs ' + 'Redis denied subscribing: %s. Check SUBSCRIBE/UNSUBSCRIBE/PING permissions and channel ACLs ' '(&%s and &%s#*). To disable room channels, set ' 'WEBSOCKET_REDIS_ROOM_CHANNELS=False on every instance and fully restart the fleet.', exc, @@ -84,18 +151,12 @@ class AsyncRedisRoomChannelManager(AsyncRedisManager): f'Cannot receive from redis... retrying in {retry_sleep} secs', extra={'redis_exception': str(exc)}, ) - subscribed = False await asyncio.sleep(retry_sleep) retry_sleep *= 2 if retry_sleep > 60: retry_sleep = 60 async def _listen(self): - main_channel = self.channel.encode() async for message in self._redis_listen_with_retries(): - if 'data' not in message: - continue - if (message['type'] == 'message' and message['channel'] == main_channel) or ( - message['type'] == 'pmessage' and message['channel'] in self._local_room_channels - ): + if message.get('type') == 'message': yield message['data']