Implement Centralized Aria Websocket Listener

Signed-off-by: lzzy12 <jhashivam2020@gmail.com>

bot/helper/ext_utils/bot_utils.py

@@ -61,7 +61,7 @@ def get_readable_file_size(size_in_bytes) -> str:
 def getDownloadByGid(gid):
     with download_dict_lock:
         for dl in download_dict.values():
-            if dl.status() == MirrorStatus.STATUS_DOWNLOADING or dl.status() == MirrorStatus.STATUS_WAITING:
+            if dl.status() != MirrorStatus.STATUS_UPLOADING and dl.status() != MirrorStatus.STATUS_ARCHIVING:
                 if dl.gid() == gid:
                     return dl
     return None
@@ -135,3 +135,13 @@ def is_magnet(url: str):
     if magnet:
         return True
     return False
+
+def new_thread(fn):
+    """To use as decorator to make a function call threaded.
+    Needs import
+    from threading import Thread"""
+    def wrapper(*args, **kwargs):
+        thread = threading.Thread(target=fn, args=args, kwargs=kwargs)
+        thread.start()
+        return thread
+    return wrapper

bot/helper/mirror_utils/download_utils/aria2_download.py

@@ -1,85 +1,76 @@
-from bot import aria2
+from bot import aria2, download_dict_lock
 from bot.helper.ext_utils.bot_utils import *
 from .download_helper import DownloadHelper
 from bot.helper.mirror_utils.status_utils.aria_download_status import AriaDownloadStatus
 from bot.helper.telegram_helper.message_utils import *
 import threading
 from aria2p import API
+from time import sleep
 
 
 class AriaDownloadHelper(DownloadHelper):
 
-    def __init__(self, listener):
+    def __init__(self):
         super().__init__()
-        self.gid = None
-        self.__listener = listener
-        self._resource_lock = threading.RLock()
 
+    @new_thread
     def __onDownloadStarted(self, api, gid):
-        with self._resource_lock:
+        LOGGER.info(f"onDownloadStart: {gid}")
-            if self.gid == gid:
+        update_all_messages()
-                download = api.get_download(gid)
-                self.name = download.name
-                update_all_messages()
 
     def __onDownloadComplete(self, api: API, gid):
-        with self._resource_lock:
+        LOGGER.info(f"onDownloadComplete: {gid}")
-            if self.gid == gid:
+        dl = getDownloadByGid(gid)
-                download = api.get_download(gid)
+        download = api.get_download(gid)
-                if download.followed_by_ids:
+        if download.followed_by_ids:
-                    self.gid = download.followed_by_ids[0]
+            new_gid = download.followed_by_ids[0]
-                    with download_dict_lock:
+            with download_dict_lock:
-                        download_dict[self.__listener.uid] = AriaDownloadStatus(self, self.__listener)
+                download_dict[dl.uid()] = AriaDownloadStatus(new_gid, dl.getListener())
-                    if download.is_torrent:
+                if download.is_torrent:
-                        download_dict[self.__listener.uid].is_torrent = True
+                    download_dict[dl.uid()].is_torrent = True
-                    update_all_messages()
+            update_all_messages()
-                    LOGGER.info(f'Changed gid from {gid} to {self.gid}')
+            LOGGER.info(f'Changed gid from {gid} to {new_gid}')
-                else:
-                    self.__listener.onDownloadComplete()
-
-    def __onDownloadPause(self, api, gid):
-        if self.gid == gid:
-            LOGGER.info("Called onDownloadPause")
-            self.__listener.onDownloadError('Download stopped by user!')
-
-    def __onDownloadStopped(self, api, gid):
-        if self.gid == gid:
-            LOGGER.info("Called on_download_stop")
-            self.__listener.onDownloadError('Download stopped by user!')
-
-    def __onDownloadError(self, api, gid):
-        with self._resource_lock:
-            if self.gid == gid:
-                download = api.get_download(gid)
-                error = download.error_message
-                LOGGER.info(f"Download Error: {error}")
-                self.__listener.onDownloadError(error)
-
-    def add_download(self, link: str, path):
-        if is_magnet(link):
-            download = aria2.add_magnet(link, {'dir': path})
         else:
-            download = aria2.add_uris([link], {'dir': path})
+            if dl: threading.Thread(target=dl.getListener().onDownloadComplete).start()
-        self.gid = download.gid
+
-        with download_dict_lock:
+    @new_thread
-            download_dict[self.__listener.uid] = AriaDownloadStatus(self, self.__listener)
+    def __onDownloadPause(self, api, gid):
-        if download.error_message:
+        LOGGER.info(f"onDownloadPause: {gid}")
-            self.__listener.onDownloadError(download.error_message)
+        dl = getDownloadByGid(gid)
-            return
+        dl.getListener().onDownloadError('Download stopped by user!')
-        LOGGER.info(f"Started: {self.gid} DIR:{download.dir} ")
+
+    @new_thread
+    def __onDownloadStopped(self, api, gid):
+        LOGGER.info(f"onDownloadStop: {gid}")
+        dl = getDownloadByGid(gid)
+        if dl: dl.getListener().onDownloadError('Download stopped by user!')
+
+    @new_thread
+    def __onDownloadError(self, api, gid):
+        sleep(0.5)  # sleep for split second to ensure proper dl gid update from onDownloadComplete
+        LOGGER.info(f"onDownloadError: {gid}")
+        dl = getDownloadByGid(gid)
+        download = api.get_download(gid)
+        error = download.error_message
+        LOGGER.info(f"Download Error: {error}")
+        if dl: dl.getListener().onDownloadError(error)
+
+    def start_listener(self):
         aria2.listen_to_notifications(threaded=True, on_download_start=self.__onDownloadStarted,
                                       on_download_error=self.__onDownloadError,
                                       on_download_pause=self.__onDownloadPause,
                                       on_download_stop=self.__onDownloadStopped,
                                       on_download_complete=self.__onDownloadComplete)
 
-    def cancel_download(self):
+    def add_download(self, link: str, path, listener):
-        download = aria2.get_download(self.gid)
+        if is_magnet(link):
-        if download.is_waiting:
+            download = aria2.add_magnet(link, {'dir': path})
-            aria2.remove([download])
+        else:
-            self.__listener.onDownloadError("Cancelled by user")
+            download = aria2.add_uris([link], {'dir': path})
+
+        if download.error_message:  # no need to proceed further at this point
+            listener.onDownloadError(download.error_message)
             return
-        if len(download.followed_by_ids) != 0:
+        with download_dict_lock:
-            downloads = aria2.get_downloads(download.followed_by_ids)
+            download_dict[listener.uid] = AriaDownloadStatus(download.gid, listener)
-            aria2.pause(downloads)
+            LOGGER.info(f"Started: {download.gid} DIR:{download.dir} ")
-        aria2.pause([download])

bot/helper/mirror_utils/status_utils/aria_download_status.py

@@ -1,4 +1,4 @@
-from bot import aria2, DOWNLOAD_DIR
+from bot import aria2, DOWNLOAD_DIR, LOGGER
 from bot.helper.ext_utils.bot_utils import MirrorStatus
 from .status import Status
 
@@ -9,13 +9,12 @@ def get_download(gid):
 
 class AriaDownloadStatus(Status):
 
-    def __init__(self, obj, listener):
+    def __init__(self, gid, listener):
         super().__init__()
         self.upload_name = None
         self.is_archiving = False
-        self.obj = obj
+        self.__gid = gid
-        self.__gid = obj.gid
+        self.__download = get_download(self.__gid)
-        self.__download = get_download(obj.gid)
         self.__uid = listener.uid
         self.__listener = listener
         self.message = listener.message
@@ -75,7 +74,16 @@ class AriaDownloadStatus(Status):
         return self.__download
 
     def download(self):
-        return self.obj
+        return self
+
+    def updateName(self,name):
+        self.__name = name
+
+    def updateGid(self,gid):
+        self.__gid = gid
+
+    def getListener(self):
+        return self.__listener
     
     def uid(self):
         return self.__uid
@@ -83,3 +91,16 @@ class AriaDownloadStatus(Status):
     def gid(self):
         self.__update()
         return self.__gid
+
+    def cancel_download(self):
+        LOGGER.info(f"Cancelling Download: {self.name()}")
+        download = self.aria_download()
+        if download.is_waiting:
+            aria2.remove([download])
+            self.__listener.onDownloadError("Cancelled by user")
+            return
+        if len(download.followed_by_ids) != 0:
+            downloads = aria2.get_downloads(download.followed_by_ids)
+            aria2.pause(downloads)
+        aria2.pause([download])
+

bot/helper/telegram_helper/message_utils.py

@@ -66,7 +66,7 @@ def update_all_messages():
     msg = get_readable_message()
     with status_reply_dict_lock:
         for chat_id in list(status_reply_dict.keys()):
-            if msg != status_reply_dict[chat_id].text:
+            if status_reply_dict[chat_id] and msg != status_reply_dict[chat_id].text:
                 try:
                     editMessage(msg, status_reply_dict[chat_id])
                 except BadRequest as e:

bot/modules/mirror.py

@@ -6,7 +6,7 @@ from bot import dispatcher, DOWNLOAD_DIR, DOWNLOAD_STATUS_UPDATE_INTERVAL, downl
 from bot.helper.ext_utils import fs_utils, bot_utils
 from bot.helper.ext_utils.bot_utils import setInterval
 from bot.helper.ext_utils.exceptions import DirectDownloadLinkException
-from bot.helper.mirror_utils.download_utils import aria2_download
+from bot.helper.mirror_utils.download_utils.aria2_download import AriaDownloadHelper
 from bot.helper.mirror_utils.download_utils.direct_link_generator import direct_link_generator
 from bot.helper.mirror_utils.download_utils.telegram_downloader import TelegramDownloadHelper
 from bot.helper.mirror_utils.status_utils import listeners
@@ -20,6 +20,9 @@ from bot.helper.mirror_utils.download_utils.youtube_dl_download_helper import Yo
 import pathlib
 import os
 
+ariaDlManager = AriaDownloadHelper()
+ariaDlManager.start_listener()
+
 
 class MirrorListener(listeners.MirrorListeners):
 
@@ -83,7 +86,6 @@ class MirrorListener(listeners.MirrorListeners):
                 del download_dict[self.uid]
                 LOGGER.info(f"Deleting folder: {download.path()}")
                 fs_utils.clean_download(download.path())
-                LOGGER.info(f"Deleting {download.name()} from download_dict.")
                 LOGGER.info(str(download_dict))
             except Exception as e:
                 LOGGER.error(str(e))
@@ -186,8 +188,7 @@ def _mirror(bot, update, isTar=False):
     except DirectDownloadLinkException as e:
         LOGGER.info(f'{link}: {e}')
     listener = MirrorListener(bot, update, isTar, tag)
-    aria = aria2_download.AriaDownloadHelper(listener)
+    ariaDlManager.add_download(link, f'{DOWNLOAD_DIR}/{listener.uid}/',listener)
-    aria.add_download(link, f'{DOWNLOAD_DIR}/{listener.uid}/')
     sendStatusMessage(update, bot)
     if len(Interval) == 0:
         Interval.append(setInterval(DOWNLOAD_STATUS_UPDATE_INTERVAL, update_all_messages))