From c9634ba10aeb41a2649e9ee2ee7ebe72f1f12ae9 Mon Sep 17 00:00:00 2001
From: Sascha Bischoff <sascha.bischoff@arm.com>
Date: Thu, 24 Feb 2022 13:19:50 +0000
Subject: [PATCH] Periodically check for websocket disconnect and re-init

Both network issues and restarts of Unifi Protect can cause the
websocket to disconnect. Once this happens, no more events are
recieved, and hence no events are stored via rclone.

We add a task which checks that the websocket is connected every
minute. If the websocket is not connected, the connection is totally
reset. For a simple network issue, is should be sufficient to just
call pyunifiprotect's update(), but this doesn't work when protect has
been restarted. Given that this is a tool that should always be
running, we opt for the most extreme option of totally resetting the
connection, and re-establishing it from scratch.
---
 unifi_protect_backup/unifi_protect_backup.py | 69 ++++++++++++++++++--
 1 file changed, 64 insertions(+), 5 deletions(-)

diff --git a/unifi_protect_backup/unifi_protect_backup.py b/unifi_protect_backup/unifi_protect_backup.py
index a2201a9..a0c6414 100644
--- a/unifi_protect_backup/unifi_protect_backup.py
+++ b/unifi_protect_backup/unifi_protect_backup.py
@@ -230,12 +230,18 @@ class UnifiProtectBackup:
         self.retention = retention
         self.rclone_args = rclone_args
 
+        self.address = address
+        self.port = port
+        self.username = username
+        self.password = password
+        self.verify_ssl = verify_ssl
+
         self._protect = ProtectApiClient(
-            address,
-            port,
-            username,
-            password,
-            verify_ssl=verify_ssl,
+            self.address,
+            self.port,
+            self.username,
+            self.password,
+            verify_ssl=self.verify_ssl,
             subscribed_models={ModelType.EVENT},
         )
         self.ignore_cameras = ignore_cameras
@@ -289,6 +295,59 @@ class UnifiProtectBackup:
                 logger.warn(f"stdout:\n{stdout.decode()}")
                 logger.warn(f"stderr:\n{stderr.decode()}")
 
+        # We need to catch websocket disconnect and trigger a reconnect.
+        @aiocron.crontab("* * * * *")
+        async def check_websocket_and_reconnect():
+            logger.debug("Checking the status of the websocket...")
+            if self._protect.check_ws():
+                logger.debug("Websocket is connected.")
+            else:
+                logger.warn("Lost connection to Unifi Protect.")
+
+                # Unsubscribe, close the session.
+                self._unsub()
+                await self._protect.close_session()
+
+                while True:
+                    logger.warn("Attempting reconnect...")
+
+                    try:
+                        # Start again from scratch. In principle if Unifi
+                        # Protect has not been restarted we should just be able
+                        # to call self._protect.update() to reconnect to the
+                        # websocket. However, if the server has been restarted a
+                        # call to self._protect.check_ws() returns true and some
+                        # seconds later pyunifiprotect detects the websocket as
+                        # disconnected again. Therefore, kill it all and try
+                        # again!
+                        replacement_protect = ProtectApiClient(
+                            self.address,
+                            self.port,
+                            self.username,
+                            self.password,
+                            verify_ssl=self.verify_ssl,
+                            subscribed_models={ModelType.EVENT},
+                        )
+                        # Start the pyunifiprotect connection by calling `update`
+                        await replacement_protect.update()
+                        if replacement_protect.check_ws():
+                            self._protect = replacement_protect
+                            self._unsub = self._protect.subscribe_websocket(self._websocket_callback)
+                            break
+                        else:
+                            logger.warn("Unable to establish connection to Unifi Protect")
+                    except Exception as e:
+                        logger.warn("Unexpected exception occurred while trying to reconnect:")
+                        logger.exception(e)
+                    finally:
+                        # If we get here we need to close the replacement session again
+                        await replacement_protect.close_session()
+
+                        # Back off for a little while
+                        await asyncio.sleep(10)
+
+                logger.info("Re-established connection to Unifi Protect and to the websocket.")
+
         # Launches the main loop
         logger.info("Listening for events...")
         await self._backup_events()