From 5128f6310bd3db46833440b6479542e7c8ce8cfe Mon Sep 17 00:00:00 2001
From: Dan <14043624+delivrance@users.noreply.github.com>
Date: Sat, 23 Dec 2017 23:43:20 +0100
Subject: [PATCH] Wait for threads to stop when closing the session

---
 pyrogram/session/session.py | 18 ++++++++----------
 1 file changed, 8 insertions(+), 10 deletions(-)

diff --git a/pyrogram/session/session.py b/pyrogram/session/session.py
index aad0928a..d27c6a83 100644
--- a/pyrogram/session/session.py
+++ b/pyrogram/session/session.py
@@ -124,11 +124,6 @@ class Session:
                 self.current_salt = FutureSalt(0, 0, self._send(functions.Ping(0)).new_server_salt)
                 self.current_salt = self._send(functions.GetFutureSalts(1)).salts[0]
 
-                if self.next_salt_thread is not None:
-                    self.next_salt_thread.join()
-
-                self.next_salt_thread_event.clear()
-
                 self.next_salt_thread = Thread(target=self.next_salt, name="NextSaltThread")
                 self.next_salt_thread.start()
 
@@ -147,11 +142,6 @@ class Session:
                         )
                     ).text
 
-                if self.ping_thread is not None:
-                    self.ping_thread.join()
-
-                self.ping_thread_event.clear()
-
                 self.ping_thread = Thread(target=self.ping, name="PingThread")
                 self.ping_thread.start()
 
@@ -170,8 +160,16 @@ class Session:
 
     def stop(self):
         self.is_connected.clear()
+
         self.ping_thread_event.set()
         self.next_salt_thread_event.set()
+
+        self.ping_thread.join()
+        self.next_salt_thread.join()
+
+        self.ping_thread_event.clear()
+        self.next_salt_thread_event.clear()
+
         self.connection.close()
 
         for i in range(self.WORKERS):