locust-cloud 1.11.7__py3-none-any.whl → 1.12.0__py3-none-any.whl

locust_cloud/__init__.py

@@ -3,6 +3,15 @@ import os
 import sys
 
 os.environ["LOCUST_SKIP_MONKEY_PATCH"] = "1"
+
+from locust_cloud.socket_logging import setup_socket_logging
+
+if os.environ.get("LOCUST_MODE_MASTER") == "1":
+    major, minor, *rest = os.environ["LOCUSTCLOUD_CLIENT_VERSION"].split(".")
+
+    if int(major) > 1 or int(major) == 1 and int(minor) >= 12:
+        setup_socket_logging()
+
 __version__ = importlib.metadata.version("locust-cloud")
 
 import logging
@@ -83,6 +92,7 @@ def add_arguments(parser: LocustArgumentParser):
         type=str,
         env_var="LOCUSTCLOUD_PROFILE",
         help=configargparse.SUPPRESS,
+        default=None,
     )
 
 
@@ -118,6 +128,7 @@ def on_locust_init(environment: locust.env.Environment, **_args):
     if environment.web_ui:
         environment.web_ui.template_args["locustVersion"] = locust.__version__
         environment.web_ui.template_args["locustCloudVersion"] = __version__
+        environment.web_ui.template_args["webBasePath"] = environment.parsed_options.web_base_path
 
         if environment.parsed_options.graph_viewer:
             environment.web_ui.template_args["isGraphViewer"] = True

locust_cloud/auth.py

@@ -7,7 +7,7 @@ import locust.env
 import requests
 import werkzeug
 from flask import Blueprint, redirect, request, session, url_for
-from flask_login import UserMixin, login_user
+from flask_login import UserMixin, login_required, login_user, logout_user
 from locust.html import render_template_from
 from locust_cloud import __version__
 
@@ -95,6 +95,9 @@ def register_auth(environment: locust.env.Environment):
             if credentials.get("challenge_session"):
                 session["challenge_session"] = credentials.get("challenge_session")
                 session["username"] = username
+
+                session["auth_error"] = ""
+
                 return redirect(url_for("locust_cloud_auth.password_reset"))
             if os.getenv("CUSTOMER_ID", "") and credentials.get("customer_id") != os.getenv("CUSTOMER_ID", ""):
                 session["auth_error"] = "Invalid login for this deployment"
@@ -405,12 +408,14 @@ def register_auth(environment: locust.env.Environment):
             auth_response.raise_for_status()
 
             session["username"] = ""
-            session["auth_info"] = "Password reset successfully! Please login"
             session["auth_error"] = ""
 
             if session.get("challenge_session"):
                 session["challenge_session"] = ""
-                return redirect("https://docs.locust.cloud/")
+
+                return redirect(url_for("locust_cloud_auth.password_reset_success"))
+
+            session["auth_info"] = "Password reset successfully! Please login"
 
             return redirect(url_for("locust.login"))
         except requests.exceptions.HTTPError as e:
@@ -420,4 +425,19 @@ def register_auth(environment: locust.env.Environment):
 
             return redirect(url_for("locust_cloud_auth.password_reset"))
 
+    @auth_blueprint.route("/password-reset-success")
+    def password_reset_success():
+        return render_template_from(
+            "auth.html",
+            auth_args={
+                "info": "Password successfully set! Please review the [documentation](https://docs.locust.cloud/) and start your first testrun! If you have already ran some tests, you may also [login](/login)"
+            },
+        )
+
+    @auth_blueprint.route("/logout", methods=["POST"])
+    @login_required
+    def logout():
+        logout_user()
+        return redirect(url_for("locust.login"))
+
     environment.web_ui.app.register_blueprint(auth_blueprint)

locust_cloud/cloud.py

@@ -1,19 +1,19 @@
 import base64
 import gzip
-import json
 import logging
 import os
 import sys
 import time
 import tomllib
+import urllib.parse
 from argparse import Namespace
 from collections import OrderedDict
-from datetime import UTC, datetime, timedelta
 from typing import IO, Any
 
 import configargparse
 import requests
-from botocore.exceptions import ClientError
+import socketio
+import socketio.exceptions
 from locust_cloud import __version__
 from locust_cloud.credential_manager import CredentialError, CredentialManager
 
@@ -200,8 +200,6 @@ def main() -> None:
     if options.region:
         os.environ["AWS_DEFAULT_REGION"] = options.region
 
-    deployments: list[Any] = []
-
     if not ((options.username and options.password) or (options.aws_access_key_id and options.aws_secret_access_key)):
         logger.error(
             "Authentication is required to use Locust Cloud. Please ensure the LOCUSTCLOUD_USERNAME and LOCUSTCLOUD_PASSWORD environment variables are set."
@@ -297,7 +295,7 @@ def main() -> None:
             sys.exit(1)
 
         if response.status_code == 200:
-            deployments = response.json().get("deployments", [])
+            log_ws_url = response.json()["log_ws_url"]
         else:
             try:
                 logger.error(f"{response.json()['Message']} (HTTP {response.status_code}/{response.reason})")
@@ -313,80 +311,71 @@ def main() -> None:
         logger.debug("Interrupted by user")
         sys.exit(0)
 
-    log_group_name = "/eks/dmdb-default" if options.region == "us-east-1" else "/eks/locust-default"
-    master_pod_name = next((deployment for deployment in deployments if "master" in deployment), None)
-
-    if not master_pod_name:
-        logger.error(
-            "Master pod not found among deployed pods. Something went wrong during the load generator deployment, please try again or contact an administrator for assistance"
-        )
-        sys.exit(1)
-
     logger.debug("Load generators deployed successfully!")
+    logger.info("Waiting for pods to be ready...")
 
     try:
-        logger.info("Waiting for pods to be ready...")
+        ws_connection_info = urllib.parse.urlparse(log_ws_url)
+        sio = socketio.Client(handle_sigint=False)
 
-        log_stream: str | None = None
-        while log_stream is None:
-            try:
-                client = credential_manager.session.client("logs")
-                response = client.describe_log_streams(
-                    logGroupName=log_group_name,
-                    logStreamNamePrefix=f"from-fluent-bit-kube.var.log.containers.{master_pod_name}",
-                )
-                all_streams = response.get("logStreams", [])
-                if all_streams:
-                    log_stream = all_streams[0].get("logStreamName")
-                else:
-                    time.sleep(1)
-            except ClientError as e:
-                logger.error(f"Error describing log streams: {e}")
-                time.sleep(5)
-        logger.debug("Pods are ready, switching to Locust logs")
-
-        timestamp = int((datetime.now(UTC) - timedelta(minutes=5)).timestamp() * 1000)
-
-        while True:
+        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():
+            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)
+
+        @sio.event
+        def shutdown(message):
+            logger.debug("Got shutdown from locust master")
+            if message:
+                print(message)
+
+            nonlocal run
+            run = False
+
+        for _ in range(5 * 60):  # try for 5 minutes
             try:
-                client = credential_manager.session.client("logs")
-                response = client.get_log_events(
-                    logGroupName=log_group_name,
-                    logStreamName=log_stream,
-                    startTime=timestamp,
-                    startFromHead=True,
+                sio.connect(
+                    f"{ws_connection_info.scheme}://{ws_connection_info.netloc}", socketio_path=ws_connection_info.path
                 )
-                locust_shutdown = False
-                for event in response.get("events", []):
-                    message = event.get("message", "")
-                    event_timestamp = event.get("timestamp", timestamp) + 1
-                    try:
-                        message_json = json.loads(message)
-                        if "log" in message_json:
-                            print(message_json["log"])
-
-                            if "Shutting down (exit code" in message_json["log"]:
-                                locust_shutdown = True
-
-                    except json.JSONDecodeError:
-                        print(message)
-                    timestamp = event_timestamp
-
-                if locust_shutdown:
-                    break
-
-                time.sleep(5)
-            except ClientError as e:
-                error_code = e.response.get("Error", {}).get("Code", "")
-                if error_code == "ExpiredTokenException":
-                    logger.debug("AWS session token expired during log streaming. Refreshing credentials.")
-                    time.sleep(5)
+                break
+            except socketio.exceptions.ConnectionError:
+                time.sleep(1)
+
+        else:  # no break
+            raise Exception("Failed to obtain socket connection")
+
+        wait()
+
     except KeyboardInterrupt:
         logger.debug("Interrupted by user")
+        delete(credential_manager)
+        wait()
     except Exception as e:
         logger.exception(e)
+        delete(credential_manager)
         sys.exit(1)
-    finally:
+    else:
         delete(credential_manager)
 
 

locust_cloud/socket_logging.py

@@ -0,0 +1,103 @@
+import atexit
+import logging
+import os
+import sys
+from collections import deque
+from contextlib import suppress
+
+import flask
+import gevent
+import gevent.pywsgi
+import geventwebsocket.handler
+import socketio
+
+
+def setup_socket_logging():
+    """
+    Set up a separate server listening for incomming websocket connections.
+    Because it listens on a separate port from the webui and the ALB only
+    accepts connections on port 443, this will be exposed on a different
+    path (/<customer-id>/socket-logs) with a separate target group.
+    The targetgroup will want to do health checks before exposing the server
+    so a small flask app exposing the health check endpoint is added to the
+    server in addition to the websocket stuff.
+    """
+
+    logger = logging.getLogger(__name__)
+    logger.propagate = False
+    logger.addHandler(logging.NullHandler())
+
+    # This app will use a logger with the same name as the application
+    # which means it will pick up the one set up above.
+    healthcheck_app = flask.Flask(__name__)
+
+    # /login is the health check endpoint currently configured for the
+    # ALB controller in kubernetes.
+    # See cloud-onboarder/kubernetes/alb-load-balancer-ingress.yaml
+    @healthcheck_app.route("/login")
+    def healthcheck():
+        return ""
+
+    sio = socketio.Server(async_handlers=True, always_connect=True, async_mode="gevent", cors_allowed_origins="*")
+    sio_app = socketio.WSGIApp(sio, healthcheck_app, socketio_path=f"/{os.environ['CUSTOMER_ID']}/socket-logs")
+    message_queue = deque(maxlen=500)
+
+    class QueueCopyStream:
+        def __init__(self, name, original):
+            self.name = name
+            self.original = original
+
+        def write(self, message):
+            self.original.write(message)
+            message_queue.append((self.name, message))
+
+        def flush(self):
+            self.original.flush()
+
+    def connected_sid():
+        for rooms in sio.manager.rooms.values():
+            for clients in rooms.values():
+                for sid in clients.keys():
+                    return sid
+
+    def emitter():
+        while True:
+            while message_queue and (sid := connected_sid()):
+                name, message = message_queue[0]
+
+                with suppress(TimeoutError):
+                    sio.call(name, message, to=sid, timeout=5)
+                    message_queue.popleft()
+
+                    if name == "shutdown":
+                        return
+
+            gevent.sleep(1)
+
+    emitter_greenlet = gevent.spawn(emitter)
+
+    sys.stderr = QueueCopyStream("stderr", sys.stderr)
+    sys.stdout = QueueCopyStream("stdout", sys.stdout)
+
+    @atexit.register
+    def notify_shutdown(*args, **kwargs):
+        message_queue.append(("shutdown", ""))
+        emitter_greenlet.join(timeout=30)
+
+    class WebSocketHandlerWithoutLogging(geventwebsocket.handler.WebSocketHandler):
+        """
+        Subclassing WebSocketHandler so it doesn't set a logger on
+        the server that I've explicitly configured not to have one.
+        """
+
+        @property
+        def logger(self):
+            return logger
+
+    @gevent.spawn
+    def start_websocket_server():
+        server = gevent.pywsgi.WSGIServer(
+            ("", 1095), sio_app, log=None, error_log=None, handler_class=WebSocketHandlerWithoutLogging
+        )
+        server.serve_forever()
+        gevent.get_hub().join()

locust_cloud/timescale/exporter.py

@@ -147,7 +147,7 @@ class Exporter:
 
     def on_quit(self, exit_code, **kwargs):
         self._finished = True
-        atexit._clear()  # make sure we dont capture additional ctrl-c:s
+        atexit.unregister(self.log_stop_test_run)  # make sure we dont capture additional ctrl-c:s
         self._background.join(timeout=10)
         if getattr(self, "_update_end_time_task", False):
             self._update_end_time_task.kill()

locust_cloud/timescale/queries.py

@@ -293,7 +293,7 @@ WHERE id = current_user
 profiles = """
 SELECT DISTINCT
 CASE
-    WHEN profile IS NOT NULL THEN profile
+    WHEN profile IS NOT NULL AND profile != '' THEN profile
     ELSE locustfile
 END AS profile
 FROM testruns

locust_cloud/timescale/query.py

@@ -1,4 +1,5 @@
 import logging
+from datetime import UTC
 
 from flask import Blueprint, make_response, request
 from flask_login import login_required
@@ -22,6 +23,7 @@ def register_query(environment, pool):
         results = []
         try:
             if query and queries[query]:
+                sql = queries[query]
                 # start_time = time.perf_counter()
                 with pool.connection() as conn:
                     # get_conn_time = (time.perf_counter() - start_time) * 1000
@@ -38,8 +40,11 @@ def register_query(environment, pool):
                                 f"UI asked for too long time interval. Start was {sql_params['start']}, end was {sql_params['end']}"
                             )
                             return []
+                        if start_time >= datetime(2024, 10, 30, 11, tzinfo=UTC):
+                            # when runs before this go out of scope, we can just update the query
+                            sql = sql.replace("FROM requests_summary_view", "FROM requests_summary_view_v1")
 
-                    cursor = conn.execute(queries[query], sql_params)
+                    cursor = conn.execute(sql, sql_params)
                     # exec_time = (time.perf_counter() - start_time) * 1000
                     assert cursor
                     # start_time = time.perf_counter()