Add --expect-workers-max-wait parameter.

locustio · Nov 1, 2021 · e5de1c5 · e5de1c5
1 parent 5186af4
commit e5de1c5
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 0 deletions.
diff --git a/locust/argument_parser.py b/locust/argument_parser.py
@@ -292,6 +292,14 @@ def setup_parser_arguments(parser):
         help="How many workers master should expect to connect before starting the test (only when --headless/autostart is used).",
         env_var="LOCUST_EXPECT_WORKERS",
     )
+    master_group.add_argument(
+        "--expect-workers-max-wait",
+        type=int,
+        default=0,
+        help="How long should the master wait for workers to connect before giving up. Defaults to wait forever",
+        env_var="LOCUST_EXPECT_WORKERS_MAX_WAIT",
+    )
+
     master_group.add_argument(
         "--expect-slaves",
         action="store_true",

diff --git a/locust/main.py b/locust/main.py
@@ -346,7 +346,11 @@ def spawn_run_time_quit_greenlet():
     def start_automatic_run():
         if options.master:
             # wait for worker nodes to connect
+            start_time = time.monotonic()
             while len(runner.clients.ready) < options.expect_workers:
+                if options.expect_workers_max_wait and options.expect_workers_max_wait < time.monotonic() - start_time:
+                    logger.error("Gave up waiting for workers to connect.")
+                    sys.exit(1)
                 logging.info(
                     "Waiting for workers to be ready, %s of %s connected",
                     len(runner.clients.ready),