locust-cloud 1.13.0__py3-none-any.whl → 1.14.1__py3-none-any.whl

locust_cloud/cloud.py

@@ -167,8 +167,6 @@ logging.basicConfig(
 )
 logger = logging.getLogger(__name__)
 # Restore log level for other libs. Yes, this can be done more nicely
-logging.getLogger("botocore").setLevel(logging.INFO)
-logging.getLogger("boto3").setLevel(logging.INFO)
 logging.getLogger("requests").setLevel(logging.INFO)
 logging.getLogger("urllib3").setLevel(logging.INFO)
 
@@ -368,6 +366,204 @@ class ApiSession(requests.Session):
         return super().request(method, f"{self.api_url}{url}", *args, **kwargs)
 
 
+class SessionMismatchError(Exception):
+    pass
+
+
+class WebsocketTimeout(Exception):
+    pass
+
+
+class Websocket:
+    def __init__(self) -> None:
+        """
+        This class was created to encapsulate all the logic involved in the websocket implementation.
+        The behaviour of the socketio client once a connection has been established
+        is to try to reconnect forever if the connection is lost.
+        The way this can be canceled is by setting the _reconnect_abort (threading.Event) on the client
+        in which case it will simply proceed with shutting down without giving any indication of an error.
+        This class handles timeouts for connection attempts as well as some logic around when the
+        socket can be shut down. See descriptions on the methods for further details.
+        """
+        self.__shutdown_allowed = threading.Event()
+        self.__timeout_on_disconnect = True
+        self.initial_connect_timeout = 120
+        self.reconnect_timeout = 10
+        self.wait_timeout = 0
+        self.exception: None | Exception = None
+
+        self.sio = socketio.Client(handle_sigint=False)
+        self.sio._reconnect_abort = threading.Event()
+        # 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.
+
+        self.sio.on("connect", self.__on_connect)
+        self.sio.on("disconnect", self.__on_disconnect)
+        self.sio.on("connect_error", self.__on_connect_error)
+        self.sio.on("events", self.__on_events)
+
+        self.__processed_events: set[int] = set()
+
+    def __set_connection_timeout(self, timeout) -> None:
+        """
+        Start a threading.Timer that will set the threading.Event on the socketio client
+        that aborts any further attempts to reconnect, sets an exception on the websocket
+        that will be raised from the wait method and the threading.Event __shutdown_allowed
+        on the websocket that tells the wait method that it should stop blocking.
+        """
+
+        def _timeout():
+            logger.debug(f"Websocket connection timed out after {timeout} seconds")
+            self.sio._reconnect_abort.set()
+            self.exception = WebsocketTimeout("Timed out connecting to locust master")
+            self.__shutdown_allowed.set()
+
+        self.__connect_timeout_timer = threading.Timer(timeout, _timeout)
+        self.__connect_timeout_timer.daemon = True
+        logger.debug(f"Setting websocket connection timeout to {timeout} seconds")
+        self.__connect_timeout_timer.start()
+
+    def connect(self, url, *, auth) -> None:
+        """
+        Send along retry=True when initiating the socketio client connection
+        to make it use it's builtin logic for retrying failed connections that
+        is usually used for reconnections. This will retry forever.
+        When connecting start a timer to trigger disabling the retry logic and
+        raise a WebsocketTimeout exception.
+        """
+        ws_connection_info = urllib.parse.urlparse(url)
+        self.__set_connection_timeout(self.initial_connect_timeout)
+        try:
+            self.sio.connect(
+                f"{ws_connection_info.scheme}://{ws_connection_info.netloc}",
+                auth=auth,
+                retry=True,
+                **{"socketio_path": ws_connection_info.path} if ws_connection_info.path else {},
+            )
+        except socketio.exceptions.ConnectionError:
+            if self.exception:
+                raise self.exception
+
+            raise
+
+    def shutdown(self) -> None:
+        """
+        When shutting down the socketio client a disconnect event will fire.
+        Before doing so disable the behaviour of starting a threading.Timer
+        to handle timeouts on attempts to reconnect since no further such attempts
+        will be made.
+        If such a timer is already running, cancel it since the client is being shutdown.
+        """
+        self.__timeout_on_disconnect = False
+        if hasattr(self, "__connect_timeout_timer"):
+            self.__connect_timeout_timer.cancel()
+        self.sio.shutdown()
+
+    def wait(self, timeout=False) -> bool:
+        """
+        Block until the threading.Event __shutdown_allowed is set, with a timeout if indicated.
+        If an exception has been set on the websocket (from a connection timeout timer or the
+        __on_connect_error method), raise it.
+        """
+        timeout = self.wait_timeout if timeout else None
+        logger.debug(f"Waiting for shutdown for {str(timeout)+'s' if timeout else 'ever'}")
+        res = self.__shutdown_allowed.wait(timeout)
+        if self.exception:
+            raise self.exception
+        return res
+
+    def __on_connect(self) -> None:
+        """
+        This gets events whenever a connection is successfully established.
+        When this happens, cancel the running threading.Timer that would
+        abort reconnect attempts and raise a WebsocketTimeout exception.
+        The wait_timeout is originally set to zero when creating the websocket
+        but once a connection has been established this is raised to ensure
+        that the server is given the chance to send all the logs and an
+        official shutdown event.
+        """
+        self.__connect_timeout_timer.cancel()
+        self.wait_timeout = 90
+        logger.debug("Websocket connected")
+
+    def __on_disconnect(self) -> None:
+        """
+        This gets events whenever a connection is lost.
+        The socketio client will try to reconnect forever so,
+        unless the behaviour has been disabled, a threading.Timer
+        is started that will abort reconnect attempts and raise a
+        WebsocketTimeout exception.
+        """
+        if self.__timeout_on_disconnect:
+            self.__set_connection_timeout(self.reconnect_timeout)
+        logger.debug("Websocket disconnected")
+
+    def __on_events(self, data):
+        """
+        This gets events explicitly sent by the websocket server.
+        This will either be messages to print on stdout/stderr or
+        an indication that the CLI can shut down in which case the
+        threading.Event __shutdown_allowed gets set on the websocket
+        that tells the wait method that it should stop blocking.
+        """
+        shutdown = False
+        shutdown_message = ""
+
+        if data["id"] in self.__processed_events:
+            logger.debug(f"Got duplicate data on websocket, id {data['id']}")
+            return
+
+        self.__processed_events.add(data["id"])
+
+        for event in data["events"]:
+            type = event["type"]
+
+            if type == "shutdown":
+                shutdown = True
+                shutdown_message = event["message"]
+            elif type == "stdout":
+                sys.stdout.write(event["message"])
+            elif type == "stderr":
+                sys.stderr.write(event["message"])
+            else:
+                raise Exception("Unexpected event type")
+
+        if shutdown:
+            logger.debug("Got shutdown from locust master")
+            if shutdown_message:
+                print(shutdown_message)
+
+            self.__shutdown_allowed.set()
+
+    def __on_connect_error(self, data) -> None:
+        """
+        This gets events whenever there's an error during connection attempts.
+        The specific case that is handled below is triggered when the connection
+        is made with the auth parameter not matching the session ID on the server.
+        If this error occurs it's because the connection is attempted towards an
+        instance of locust not started by this CLI.
+
+        In that case:
+        Cancel the running threading.Timer that would abort reconnect attempts
+        and raise a WebsocketTimeout exception.
+        Set an exception on the websocket that will be raised from the wait method.
+        Cancel further reconnect attempts.
+        Set the threading.Event __shutdown_allowed on the websocket that tells the
+        wait method that it should stop blocking.
+        """
+        # Do nothing if it's not the specific case we know how to deal with
+        if not (isinstance(data, dict) and data.get("message") == "Session mismatch"):
+            return
+
+        self.__connect_timeout_timer.cancel()
+        self.exception = SessionMismatchError(
+            "The session from this run of locust-cloud did not match the one on the server"
+        )
+        self.sio._reconnect_abort.set()
+        self.__shutdown_allowed.set()
+
+
 def main() -> None:
     if options.version:
         print(f"locust-cloud version {__version__}")
@@ -384,12 +580,13 @@ def main() -> None:
         sys.exit()
 
     session = ApiSession()
+    websocket = Websocket()
 
-    try:
-        if options.delete:
-            delete(session)
-            return
+    if options.delete:
+        delete(session)
+        sys.exit()
 
+    try:
         try:
             with open(options.locustfile, "rb") as f:
                 locustfile_data = base64.b64encode(gzip.compress(f.read())).decode()
@@ -448,9 +645,7 @@ def main() -> None:
             logger.error(f"Failed to deploy the load generators: {e}")
             sys.exit(1)
 
-        if response.status_code == 200:
-            log_ws_url = response.json()["log_ws_url"]
-        else:
+        if response.status_code != 200:
             try:
                 logger.error(f"{response.json()['Message']} (HTTP {response.status_code}/{response.reason})")
             except Exception:
@@ -459,66 +654,33 @@ def main() -> None:
                 )
             sys.exit(1)
 
-    except KeyboardInterrupt:
-        # TODO: This would potentially leave a deployment running, combine with try-catch below?
-        logger.debug("Interrupted by user")
-        sys.exit(0)
-
-    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)
-    sio = socketio.Client(handle_sigint=False)
-
-    try:
-        ws_connection_info = urllib.parse.urlparse(log_ws_url)
-
-        @sio.event
-        def connect():
-            shutdown_allowed.clear()
-            connect_timeout.cancel()
-            logger.debug("Websocket connection established, switching to Locust logs")
-
-        @sio.event
-        def disconnect():
-            logger.debug("Websocket disconnected")
-
-        @sio.event
-        def stderr(message):
-            sys.stderr.write(message)
-
-        @sio.event
-        def stdout(message):
-            sys.stdout.write(message)
+        log_ws_url = response.json()["log_ws_url"]
+        session_id = response.json()["session_id"]
+        logger.debug(f"Session ID is {session_id}")
 
-        @sio.event
-        def shutdown(message):
-            logger.debug("Got shutdown from locust master")
-            if message:
-                print(message)
-
-            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.info("Waiting for load generators to be ready...")
+        websocket.connect(
+            log_ws_url,
+            auth=session_id,
         )
-        logger.debug("Waiting for shutdown")
-        shutdown_allowed.wait()
+        websocket.wait()
 
     except KeyboardInterrupt:
         logger.debug("Interrupted by user")
         delete(session)
-        shutdown_allowed.wait(timeout=90)
+        try:
+            websocket.wait(timeout=True)
+        except (WebsocketTimeout, SessionMismatchError) as e:
+            logger.error(str(e))
+            sys.exit(1)
+    except WebsocketTimeout as e:
+        logger.error(str(e))
+        delete(session)
+        sys.exit(1)
+    except SessionMismatchError as e:
+        # In this case we do not trigger the teardown since the running instance is not ours
+        logger.error(str(e))
+        sys.exit(1)
     except Exception as e:
         logger.exception(e)
         delete(session)
@@ -526,7 +688,8 @@ def main() -> None:
     else:
         delete(session)
     finally:
-        sio.shutdown()
+        logger.debug("Shutting down websocket")
+        websocket.shutdown()
 
 
 def delete(session):
@@ -545,8 +708,6 @@ def delete(session):
     except Exception as e:
         logger.error(f"Could not automatically tear down Locust Cloud: {e.__class__.__name__}:{e}")
 
-    logger.info("Done! ✨")  # FIXME: Should probably not say it's done since at this point it could still be running
-
 
 if __name__ == "__main__":
     main()

{locust_cloud-1.13.0.dist-info → locust_cloud-1.14.1.dist-info}/METADATA

@@ -1,14 +1,14 @@
 Metadata-Version: 2.4
 Name: locust-cloud
-Version: 1.13.0
+Version: 1.14.1
 Summary: Locust Cloud
 Project-URL: Homepage, https://locust.cloud
 Requires-Python: >=3.11
-Requires-Dist: boto3==1.34.125
 Requires-Dist: configargparse==1.7
 Requires-Dist: platformdirs>=4.3.6
 Requires-Dist: pyjwt>=2.0
 Requires-Dist: python-socketio[client]==5.11.4
+Requires-Dist: requests==2.32.3
 Description-Content-Type: text/markdown
 
 # Locust Cloud

locust_cloud-1.14.1.dist-info/RECORD

@@ -0,0 +1,5 @@
+locust_cloud/cloud.py,sha256=2zRfdvnrSqEBzCspEtL3H2h0H-Y2gIHs4GmZUX-P9bw,25357
+locust_cloud-1.14.1.dist-info/METADATA,sha256=wJHNidr7_UQ39nTxdPml1Dd7lcL-X4vfnIxbjGMn9i0,497
+locust_cloud-1.14.1.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+locust_cloud-1.14.1.dist-info/entry_points.txt,sha256=PGyAb4e3aTsGS3N3VGShDl6VzJaXy7QwsEgsLOC7V00,57
+locust_cloud-1.14.1.dist-info/RECORD,,

locust_cloud-1.13.0.dist-info/RECORD

@@ -1,5 +0,0 @@
-locust_cloud/cloud.py,sha256=r0lWoP1QiT1Klkrecsr6LyPD_l1GMW2BFGRCRVkLowM,18124
-locust_cloud-1.13.0.dist-info/METADATA,sha256=STWxyZOw0qE003r-Pf-8RU3KPFK3yjH_NsSo5ffoHh0,496
-locust_cloud-1.13.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-locust_cloud-1.13.0.dist-info/entry_points.txt,sha256=PGyAb4e3aTsGS3N3VGShDl6VzJaXy7QwsEgsLOC7V00,57
-locust_cloud-1.13.0.dist-info/RECORD,,