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:
@max-rocket-internet Do you have a minute to look at this?
This issue is stale because it has been open 30 days with no activity. Remove stale label or comment or this will be closed in 10 days.
This issue was closed because it has been stalled for 10 days with no activity. This does not necessarily mean that the issue is bad, but it most likely means that nobody is willing to take the time to fix it. If you have found Locust useful, then consider contributing a fix yourself!
I believe this issue still persists (using 2.14.2). From the error trace we see that it fails in
`environment.runner.shape_greenlet.join()`

It's dirty, but we could catch it like this to let it end gracefully: 

```
# start the test
            if environment.shape_class:
                if options.run_time:
                    sys.stderr.write("It makes no sense to combine --run-time and LoadShapes. Bailing out.\n")
                    sys.exit(1)
                try:
                    logger.info("Modified Locust Main is running")
                    environment.runner.start_shape()
                    environment.runner.shape_greenlet.join()
                except KeyboardInterrupt:
                    logging.info("Exiting due to CTRL+C interruption")
                finally:
                    stop_and_optionally_quit()
```

Thoughts? Locust's inability to gracefully exit means `@events.test_stop.add_listener` doesn't work. 
LGTM, PR very welcome. PR with unit test even more so :) 

I'd skip the "modified locust main is running" logging though (unclear what that even means)

Resolving patch:
diff --git a/locust/main.py b/locust/main.py
--- a/locust/main.py
+++ b/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)


Test patch:
diff --git a/locust/test/test_main.py b/locust/test/test_main.py
--- a/locust/test/test_main.py
+++ b/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):