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

locust_cloud/__init__.py

@@ -105,21 +105,17 @@ def on_locust_init(environment: locust.env.Environment, **_args):
     if not (os.environ.get("PGHOST")):
         return
 
-    try:
-        conninfo = make_conninfo(
-            sslmode="require",
-        )
-        pool = ConnectionPool(
-            conninfo,
-            min_size=1,
-            max_size=20,
-            configure=set_autocommit,
-            check=ConnectionPool.check_connection,
-        )
-        pool.wait()
-    except Exception as e:
-        logger.exception(e)
-        raise
+    conninfo = make_conninfo(
+        sslmode="require",
+    )
+    pool = ConnectionPool(
+        conninfo,
+        min_size=1,
+        max_size=20,
+        configure=set_autocommit,
+        check=ConnectionPool.check_connection,
+    )
+    pool.wait(timeout=10)
 
     if not environment.parsed_options.graph_viewer:
         IdleExit(environment)

locust_cloud/cloud.py

@@ -401,11 +401,12 @@ def delete(credential_manager):
             headers=headers,
         )
 
-        if response.status_code != 200:
+        if response.status_code == 200:
+            logger.debug(response.json()["message"])
+        else:
             logger.info(
                 f"Could not automatically tear down Locust Cloud: HTTP {response.status_code}/{response.reason} - Response: {response.text} - URL: {response.request.url}"
             )
-        logger.debug(response.json()["message"])
     except Exception as e:
         logger.error(f"Could not automatically tear down Locust Cloud: {e.__class__.__name__}:{e}")
 

locust_cloud/socket_logging.py

@@ -3,13 +3,13 @@ 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
+import socketio.exceptions
 
 
 def setup_socket_logging():
@@ -23,13 +23,13 @@ def setup_socket_logging():
     server in addition to the websocket stuff.
     """
 
-    logger = logging.getLogger(__name__)
-    logger.propagate = False
-    logger.addHandler(logging.NullHandler())
+    quiet_logger = logging.getLogger("be_quiet")
+    quiet_logger.propagate = False
+    quiet_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__)
+    healthcheck_app = flask.Flask("be_quiet")
 
     # /login is the health check endpoint currently configured for the
     # ALB controller in kubernetes.
@@ -38,18 +38,33 @@ def setup_socket_logging():
     def healthcheck():
         return ""
 
+    logger = logging.getLogger(__name__)
+    socketio_path = f"/{os.environ['CUSTOMER_ID']}/socket-logs"
     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")
+    sio_app = socketio.WSGIApp(sio, healthcheck_app, socketio_path=socketio_path)
     message_queue = deque(maxlen=500)
 
+    @sio.event
+    def connect(sid, environ, auth):  # noqa: ARG001
+        logger.debug("Client connected to socketio server")
+
+    @sio.event
+    def disconnect(sid):  # noqa: ARG001
+        logger.debug("Client disconnected from socketio server")
+
     class QueueCopyStream:
         def __init__(self, name, original):
             self.name = name
             self.original = original
 
         def write(self, message):
-            self.original.write(message)
+            """
+            Writes to the queue first since when running in gevent the write to
+            stdout/stderr can yield and we want to ensure the same ordering in
+            the queue as for the written messages.
+            """
             message_queue.append((self.name, message))
+            self.original.write(message)
 
         def flush(self):
             self.original.flush()
@@ -65,13 +80,20 @@ def setup_socket_logging():
             while message_queue and (sid := connected_sid()):
                 name, message = message_queue[0]
 
-                with suppress(TimeoutError):
+                try:
+                    if name == "shutdown":
+                        logger.debug("Sending websocket shutdown event")
+
                     sio.call(name, message, to=sid, timeout=5)
                     message_queue.popleft()
 
                     if name == "shutdown":
+                        logger.debug("Websocket shutdown event acknowledged by client")
                         return
 
+                except socketio.exceptions.TimeoutError:
+                    logger.debug("Timed out waiting for client to aknowledge websocket message")
+
             gevent.sleep(1)
 
     emitter_greenlet = gevent.spawn(emitter)
@@ -81,6 +103,7 @@ def setup_socket_logging():
 
     @atexit.register
     def notify_shutdown(*args, **kwargs):
+        logger.debug("Adding shutdown event to websocket queue")
         message_queue.append(("shutdown", ""))
         emitter_greenlet.join(timeout=30)
 
@@ -92,10 +115,11 @@ def setup_socket_logging():
 
         @property
         def logger(self):
-            return logger
+            return quiet_logger
 
     @gevent.spawn
     def start_websocket_server():
+        logger.debug(f"Starting socketio server on port 1095 with path {socketio_path}")
         server = gevent.pywsgi.WSGIServer(
             ("", 1095), sio_app, log=None, error_log=None, handler_class=WebSocketHandlerWithoutLogging
         )

locust_cloud/timescale/queries.py

@@ -4,10 +4,8 @@ requests_query = """
 SELECT
 	name,
   request_type as method,
-	SUM(average * count) / SUM(count) as average,
 	SUM(count) as requests,
 	SUM(failed_count) as failed,
-	MIN(min),
 	MAX(max),
 	SUM(failed_count) / SUM(count) * 100 as "errorPercentage"
 FROM requests_summary_view
@@ -200,7 +198,6 @@ SELECT
  fail_ratio as "failRatio",
  requests,
  date_trunc('second', end_time - id) AS "runTime",
- description,
  exit_code as "exitCode",
  username,
  worker_count as "workerCount",