Fix #1910 KeyboardInterrupt with load shapes

locust/main.py

@@ -378,9 +378,13 @@ def start_automatic_run():
                 if options.run_time:
                     sys.stderr.write("It makes no sense to combine --run-time and LoadShapes. Bailing out.\n")
                     sys.exit(1)
-                environment.runner.start_shape()
-                environment.runner.shape_greenlet.join()
-                stop_and_optionally_quit()
+                try:
+                    environment.runner.start_shape()
+                    environment.runner.shape_greenlet.join()
+                except KeyboardInterrupt:
+                    logging.info("Exiting due to CTRL+C interruption")
+                finally:
+                    stop_and_optionally_quit()
             else:
                 headless_master_greenlet = gevent.spawn(runner.start, options.num_users, options.spawn_rate)
                 headless_master_greenlet.link_exception(greenlet_exception_handler)

locust/test/test_main.py

@@ -1185,6 +1185,51 @@ def task1(self):
             self.assertIn("No tasks defined on MyUser", stderr)
             self.assertEqual(1, proc.returncode)
 
+    def test_graceful_exit_when_keyboard_interrupt(self):
+        with temporary_file(
+            content=textwrap.dedent(
+                """
+                from locust import User, events,task, between, LoadTestShape
+                @events.test_stop.add_listener
+                def on_test_stop(environment, **kwargs) -> None:
+                    print("Test Stopped")
+
+                class LoadTestShape(LoadTestShape):
+                    def tick(self):
+                        run_time = self.get_run_time()
+                        if run_time < 2:
+                            return (10, 1)
+
+                        return None
+
+                class TestUser(User):
+                    @task
+                    def my_task(self):
+                        print("running my_task()")
+            """
+            )
+        ) as mocked:
+            proc = subprocess.Popen(
+                [
+                    "locust",
+                    "-f",
+                    mocked,
+                    "--headless",
+                ],
+                stdout=PIPE,
+                stderr=PIPE,
+                text=True,
+            )
+            gevent.sleep(1.9)
+            proc.send_signal(signal.SIGINT)
+            stdout, stderr = proc.communicate()
+            print(stderr, stdout)
+            self.assertIn("Shape test starting", stderr)
+            self.assertIn("Exiting due to CTRL+C interruption", stderr)
+            self.assertIn("Test Stopped", stdout)
+            # ensure stats printer printed at least one report before shutting down and that there was a final report printed as well
+            self.assertRegex(stderr, r".*Aggregated[\S\s]*Shutting down[\S\s]*Aggregated.*")
+
 
 class DistributedIntegrationTests(ProcessIntegrationTest):
     def test_expect_workers(self):