Improve worker resilience

Author: alexander-akhmetov

docs/source/changelog.rst

@@ -8,6 +8,8 @@ Changelog
 
 - Python versions 3.7 and 3.8 are no longer supported.
 - tdlib 1.8.31.
+- Fix: Handle errors during updates processing
+- Fix: Handle queue full errors
 
 [0.18.0] - 2023-03-13
 

telegram/client.py

@@ -559,7 +559,14 @@ def _run_handlers(self, update: Dict[Any, Any]) -> None:
         update_type: str = update.get("@type", "unknown")
 
         for handler in self._update_handlers[update_type]:
-            self._workers_queue.put((handler, update), timeout=self._queue_put_timeout)
+            try:
+                self._workers_queue.put((handler, update), timeout=self._queue_put_timeout)
+            except queue.Full:
+                logger.error(
+                    "Queue is full, update %s dropped for handler %s",
+                    update_type,
+                    handler.__name__ if hasattr(handler, "__name__") else handler,
+                )
 
     def remove_update_handler(self, handler_type: str, func: Callable) -> None:
         """

telegram/worker.py

@@ -41,8 +41,12 @@ def _run_thread(self) -> None:
             except Empty:
                 continue
 
-            handler(update)
-            self._queue.task_done()
+            try:
+                handler(update)
+            except Exception:
+                logger.exception("Handler raised an exception")
+            finally:
+                self._queue.task_done()
 
     def stop(self) -> None:
         self._is_enabled = False

tests/test_telegram_methods.py

@@ -1,11 +1,14 @@
 import pytest
+import queue
+import time
 
 from unittest.mock import patch
 
 from telegram import VERSION
 from telegram.utils import AsyncResult
 from telegram.client import Telegram, MESSAGE_HANDLER_TYPE, AuthorizationState
 from telegram.text import Spoiler
+from telegram.worker import SimpleWorker
 
 API_ID = 1
 API_HASH = "hash"
@@ -90,7 +93,10 @@ def test_parse_text_entities(self, telegram):
 
     def test_send_phone_number_or_bot_token(self, telegram):
         # check that the dunction calls _send_phone_number or _send_bot_token
-        with patch.object(telegram, "_send_phone_number"), patch.object(telegram, "_send_bot_token"):
+        with (
+            patch.object(telegram, "_send_phone_number"),
+            patch.object(telegram, "_send_bot_token"),
+        ):
             telegram.phone = "123"
             telegram.bot_token = None
 
@@ -112,7 +118,9 @@ def test_send_bot_token(self, telegram):
             telegram._send_bot_token()
 
             exp_data = {"@type": "checkAuthenticationBotToken", "token": "some-token"}
-            telegram._send_data.assert_called_once_with(exp_data, result_id="updateAuthorizationState")
+            telegram._send_data.assert_called_once_with(
+                exp_data, result_id="updateAuthorizationState"
+            )
 
     def test_add_message_handler(self, telegram):
         # check that add_message_handler
@@ -201,7 +209,9 @@ def test_get_web_page_instant_view(self, telegram):
         url = "https://yandex.ru/"
         force_full = False
 
-        async_result = telegram.get_web_page_instant_view(url=url, force_full=force_full)
+        async_result = telegram.get_web_page_instant_view(
+            url=url, force_full=force_full
+        )
 
         exp_data = {
             "@type": "getWebPageInstantView",
@@ -263,7 +273,9 @@ def test_get_chats(self, telegram):
         offset_chat_id = 1
         limit = 100
 
-        async_result = telegram.get_chats(offset_order=offset_order, offset_chat_id=offset_chat_id, limit=limit)
+        async_result = telegram.get_chats(
+            offset_order=offset_order, offset_chat_id=offset_chat_id, limit=limit
+        )
 
         exp_data = {
             "@type": "getChats",
@@ -366,7 +378,9 @@ def test_update_async_result_returns_async_result_with_same_id(self, telegram):
         assert async_result == new_async_result
 
     def test_result_id_should_be_replaced_if_it_is_auth_process(self, telegram):
-        async_result = AsyncResult(client=telegram, result_id="updateAuthorizationState")
+        async_result = AsyncResult(
+            client=telegram, result_id="updateAuthorizationState"
+        )
         telegram._results["updateAuthorizationState"] = async_result
 
         update = {
@@ -404,7 +418,9 @@ def _get_async_result(data, request_id=None):
         )
 
         telegram._set_initial_params = lambda: _get_async_result(
-            data={"authorization_state": {"@type": "authorizationStateWaitEncryptionKey"}}
+            data={
+                "authorization_state": {"@type": "authorizationStateWaitEncryptionKey"}
+            }
         )
         telegram._send_encryption_key = lambda: _get_async_result(
             data={"authorization_state": {"@type": "authorizationStateWaitPhoneNumber"}}
@@ -444,7 +460,9 @@ def _get_async_result(data, request_id=None):
         )
 
         telegram._set_initial_params = lambda: _get_async_result(
-            data={"authorization_state": {"@type": "authorizationStateWaitEncryptionKey"}}
+            data={
+                "authorization_state": {"@type": "authorizationStateWaitEncryptionKey"}
+            }
         )
         telegram._send_encryption_key = lambda: _get_async_result(
             data={"authorization_state": {"@type": "authorizationStateWaitPhoneNumber"}}
@@ -453,7 +471,9 @@ def _get_async_result(data, request_id=None):
             data={"authorization_state": {"@type": "authorizationStateWaitCode"}}
         )
         telegram._send_telegram_code = lambda _: _get_async_result(
-            data={"authorization_state": {"@type": "authorizationStateWaitRegistration"}}
+            data={
+                "authorization_state": {"@type": "authorizationStateWaitRegistration"}
+            }
         )
         telegram._register_user = lambda _, __: _get_async_result(
             data={"authorization_state": {"@type": "authorizationStateWaitPassword"}}
@@ -478,3 +498,44 @@ def _get_async_result(data, request_id=None):
         assert state == telegram.authorization_state == AuthorizationState.READY
 
         assert telegram._tdjson.send.call_count == 0
+
+
+class TestWorker:
+    def test_worker_continues_after_handler_exception(self):
+        """Handler exceptions should not kill the worker thread and task_done must be called"""
+        q = queue.Queue()
+        worker = SimpleWorker(q)
+        worker.run()
+
+        results = []
+
+        def failing_handler(update):
+            raise ValueError("Handler failed")
+
+        def working_handler(update):
+            results.append(update)
+
+        # Put two items: one with a failing handler, one with a working handler
+        q.put((failing_handler, {"id": 1}))
+        q.put((working_handler, {"id": 2}))
+
+        # Give the worker time to process both items.
+        # Can't use join when the test fails.
+        time.sleep(1)
+
+        worker.stop()
+
+        assert results == [{"id": 2}]
+
+    def test_run_handlers_continues_on_queue_full(self, telegram):
+        """queue.Full should not crash the listener"""
+
+        def my_handler():
+            pass
+
+        telegram.add_message_handler(my_handler)
+
+        # Mock the queue to always raise queue.Full
+        with patch.object(telegram._workers_queue, "put", side_effect=queue.Full):
+            # This should not raise an exception
+            telegram._run_handlers({"@type": MESSAGE_HANDLER_TYPE})