Add --processes parameter to automatically fork worker processes

.github/workflows/tests.yml

@@ -14,12 +14,13 @@ jobs:
       fail-fast: false
       matrix:
         include:
-          #- {name: Linux, python: '3.9', os: ubuntu-latest, tox: py39}
           #- {name: Windows, python: '3.9', os: windows-latest, tox: py39}
           #- {name: Mac, python: '3.9', os: macos-latest, tox: py39}
           - { name: "flake8", python: "3.11", os: ubuntu-latest, tox: "flake8" }
           - { name: "black", python: "3.11", os: ubuntu-latest, tox: "black" }
           - { name: "mypy", python: "3.10", os: ubuntu-latest, tox: "mypy" }
+          # run some integration tests and abort immediately if they fail, for faster feedback
+          - { name: "fail_fast_test_main", python: "3.12", os: ubuntu-latest, tox: fail_fast_test_main }
           - { name: "3.12", python: "3.12", os: ubuntu-latest, tox: py312 }
           - { name: "3.11", python: "3.11", os: ubuntu-latest, tox: py311 }
           - { name: "3.10", python: "3.10", os: ubuntu-latest, tox: py310 }

locust/argument_parser.py

@@ -423,6 +423,12 @@ def setup_parser_arguments(parser):
         help="Set locust to run in distributed mode with this process as worker",
         env_var="LOCUST_MODE_WORKER",
     )
+    worker_group.add_argument(
+        "--processes",
+        type=int,
+        help="Number of times to fork the locust process, to enable using multiple CPU cores. Use -1 to launch one process per CPU core in your system. Combine with --worker flag or let it automatically set --worker and --master flags for an all-in-one-solution. Not available on Windows. Experimental.",
+        env_var="LOCUST_PROCESSES",
+    )
     worker_group.add_argument(
         "--slave",
         action="store_true",

locust/main.py

@@ -28,6 +28,7 @@
 from .input_events import input_listener
 from .html import get_html_report
 from .util.load_locustfile import load_locustfile
+import traceback
 
 try:
     # import locust_plugins if it is installed, to allow it to register custom arguments etc
@@ -37,7 +38,6 @@
 
 version = locust.__version__
 
-
 # Options to ignore when using a custom shape class without `use_common_options=True`
 # See: https://docs.locust.io/en/stable/custom-load-shape.html#use-common-options
 COMMON_OPTIONS = {
@@ -162,6 +162,84 @@ def is_valid_percentile(parameter):
             sys.stderr.write("Invalid --loglevel. Valid values are: DEBUG/INFO/WARNING/ERROR/CRITICAL\n")
             sys.exit(1)
 
+    children = []
+
+    if options.processes:
+        if os.name == "nt":
+            raise Exception("--processes is not supported in Windows (except in WSL)")
+        if options.processes == -1:
+            options.processes = os.cpu_count()
+            assert options.processes, "couldnt detect number of cpus!?"
+        elif options.processes < -1:
+            raise Exception(f"invalid processes count {options.processes}")
+        for _ in range(options.processes):
+            child_pid = gevent.fork()
+            if child_pid:
+                children.append(child_pid)
+                logging.debug(f"Started child worker with pid #{child_pid}")
+            else:
+                # child is always a worker, even when it wasnt set on command line
+                options.worker = True
+                # remove options that dont make sense on worker
+                options.run_time = None
+                options.autostart = None
+                break
+        else:
+            # we're in the parent process
+            if options.worker:
+                # ignore the first sigint in parent, and wait for the children to handle sigint
+                def sigint_handler(_signal, _frame):
+                    if getattr(sigint_handler, "has_run", False):
+                        # if parent gets repeated sigint, we kill the children hard
+                        for child_pid in children:
+                            try:
+                                logging.debug(f"Sending SIGKILL to child with pid {child_pid}")
+                                os.kill(child_pid, signal.SIGKILL)
+                            except ProcessLookupError:
+                                pass  # process already dead
+                            except Exception:
+                                logging.error(traceback.format_exc())
+                        sys.exit(1)
+                    sigint_handler.has_run = True
+
+                signal.signal(signal.SIGINT, sigint_handler)
+                exit_code = 0
+                # nothing more to do, just wait for the children to exit
+                for child_pid in children:
+                    _, child_status = os.waitpid(child_pid, 0)
+                    try:
+                        if sys.version_info > (3, 8):
+                            child_exit_code = os.waitstatus_to_exitcode(child_status)
+                            exit_code = max(exit_code, child_exit_code)
+                    except AttributeError:
+                        pass  # dammit python 3.8...
+                sys.exit(exit_code)
+            else:
+                options.master = True
+                options.expect_workers = options.processes
+
+                def kill_workers(children):
+                    exit_code = 0
+                    logging.debug("Sending SIGINT to children")
+                    for child_pid in children:
+                        try:
+                            os.kill(child_pid, signal.SIGINT)
+                        except ProcessLookupError:
+                            pass  # never mind, process was already dead
+                    logging.debug("waiting for children to terminate")
+                    for child_pid in children:
+                        _, child_status = os.waitpid(child_pid, 0)
+                        try:
+                            if sys.version_info > (3, 8):
+                                child_exit_code = os.waitstatus_to_exitcode(child_status)
+                                exit_code = max(exit_code, child_exit_code)
+                        except AttributeError:
+                            pass  # dammit python 3.8...
+                    if exit_code > 1:
+                        logging.error(f"bad response code from worker children: {exit_code}")
+
+                atexit.register(kill_workers, children)
+
     logger = logging.getLogger(__name__)
     greenlet_exception_handler = greenlet_exception_logger(logger)
 
@@ -486,7 +564,6 @@ def shutdown():
             print_stats(runner.stats, current=False)
             print_percentile_stats(runner.stats)
             print_error_report(runner.stats)
-
         environment.events.quit.fire(exit_code=code)
         sys.exit(code)
 

locust/runners.py

@@ -1360,7 +1360,11 @@ def _send_stats(self) -> None:
     def connect_to_master(self):
         self.retry += 1
         self.client.send(Message("client_ready", __version__, self.client_id))
-        success = self.connection_event.wait(timeout=CONNECT_TIMEOUT)
+        try:
+            success = self.connection_event.wait(timeout=CONNECT_TIMEOUT)
+        except KeyboardInterrupt:
+            # dont complain about getting CTRL-C
+            sys.exit(1)
         if not success:
             if self.retry < 3:
                 logger.debug(

locust/test/test_main.py

@@ -5,7 +5,7 @@
 import sys
 import platform
 import unittest
-
+import socket
 import pty
 import signal
 import subprocess
@@ -20,6 +20,11 @@
 from .util import temporary_file, get_free_tcp_port, patch_env
 
 
+def is_port_in_use(port: int) -> bool:
+    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+        return s.connect_ex(("localhost", port)) == 0
+
+
 MOCK_LOCUSTFILE_CONTENT_A = textwrap.dedent(
     """
     from locust import User, task, constant, events
@@ -52,6 +57,13 @@ def tearDown(self):
         self.timeout.cancel()
         super().tearDown()
 
+    def assert_run(self, cmd: list[str], timeout: int = 5) -> subprocess.CompletedProcess[str]:
+        out = subprocess.run(cmd, capture_output=True, text=True, timeout=timeout)
+        self.assertEqual(0, out.returncode, f"locust run failed with exit code {out.returncode}:\n{out.stderr}")
+        return out
+
+
+class StandaloneIntegrationTests(ProcessIntegrationTest):
     def test_help_arg(self):
         output = subprocess.check_output(
             ["locust", "--help"],
@@ -65,13 +77,6 @@ def test_help_arg(self):
         self.assertIn("Logging options:", output)
         self.assertIn("--skip-log-setup      Disable Locust's logging setup.", output)
 
-    def assert_run(self, cmd: list[str], timeout: int = 5) -> subprocess.CompletedProcess[str]:
-        out = subprocess.run(cmd, capture_output=True, text=True, timeout=timeout)
-        self.assertEqual(0, out.returncode, f"locust run failed with exit code {out.returncode}:\n{out.stderr}")
-        return out
-
-
-class StandaloneIntegrationTests(ProcessIntegrationTest):
     def test_custom_arguments(self):
         port = get_free_tcp_port()
         with temporary_file(
@@ -1322,6 +1327,22 @@ def my_task(self):
 
 
 class DistributedIntegrationTests(ProcessIntegrationTest):
+    failed_port_check = False
+
+    def setUp(self):
+        if self.failed_port_check:
+            # fail immediately
+            raise Exception("Port 5557 was (still) busy when starting a new test case")
+        for _ in range(5):
+            if not is_port_in_use(5557):
+                break
+            else:
+                gevent.sleep(1)
+        else:
+            self.failed_port_check = True
+            raise Exception("Port 5557 was (still) busy when starting a new test case")
+        super().setUp()
+
     def test_expect_workers(self):
         with mock_locustfile() as mocked:
             proc = subprocess.Popen(
@@ -1727,3 +1748,115 @@ def my_task(self):
             for i in range(2):
                 if found[i] != i:
                     raise Exception(f"expected index {i} but got", found[i])
+
+    def test_processes(self):
+        with mock_locustfile() as mocked:
+            command = f"locust -f {mocked.file_path} --processes 4 --headless --run-time 1 --exit-code-on-error 0"
+            proc = subprocess.Popen(
+                command,
+                shell=True,
+                stdout=PIPE,
+                stderr=PIPE,
+                text=True,
+            )
+            try:
+                _, stderr = proc.communicate(timeout=9)
+            except Exception:
+                proc.kill()
+                assert False, f"locust process never finished: {command}"
+            self.assertNotIn("Traceback", stderr)
+            self.assertIn("(index 3) reported as ready", stderr)
+            self.assertIn("Shutting down (exit code 0)", stderr)
+
+    def test_processes_autodetect(self):
+        with mock_locustfile() as mocked:
+            command = f"locust -f {mocked.file_path} --processes -1 --headless --run-time 1 --exit-code-on-error 0"
+            proc = subprocess.Popen(
+                command,
+                shell=True,
+                stdout=PIPE,
+                stderr=PIPE,
+                text=True,
+            )
+            try:
+                _, stderr = proc.communicate(timeout=9)
+            except Exception:
+                proc.kill()
+                assert False, f"locust process never finished: {command}"
+            self.assertNotIn("Traceback", stderr)
+            self.assertIn("(index 0) reported as ready", stderr)
+            self.assertIn("Shutting down (exit code 0)", stderr)
+
+    def test_processes_separate_worker(self):
+        with mock_locustfile() as mocked:
+            master_proc = subprocess.Popen(
+                f"locust -f {mocked.file_path} --master --headless --run-time 1 --exit-code-on-error 0 --expect-workers-max-wait 2",
+                shell=True,
+                stdout=PIPE,
+                stderr=PIPE,
+                text=True,
+            )
+
+            worker_parent_proc = subprocess.Popen(
+                f"locust -f {mocked.file_path} --processes 4 --worker",
+                shell=True,
+                stdout=PIPE,
+                stderr=PIPE,
+                text=True,
+            )
+
+            try:
+                _, worker_stderr = worker_parent_proc.communicate(timeout=9)
+            except Exception:
+                master_proc.kill()
+                worker_parent_proc.kill()
+                _, worker_stderr = worker_parent_proc.communicate()
+                _, master_stderr = master_proc.communicate()
+                assert False, f"worker never finished: {worker_stderr}"
+
+            try:
+                _, master_stderr = master_proc.communicate(timeout=9)
+            except Exception:
+                master_proc.kill()
+                worker_parent_proc.kill()
+                _, worker_stderr = worker_parent_proc.communicate()
+                _, master_stderr = master_proc.communicate()
+                assert False, f"master never finished: {master_stderr}"
+
+            _, worker_stderr = worker_parent_proc.communicate()
+            _, master_stderr = master_proc.communicate()
+            self.assertNotIn("Traceback", worker_stderr)
+            self.assertNotIn("Traceback", master_stderr)
+            self.assertNotIn("Gave up waiting for workers to connect", master_stderr)
+            self.assertIn("(index 3) reported as ready", master_stderr)
+            self.assertIn("Shutting down (exit code 0)", master_stderr)
+
+    def test_processes_ctrl_c(self):
+        with mock_locustfile() as mocked:
+            proc = subprocess.Popen(
+                [
+                    "locust",
+                    "-f",
+                    mocked.file_path,
+                    "--processes",
+                    "4",
+                    "--headless",
+                    "-L",
+                    "DEBUG",
+                ],
+                stdout=PIPE,
+                stderr=PIPE,
+                text=True,
+            )
+            gevent.sleep(3)
+            proc.send_signal(signal.SIGINT)
+            try:
+                _, stderr = proc.communicate(timeout=3)
+            except Exception:
+                proc.kill()
+                _, stderr = proc.communicate()
+                assert False, f"locust process never finished: {stderr}"
+            self.assertNotIn("Traceback", stderr)
+            self.assertIn("(index 3) reported as ready", stderr)
+            self.assertIn("The last worker quit, stopping test", stderr)
+            self.assertIn("Shutting down (exit code 0)", stderr)

tox.ini

@@ -5,6 +5,7 @@ envlist =
     black
     mypy
     py{37,38,39,310,311,312}
+    fail_fast_test_main
 
 [flake8]
 extend-exclude = build,examples/issue_*.py,src/readthedocs-sphinx-search/
@@ -34,6 +35,11 @@ commands =
     ; bash -ec '! grep . {temp_dir}/err.txt' # should be empty
     bash -ec 'PYTHONUNBUFFERED=1 python3 examples/debugging_advanced.py | grep done'
 
+[testenv:fail_fast_test_main]
+commands = 
+    python3 -m pip install .
+    python3 -m unittest -f locust.test.test_main
+
 [testenv:black]
 deps = black==23.10.1
 commands = black --check .