diff --git a/backend/open_webui/env.py b/backend/open_webui/env.py index fe5f8efb21..dc160e1592 100644 --- a/backend/open_webui/env.py +++ b/backend/open_webui/env.py @@ -493,6 +493,12 @@ else: WEBSOCKET_REDIS_URL = os.getenv('WEBSOCKET_REDIS_URL', REDIS_URL) WEBSOCKET_REDIS_CLUSTER = os.getenv('WEBSOCKET_REDIS_CLUSTER', str(REDIS_CLUSTER)).lower() == 'true' +# publishes room-targeted emits on per-room redis channels so instances skip +# messages for rooms without local members; must be identical across the fleet +# (toggle with a full restart, not a rolling one), set false for the previous +# shared-channel-only delivery +WEBSOCKET_REDIS_ROOM_CHANNELS = os.getenv('WEBSOCKET_REDIS_ROOM_CHANNELS', 'True').lower() == 'true' + websocket_redis_lock_timeout = os.getenv('WEBSOCKET_REDIS_LOCK_TIMEOUT', '60') try: diff --git a/backend/open_webui/socket/main.py b/backend/open_webui/socket/main.py index 659dc3da49..796c08b9cf 100644 --- a/backend/open_webui/socket/main.py +++ b/backend/open_webui/socket/main.py @@ -24,6 +24,7 @@ from open_webui.env import ( WEBSOCKET_REDIS_CLUSTER, WEBSOCKET_REDIS_LOCK_TIMEOUT, WEBSOCKET_REDIS_OPTIONS, + WEBSOCKET_REDIS_ROOM_CHANNELS, WEBSOCKET_REDIS_URL, WEBSOCKET_SENTINEL_HOSTS, WEBSOCKET_SENTINEL_PORT, @@ -38,6 +39,7 @@ from open_webui.models.chats import Chats from open_webui.models.folders import Folders from open_webui.models.notes import Notes, NoteUpdateForm from open_webui.models.users import UserNameResponse, Users +from open_webui.socket.redis_room_channels import AsyncRedisRoomChannelManager from open_webui.socket.utils import CachedRedisDict, RedisDict, RedisLock, YdocManager from open_webui.tasks import ( REDIS_PUBSUB_MAX_RECONNECT_INTERVAL, @@ -93,7 +95,8 @@ if WEBSOCKET_MANAGER == 'redis': if sentinel_hosts else WEBSOCKET_REDIS_URL ) - redis_manager = socketio.AsyncRedisManager(ws_redis_url, redis_options=WEBSOCKET_REDIS_OPTIONS, json=SOCKETIO_JSON) + manager_class = AsyncRedisRoomChannelManager if WEBSOCKET_REDIS_ROOM_CHANNELS else socketio.AsyncRedisManager + redis_manager = manager_class(ws_redis_url, redis_options=WEBSOCKET_REDIS_OPTIONS, json=SOCKETIO_JSON) sio = socketio.AsyncServer( cors_allowed_origins=SOCKETIO_CORS_ORIGINS, async_mode='asgi', diff --git a/backend/open_webui/socket/redis_room_channels.py b/backend/open_webui/socket/redis_room_channels.py new file mode 100644 index 0000000000..715ccc1ea9 --- /dev/null +++ b/backend/open_webui/socket/redis_room_channels.py @@ -0,0 +1,82 @@ +"""Per-room redis channels let instances skip the decode and packet encode for rooms with no local members.""" + +import asyncio + +from socketio import AsyncRedisManager + + +class AsyncRedisRoomChannelManager(AsyncRedisManager): + name = 'aioredisroomchannel' + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self._local_room_channels = set() + + # collision-free while namespaces contain no '#' (socket.io default '/'); rooms may contain '#' + def _room_channel(self, namespace, room): + return f'{self.channel}#{namespace}#{room}'.encode() + + 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)) + + 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)) + + async def _publish(self, data): + if data.get('method') == 'emit' and isinstance(data.get('room'), str): + channel = self._room_channel(data['namespace'], data['room']) + else: + channel = self.channel + _, error = self._get_redis_module_and_error() + for retries_left in range(1, -1, -1): # 2 attempts + try: + if not self.connected: + self._redis_connect() + return await self.redis.publish(channel, self.json.dumps(data)) + except error as exc: + if retries_left > 0: + self._get_logger().error('Cannot publish to redis... retrying', extra={'redis_exception': str(exc)}) + self.connected = False + else: + self._get_logger().error( + 'Cannot publish to redis... giving up', extra={'redis_exception': str(exc)} + ) + break + + 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(): + yield message + except error as exc: + self._get_logger().error( + 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 + ): + yield message['data']