locust 2.19.2.dev24__tar.gz → 2.19.2.dev41__tar.gz

{locust-2.19.2.dev24 → locust-2.19.2.dev41}/.vscode/settings.json

@@ -35,9 +35,6 @@
         "**/.env": true
     },
     "restructuredtext.confPath": "${workspaceFolder}/docs",
-    "python.linting.pylintEnabled": false,
-    "python.linting.flake8Enabled": true,
-    "python.linting.enabled": true,
     "[python]": {
         "editor.defaultFormatter": "ms-python.black-formatter"
     },

{locust-2.19.2.dev24/locust.egg-info → locust-2.19.2.dev41}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: locust
-Version: 2.19.2.dev24
+Version: 2.19.2.dev41
 Summary: Developer friendly load testing framework
 License: MIT
 Project-URL: Homepage, https://locust.io/

{locust-2.19.2.dev24 → locust-2.19.2.dev41}/locust/_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '2.19.2.dev24'
-__version_tuple__ = version_tuple = (2, 19, 2, 'dev24')
+__version__ = version = '2.19.2.dev41'
+__version_tuple__ = version_tuple = (2, 19, 2, 'dev41')

locust-2.19.2.dev41/locust/rpc/protocol.py

@@ -0,0 +1,45 @@
+import msgpack
+import datetime
+from typing import Type
+
+try:
+    from bson import ObjectId  # type: ignore
+except ImportError:
+
+    class ObjectId:  # type: ignore
+        def __init__(self, s):
+            raise Exception("You need to install pymongo or at least bson to be able to send/receive ObjectIds")
+
+
+def decode(obj):
+    if "__datetime__" in obj:
+        obj = datetime.datetime.strptime(obj["as_str"], "%Y%m%dT%H:%M:%S.%f")
+    elif "__ObjectId__" in obj:
+        obj = ObjectId(obj["as_str"])
+    return obj
+
+
+def encode(obj):
+    if isinstance(obj, datetime.datetime):
+        return {"__datetime__": True, "as_str": obj.strftime("%Y%m%dT%H:%M:%S.%f")}
+    elif isinstance(obj, ObjectId):
+        return {"__ObjectId__": True, "as_str": str(obj)}
+    return obj
+
+
+class Message:
+    def __init__(self, message_type, data, node_id):
+        self.type = message_type
+        self.data = data
+        self.node_id = node_id
+
+    def __repr__(self):
+        return f"<Message {self.type}:{self.node_id}>"
+
+    def serialize(self):
+        return msgpack.dumps((self.type, self.data, self.node_id), default=encode)
+
+    @classmethod
+    def unserialize(cls, data):
+        msg = cls(*msgpack.loads(data, raw=False, strict_map_key=False, object_hook=decode))
+        return msg

{locust-2.19.2.dev24 → locust-2.19.2.dev41}/locust/runners.py

@@ -37,6 +37,7 @@ import greenlet
 import psutil
 from gevent.event import Event
 from gevent.pool import Group
+import inspect
 
 from . import User
 from locust import __version__
@@ -380,7 +381,11 @@ class Runner:
         """
         if self.state == STATE_STOPPED:
             return
-        logger.debug("Stopping all users")
+        try:
+            caller = inspect.getframeinfo(inspect.stack()[1][0])
+            logger.debug(f"Stopping all users (called from {caller.filename}:{caller.lineno})")
+        except Exception:
+            logger.debug("Stopping all users (couldnt determine where stop() was called from)")
         self.environment.events.test_stopping.fire(environment=self.environment)
         self.final_user_classes_count = {**self.user_classes_count}
         self.update_state(STATE_CLEANUP)
@@ -998,6 +1003,10 @@ class MasterRunner(DistributedRunner):
                 self.connection_broken = True
                 gevent.sleep(FALLBACK_INTERVAL)
                 continue
+            except KeyboardInterrupt:
+                logging.debug(
+                    "Got KeyboardInterrupt in client_listener. Other greenlets should catch this and shut down."
+                )
             if msg.type == "client_ready":
                 if not msg.data:
                     logger.error(f"An old (pre 2.0) worker tried to connect ({client_id}). That's not going to work.")
@@ -1100,7 +1109,11 @@ class MasterRunner(DistributedRunner):
                 logger.debug(
                     f"Received {msg.type} message from worker {msg.node_id} (index {self.get_worker_index(msg.node_id)})"
                 )
-                self.custom_messages[msg.type](environment=self.environment, msg=msg)
+                try:
+                    self.custom_messages[msg.type](environment=self.environment, msg=msg)
+                except Exception:
+                    logging.error(f"Uncaught exception in handler for {msg.type}\n{traceback.format_exc()}")
+
             else:
                 logger.warning(
                     f"Unknown message type received from worker {msg.node_id} (index {self.get_worker_index(msg.node_id)}): {msg.type}"