Fix: Improve PubSub resilience and exception handling

This commit introduces several improvements to the Pub/Sub system to make it more resilient to connection issues and to handle exceptions more gracefully.

The following changes are included:

- A `health_check_interval` is added to the Redis connection to keep it alive and prevent timeouts.
- Reconnection logic is implemented in the `listen` method. If the connection to Redis is lost, the client will now attempt to reconnect and re-subscribe to the channel automatically.
- The exception handling in the `listen` method is made more specific, catching `ConnectionError` for reconnection and other `RedisError` exceptions for graceful exit.

These changes improve the overall stability and robustness of the Pub/Sub system, especially in environments where the Redis server might be restarted or the network connection is not perfectly stable.

Signed-off-by: Denys Fedoryshchenko <[email protected]>

Author: nuclearcat

api/pubsub.py

@@ -39,7 +39,9 @@ def __init__(self, host=None, db_number=None):
             host = self._settings.redis_host
         if db_number is None:
             db_number = self._settings.redis_db_number
-        self._redis = aioredis.from_url(f'redis://{host}/{db_number}')
+        self._redis = aioredis.from_url(
+            f'redis://{host}/{db_number}', health_check_interval=30
+        )
         # self._subscriptions is a dict that matches a subscription id
         # (key) with a Subscription object ('sub') and a redis
         # PubSub object ('redis_sub'). For instance:
@@ -135,9 +137,22 @@ async def listen(self, sub_id, user=None):
                                f"not owned by {user}")
         while True:
             self._subscriptions[sub_id]['last_poll'] = datetime.utcnow()
-            msg = await sub['redis_sub'].get_message(
-                ignore_subscribe_messages=True, timeout=1.0
-            )
+            msg = None
+            try:
+                msg = await sub['redis_sub'].get_message(
+                    ignore_subscribe_messages=True, timeout=1.0
+                )
+            except aioredis.ConnectionError:
+                async with self._lock:
+                    channel = self._subscriptions[sub_id]['sub'].channel
+                    new_redis_sub = self._redis.pubsub()
+                    await new_redis_sub.subscribe(channel)
+                    self._subscriptions[sub_id]['redis_sub'] = new_redis_sub
+                    sub['redis_sub'] = new_redis_sub
+                continue
+            except aioredis.RedisError:
+                return None  # Handle any exceptions gracefully
+
             if msg is None:
                 continue
             msg_data = json.loads(msg['data'])