locust-cloud 0.1.4__py3-none-any.whl

locust_cloud/__init__.py

@@ -0,0 +1,50 @@
+from locust_cloud.timescale.exporter import Timescale
+
+import os
+
+from locust import events
+from locust.argument_parser import LocustArgumentParser
+
+if "LOCUST_BUILD_PATH" not in os.environ:
+    os.environ["LOCUST_BUILD_PATH"] = os.path.join(os.path.dirname(__file__), "webui/dist")
+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)
+
+
+@events.init_command_line_parser.add_listener
+def add_arguments(parser: LocustArgumentParser):
+    locust_cloud = parser.add_argument_group(
+        "locust-cloud",
+        "Arguments for use with Locust cloud!",
+    )
+
+    locust_cloud.add_argument(
+        "--exporter",
+        default=True,
+        action="store_true",
+        env_var="LOCUST_EXPORTER",
+        help="Exports Locust stats to Timescale",
+    )
+    locust_cloud.add_argument(
+        "--description",
+        type=str,
+        env_var="LOCUST_DESCRIPTION",
+        default="",
+        help="Description of the test being run",
+    )
+
+
+@events.init.add_listener
+def on_locust_init(environment, **args):
+    if environment.parsed_options.exporter:
+        Timescale(
+            environment,
+            pg_user=PG_USER,
+            pg_host=PG_HOST,
+            pg_password=PG_PASSWORD,
+            pg_database=PG_DATABASE,
+            pg_port=PG_PORT,
+        )

locust_cloud/timescale/exporter.py

@@ -0,0 +1,336 @@
+import json
+import logging
+import os
+import socket
+import sys
+from contextlib import contextmanager
+
+import gevent
+import locust.env
+from gevent.lock import Semaphore
+from locust.exception import CatchResponseError
+
+try:
+    import psycogreen.gevent
+except ModuleNotFoundError as e:
+    logging.error(f"'{e}', you need to install it using 'pip install psycogreen'")
+    sys.exit(1)
+
+psycogreen.gevent.patch_psycopg()
+import atexit
+from datetime import UTC, datetime, timedelta
+
+import greenlet
+import psycopg2
+import psycopg2.extras
+
+
+def safe_serialize(obj):
+    def default(o):
+        return f"<<non-serializable: {type(o).__qualname__}>>"
+
+    return json.dumps(obj, default=default)
+
+
+def print_t(s):
+    print(str(s), end="\t")
+
+
+class Timescale:
+    """
+    See timescale_listener_ex.py for documentation
+    """
+
+    dblock = Semaphore()
+    first_instance = True
+
+    def __init__(self, environment: locust.env.Environment, pg_user, pg_host, pg_password, pg_database, pg_port):
+        if not Timescale.first_instance:
+            # we should refactor this into a module as it is much more pythonic
+            raise Exception(
+                "You tried to initialize the Timescale listener twice, maybe both in your locustfile and using command line --timescale? Ignoring second initialization."
+            )
+        Timescale.first_instance = False
+        self.env = environment
+        self._run_id = ""
+        self.dbconn = None
+        self._samples: list[dict] = []
+        self._background = gevent.spawn(self._run)
+        self._hostname = socket.gethostname()
+        self._finished = False
+        self._pid = os.getpid()
+
+        self.pg_user = pg_user
+        self.pg_host = pg_host
+        self.pg_password = pg_password
+        self.pg_database = pg_database
+        self.pg_port = pg_port
+
+        events = self.env.events
+        events.test_start.add_listener(self.on_test_start)
+        events.test_stop.add_listener(self.on_test_stop)
+        events.request.add_listener(self.on_request)
+        events.cpu_warning.add_listener(self.on_cpu_warning)
+        events.quit.add_listener(self.on_quit)
+        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)
+
+    @contextmanager
+    def dbcursor(self):
+        with self.dblock:
+            try:
+                if self.dbconn:
+                    if self.dbconn.closed:
+                        self.dbconn = self._dbconn()
+                    yield self.dbconn.cursor()
+            except psycopg2.Error:
+                try:
+                    # try to recreate connection
+                    self.dbconn = self._dbconn()
+                except Exception:
+                    pass
+                raise
+
+    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:
+            timestamp = datetime.now(UTC).isoformat()
+        if not message:
+            message = f"High CPU usage ({cpu_usage}%)"
+        with self.dbcursor() as cur:
+            cur.execute(
+                "INSERT INTO events (time, text, run_id) VALUES (%s, %s, %s)", (timestamp, message, self._run_id)
+            )
+
+    def on_test_start(self, environment: locust.env.Environment):
+        try:
+            self.dbconn = self._dbconn()
+        except psycopg2.OperationalError as e:
+            logging.error(e)
+            sys.exit(1)
+
+        if not self.env.parsed_options or not self.env.parsed_options.worker:
+            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.log_start_testrun()
+            self._user_count_logger = gevent.spawn(self._log_user_count)
+
+    def _dbconn(self) -> psycopg2.extensions.connection:
+        try:
+            conn = psycopg2.connect(
+                host=self.pg_host,
+                user=self.pg_user,
+                password=self.pg_password,
+                database=self.pg_database,
+                port=self.pg_port,
+                keepalives_idle=120,
+                keepalives_interval=20,
+                keepalives_count=6,
+            )
+
+            conn.autocommit = True
+        except Exception:
+            sys.stderr.write(f"Could not connect to postgres ({self.pg_user}@{self.pg_host}:{self.pg_port}).")
+            sys.exit(1)
+
+        return conn
+
+    def _log_user_count(self):
+        while True:
+            if self.env.runner is None:
+                return  # there is no runner, so nothing to log...
+            try:
+                with self.dbcursor() as cur:
+                    cur.execute(
+                        """INSERT INTO number_of_users(time, run_id, user_count) VALUES (%s, %s, %s)""",
+                        (datetime.now(UTC).isoformat(), self._run_id, self.env.runner.user_count),
+                    )
+            except psycopg2.Error as error:
+                logging.error("Failed to write user count to Postgresql: " + repr(error))
+                try:
+                    # try to recreate connection
+                    self.user_conn = self._dbconn()
+                except Exception:
+                    pass
+            gevent.sleep(2.0)
+
+    def _run(self):
+        while True:
+            if self._samples:
+                # Buffer samples, so that a locust greenlet will write to the new list
+                # instead of the one that has been sent into postgres client
+                samples_buffer = self._samples
+                self._samples = []
+                self.write_samples_to_db(samples_buffer)
+            else:
+                if self._finished:
+                    break
+            gevent.sleep(0.5)
+
+    def write_samples_to_db(self, samples):
+        try:
+            with self.dbcursor() as cur:
+                psycopg2.extras.execute_values(
+                    cur,
+                    """INSERT INTO requests(time,run_id,greenlet_id,loadgen,name,request_type,response_time,success,response_length,exception,pid,url,context) VALUES %s""",
+                    samples,
+                    template="(%(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)",
+                )
+
+        except psycopg2.Error as error:
+            logging.error("Failed to write samples to Postgresql timescale database: " + repr(error))
+            sys.exit(1)
+
+    def on_test_stop(self, environment):
+        if getattr(self, "_user_count_logger", False):
+            self._user_count_logger.kill()
+        self.log_stop_test_run()
+
+    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)
+        if getattr(self, "_user_count_logger", False):
+            self._user_count_logger.kill()
+        self.log_stop_test_run(exit_code)
+
+    def on_request(
+        self,
+        request_type,
+        name,
+        response_time,
+        response_length,
+        exception,
+        context,
+        start_time=None,
+        url=None,
+        **kwargs,
+    ):
+        success = 0 if exception else 1
+        if start_time:
+            time = datetime.fromtimestamp(start_time, tz=UTC)
+        else:
+            # some users may not send start_time, so we just make an educated guess
+            # (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": psycopg2.extras.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)
+            else:
+                try:
+                    sample["exception"] = repr(exception)
+                except AttributeError:
+                    sample["exception"] = f"{exception.__class__} (and it has no string representation)"
+        else:
+            sample["exception"] = None
+
+        self._samples.append(sample)
+
+    def log_start_testrun(self):
+        cmd = sys.argv[1:]
+        with self.dbcursor() as cur:
+            cur.execute(
+                "INSERT INTO testruns (id, num_users, description, arguments) VALUES (%s,%s,%s,%s)",
+                (
+                    self._run_id,
+                    self.env.parsed_options.num_users or self.env.runner.user_count or 0,
+                    "self.env.parsed_options.description",
+                    " ".join(cmd),
+                ),
+            )
+            cur.execute(
+                "INSERT INTO events (time, text, run_id) VALUES (%s, %s, %s)",
+                (datetime.now(UTC).isoformat(), "Test run started", self._run_id),
+            )
+
+    def spawning_complete(self, user_count):
+        if not self.env.parsed_options.worker:  # only log for master/standalone
+            end_time = datetime.now(UTC)
+            try:
+                with self.dbcursor() as cur:
+                    cur.execute(
+                        "INSERT INTO events (time, text, run_id) VALUES (%s, %s, %s)",
+                        (end_time, f"Rampup complete, {user_count} users spawned", self._run_id),
+                    )
+            except psycopg2.Error as error:
+                logging.error(
+                    "Failed to insert rampup complete event time to Postgresql timescale database: " + repr(error)
+                )
+
+    def log_stop_test_run(self, exit_code=None):
+        logging.debug(f"Test run id {self._run_id} stopping")
+        if self.env.parsed_options.worker:
+            return  # only run on master or standalone
+        if getattr(self, "dbconn", None) is None:
+            return  # test_start never ran, so there's not much for us to do
+        end_time = datetime.now(UTC)
+        try:
+            with self.dbcursor() as cur:
+                cur.execute(
+                    "UPDATE testruns SET end_time = %s, exit_code = %s where id = %s",
+                    (end_time, exit_code, self._run_id),
+                )
+                cur.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:
+                    cur.execute(
+                        """
+UPDATE testruns
+SET (requests, resp_time_avg, rps_avg, fail_ratio) =
+(SELECT reqs, resp_time, reqs / GREATEST(duration, 1), fails / reqs) FROM
+(SELECT
+ COUNT(*)::numeric AS reqs,
+ AVG(response_time)::numeric as resp_time
+ FROM requests WHERE run_id = %s AND time > %s) AS _,
+(SELECT
+ EXTRACT(epoch FROM (SELECT MAX(time)-MIN(time) FROM requests WHERE run_id = %s AND time > %s))::numeric AS duration) AS __,
+(SELECT
+ COUNT(*)::numeric AS fails
+ FROM requests WHERE run_id = %s AND time > %s AND success = 0) AS ___
+WHERE id = %s""",
+                        [self._run_id] * 7,
+                    )
+                except psycopg2.errors.DivisionByZero:
+                    logging.info(
+                        "Got DivisionByZero error when trying to update testruns, most likely because there were no requests logged"
+                    )
+        except psycopg2.Error as error:
+            logging.error(
+                "Failed to update testruns record (or events) with end time to Postgresql timescale database: "
+                + repr(error)
+            )

locust_cloud/webui/.eslintrc

@@ -0,0 +1,41 @@
+{
+  "parser": "@typescript-eslint/parser",
+  "extends": ["plugin:@typescript-eslint/recommended"],
+  "plugins": ["react", "react-hooks", "@typescript-eslint", "prettier", "unused-imports", "import"],
+  "rules": {
+    "react/display-name": "off",
+    "@typescript-eslint/no-explicit-any": "off",
+    "no-console": "error",
+    "react/jsx-sort-props": 2,
+    "react/sort-prop-types": 2,
+    "import/order": [
+      "error",
+      {
+        "groups": ["external", "internal"],
+        "newlines-between": "always",
+        "alphabetize": { "order": "asc", "caseInsensitive": true },
+        "pathGroups": [
+          {
+            "pattern": "react",
+            "group": "external",
+            "position": "before",
+          },
+          {
+            "pattern": "App",
+            "group": "internal",
+          },
+          {
+            "pattern": "Report",
+            "group": "internal",
+          },
+          {
+            "pattern": "{components,hooks,redux,utils}/**",
+            "group": "internal",
+          },
+        ],
+        "distinctGroup": false,
+        "pathGroupsExcludedImportTypes": ["internal"],
+      },
+    ],
+  },
+}

locust_cloud/webui/.gitignore

@@ -0,0 +1,4 @@
+__pycache__
+node_modules
+tsconfig.tsbuildinfo
+dist/**

locust_cloud/webui/.prettierrc

@@ -0,0 +1,9 @@
+{
+  "arrowParens": "avoid",
+  "semi": true,
+  "singleQuote": true,
+  "jsxSingleQuote": true,
+  "trailingComma": "all",
+  "bracketSpacing": true,
+  "printWidth": 100
+}