Improve gdrive progress tracking

* Still bugged

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

bot/helper/ext_utils/bot_utils.py

@@ -42,7 +42,7 @@ def get_download_status_list():
 
 def get_progress_bar_string(status):
     if status.status() == MirrorStatus.STATUS_UPLOADING:
-        completed = status.uploaded_bytes / 8
+        completed = status.upload_helper.uploaded_bytes / 8
     else:
         completed = status.download().completed_length / 8
     total = status.download().total_length / 8

bot/helper/mirror_utils/download_status.py

@@ -14,8 +14,7 @@ class DownloadStatus:
         self.__gid = gid
         self.__download = get_download(gid)
         self.__uid = message_id
-        self.uploaded_bytes = 0
+        self.upload_helper = None
-        self.upload_time = 0
 
     def __update(self):
         self.__download = get_download(self.__gid)
@@ -26,12 +25,12 @@ class DownloadStatus:
         :return: returns progress in percentage
         """
         self.__update()
-        if self.status() == MirrorStatus.STATUS_UPLOADING:
+        if self.upload_helper is not None:
             return f'{round(self.upload_progress(), 2)}%'
         return self.__download.progress_string()
 
     def upload_progress(self):
-        return self.uploaded_bytes / self.download().total_length * 100
+        return self.upload_helper.uploaded_bytes / self.download().total_length * 100
 
     def __size(self):
         """
@@ -42,17 +41,13 @@ class DownloadStatus:
 
     def __upload_speed(self):
         """
-        Calculates upload speed in bytes/second
         :return: Upload speed in Bytes/Seconds
         """
-        try:
+        return self.upload_helper.speed()
-            return self.uploaded_bytes / self.upload_time
-        except ZeroDivisionError:
-            return 0
 
     def speed(self):
         self.__update()
-        if self.status() == MirrorStatus.STATUS_UPLOADING:
+        if self.upload_helper is not None:
             return f'{get_readable_file_size(self.__upload_speed())}/s'
         return self.__download.download_speed_string()
 
@@ -69,9 +64,9 @@ class DownloadStatus:
 
     def eta(self):
         self.__update()
-        if self.status() == MirrorStatus.STATUS_UPLOADING:
+        if self.upload_helper is not None:
             try:
-                seconds = round((self.__size() - self.uploaded_bytes) / self.__upload_speed(), 2)
+                seconds = (self.__size() - self.upload_helper.uploaded_bytes) / self.__upload_speed()
                 return f'{get_readable_time(seconds)}'
             except ZeroDivisionError:
                 return '-'
@@ -79,20 +74,17 @@ class DownloadStatus:
 
     def status(self):
         self.__update()
-        status = None
         if self.is_archiving:
             status = MirrorStatus.STATUS_ARCHIVING
-        if self.__download.is_waiting:
+        elif self.download().is_waiting:
             status = MirrorStatus.STATUS_WAITING
         elif self.download().is_paused:
             status = MirrorStatus.STATUS_CANCELLED
-        elif self.__download.is_complete:
+        elif self.upload_helper is not None:
-            # If download exists and is complete the it must be uploading
-            # otherwise the gid would have been removed from the download_list
             status = MirrorStatus.STATUS_UPLOADING
         elif self.__download.has_failed:
             status = MirrorStatus.STATUS_FAILED
-        elif self.__download.is_active:
+        else:
             status = MirrorStatus.STATUS_DOWNLOADING
         return status
 

bot/helper/mirror_utils/gdriveTools.py

@@ -25,33 +25,49 @@ class GoogleDriveHelper:
         self.__G_DRIVE_BASE_DOWNLOAD_URL = "https://drive.google.com/uc?id={}&export=download"
         self.__listener = listener
         self.__service = self.authorize()
-        self.uploadedBytes = 0
+        self._file_uploaded_bytes = 0
+        self.uploaded_bytes = 0
         self.start_time = 0
+        self.total_time = 0
         self._should_update = True
-        self._do_progress_update = True
+        self.is_uploading = True
+        self.is_cancelled = False
         self.status = None
 
+    def cancel(self):
+        self.is_cancelled = True
+        self.is_uploading = False
+
+    def speed(self):
+        """
+        It calculates the average upload speed and returns it in bytes/seconds unit
+        :return: Upload speed in bytes/second
+        """
+        try:
+            return self.uploaded_bytes / self.total_time
+        except ZeroDivisionError:
+            return 0
+
     def _on_upload_progress(self):
-        while self._do_progress_update:
+        while self.is_uploading:
             if self.status is not None:
-                time_lapsed = time.time() - self.start_time
+                chunk_size = self.status.total_size * self.status.progress() - self._file_uploaded_bytes
-                # Update the message only if status is not null and loop_count is multiple of 50
+                self._file_uploaded_bytes = self.status.total_size * self.status.progress()
-                chunk_size = self.status.total_size * self.status.progress() - self.uploadedBytes
+                LOGGER.info(f'Chunk size: {get_readable_file_size(chunk_size)}')
-                self.uploadedBytes = self.status.total_size * self.status.progress()
+                self.uploaded_bytes += chunk_size
-
+                self.total_time += DOWNLOAD_STATUS_UPDATE_INTERVAL
-                # LOGGER.info(f'{file_name}: {status.progress() * 100}')
-
-                with download_dict_lock:
-                    download_dict[self.__listener.uid].uploaded_bytes += chunk_size
-                    download_dict[self.__listener.uid].upload_time = time_lapsed
 
                 if self._should_update:
                     try:
+                        LOGGER.info('Updating messages')
                         _list = get_download_status_list()
                         index = get_download_index(_list, get_download(self.__listener.message.message_id).gid)
                         self.__listener.onUploadProgress(_list, index)
-                    except KillThreadException:
+                    except KillThreadException as e:
+                        LOGGER.info(f'Stopped calling onDownloadProgress(): {str(e)}')
                         self._should_update = False
+            else:
+                LOGGER.info('status: None')
             time.sleep(DOWNLOAD_STATUS_UPDATE_INTERVAL)
 
     def upload_file(self, file_path, file_name, mime_type, parent_id):
@@ -78,10 +94,12 @@ class GoogleDriveHelper:
         # Insert a file
         drive_file = self.__service.files().create(body=file_metadata, media_body=media_body)
         response = None
-        threading.Thread(target=self._on_upload_progress).start()
+        file_start_time = time.time()
         while response is None:
+            if self.is_cancelled:
+                return None
             self.status, response = drive_file.next_chunk()
-        self._do_progress_update = False
+        self._file_uploaded_bytes = 0
         # Insert new permissions
         self.__service.permissions().create(fileId=response['id'], body=permissions).execute()
         # Define file instance and get url for download
@@ -97,22 +115,28 @@ class GoogleDriveHelper:
         file_path = f"{file_dir}/{file_name}"
         LOGGER.info("Uploading File: " + file_name)
         self.start_time = time.time()
+        threading.Thread(target=self._on_upload_progress).start()
         if os.path.isfile(file_path):
             try:
                 mime_type = get_mime_type(file_path)
-                g_drive_link = self.upload_file(file_path, file_name, mime_type, parent_id)
+                link = self.upload_file(file_path, file_name, mime_type, parent_id)
+                if link is None:
+                    raise Exception('Upload has been manually cancelled')
                 LOGGER.info("Uploaded To G-Drive: " + file_path)
-                link = g_drive_link
             except Exception as e:
                 LOGGER.error(str(e))
                 e_str = str(e).replace('<', '')
                 e_str = e_str.replace('>', '')
                 self.__listener.onUploadError(e_str, _list, index)
                 return
+            finally:
+                self.is_uploading = False
         else:
             try:
                 dir_id = self.create_directory(os.path.basename(os.path.abspath(file_name)), parent_id)
-                self.upload_dir(file_path, dir_id)
+                result = self.upload_dir(file_path, dir_id)
+                if result is None:
+                    raise Exception('Upload has been manually cancelled!')
                 LOGGER.info("Uploaded To G-Drive: " + file_name)
                 link = f"https://drive.google.com/folderview?id={dir_id}"
             except Exception as e:
@@ -121,6 +145,8 @@ class GoogleDriveHelper:
                 e_str = e_str.replace('>', '')
                 self.__listener.onUploadError(e_str, _list, index)
                 return
+            finally:
+                self.is_uploading = False
         LOGGER.info(download_dict)
         self.__listener.onUploadComplete(link, _list, index)
         LOGGER.info("Deleting downloaded file/folder..")
@@ -152,6 +178,8 @@ class GoogleDriveHelper:
         new_id = None
         for item in list_dirs:
             current_file_name = os.path.join(input_directory, item)
+            if self.is_cancelled:
+                return None
             if os.path.isdir(current_file_name):
                 current_dir_id = self.create_directory(item, parent_id)
                 new_id = self.upload_dir(current_file_name, current_dir_id)

bot/modules/cancel_mirror.py

@@ -24,6 +24,10 @@ def cancel_mirror(update: Update, context):
         downloads = aria2.get_downloads(download.followed_by_ids)
         aria2.pause(downloads)
     aria2.pause([download])
+    with download_dict_lock:
+        upload_helper = download_dict[mirror_message.message_id].upload_helper
+    if upload_helper is not None:
+        upload_helper.cancel()
     clean_download(f'{DOWNLOAD_DIR}{mirror_message.message_id}/')
 
 

bot/modules/mirror.py

@@ -1,5 +1,5 @@
 from telegram.ext import CommandHandler, run_async
-from telegram.error import BadRequest
+from telegram.error import BadRequest, TimedOut
 from bot.helper.mirror_utils import download_tools, gdriveTools, listeners
 from bot import LOGGER, dispatcher, DOWNLOAD_DIR
 from bot.helper.ext_utils import fs_utils, bot_utils
@@ -49,6 +49,8 @@ class MirrorListener(listeners.MirrorListeners):
             download_dict[self.uid].is_archiving = False
             download_dict[self.uid].upload_name = name
         gdrive = gdriveTools.GoogleDriveHelper(self)
+        with download_dict_lock:
+            download_dict[self.uid].upload_helper = gdrive
         gdrive.upload(name)
 
     def onDownloadError(self, error, progress_status_list: list, index: int):
@@ -109,8 +111,10 @@ class MirrorListener(listeners.MirrorListeners):
         msg = get_readable_message(progress)
         try:
             editMessage(msg, self.context, self.reply_message)
-        except BadRequest:
+        except BadRequest as e:
-            raise KillThreadException('Message deleted. Do not call this method from the thread')
+            raise KillThreadException(str(e))
+        except TimedOut:
+            pass
 
 
 def _mirror(update, context, isTar=False):