Fix dead locks due to while loop in multiple threads trying to access a shared state

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

bot/__main__.py

@@ -1,6 +1,6 @@
 from telegram.ext import CommandHandler, run_async
 from bot import dispatcher, LOGGER, updater, aria2
-import bot.mirror,bot.list
+import bot.mirror, bot.list, bot.mirror_status
 from bot.helper import fs_utils
 import signal
 

bot/helper/bot_utils.py

@@ -1,3 +1,4 @@
+
 from bot import download_dict
 
 
@@ -35,3 +36,10 @@ def get_readable_message(progress_list: list = download_dict.values()):
                                            status.progress(), status.size(),
                                            status.speed(), status.eta())
     return msg
+
+
+# Custom Exception class for killing thread as soon as they aren't needed
+class KillThreadException(Exception):
+    def __init__(self, message, error=None):
+        super().__init__(message)
+        self.error = error

bot/helper/download_tools.py

@@ -55,7 +55,10 @@ class DownloadHelper:
                 # Check every few seconds
                 sleep(DOWNLOAD_STATUS_UPDATE_INTERVAL)
                 new_gid = self.__get_followed_download_gid()
-                self.__listener.onDownloadProgress(get_download_status_list(), index)
+                try:
+                    self.__listener.onDownloadProgress(get_download_status_list(), index)
+                except KillThreadException:
+                    return
             download_dict[self.__listener.update.update_id] = DownloadStatus(new_gid, self.__listener.update.update_id)
 
         # Start tracking the actual download
@@ -70,7 +73,10 @@ class DownloadHelper:
             # TODO: Find a better way to differentiate between 2 list of objects
             progress_str_list = get_download_str()
             if progress_str_list != previous:
-                self.__listener.onDownloadProgress(status_list, index)
+                try:
+                    self.__listener.onDownloadProgress(status_list, index)
+                except KillThreadException:
+                    return
                 previous = progress_str_list
             download = self.__get_download()
             sleep(DOWNLOAD_STATUS_UPDATE_INTERVAL)

bot/helper/message_utils.py

@@ -1,7 +1,7 @@
 from telegram.error import BadRequest
 from telegram.message import Message
 from telegram.update import Update
-
+from bot import LOGGER
 
 def sendMessage(text: str, context, update: Update):
     return context.bot.send_message(update.message.chat_id,
@@ -10,12 +10,9 @@ def sendMessage(text: str, context, update: Update):
 
 
 def editMessage(text: str, context, message: Message):
-    try:
         context.bot.edit_message_text(text=text, message_id=message.message_id,
                                       chat_id=message.chat.id,
                                       parse_mode='HTMl')
-    except BadRequest:
-        pass
 
 
 def deleteMessage(context, message: Message):

bot/mirror.py

@@ -5,7 +5,7 @@ from bot.helper import fs_utils
 from bot import download_dict, status_reply_dict, DOWNLOAD_STATUS_UPDATE_INTERVAL
 from bot.helper.message_utils import *
 from time import sleep
-from bot.helper.bot_utils import get_readable_message
+from bot.helper.bot_utils import get_readable_message, KillThreadException
 LOGGER.info('mirror.py')
 
 
@@ -19,7 +19,10 @@ class MirrorListener(listeners.MirrorListeners):
     def onDownloadProgress(self, progress_status_list: list, index: int):
         msg = get_readable_message(progress_status_list)
         LOGGER.info("Editing message")
-        editMessage(msg, self.context, self.reply_message)
+        try:
+            editMessage(msg, self.context, self.reply_message)
+        except BadRequest:
+            raise KillThreadException('Message deleted. Terminate thread')
 
     def onDownloadComplete(self, progress_status_list, index: int):
         msg = get_readable_message(progress_status_list)
@@ -70,28 +73,6 @@ def mirror(update, context):
     aria.add_download(link)
 
 
-@run_async
-def mirror_status(update: Update, context):
-    message = get_readable_message()
-    if len(message) == 0:
-        message = "No active downloads"
-        sendMessage(message, context, update)
-        return
-    index = update.effective_chat.id
-    if index in status_reply_dict.keys():
-        deleteMessage(context, status_reply_dict[index])
-        del status_reply_dict[index]
-    while True:
-        message = get_readable_message()
-        if index in status_reply_dict.keys():
-            if len(message) == 0:
-                message = "No active downloads"
-                editMessage(message, context, status_reply_dict[index])
-                break
-            editMessage(message, context, status_reply_dict[index])
-        else:
-            status_reply_dict[index] = sendMessage(message, context, update)
-        sleep(DOWNLOAD_STATUS_UPDATE_INTERVAL)
 
 
 @run_async
@@ -100,6 +81,4 @@ def cancel_mirror(update, context):
 
 
 mirror_handler = CommandHandler('mirror', mirror)
-mirror_status_handler = CommandHandler('status', mirror_status)
 dispatcher.add_handler(mirror_handler)
-dispatcher.add_handler(mirror_status_handler)

bot/mirror_status.py

@@ -0,0 +1,43 @@
+from telegram.ext import CommandHandler, run_async
+from bot import dispatcher, status_reply_dict, DOWNLOAD_STATUS_UPDATE_INTERVAL
+from bot.helper.message_utils import *
+from time import sleep
+from bot.helper.bot_utils import get_readable_message
+
+
+@run_async
+def mirror_status(update: Update, context):
+    message = get_readable_message()
+    if len(message) == 0:
+        message = "No active downloads"
+        sendMessage(message, context, update)
+        return
+    index = update.effective_chat.id
+    if index in status_reply_dict.keys():
+        deleteMessage(context, status_reply_dict[index])
+        del status_reply_dict[index]
+    kill_thread = False
+    while len(message) != 0:
+        message = get_readable_message()
+        if index in status_reply_dict.keys():
+            if len(message) == 0:
+                message = "No active downloads"
+                editMessage(message, context, status_reply_dict[index])
+                break
+            try:
+                editMessage(message, context, status_reply_dict[index])
+            except BadRequest:
+                break
+        else:
+            # If the loop returns here 2nd time, it means the message
+            # has been replaced by a new message due to a second /status command in the chat.
+            # So we kill the thread by simply breaking the loop
+            if kill_thread:
+                break
+            status_reply_dict[index] = sendMessage(message, context, update)
+            kill_thread = True
+        sleep(DOWNLOAD_STATUS_UPDATE_INTERVAL)
+
+
+mirror_status_handler = CommandHandler('status', mirror_status)
+dispatcher.add_handler(mirror_status_handler)