locust 2.22.1.dev45__py3-none-any.whl → 2.22.1.dev67__py3-none-any.whl

locust/_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '2.22.1.dev45'
-__version_tuple__ = version_tuple = (2, 22, 1, 'dev45')
+__version__ = version = '2.22.1.dev67'
+__version_tuple__ = version_tuple = (2, 22, 1, 'dev67')

locust/argument_parser.py

@@ -1,14 +1,20 @@
 from __future__ import annotations
 
 import locust
+from locust import runners
+from locust.rpc import Message, zmqrpc
 
 import os
 import platform
+import socket
 import sys
 import textwrap
 from typing import Any, NamedTuple
+from uuid import uuid4
 
 import configargparse
+import gevent
+from gevent.event import Event
 
 version = locust.__version__
 
@@ -175,6 +181,46 @@ See documentation for more details, including how to set options using a file or
     return parser
 
 
+def download_locustfile_from_master(master_host: str, master_port: int) -> str:
+    client_id = socket.gethostname() + "_download_locustfile_" + uuid4().hex
+    tempclient = zmqrpc.Client(master_host, master_port, client_id)
+    got_reply = False
+
+    def ask_for_locustfile():
+        while not got_reply:
+            tempclient.send(Message("locustfile", None, client_id))
+            gevent.sleep(1)
+
+    def wait_for_reply():
+        return tempclient.recv()
+
+    gevent.spawn(ask_for_locustfile)
+    try:
+        # wait same time as for client_ready ack. not that it is really relevant...
+        msg = gevent.spawn(wait_for_reply).get(timeout=runners.CONNECT_TIMEOUT * runners.CONNECT_RETRY_COUNT)
+        got_reply = True
+    except gevent.Timeout:
+        sys.stderr.write(
+            f"Got no locustfile response from master, gave up after {runners.CONNECT_TIMEOUT * runners.CONNECT_RETRY_COUNT}s\n"
+        )
+        sys.exit(1)
+
+    if msg.type != "locustfile":
+        sys.stderr.write(f"Got wrong message type from master {msg.type}\n")
+        sys.exit(1)
+
+    if "error" in msg.data:
+        sys.stderr.write(f"Got error from master: {msg.data['error']}\n")
+        sys.exit(1)
+
+    filename = msg.data["filename"]
+    with open(filename, "w") as local_file:
+        local_file.write(msg.data["contents"])
+
+    tempclient.close()
+    return filename
+
+
 def parse_locustfile_option(args=None) -> list[str]:
     """
     Construct a command line parser that is only used to parse the -f argument so that we can
@@ -197,9 +243,41 @@ def parse_locustfile_option(args=None) -> list[str]:
         action="store_true",
         default=False,
     )
+    # the following arguments are only used for downloading the locustfile from master
+    parser.add_argument(
+        "--worker",
+        action="store_true",
+        env_var="LOCUST_MODE_WORKER",
+    )
+    parser.add_argument(
+        "--master",  # this is just here to prevent argparse from giving the dreaded "ambiguous option: --master could match --master-host, --master-port"
+        action="store_true",
+        env_var="LOCUST_MODE_MASTER",
+    )
+    parser.add_argument(
+        "--master-host",
+        default="127.0.0.1",
+        env_var="LOCUST_MASTER_NODE_HOST",
+    )
+    parser.add_argument(
+        "--master-port",
+        type=int,
+        default=5557,
+        env_var="LOCUST_MASTER_NODE_PORT",
+    )
 
     options, _ = parser.parse_known_args(args=args)
 
+    if options.locustfile == "-":
+        if not options.worker:
+            sys.stderr.write(
+                "locustfile was set to '-' (meaning to download from master) but --worker was not specified.\n"
+            )
+            sys.exit(1)
+        # having this in argument_parser module is a bit weird, but it needs to be done early
+        filename = download_locustfile_from_master(options.master_host, options.master_port)
+        return [filename]
+
     # Comma separated string to list
     locustfile_as_list = [locustfile.strip() for locustfile in options.locustfile.split(",")]
 
@@ -457,7 +535,7 @@ Typically ONLY these options (and --locustfile) need to be specified on workers,
     worker_group.add_argument(
         "--worker",
         action="store_true",
-        help="Set locust to run in distributed mode with this process as worker",
+        help="Set locust to run in distributed mode with this process as worker. Can be combined with setting --locustfile to '-' to download it from master.",
         env_var="LOCUST_MODE_WORKER",
     )
     worker_group.add_argument(

locust/runners.py

@@ -1033,6 +1033,40 @@ class MasterRunner(DistributedRunner):
                 # emit a warning if the worker's clock seem to be out of sync with our clock
                 # if abs(time() - msg.data["time"]) > 5.0:
                 #    warnings.warn("The worker node's clock seem to be out of sync. For the statistics to be correct the different locust servers need to have synchronized clocks.")
+            elif msg.type == "locustfile":
+                logging.debug("Worker requested locust file")
+                assert self.environment.parsed_options
+                filename = (
+                    "locustfile.py"
+                    if self.environment.parsed_options.locustfile == "locustfile"
+                    else self.environment.parsed_options.locustfile
+                )
+                try:
+                    with open(filename) as f:
+                        file_contents = f.read()
+                except Exception as e:
+                    logger.error(
+                        f"--locustfile must be a plain filename (not a module name) for file distribution to work {e}"
+                    )
+                    self.send_message(
+                        "locustfile",
+                        client_id=client_id,
+                        data={
+                            "error": f"locustfile parameter on master must be a plain filename (not a module name) (was '{filename}')"
+                        },
+                    )
+                else:
+                    if getattr(self, "_old_file_contents", file_contents) != file_contents:
+                        logger.warning(
+                            "Locustfile contents changed on disk after first worker requested locustfile, sending new content. If you make any major changes (like changing User class names) you need to restart master."
+                        )
+                    self._old_file_contents = file_contents
+                    self.send_message(
+                        "locustfile",
+                        client_id=client_id,
+                        data={"filename": filename, "contents": file_contents},
+                    )
+                continue
             elif msg.type == "client_stopped":
                 if msg.node_id not in self.clients:
                     logger.warning(f"Received {msg.type} message from an unknown worker: {msg.node_id}.")

locust/test/test_main.py

@@ -21,6 +21,8 @@ from pyquery import PyQuery as pq
 from .mock_locustfile import MOCK_LOCUSTFILE_CONTENT, mock_locustfile
 from .util import get_free_tcp_port, patch_env, temporary_file
 
+SHORT_SLEEP = 2 if sys.platform == "darwin" else 1  # macOS is slow on GH, give it some extra time
+
 
 def is_port_in_use(port: int) -> bool:
     with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
@@ -202,7 +204,7 @@ class StandaloneIntegrationTests(ProcessIntegrationTest):
             )
         ) as file_path:
             proc = subprocess.Popen(["locust", "-f", file_path], stdout=PIPE, stderr=PIPE, text=True)
-            gevent.sleep(1)
+            gevent.sleep(SHORT_SLEEP)
             proc.send_signal(signal.SIGTERM)
             stdout, stderr = proc.communicate()
             self.assertIn("Starting web interface at", stderr)
@@ -295,7 +297,7 @@ class StandaloneIntegrationTests(ProcessIntegrationTest):
                 proc = subprocess.Popen(
                     ["locust", "-f", f"{mocked1.file_path},{mocked2.file_path}"], stdout=PIPE, stderr=PIPE, text=True
                 )
-                gevent.sleep(1)
+                gevent.sleep(SHORT_SLEEP)
                 proc.send_signal(signal.SIGTERM)
                 stdout, stderr = proc.communicate()
                 self.assertIn("Starting web interface at", stderr)
@@ -310,7 +312,7 @@ class StandaloneIntegrationTests(ProcessIntegrationTest):
             with mock_locustfile(content=MOCK_LOCUSTFILE_CONTENT_A, dir=temp_dir):
                 with mock_locustfile(content=MOCK_LOCUSTFILE_CONTENT_B, dir=temp_dir):
                     proc = subprocess.Popen(["locust", "-f", temp_dir], stdout=PIPE, stderr=PIPE, text=True)
-                    gevent.sleep(1)
+                    gevent.sleep(SHORT_SLEEP)
                     proc.send_signal(signal.SIGTERM)
                     stdout, stderr = proc.communicate()
                     self.assertIn("Starting web interface at", stderr)
@@ -355,7 +357,7 @@ class StandaloneIntegrationTests(ProcessIntegrationTest):
                 proc = subprocess.Popen(
                     ["locust", "-f", f"{mocked1.file_path},{mocked2}"], stdout=PIPE, stderr=PIPE, text=True
                 )
-                gevent.sleep(1)
+                gevent.sleep(SHORT_SLEEP)
                 proc.send_signal(signal.SIGTERM)
                 stdout, stderr = proc.communicate()
                 self.assertIn("Starting web interface at", stderr)
@@ -1660,6 +1662,171 @@ class SecondUser(HttpUser):
             self.assertEqual(0, proc.returncode)
             self.assertEqual(0, proc_worker.returncode)
 
+    def test_locustfile_distribution(self):
+        LOCUSTFILE_CONTENT = textwrap.dedent(
+            """
+            from locust import User, task, constant
+
+            class User1(User):
+                wait_time = constant(1)
+
+                @task
+                def t(self):
+                    pass
+            """
+        )
+        with mock_locustfile(content=LOCUSTFILE_CONTENT) as mocked:
+            proc = subprocess.Popen(
+                [
+                    "locust",
+                    "-f",
+                    mocked.file_path,
+                    "--headless",
+                    "--master",
+                    "--expect-workers",
+                    "2",
+                    "-t",
+                    "1s",
+                ],
+                stderr=STDOUT,
+                stdout=PIPE,
+                text=True,
+            )
+            proc_worker = subprocess.Popen(
+                [
+                    "locust",
+                    "-f",
+                    "-",
+                    "--worker",
+                ],
+                stderr=STDOUT,
+                stdout=PIPE,
+                text=True,
+            )
+            gevent.sleep(2)
+            # modify the locustfile to trigger warning about file change when the second worker connects
+            with open(mocked.file_path, "w") as locustfile:
+                locustfile.write(LOCUSTFILE_CONTENT)
+                locustfile.write("\n# New comment\n")
+            gevent.sleep(2)
+            proc_worker2 = subprocess.Popen(
+                [
+                    "locust",
+                    "-f",
+                    "-",
+                    "--worker",
+                ],
+                stderr=STDOUT,
+                stdout=PIPE,
+                text=True,
+            )
+            stdout = proc.communicate()[0]
+            proc_worker2.communicate()
+            proc_worker.communicate()
+
+            self.assertIn('All users spawned: {"User1": 1} (1 total users)', stdout)
+            self.assertIn("Locustfile contents changed on disk after first worker requested locustfile", stdout)
+            self.assertIn("Shutting down (exit code 0)", stdout)
+
+            self.assertEqual(0, proc.returncode)
+            self.assertEqual(0, proc_worker.returncode)
+
+    def test_locustfile_distribution_with_workers_started_first(self):
+        LOCUSTFILE_CONTENT = textwrap.dedent(
+            """
+            from locust import User, task, constant
+
+            class User1(User):
+                wait_time = constant(1)
+
+                @task
+                def t(self):
+                    print("hello")
+            """
+        )
+        with mock_locustfile(content=LOCUSTFILE_CONTENT) as mocked:
+            proc_worker = subprocess.Popen(
+                [
+                    "locust",
+                    "-f",
+                    "-",
+                    "--worker",
+                ],
+                stderr=STDOUT,
+                stdout=PIPE,
+                text=True,
+            )
+            gevent.sleep(2)
+            proc = subprocess.Popen(
+                [
+                    "locust",
+                    "-f",
+                    mocked.file_path,
+                    "--headless",
+                    "--master",
+                    "--expect-workers",
+                    "1",
+                    "-t",
+                    "1",
+                ],
+                stderr=STDOUT,
+                stdout=PIPE,
+                text=True,
+            )
+
+            stdout = proc.communicate()[0]
+            worker_stdout = proc_worker.communicate()[0]
+
+            self.assertIn('All users spawned: {"User1": ', stdout)
+            self.assertIn("Shutting down (exit code 0)", stdout)
+
+            self.assertEqual(0, proc.returncode)
+            self.assertEqual(0, proc_worker.returncode)
+            self.assertIn("hello", worker_stdout)
+
+    def test_distributed_with_locustfile_distribution_not_plain_filename(self):
+        LOCUSTFILE_CONTENT = textwrap.dedent(
+            """
+            from locust import User, task, constant
+
+            class User1(User):
+                wait_time = constant(1)
+
+                @task
+                def t(self):
+                    pass
+            """
+        )
+        with mock_locustfile(content=LOCUSTFILE_CONTENT) as mocked:
+            proc = subprocess.Popen(
+                [
+                    "locust",
+                    "-f",
+                    mocked.file_path[:-3],  # remove ".py"
+                    "--headless",
+                    "--master",
+                ],
+                stderr=STDOUT,
+                stdout=PIPE,
+                text=True,
+            )
+            proc_worker = subprocess.Popen(
+                [
+                    "locust",
+                    "-f",
+                    "-",
+                    "--worker",
+                ],
+                stderr=STDOUT,
+                stdout=PIPE,
+                text=True,
+            )
+            stdout = proc_worker.communicate()[0]
+            self.assertIn("Got error from master: locustfile parameter on master must be a plain filename", stdout)
+            proc.kill()
+            master_stdout = proc.communicate()[0]
+            self.assertIn("--locustfile must be a plain filename (not a module name) for file distribut", master_stdout)
+
     def test_json_schema(self):
         LOCUSTFILE_CONTENT = textwrap.dedent(
             """
@@ -1972,18 +2139,19 @@ class AnyUser(HttpUser):
                 text=True,
                 start_new_session=True,
             )
-            gevent.sleep(1)
+            gevent.sleep(2)
             master_proc.kill()
             master_proc.wait()
             try:
-                _, worker_stderr = worker_parent_proc.communicate(timeout=7)
+                worker_stdout, worker_stderr = worker_parent_proc.communicate(timeout=7)
             except Exception:
                 os.killpg(worker_parent_proc.pid, signal.SIGTERM)
-                _, worker_stderr = worker_parent_proc.communicate()
-                assert False, f"worker never finished: {worker_stderr}"
+                worker_stdout, worker_stderr = worker_parent_proc.communicate()
+                assert False, f"worker never finished: {worker_stdout} / {worker_stderr}"
 
             self.assertNotIn("Traceback", worker_stderr)
             self.assertIn("Didn't get heartbeat from master in over ", worker_stderr)
+            self.assertIn("worker index:", worker_stdout)
 
     @unittest.skipIf(os.name == "nt", reason="--processes doesnt work on windows")
     def test_processes_error_doesnt_blow_up_completely(self):
@@ -2010,6 +2178,7 @@ class AnyUser(HttpUser):
             self.assertNotIn("Traceback", stderr)
 
     @unittest.skipIf(os.name == "nt", reason="--processes doesnt work on windows")
+    @unittest.skipIf(sys.platform == "darwin", reason="Flaky on macOS :-/")
     def test_processes_workers_quit_unexpected(self):
         content = """
 from locust import runners, events, User, task