Repository: locustio/locust

Issue:
KeyboardInterrupt does not gracefully exit custom load shape

### Describe the bug
KeyboardInterrupt does not gracefully exit custom load shape.

### Expected behavior
A report should be generated after Ctrl-C, like with a normal test case.

### Actual behavior
The program crashes:

```
[2021-10-21 10:07:23,218] user-vm/INFO/locust.runners: Shape worker starting
[2021-10-21 10:07:23,218] user-vm/INFO/locust.runners: Shape test updating to 10 users at 10.00 spawn rate
[2021-10-21 10:07:23,218] user-vm/INFO/locust.runners: Ramping to 10 users at a rate of 10.00 per second
[2021-10-21 10:07:23,219] user-vm/INFO/locust.runners: All users spawned: {"WebsiteUser": 10} (10 total users)
^CKeyboardInterrupt
2021-10-21T08:07:24Z
Traceback (most recent call last):
  File "/home/user/.local/bin/locust", line 8, in <module>
    sys.exit(main())
  File "/home/user/.local/lib/python3.7/site-packages/locust/main.py", line 383, in main
    start_automatic_run()
  File "/home/user/.local/lib/python3.7/site-packages/locust/main.py", line 370, in start_automatic_run
    environment.runner.shape_greenlet.join()
  File "src/gevent/greenlet.py", line 831, in gevent._gevent_cgreenlet.Greenlet.join
  File "src/gevent/greenlet.py", line 857, in gevent._gevent_cgreenlet.Greenlet.join
  File "src/gevent/greenlet.py", line 846, in gevent._gevent_cgreenlet.Greenlet.join
  File "src/gevent/_greenlet_primitives.py", line 61, in gevent._gevent_c_greenlet_primitives.SwitchOutGreenletWithLoop.switch
  File "src/gevent/_greenlet_primitives.py", line 61, in gevent._gevent_c_greenlet_primitives.SwitchOutGreenletWithLoop.switch
  File "src/gevent/_greenlet_primitives.py", line 65, in gevent._gevent_c_greenlet_primitives.SwitchOutGreenletWithLoop.switch
  File "src/gevent/_gevent_c_greenlet_primitives.pxd", line 35, in gevent._gevent_c_greenlet_primitives._greenlet_switch
KeyboardInterrupt

```

### Steps to reproduce
Use [this example](https://github.com/locustio/locust/blob/master/examples/custom_shape/step_load.py) from this very repo, launched as `locust --headless -f step_load.py --host https://localhost`:

### Environment

- OS: Linux
- Python version: 3.7.3
- Locust version: 2.4.1
- Locust command line that you ran:`locust --headless -f step_load.py --host https://localhost`
- Locust file contents (anonymized if necessary):  https://github.com/locustio/locust/blob/master/examples/custom_shape/step_load.py

Discussion:
JevonCowell: Please let me know if any other tests need to be written. This is my first Pull Request on an Open Source Project 😃.
cyberw: Thanks! I dont think the flake8 issue is your fault so you can ignore it. In the test, I think you should send the interrupt signal from the test rather than throwing an exception from inside the locustfile. Can you change that? There are some other tests like that in test_main.py
cyberw: Nice! For future contributions, use `#issuenumber` instead of `ISSUE-issuenumber:` in commit messages to connect them to a specific issue (or leave it out completely, its enough that the PR name and possibly branch name reference it). 

I'll fix the flake8 issue in master.

Patch:
diff --git a/locust/main.py b/locust/main.py
index df6caa3b..cab7824e 100644
--- a/locust/main.py
+++ b/locust/main.py
@@ -378,9 +378,13 @@ See https://github.com/locustio/locust/wiki/Installation#increasing-maximum-numb
                 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)
diff --git a/tox.ini b/tox.ini
index 94f43958..b85cc8a2 100644
--- a/tox.ini
+++ b/tox.ini
@@ -39,7 +39,7 @@ deps = black==23.1.0
 commands = black --check .
 
 [testenv:flake8]
-deps = flake8
+deps = flake8==6.0.0
 commands = flake8 . --count --show-source --statistics
 
 [testenv:mypy]


Test patch:
diff --git a/locust/test/test_main.py b/locust/test/test_main.py
index cf76cd93..6048f3d1 100644
--- a/locust/test/test_main.py
+++ b/locust/test/test_main.py
@@ -1185,6 +1185,51 @@ class MyUser(HttpUser):
             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):