From cdefe1234800e9a8b36ca5f55a51f73f14696293 Mon Sep 17 00:00:00 2001 From: jaskaranSM Date: Fri, 15 May 2020 13:34:09 +0530 Subject: [PATCH] Implement Centralized Aria Websocket Listener --- bot/helper/ext_utils/bot_utils.py | 12 ++- .../download_utils/aria2_download.py | 97 +++++++++---------- .../status_utils/aria_download_status.py | 33 +++++-- bot/helper/telegram_helper/message_utils.py | 2 +- bot/modules/mirror.py | 9 +- 5 files changed, 89 insertions(+), 64 deletions(-) diff --git a/bot/helper/ext_utils/bot_utils.py b/bot/helper/ext_utils/bot_utils.py index 8d0e6ce..5998e53 100644 --- a/bot/helper/ext_utils/bot_utils.py +++ b/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 \ No newline at end of file diff --git a/bot/helper/mirror_utils/download_utils/aria2_download.py b/bot/helper/mirror_utils/download_utils/aria2_download.py index bd5c783..e30a0a9 100644 --- a/bot/helper/mirror_utils/download_utils/aria2_download.py +++ b/bot/helper/mirror_utils/download_utils/aria2_download.py @@ -1,85 +1,78 @@ -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: - download = api.get_download(gid) - self.name = download.name - update_all_messages() + LOGGER.info(f"onDownloadStart: {gid}") + update_all_messages() def __onDownloadComplete(self, api: API, gid): - with self._resource_lock: - LOGGER.info(f"onDownloadComplete: {gid}") - if self.gid == gid: - download = api.get_download(gid) - if download.followed_by_ids: - self.gid = download.followed_by_ids[0] - with download_dict_lock: - download_dict[self.__listener.uid] = AriaDownloadStatus(self, self.__listener) - if download.is_torrent: - download_dict[self.__listener.uid].is_torrent = True - update_all_messages() - LOGGER.info(f'Changed gid from {gid} to {self.gid}') - else: - self.__listener.onDownloadComplete() + LOGGER.info(f"onDownloadComplete: {gid}") + dl = getDownloadByGid(gid) + download = api.get_download(gid) + if download.followed_by_ids: + new_gid = download.followed_by_ids[0] + with download_dict_lock: + download_dict[dl.uid()] = AriaDownloadStatus(new_gid,dl.getListener()) + if download.is_torrent: + download_dict[dl.uid()].is_torrent = True + update_all_messages() + LOGGER.info(f'Changed gid from {gid} to {new_gid}') + else: + if dl: threading.Thread(target=dl.getListener().onDownloadComplete).start() + @new_thread def __onDownloadPause(self, api, gid): - if self.gid == gid: - LOGGER.info("Called onDownloadPause") - self.__listener.onDownloadError('Download stopped by user!') + LOGGER.info(f"onDownloadPause: {gid}") + dl = getDownloadByGid(gid) + dl.getListener().onDownloadError('Download stopped by user!') + @new_thread def __onDownloadStopped(self, api, gid): - if self.gid == gid: - LOGGER.info("Called on_download_stop") - self.__listener.onDownloadError('Download stopped by user!') + LOGGER.info(f"onDownloadStop: {gid}") + dl = getDownloadByGid(gid) + if dl: dl.getListener().onDownloadError('Download stopped by user!') + @new_thread 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) + 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 add_download(self, link: str, path): + 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 add_download(self, link: str, path,listener): if is_magnet(link): download = aria2.add_magnet(link, {'dir': path}) else: download = aria2.add_uris([link], {'dir': path}) - with self._resource_lock: - self.gid = download.gid + if download.error_message: #no need to proceed further at this point + listener.onDownloadError(download.error_message) + return with download_dict_lock: - download_dict[self.__listener.uid] = AriaDownloadStatus(self, self.__listener) - LOGGER.info(f"Started: {self.gid} DIR:{download.dir} ") + download_dict[listener.uid] = AriaDownloadStatus(download.gid,listener) + LOGGER.info(f"Started: {download.gid} DIR:{download.dir} ") + - def cancel_download(self): - download = aria2.get_download(self.gid) - 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]) diff --git a/bot/helper/mirror_utils/status_utils/aria_download_status.py b/bot/helper/mirror_utils/status_utils/aria_download_status.py index 502aaf0..c3199f7 100644 --- a/bot/helper/mirror_utils/status_utils/aria_download_status.py +++ b/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 = obj.gid - self.__download = get_download(obj.gid) + self.__gid = gid + self.__download = get_download(self.__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]) + diff --git a/bot/helper/telegram_helper/message_utils.py b/bot/helper/telegram_helper/message_utils.py index 6980a23..9d3eb02 100644 --- a/bot/helper/telegram_helper/message_utils.py +++ b/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 Exception as e: diff --git a/bot/modules/mirror.py b/bot/modules/mirror.py index 0c613c5..f6831fb 100644 --- a/bot/modules/mirror.py +++ b/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 @@ -19,6 +19,9 @@ from bot.helper.telegram_helper.message_utils import * import pathlib import os +ariaDlManager = AriaDownloadHelper() +ariaDlManager.start_listener() + class MirrorListener(listeners.MirrorListeners): @@ -82,7 +85,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)) @@ -185,8 +187,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) - aria.add_download(link, f'{DOWNLOAD_DIR}/{listener.uid}/') + ariaDlManager.add_download(link, f'{DOWNLOAD_DIR}/{listener.uid}/',listener) sendStatusMessage(update, bot) if len(Interval) == 0: Interval.append(setInterval(DOWNLOAD_STATUS_UPDATE_INTERVAL, update_all_messages))