From 98615146a67c7370007b781e2358d48055a1ab44 Mon Sep 17 00:00:00 2001
From: Aljoscha Rittner <aljoscha.rittner@ipoint-systems.de>
Date: Sat, 25 Jun 2022 22:17:17 +0200
Subject: [PATCH] Resilient stream.close handling and early close() while retry
 sleeps

fixes #212
---
 mastodon/Mastodon.py | 35 +++++++++++++++++++++++++++--------
 1 file changed, 27 insertions(+), 8 deletions(-)

diff --git a/mastodon/Mastodon.py b/mastodon/Mastodon.py
index 0922a6f..64450b6 100644
--- a/mastodon/Mastodon.py
+++ b/mastodon/Mastodon.py
@@ -3610,7 +3610,8 @@ class Mastodon:
                 
             def close(self):
                 self.closed = True
-                self.connection.close()
+                if not self.connection is None:
+                    self.connection.close()
 
             def is_alive(self):
                 return self._thread.is_alive()
@@ -3621,6 +3622,14 @@ class Mastodon:
                 else:
                     return True
 
+            def _sleep_attentive(self):
+                if self._thread != threading.current_thread():
+                    raise RuntimeError ("Illegal call from outside the stream_handle thread")
+                time_remaining = self.reconnect_async_wait_sec
+                while time_remaining>0 and not self.closed:
+                    time.sleep(0.5)
+                    time_remaining -= 0.5
+
             def _threadproc(self):
                 self._thread = threading.current_thread()
                 
@@ -3642,16 +3651,26 @@ class Mastodon:
                         self.reconnecting = True
                         connect_success = False
                         while not connect_success:
-                            connect_success = True
+                            if self.closed:
+                                # Someone from outside stopped the streaming
+                                self.running = False
+                                break
                             try:
-                                self.connection = self.connect_func()
-                                if self.connection.status_code != 200:
-                                    time.sleep(self.reconnect_async_wait_sec)
-                                    connect_success = False
-                                    exception = MastodonNetworkError("Could not connect to server.")
+                                the_connection = self.connect_func()
+                                if the_connection.status_code != 200:
+                                    exception = MastodonNetworkError(f"Could not connect to server. "
+                                                                     f"HTTP status: {the_connection.status_code}")
                                     listener.on_abort(exception)
+                                    self._sleep_attentive()
+                                if self.closed:
+                                    # Here we have maybe a rare race condition. Exactly on connect, someone
+                                    # stopped the streaming before. We close the previous established connection:
+                                    the_connection.close()
+                                else:
+                                    self.connection = the_connection
+                                    connect_success = True
                             except:
-                                time.sleep(self.reconnect_async_wait_sec)
+                                self._sleep_attentive()
                                 connect_success = False
                         self.reconnecting = False
                     else: