fix(dispatcher): mark tasks as done after taking elements off updates_queue

d-Rickyy-b · d-Rickyy-b · commit 0e5770105a54 · 2023-07-16T22:58:49.000+02:00
This commit fixes an issue where the `updates_queue.join()` method would never return. When waiting for the `updates_queue` to get emptied, the asyncio queue implementation relies on an internal counter. This counter is incremented on each `put()` but is not automatically decremented on each `get()`. In order to decrement the counter we need to use `updates_queue.task_done()`, as per the [docs](https://docs.python.org/3/library/asyncio-queue.html#asyncio.Queue.task_done). Since this wasn't done in pyrogram yet, there was no proper way to use the `join()` method on the updates_queue and hence wait until all updates got processed. A workaround would be to use a while loop and check if `updates_queue.empty()` returns true. But using `join()` is just a lot cleaner and does not require a loop on our end.
diff --git a/pyrogram/dispatcher.py b/pyrogram/dispatcher.py
@@ -257,3 +257,5 @@ async def handler_worker(self, lock):
                 pass
             except Exception as e:
                 log.exception(e)
+            finally:
+                self.updates_queue.task_done()
