locust 2.19.2.dev24__py3-none-any.whl → 2.19.2.dev41__py3-none-any.whl

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/rpc/protocol.py

@@ -1,4 +1,30 @@
 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:
@@ -11,9 +37,9 @@ class Message:
         return f"<Message {self.type}:{self.node_id}>"
 
     def serialize(self):
-        return msgpack.dumps((self.type, self.data, self.node_id))
+        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))
+        msg = cls(*msgpack.loads(data, raw=False, strict_map_key=False, object_hook=decode))
         return msg

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}"