locust-cloud 1.6.0__py3-none-any.whl → 1.8.0__py3-none-any.whl

locust_cloud/__init__.py

@@ -6,34 +6,29 @@ __version__ = importlib.metadata.version("locust-cloud")
 
 import argparse
 import logging
-import sys
 
+import configargparse
 import locust.env
 import psycopg
 from locust import events
 from locust.argument_parser import LocustArgumentParser
 from locust_cloud.auth import register_auth
+from locust_cloud.idle_exit import IdleExit
 from locust_cloud.timescale.exporter import Exporter
 from locust_cloud.timescale.query import register_query
 from psycopg.conninfo import make_conninfo
 from psycopg_pool import ConnectionPool
 
-PG_USER = os.environ.get("PG_USER")
-PG_HOST = os.environ.get("PG_HOST")
-PG_PASSWORD = os.environ.get("PG_PASSWORD")
-PG_DATABASE = os.environ.get("PG_DATABASE")
-PG_PORT = os.environ.get("PG_PORT", 5432)
 GRAPH_VIEWER = os.environ.get("GRAPH_VIEWER")
-MAX_USER_COUNT = os.environ.get("MAX_USER_COUNT")
 logger = logging.getLogger(__name__)
 
 
 @events.init_command_line_parser.add_listener
 def add_arguments(parser: LocustArgumentParser):
-    if not (PG_HOST or GRAPH_VIEWER):
+    if not (os.environ.get("PGHOST") or GRAPH_VIEWER):
         parser.add_argument_group(
             "locust-cloud",
-            "locust-cloud disabled, because PG_HOST was not set - this is normal for local runs",
+            "locust-cloud disabled, because PGHOST was not set - this is normal for local runs",
         )
         return
 
@@ -56,61 +51,48 @@ def add_arguments(parser: LocustArgumentParser):
         default="",
         help="Description of the test being run",
     )
+    # do not set
+    # used for sending the run id from master to workers
+    locust_cloud.add_argument(
+        "--run-id",
+        type=str,
+        env_var="LOCUSTCLOUD_RUN_ID",
+        help=configargparse.SUPPRESS,
+    )
 
 
 def set_autocommit(conn: psycopg.Connection):
     conn.autocommit = True
 
 
-def create_connection_pool(
-    pg_user: str, pg_host: str, pg_password: str, pg_database: str, pg_port: str | int
-) -> ConnectionPool:
+@events.init.add_listener
+def on_locust_init(environment: locust.env.Environment, **_args):
+    if not (os.environ.get("PGHOST")):
+        return
+
     try:
         conninfo = make_conninfo(
-            dbname=pg_database,
-            user=pg_user,
-            port=pg_port,
-            password=pg_password,
-            host=pg_host,
             sslmode="require",
-            options="-c statement_timeout=55000",
         )
-        return ConnectionPool(
+        pool = ConnectionPool(
             conninfo,
             min_size=1,
-            max_size=10,
+            max_size=20,
             configure=set_autocommit,
-        )
-    except Exception:
-        sys.stderr.write(f"Could not connect to postgres ({pg_user}@{pg_host}:{pg_port}).")
-        sys.exit(1)
-
-
-@events.init.add_listener
-def on_locust_init(environment: locust.env.Environment, **_args):
-    if not (PG_HOST and PG_USER and PG_PASSWORD and PG_DATABASE and PG_PORT):
-        return
-
-    try:
-        pool = create_connection_pool(
-            pg_user=PG_USER,
-            pg_host=PG_HOST,
-            pg_password=PG_PASSWORD,
-            pg_database=PG_DATABASE,
-            pg_port=PG_PORT,
+            check=ConnectionPool.check_connection,
         )
         pool.wait()
     except Exception as e:
         logger.exception(e)
-        logger.error(f"{PG_HOST=}")
         raise
 
+    if not GRAPH_VIEWER:
+        IdleExit(environment)
+
     if not GRAPH_VIEWER and environment.parsed_options and environment.parsed_options.exporter:
         Exporter(environment, pool)
 
     if environment.web_ui:
-        environment.web_ui.template_args["maxUserCount"] = MAX_USER_COUNT
-
         if GRAPH_VIEWER:
             environment.web_ui.template_args["isGraphViewer"] = True
 

locust_cloud/auth.py

@@ -1,3 +1,4 @@
+import logging
 import os
 from datetime import UTC, datetime, timedelta
 from typing import TypedDict
@@ -12,6 +13,8 @@ from locust_cloud.constants import DEFAULT_DEPLOYER_URL
 
 DEPLOYER_URL = os.environ.get("LOCUSTCLOUD_DEPLOYER_URL", DEFAULT_DEPLOYER_URL)
 
+logger = logging.getLogger(__name__)
+
 
 class Credentials(TypedDict):
     user_sub_id: str
@@ -71,21 +74,22 @@ def register_auth(environment: locust.env.Environment):
                 headers={"X-Client-Version": __version__},
             )
 
-            if auth_response.status_code == 200:
-                credentials = auth_response.json()
-                response = redirect(url_for("index"))
-                response = set_credentials(username, credentials, response)
-                login_user(AuthUser(credentials["user_sub_id"]))
+            auth_response.raise_for_status()
 
-                return response
+            credentials = auth_response.json()
+            response = redirect(url_for("index"))
+            response = set_credentials(username, credentials, response)
+            login_user(AuthUser(credentials["user_sub_id"]))
 
-            environment.web_ui.auth_args = {**environment.web_ui.auth_args, "error": "Invalid username or password"}
+            return response
+        except requests.exceptions.HTTPError as e:
+            if e.response.status_code == 401:
+                environment.web_ui.auth_args["error"] = "Invalid username or password"
+            else:
+                logger.error(f"Unknown response from auth: {e.response.status_code} {e.response.text}")
 
-            return redirect(url_for("login"))
-        except Exception:
-            environment.web_ui.auth_args = {
-                **environment.web_ui.auth_args,
-                "error": "An unknown error occured, please try again",
-            }
+                environment.web_ui.auth_args["error"] = (
+                    "Unknown error during authentication, check logs and/or contact support"
+                )
 
             return redirect(url_for("login"))

locust_cloud/cloud.py

@@ -199,12 +199,7 @@ def main() -> None:
     deployments: list[Any] = []
     worker_count: int = max(options.workers or math.ceil(options.users / USERS_PER_WORKER), 2)
     os.environ["AWS_DEFAULT_REGION"] = options.region
-    if options.users > 5000000:
-        logger.error("You asked for more than 5000000 Users, that isn't allowed.")
-        sys.exit(1)
-    if worker_count > 1000:
-        logger.error("You asked for more than 20 workers, that isn't allowed.")
-        sys.exit(1)
+
     try:
         if not (
             (options.username and options.password) or (options.aws_access_key_id and options.aws_secret_access_key)
@@ -295,6 +290,7 @@ def main() -> None:
                 *locust_env_variables,
             ],
             "worker_count": worker_count,
+            "user_count": options.users,
             "image_tag": options.image_tag,
         }
         headers = {
@@ -372,6 +368,7 @@ def main() -> None:
                     startTime=timestamp,
                     startFromHead=True,
                 )
+                locust_shutdown = False
                 for event in response.get("events", []):
                     message = event.get("message", "")
                     event_timestamp = event.get("timestamp", timestamp) + 1
@@ -379,9 +376,17 @@ def main() -> None:
                         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", "")

locust_cloud/idle_exit.py

@@ -0,0 +1,38 @@
+import logging
+import sys
+
+import gevent
+import locust.env
+from locust import events
+
+logger = logging.getLogger(__name__)
+
+
+class IdleExit:
+    def __init__(self, environment: locust.env.Environment):
+        self.environment = environment
+        self._destroy_task: gevent.Greenlet | None = None
+        events.test_start.add_listener(self.on_locust_state_change)
+        events.test_stop.add_listener(self.on_test_stop)
+        events.quit.add_listener(self.on_locust_state_change)
+
+        if not self.environment.parsed_options.autostart:
+            self._destroy_task = gevent.spawn(self._destroy)
+
+    def _destroy(self):
+        gevent.sleep(1800)
+        logger.info("Locust was detected as idle (no test running) for more than 30 minutes")
+        self.environment.runner.quit()
+
+        if self.environment.web_ui:
+            self.environment.web_ui.greenlet.kill(timeout=5)
+
+            if self.environment.web_ui.greenlet.started:
+                sys.exit(1)
+
+    def on_test_stop(self, **kwargs):
+        self._destroy_task = gevent.spawn(self._destroy)
+
+    def on_locust_state_change(self, **kwargs):
+        if self._destroy_task:
+            self._destroy_task.kill()

locust_cloud/timescale/exporter.py

@@ -22,15 +22,23 @@ def safe_serialize(obj):
     return json.dumps(obj, default=default)
 
 
+def format_datetime(d: datetime):
+    return d.strftime("%Y-%m-%d, %H:%M:%S.%f")
+
+
+def parse_datetime(s: str):
+    return datetime.strptime(s, "%Y-%m-%d, %H:%M:%S.%f").replace(tzinfo=UTC)
+
+
 class Exporter:
     def __init__(self, environment: locust.env.Environment, pool):
         self.env = environment
         self._run_id = None
-        self._samples: list[dict] = []
+        self._samples: list[tuple] = []
         self._background = gevent.spawn(self._run)
-        self._update_end_time_task = gevent.spawn(self._update_end_time)
         self._hostname = socket.gethostname()
         self._finished = False
+        self._has_logged_test_stop = False
         self._pid = os.getpid()
         self.pool = pool
 
@@ -43,13 +51,6 @@ class Exporter:
         events.spawning_complete.add_listener(self.spawning_complete)
         atexit.register(self.log_stop_test_run)
 
-        if self.env.runner is not None:
-            self.env.runner.register_message("run_id", self.set_run_id)
-
-    def set_run_id(self, environment, msg, **kwargs):
-        logging.debug(f"run id from master: {msg.data}")
-        self._run_id = datetime.strptime(msg.data, "%Y-%m-%d, %H:%M:%S.%f").replace(tzinfo=UTC)
-
     def on_cpu_warning(self, environment: locust.env.Environment, cpu_usage, message=None, timestamp=None, **kwargs):
         # passing a custom message & timestamp to the event is a haxx to allow using this event for reporting generic events
         if not timestamp:
@@ -58,18 +59,20 @@ class Exporter:
             message = f"High CPU usage ({cpu_usage}%)"
         with self.pool.connection() as conn:
             conn.execute(
-                "INSERT INTO events (time, text, run_id) VALUES (%s, %s, %s)", (timestamp, message, self._run_id)
+                "INSERT INTO events (time, text, run_id, customer) VALUES (%s, %s, %s, current_user)",
+                (timestamp, message, self._run_id),
             )
 
     def on_test_start(self, environment: locust.env.Environment):
         if not self.env.parsed_options or not self.env.parsed_options.worker:
+            self._has_logged_test_stop = False
             self._run_id = environment._run_id = datetime.now(UTC)  # type: ignore
-            msg = environment._run_id.strftime("%Y-%m-%d, %H:%M:%S.%f")  # type: ignore
-            if environment.runner is not None:
-                logging.debug(f"about to send run_id to workers: {msg}")
-                environment.runner.send_message("run_id", msg)
+            self.env.parsed_options.run_id = format_datetime(environment._run_id)  # type: ignore
             self.log_start_testrun()
             self._user_count_logger = gevent.spawn(self._log_user_count)
+            self._update_end_time_task = gevent.spawn(self._update_end_time)
+        if self.env.parsed_options.worker:
+            self._run_id = parse_datetime(self.env.parsed_options.run_id)
 
     def _log_user_count(self):
         while True:
@@ -78,7 +81,7 @@ class Exporter:
             try:
                 with self.pool.connection() as conn:
                     conn.execute(
-                        """INSERT INTO number_of_users(time, run_id, user_count) VALUES (%s, %s, %s)""",
+                        """INSERT INTO number_of_users(time, run_id, user_count, customer) VALUES (%s, %s, %s, current_user)""",
                         (datetime.now(UTC).isoformat(), self._run_id, self.env.runner.user_count),
                     )
             except psycopg.Error as error:
@@ -99,6 +102,10 @@ class Exporter:
             gevent.sleep(0.5)
 
     def _update_end_time(self):
+        # delay setting first end time
+        # so UI doesn't display temporary value
+        gevent.sleep(5)
+
         # Regularly update endtime to prevent missing endtimes when a test crashes
         while True:
             current_end_time = datetime.now(UTC)
@@ -117,35 +124,39 @@ class Exporter:
         try:
             with self.pool.connection() as conn:
                 conn: psycopg.connection.Connection
-                with conn.cursor() as cur:
-                    cur.executemany(
-                        """
-    INSERT INTO requests (time,run_id,greenlet_id,loadgen,name,request_type,response_time,success,response_length,exception,pid,url,context)
-    VALUES (%(time)s, %(run_id)s, %(greenlet_id)s, %(loadgen)s, %(name)s, %(request_type)s, %(response_time)s, %(success)s, %(response_length)s, %(exception)s, %(pid)s, %(url)s, %(context)s)
-    """,
-                        samples,
-                    )
+                with conn.cursor().copy(
+                    "COPY requests (time,run_id,greenlet_id,loadgen,name,request_type,response_time,success,response_length,exception,pid,url,context) FROM STDIN"
+                ) as copy:
+                    for sample in samples:
+                        copy.write_row(sample)
         except psycopg.Error as error:
             logging.error("Failed to write samples to Postgresql timescale database: " + repr(error))
 
     def on_test_stop(self, environment):
+        if getattr(self, "_update_end_time_task", False):
+            self._update_end_time_task.kill()
         if getattr(self, "_user_count_logger", False):
             self._user_count_logger.kill()
             with self.pool.connection() as conn:
                 conn.execute(
-                    """INSERT INTO number_of_users(time, run_id, user_count) VALUES (%s, %s, %s)""",
+                    """INSERT INTO number_of_users(time, run_id, user_count, customer) VALUES (%s, %s, %s, current_user)""",
                     (datetime.now(UTC).isoformat(), self._run_id, 0),
                 )
         self.log_stop_test_run()
+        self._has_logged_test_stop = True
 
     def on_quit(self, exit_code, **kwargs):
         self._finished = True
         atexit._clear()  # make sure we dont capture additional ctrl-c:s
         self._background.join(timeout=10)
-        self._update_end_time_task.kill()
+        if getattr(self, "_update_end_time_task", False):
+            self._update_end_time_task.kill()
         if getattr(self, "_user_count_logger", False):
             self._user_count_logger.kill()
-        self.log_stop_test_run(exit_code)
+        if not self._has_logged_test_stop:
+            self.log_stop_test_run()
+        if not self.env.parsed_options.worker:
+            self.log_exit_code(exit_code)
 
     def on_request(
         self,
@@ -159,6 +170,9 @@ class Exporter:
         url=None,
         **kwargs,
     ):
+        # handle if a worker connects after test_start
+        if not self._run_id:
+            self._run_id = parse_datetime(self.env.parsed_options.run_id)
         success = 0 if exception else 1
         if start_time:
             time = datetime.fromtimestamp(start_time, tz=UTC)
@@ -167,35 +181,35 @@ class Exporter:
             # (which will be horribly wrong if users spend a lot of time in a with/catch_response-block)
             time = datetime.now(UTC) - timedelta(milliseconds=response_time or 0)
         greenlet_id = getattr(greenlet.getcurrent(), "minimal_ident", 0)  # if we're debugging there is no greenlet
-        sample = {
-            "time": time,
-            "run_id": self._run_id,
-            "greenlet_id": greenlet_id,
-            "loadgen": self._hostname,
-            "name": name,
-            "request_type": request_type,
-            "response_time": response_time,
-            "success": success,
-            "url": url[0:255] if url else None,
-            "pid": self._pid,
-            "context": psycopg.types.json.Json(context, safe_serialize),
-        }
-
-        if response_length >= 0:
-            sample["response_length"] = response_length
-        else:
-            sample["response_length"] = None
 
         if exception:
             if isinstance(exception, CatchResponseError):
-                sample["exception"] = str(exception)
+                exception = str(exception)
             else:
                 try:
-                    sample["exception"] = repr(exception)
+                    exception = repr(exception)
                 except AttributeError:
-                    sample["exception"] = f"{exception.__class__} (and it has no string representation)"
+                    exception = f"{exception.__class__} (and it has no string representation)"
+
+            exception = exception[:300]
         else:
-            sample["exception"] = None
+            exception = None
+
+        sample = (
+            time,
+            self._run_id,
+            greenlet_id,
+            self._hostname,
+            name,
+            request_type,
+            response_time,
+            success,
+            response_length,
+            exception,
+            self._pid,
+            url[0:255] if url else None,
+            psycopg.types.json.Json(context, safe_serialize),
+        )
 
         self._samples.append(sample)
 
@@ -203,7 +217,7 @@ class Exporter:
         cmd = sys.argv[1:]
         with self.pool.connection() as conn:
             conn.execute(
-                "INSERT INTO testruns (id, num_users, worker_count, username, locustfile, description, arguments) VALUES (%s,%s,%s,%s,%s,%s,%s)",
+                "INSERT INTO testruns (id, num_users, worker_count, username, locustfile, description, arguments, customer) VALUES (%s,%s,%s,%s,%s,%s,%s,current_user)",
                 (
                     self._run_id,
                     self.env.runner.target_user_count if self.env.runner else 1,
@@ -220,7 +234,7 @@ class Exporter:
                 ),
             )
             conn.execute(
-                "INSERT INTO events (time, text, run_id) VALUES (%s, %s, %s)",
+                "INSERT INTO events (time, text, run_id, customer) VALUES (%s, %s, %s, current_user)",
                 (datetime.now(UTC).isoformat(), "Test run started", self._run_id),
             )
 
@@ -230,7 +244,7 @@ class Exporter:
             try:
                 with self.pool.connection() as conn:
                     conn.execute(
-                        "INSERT INTO events (time, text, run_id) VALUES (%s, %s, %s)",
+                        "INSERT INTO events (time, text, run_id, customer) VALUES (%s, %s, %s, current_user)",
                         (end_time, f"Rampup complete, {user_count} users spawned", self._run_id),
                     )
             except psycopg.Error as error:
@@ -238,7 +252,7 @@ class Exporter:
                     "Failed to insert rampup complete event time to Postgresql timescale database: " + repr(error)
                 )
 
-    def log_stop_test_run(self, exit_code=None):
+    def log_stop_test_run(self):
         logging.debug(f"Test run id {self._run_id} stopping")
         if self.env.parsed_options.worker:
             return  # only run on master or standalone
@@ -246,17 +260,14 @@ class Exporter:
         try:
             with self.pool.connection() as conn:
                 conn.execute(
-                    "UPDATE testruns SET end_time = %s, exit_code = %s where id = %s",
-                    (end_time, exit_code, self._run_id),
+                    "UPDATE testruns SET end_time = %s WHERE id = %s",
+                    (end_time, self._run_id),
                 )
-                conn.execute(
-                    "INSERT INTO events (time, text, run_id) VALUES (%s, %s, %s)",
-                    (end_time, f"Finished with exit code: {exit_code}", self._run_id),
-                )
-                # The AND time > run_id clause in the following statements are there to help Timescale performance
-                # We dont use start_time / end_time to calculate RPS, instead we use the time between the actual first and last request
-                # (as this is a more accurate measurement of the actual test)
+
                 try:
+                    # The AND time > run_id clause in the following statements are there to help Timescale performance
+                    # We dont use start_time / end_time to calculate RPS, instead we use the time between the actual first and last request
+                    # (as this is a more accurate measurement of the actual test)
                     conn.execute(
                         """
 UPDATE testruns
@@ -265,12 +276,12 @@ SET (requests, resp_time_avg, rps_avg, fail_ratio) =
 (SELECT
  COUNT(*)::numeric AS reqs,
  AVG(response_time)::numeric as resp_time
- FROM requests WHERE run_id = %(run_id)s AND time > %(run_id)s) AS _,
+ FROM requests_view WHERE run_id = %(run_id)s AND time > %(run_id)s) AS _,
 (SELECT
- EXTRACT(epoch FROM (SELECT MAX(time)-MIN(time) FROM requests WHERE run_id = %(run_id)s AND time > %(run_id)s))::numeric AS duration) AS __,
+ EXTRACT(epoch FROM (SELECT MAX(time)-MIN(time) FROM requests_view WHERE run_id = %(run_id)s AND time > %(run_id)s))::numeric AS duration) AS __,
 (SELECT
  COUNT(*)::numeric AS fails
- FROM requests WHERE run_id = %(run_id)s AND time > %(run_id)s AND success = 0) AS ___
+ FROM requests_view WHERE run_id = %(run_id)s AND time > %(run_id)s AND success = 0) AS ___
 WHERE id = %(run_id)s""",
                         {"run_id": self._run_id},
                     )
@@ -283,3 +294,20 @@ WHERE id = %(run_id)s""",
                 "Failed to update testruns record (or events) with end time to Postgresql timescale database: "
                 + repr(error)
             )
+
+    def log_exit_code(self, exit_code=None):
+        try:
+            with self.pool.connection() as conn:
+                conn.execute(
+                    "UPDATE testruns SET exit_code = %s WHERE id = %s",
+                    (exit_code, self._run_id),
+                )
+                conn.execute(
+                    "INSERT INTO events (time, text, run_id, customer) VALUES (%s, %s, %s, current_user)",
+                    (datetime.now(UTC).isoformat(), f"Finished with exit code: {exit_code}", self._run_id),
+                )
+        except psycopg.Error as error:
+            logging.error(
+                "Failed to update testruns record (or events) with end time to Postgresql timescale database: "
+                + repr(error)
+            )