locust-cloud 1.12.3__py3-none-any.whl → 1.12.4__py3-none-any.whl

locust_cloud/cloud.py

@@ -3,7 +3,7 @@ import gzip
 import logging
 import os
 import sys
-import time
+import threading
 import tomllib
 import urllib.parse
 from argparse import Namespace
@@ -314,22 +314,19 @@ def main() -> None:
     logger.debug("Load generators deployed successfully!")
     logger.info("Waiting for pods to be ready...")
 
+    shutdown_allowed = threading.Event()
+    shutdown_allowed.set()
+    reconnect_aborted = threading.Event()
+    connect_timeout = threading.Timer(2 * 60, reconnect_aborted.set)
+
     try:
         ws_connection_info = urllib.parse.urlparse(log_ws_url)
         sio = socketio.Client(handle_sigint=False)
 
-        run = True
-
-        def wait():
-            logger.debug("Waiting for shutdown event")
-            while run:
-                time.sleep(0.1)
-
-            logger.debug("Shutting down websocket connection")
-            sio.shutdown()
-
         @sio.event
         def connect():
+            shutdown_allowed.clear()
+            connect_timeout.cancel()
             logger.debug("Websocket connection established, switching to Locust logs")
 
         @sio.event
@@ -350,33 +347,33 @@ def main() -> None:
             if message:
                 print(message)
 
-            nonlocal run
-            run = False
-
-        for _ in range(5 * 60):  # try for 5 minutes
-            try:
-                sio.connect(
-                    f"{ws_connection_info.scheme}://{ws_connection_info.netloc}", socketio_path=ws_connection_info.path
-                )
-                break
-            except socketio.exceptions.ConnectionError:
-                time.sleep(1)
-
-        else:  # no break
-            raise Exception("Failed to obtain socket connection")
-
-        wait()
+            shutdown_allowed.set()
+
+        # The _reconnect_abort value on the socketio client will be populated with a newly created threading.Event if it's not already set.
+        # There is no way to set this by passing it in the constructor.
+        # This event is the only way to interupt the retry logic when the connection is attempted.
+        sio._reconnect_abort = reconnect_aborted
+        connect_timeout.start()
+        sio.connect(
+            f"{ws_connection_info.scheme}://{ws_connection_info.netloc}",
+            socketio_path=ws_connection_info.path,
+            retry=True,
+        )
+        logger.debug("Waiting for shutdown")
+        shutdown_allowed.wait()
 
     except KeyboardInterrupt:
         logger.debug("Interrupted by user")
         delete(credential_manager)
-        wait()
+        shutdown_allowed.wait(timeout=90)
     except Exception as e:
         logger.exception(e)
         delete(credential_manager)
         sys.exit(1)
     else:
         delete(credential_manager)
+    finally:
+        sio.shutdown()
 
 
 def delete(credential_manager):

{locust_cloud-1.12.3.dist-info → locust_cloud-1.12.4.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: locust-cloud
-Version: 1.12.3
+Version: 1.12.4
 Summary: Locust Cloud
 Project-URL: Homepage, https://locust.cloud
 Requires-Python: >=3.11

{locust_cloud-1.12.3.dist-info → locust_cloud-1.12.4.dist-info}/RECORD

@@ -1,6 +1,6 @@
 locust_cloud/__init__.py,sha256=F_-Mzw8-OpWEwpnhdIL_PWG2aDOYTXkmZSU11G5GH9s,3914
 locust_cloud/auth.py,sha256=FTQZvMadDfq7nwVlKfAITp2rI960YBpr6KO7QnXi8i0,17616
-locust_cloud/cloud.py,sha256=wiVzqcFoOZRnACg4JVrpmXggXxfJcJ9WlUSSj11abMo,13917
+locust_cloud/cloud.py,sha256=xXe7PJGVFi__93nps9wXpJ71Tt6nboYCz_m7s4apsyo,14178
 locust_cloud/credential_manager.py,sha256=MiESqS8IOxncBON04irK1-seF-VNAHHTcELcyxp1ZR8,5571
 locust_cloud/idle_exit.py,sha256=bYD6LeUa0d8Ev2b1f-tiRXVnfyolCoACcDWC8zizNPE,1189
 locust_cloud/socket_logging.py,sha256=ZQ6_a1NrSNxmLI_LWDcmS6p5RVkc8LsEnN177JoRKyg,4461
@@ -19,7 +19,7 @@ locust_cloud/webui/vitest.config.ts,sha256=Rb-lqfCVqW5EAvBCZ-U7LO-tqHbIBc_bQEvfZ
 locust_cloud/webui/yarn.lock,sha256=slIO1G7Eb2H9bFM4suyh24Mr4YGRM4bHjUhJAmSGdG4,266408
 locust_cloud/webui/dist/index.html,sha256=o7Wt2gMtE81eDJcqvbPXQ6ssS1qwfcYuMM3OzcaFzhA,664
 locust_cloud/webui/dist/assets/index-D3YieuNV.js,sha256=EtQvL7--QPGFf3JL_EFSjfdgN2I_uzr6o-hKpTqynXU,2870983
-locust_cloud-1.12.3.dist-info/METADATA,sha256=_1GlLHUWrZzjo7J6c0gMfbcBjDrc9wmaTyoIUsHLyMA,545
-locust_cloud-1.12.3.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-locust_cloud-1.12.3.dist-info/entry_points.txt,sha256=PGyAb4e3aTsGS3N3VGShDl6VzJaXy7QwsEgsLOC7V00,57
-locust_cloud-1.12.3.dist-info/RECORD,,
+locust_cloud-1.12.4.dist-info/METADATA,sha256=DAfqogr5gZYO6p68ilZsvphzWsGdEfCBz7lqjNvwR3w,545
+locust_cloud-1.12.4.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+locust_cloud-1.12.4.dist-info/entry_points.txt,sha256=PGyAb4e3aTsGS3N3VGShDl6VzJaXy7QwsEgsLOC7V00,57
+locust_cloud-1.12.4.dist-info/RECORD,,