From 17259d1b7b7ab17a40b0f7a0130543aea2f221c6 Mon Sep 17 00:00:00 2001
From: Dan <14043624+delivrance@users.noreply.github.com>
Date: Fri, 13 Apr 2018 19:09:00 +0200
Subject: [PATCH] Join threads before closing the connection

---
 pyrogram/client/client.py                | 32 ++++++++++++++++++++----
 pyrogram/client/dispatcher/dispatcher.py | 16 +++++++++---
 2 files changed, 39 insertions(+), 9 deletions(-)

diff --git a/pyrogram/client/client.py b/pyrogram/client/client.py
index 0950bd7b..51f63896 100644
--- a/pyrogram/client/client.py
+++ b/pyrogram/client/client.py
@@ -191,7 +191,9 @@ class Client:
         self.is_idle = None
 
         self.updates_queue = Queue()
+        self.updates_workers_list = []
         self.download_queue = Queue()
+        self.download_workers_list = []
 
         self.dispatcher = Dispatcher(self, workers)
         self.update_handler = None
@@ -301,10 +303,24 @@ class Client:
             self.send(functions.updates.GetState())
 
         for i in range(self.UPDATES_WORKERS):
-            Thread(target=self.updates_worker, name="UpdatesWorker#{}".format(i + 1)).start()
+            self.updates_workers_list.append(
+                Thread(
+                    target=self.updates_worker,
+                    name="UpdatesWorker#{}".format(i + 1)
+                )
+            )
+
+            self.updates_workers_list[-1].start()
 
         for i in range(self.DOWNLOAD_WORKERS):
-            Thread(target=self.download_worker, name="DownloadWorker#{}".format(i + 1)).start()
+            self.download_workers_list.append(
+                Thread(
+                    target=self.download_worker,
+                    name="DownloadWorker#{}".format(i + 1)
+                )
+            )
+
+            self.download_workers_list[-1].start()
 
         self.dispatcher.start()
 
@@ -318,17 +334,23 @@ class Client:
         if not self.is_started:
             raise ConnectionError("Client is already stopped")
 
-        self.is_started = False
-        self.session.stop()
-
         for _ in range(self.UPDATES_WORKERS):
             self.updates_queue.put(None)
 
+        for i in self.updates_workers_list:
+            i.join()
+
         for _ in range(self.DOWNLOAD_WORKERS):
             self.download_queue.put(None)
 
+        for i in self.download_workers_list:
+            i.join()
+
         self.dispatcher.stop()
 
+        self.is_started = False
+        self.session.stop()
+
         Syncer.remove(self)
 
     def authorize_bot(self):
diff --git a/pyrogram/client/dispatcher/dispatcher.py b/pyrogram/client/dispatcher/dispatcher.py
index 2710c4c0..ac61b04f 100644
--- a/pyrogram/client/dispatcher/dispatcher.py
+++ b/pyrogram/client/dispatcher/dispatcher.py
@@ -46,20 +46,28 @@ class Dispatcher:
     def __init__(self, client, workers):
         self.client = client
         self.workers = workers
+        self.workers_list = []
         self.updates = Queue()
         self.groups = OrderedDict()
 
     def start(self):
         for i in range(self.workers):
-            Thread(
-                target=self.update_worker,
-                name="UpdateWorker#{}".format(i + 1)
-            ).start()
+            self.workers_list.append(
+                Thread(
+                    target=self.update_worker,
+                    name="UpdateWorker#{}".format(i + 1)
+                )
+            )
+
+            self.workers_list[-1].start()
 
     def stop(self):
         for _ in range(self.workers):
             self.updates.put(None)
 
+        for i in self.workers_list:
+            i.join()
+
     def add_handler(self, handler, group: int):
         if group not in self.groups:
             self.groups[group] = []