Merge pull request #2239 from locustio/fix-issue-with---stop-timeout-…

…parsing-time-strings

Fix issue with --stop timeout parsing time strings

.github/workflows/tests.yml

@@ -20,7 +20,7 @@ jobs:
           - { name: "flake8", python: "3.10", os: ubuntu-latest, tox: "flake8" }
           - { name: "black", python: "3.10", os: ubuntu-latest, tox: "black" }
           - { name: "mypy", python: "3.10", os: ubuntu-latest, tox: "mypy" }
-          - { name: "3.11", python: "3.11-dev", os: ubuntu-latest, tox: py311 }
+          - { name: "3.11", python: "3.11", os: ubuntu-latest, tox: py311 }
           - { name: "3.10", python: "3.10", os: ubuntu-latest, tox: py310 }
           - { name: "3.9", python: "3.9", os: ubuntu-latest, tox: py39 }
           - { name: "3.8", python: "3.8", os: ubuntu-latest, tox: py38 }

locust/argument_parser.py

@@ -548,7 +548,7 @@ def setup_parser_arguments(parser):
         "--stop-timeout",
         action="store",
         dest="stop_timeout",
-        default=None,
+        default="0",
         help="Number of seconds to wait for a simulated user to complete any executing task before exiting. Default is to terminate immediately. This parameter only needs to be specified for the master process when running Locust distributed.",
         env_var="LOCUST_STOP_TIMEOUT",
     )

locust/env.py

@@ -78,7 +78,12 @@ def __init__(
         """Base URL of the target system"""
         self.reset_stats = reset_stats
         """Determines if stats should be reset once all simulated users have been spawned"""
-        self.stop_timeout = stop_timeout
+        if stop_timeout is not None:
+            self.stop_timeout = stop_timeout
+        elif parsed_options:
+            self.stop_timeout = float(parsed_options.stop_timeout)
+        else:
+            self.stop_timeout = 0.0
         """
         If set, the runner will try to stop the running users gracefully and wait this many seconds
         before killing them hard.

locust/main.py

@@ -44,7 +44,6 @@ def create_environment(
         events=events,
         host=options.host,
         reset_stats=options.reset_stats,
-        stop_timeout=options.stop_timeout,
         parsed_options=options,
         available_user_classes=available_user_classes,
         available_shape_classes=available_shape_classes,
@@ -129,6 +128,13 @@ def main():
     logger = logging.getLogger(__name__)
     greenlet_exception_handler = greenlet_exception_logger(logger)
 
+    if options.stop_timeout:
+        try:
+            options.stop_timeout = parse_timespan(options.stop_timeout)
+        except ValueError:
+            logger.error("Valid --stop-timeout formats are: 20, 20s, 3m, 2h, 1h20m, 3h30m10s, etc.")
+            sys.exit(1)
+
     if options.list_commands:
         print("Available Users:")
         for name in user_classes:
@@ -234,13 +240,6 @@ def main():
             logger.error("Valid --run-time formats are: 20, 20s, 3m, 2h, 1h20m, 3h30m10s, etc.")
             sys.exit(1)
 
-    if options.stop_timeout:
-        try:
-            options.stop_timeout = parse_timespan(options.stop_timeout)
-        except ValueError:
-            logger.error("Valid --stop-timeout formats are: 20, 20s, 3m, 2h, 1h20m, 3h30m10s, etc.")
-            sys.exit(1)
-
     if options.csv_prefix:
         stats_csv_writer = StatsCSVFileWriter(
             environment, stats.PERCENTILES_TO_REPORT, options.csv_prefix, options.stats_history_enabled

locust/runners.py

@@ -891,7 +891,7 @@ def stop(self, send_stop_to_client: bool = True) -> None:
                     self.server.send_to_client(Message("stop", None, client.id))
 
                 # Give an additional 60s for all workers to stop
-                timeout = gevent.Timeout(self.environment.stop_timeout or 0 + 60)
+                timeout = gevent.Timeout(self.environment.stop_timeout + 60)
                 timeout.start()
                 try:
                     while self.user_count != 0:
@@ -1296,7 +1296,7 @@ def worker(self) -> NoReturn:
                     )
                     continue
                 self.environment.host = job["host"]
-                self.environment.stop_timeout = job["stop_timeout"]
+                self.environment.stop_timeout = job["stop_timeout"] or 0.0
 
                 # receive custom arguments
                 if self.environment.parsed_options is None:
@@ -1334,7 +1334,7 @@ def worker(self) -> NoReturn:
                 # +additional_wait is just a small buffer to account for the random network latencies and/or other
                 # random delays inherent to distributed systems.
                 additional_wait = int(os.getenv("LOCUST_WORKER_ADDITIONAL_WAIT_BEFORE_READY_AFTER_STOP", 0))
-                gevent.sleep((self.environment.stop_timeout or 0) + additional_wait)
+                gevent.sleep(self.environment.stop_timeout + additional_wait)
                 self.client.send(Message("client_ready", __version__, self.client_id))
                 self.worker_state = STATE_INIT
             elif msg.type == "quit":

locust/test/test_main.py

@@ -265,7 +265,7 @@ def my_task(self):
 
     def test_default_headless_spawn_options(self):
         with mock_locustfile() as mocked:
-            output = subprocess.check_output(
+            proc = subprocess.Popen(
                 [
                     "locust",
                     "-f",
@@ -279,12 +279,38 @@ def test_default_headless_spawn_options(self):
                     "DEBUG",
                     "--exit-code-on-error",
                     "0",
+                    # just to test --stop-timeout argument parsing, doesnt actually validate its function:
+                    "--stop-timeout",
+                    "1s",
                 ],
-                stderr=subprocess.STDOUT,
-                timeout=3,
+                stdout=PIPE,
+                stderr=PIPE,
                 text=True,
-            ).strip()
-            self.assertIn('Spawning additional {"UserSubclass": 1} ({"UserSubclass": 0} already running)...', output)
+            )
+            stdout, stderr = proc.communicate(timeout=4)
+            self.assertNotIn("Traceback", stderr)
+            self.assertIn('Spawning additional {"UserSubclass": 1} ({"UserSubclass": 0} already running)...', stderr)
+            self.assertEqual(0, proc.returncode)
+
+    def test_invalid_stop_timeout_string(self):
+        with mock_locustfile() as mocked:
+            proc = subprocess.Popen(
+                [
+                    "locust",
+                    "-f",
+                    mocked.file_path,
+                    "--host",
+                    "https://test.com/",
+                    "--stop-timeout",
+                    "asdf1",
+                ],
+                stdout=PIPE,
+                stderr=PIPE,
+                text=True,
+            )
+            stdout, stderr = proc.communicate()
+            self.assertIn("ERROR/locust.main: Valid --stop-timeout formats are", stderr)
+            self.assertEqual(1, proc.returncode)
 
     def test_headless_spawn_options_wo_run_time(self):
         with mock_locustfile() as mocked:

locust/test/test_parser.py

@@ -88,7 +88,7 @@ def test_parse_options(self):
         self.assertEqual(10, options.spawn_rate)
         self.assertEqual("5m", options.run_time)
         self.assertTrue(options.reset_stats)
-        self.assertEqual(5, options.stop_timeout)
+        self.assertEqual("5", options.stop_timeout)
         self.assertEqual(["MyUserClass"], options.user_classes)
         # check default arg
         self.assertEqual(8089, options.web_port)

locust/test/test_runners.py

@@ -107,7 +107,7 @@ def __init__(self):
         self.master_bind_port = 5557
         self.heartbeat_liveness = 3
         self.heartbeat_interval = 1
-        self.stop_timeout = None
+        self.stop_timeout = 0.0
         self.connection_broken = False
 
     def reset_stats(self):
@@ -631,8 +631,7 @@ class TestUser2(User):
             def my_task(self):
                 gevent.sleep(600)
 
-        stop_timeout = 0
-        env = Environment(user_classes=[TestUser1, TestUser2], stop_timeout=stop_timeout)
+        env = Environment(user_classes=[TestUser1, TestUser2])
         local_runner = env.create_local_runner()
         web_ui = env.create_web_ui("127.0.0.1", 0)
 
@@ -778,7 +777,7 @@ class MyUser1(User):
             def my_task(self):
                 pass
 
-        env = Environment(user_classes=[MyUser1], stop_timeout=0)
+        env = Environment(user_classes=[MyUser1])
         local_runner = env.create_local_runner()
         web_ui = env.create_web_ui("127.0.0.1", 0)
 
@@ -1731,8 +1730,8 @@ def tick(self):
             "LOCUST_WORKER_ADDITIONAL_WAIT_BEFORE_READY_AFTER_STOP",
             str(locust_worker_additional_wait_before_ready_after_stop),
         ):
-            stop_timeout = 0
-            master_env = Environment(user_classes=[TestUser1], shape_class=TestShape(), stop_timeout=stop_timeout)
+            master_env = Environment(user_classes=[TestUser1], shape_class=TestShape())
+
             master_env.shape_class.reset_time()
             master = master_env.create_master_runner("*", 0)
 
@@ -1798,8 +1797,7 @@ def my_task(self):
                 gevent.sleep(600)
 
         with mock.patch("locust.runners.WORKER_REPORT_INTERVAL", new=0.3):
-            stop_timeout = 0
-            master_env = Environment(user_classes=[TestUser1, TestUser2], stop_timeout=stop_timeout)
+            master_env = Environment(user_classes=[TestUser1, TestUser2])
             master = master_env.create_master_runner("*", 0)
             web_ui = master_env.create_web_ui("127.0.0.1", 0)
 
@@ -3118,9 +3116,7 @@ def the_task(self):
                 MyTestUser._test_state = 2
 
         with mock.patch("locust.rpc.rpc.Client", mocked_rpc()) as client:
-            worker = self.get_runner(
-                environment=Environment(stop_timeout=None), user_classes=[MyTestUser], client=client
-            )
+            worker = self.get_runner(environment=Environment(), user_classes=[MyTestUser], client=client)
             self.assertEqual(1, len(client.outbox))
             self.assertEqual("client_ready", client.outbox[0].type)
             client.mocked_send(
@@ -3891,8 +3887,9 @@ class MyTestUser(User):
             tasks = [MySubTaskSet]
             wait_time = constant(3)
 
-        environment = create_environment([MyTestUser], mocked_options())
-        environment.stop_timeout = 0.3
+        options = mocked_options()
+        options.stop_timeout = 0.3
+        environment = create_environment([MyTestUser], options)
         runner = environment.create_local_runner()
         runner.start(1, 1, wait=True)
         gevent.sleep(0)

tox.ini

@@ -4,7 +4,7 @@ envlist =
     flake8
     black
     mypy
-    py{37,38,39,310,311-dev}
+    py{37,38,39,310,311}
 
 [flake8]
 extend-exclude = build,examples/issue_*.py,src/readthedocs-sphinx-search/