From e5f997d1f30f22b13f979a685a04a634b86f0c07 Mon Sep 17 00:00:00 2001 From: lzzy12 Date: Fri, 3 Jan 2020 10:41:54 +0530 Subject: [PATCH] Revert "Fixed potential race conditions in GoogleDriveHelper class" This reverts commit 20dae25dc91e554c2ac5953e2c0685669f97d7a1. Signed-off-by: lzzy12 --- .../mirror_utils/upload_utils/gdriveTools.py | 43 +++++++++++-------- bot/modules/cancel_mirror.py | 9 ++-- 2 files changed, 30 insertions(+), 22 deletions(-) diff --git a/bot/helper/mirror_utils/upload_utils/gdriveTools.py b/bot/helper/mirror_utils/upload_utils/gdriveTools.py index 8dc0172..43b9f73 100644 --- a/bot/helper/mirror_utils/upload_utils/gdriveTools.py +++ b/bot/helper/mirror_utils/upload_utils/gdriveTools.py @@ -26,35 +26,40 @@ class GoogleDriveHelper: self.__G_DRIVE_BASE_DOWNLOAD_URL = "https://drive.google.com/uc?id={}&export=download" self.__listener = listener self.__service = self.authorize() - self.__uploaded_bytes = 0 + self._file_uploaded_bytes = 0 + self.uploaded_bytes = 0 self.start_time = 0 + self.total_time = 0 + self._should_update = True self.is_uploading = True self.is_cancelled = False + self.status = None + self.updater = None self.name = name - self.resource_lock = threading.Lock() def cancel(self): - with self.resource_lock: - self.is_cancelled = True - self.is_uploading = False + self.is_cancelled = True + self.is_uploading = False - @property - def uploaded_bytes(self): - with self.resource_lock: - return self.__uploaded_bytes - def speed(self): """ It calculates the average upload speed and returns it in bytes/seconds unit :return: Upload speed in bytes/second """ try: - with self.resource_lock: - return self.__uploaded_bytes / (time.time() - self.start_time) + return self.uploaded_bytes / self.total_time except ZeroDivisionError: return 0 @retry(wait=wait_exponential(multiplier=2, min=3, max=6),stop=stop_after_attempt(5),retry=retry_if_exception_type(HttpError),before=before_log(LOGGER,logging.DEBUG)) + def _on_upload_progress(self): + if self.status is not None: + chunk_size = self.status.total_size * self.status.progress() - self._file_uploaded_bytes + self._file_uploaded_bytes = self.status.total_size * self.status.progress() + LOGGER.info(f'Chunk size: {get_readable_file_size(chunk_size)}') + self.uploaded_bytes += chunk_size + self.total_time += DOWNLOAD_STATUS_UPDATE_INTERVAL + def __upload_empty_file(self, path, file_name, mime_type, parent_id=None): media_body = MediaFileUpload(path, mimetype=mime_type, @@ -110,16 +115,11 @@ class GoogleDriveHelper: drive_file = self.__service.files().create(supportsTeamDrives=True, body=file_metadata, media_body=media_body) response = None - last_uploaded = 0 while response is None: if self.is_cancelled: return None - status, response = drive_file.next_chunk() - if status is not None: - with self.resource_lock: - chunk_size = status.total_size * status.progress() - last_uploaded - last_uploaded = status.total_size * status.progress() - self.__uploaded_bytes += chunk_size + self.status, response = drive_file.next_chunk() + self._file_uploaded_bytes = 0 # Insert new permissions if not IS_TEAM_DRIVE: self.__set_permission(response['id']) @@ -134,6 +134,7 @@ class GoogleDriveHelper: file_path = f"{file_dir}/{file_name}" LOGGER.info("Uploading File: " + file_path) self.start_time = time.time() + self.updater = setInterval(5, self._on_upload_progress) if os.path.isfile(file_path): try: mime_type = get_mime_type(file_path) @@ -146,6 +147,8 @@ class GoogleDriveHelper: LOGGER.error(e.last_attempt.exception()) self.__listener.onUploadError(e) return + finally: + self.updater.cancel() else: try: dir_id = self.create_directory(os.path.basename(os.path.abspath(file_name)), parent_id) @@ -159,6 +162,8 @@ class GoogleDriveHelper: LOGGER.error(e.last_attempt.exception()) self.__listener.onUploadError(e) return + finally: + self.updater.cancel() LOGGER.info(download_dict) self.__listener.onUploadComplete(link) LOGGER.info("Deleting downloaded file/folder..") diff --git a/bot/modules/cancel_mirror.py b/bot/modules/cancel_mirror.py index 3ec567e..b69079d 100644 --- a/bot/modules/cancel_mirror.py +++ b/bot/modules/cancel_mirror.py @@ -6,9 +6,8 @@ from bot.helper.ext_utils.fs_utils import clean_download from bot.helper.telegram_helper.bot_commands import BotCommands from time import sleep - @run_async -def cancel_mirror(bot, update): +def cancel_mirror(bot,update): mirror_message = update.message.reply_to_message with download_dict_lock: keys = download_dict.keys() @@ -32,9 +31,13 @@ def cancel_mirror(bot, update): downloads = aria2.get_downloads(download.followed_by_ids) aria2.pause(downloads) aria2.pause([download]) + + elif dl.status() == "Uploading": + sendMessage("Upload in Progress, Dont Cancel it.",bot,update) + return else: dl._listener.onDownloadError("Download stopped by user!") - sleep(1) # Wait a Second For Aria2 To free Resources. + sleep(1) #Wait a Second For Aria2 To free Resources. clean_download(f'{DOWNLOAD_DIR}{mirror_message.message_id}/')