Merge pull request #2067 from locustio/add-cpu-warning-event

Add cpu_warning event, so listeners can do some action when CPU usage is too high
locustio · Apr 6, 2022 · 482aafd · 482aafd
2 parents 6899adc + 3f5f599
commit 482aafd
Show file tree

Hide file tree

Showing 3 changed files with 22 additions and 5 deletions.
diff --git a/locust/event.py b/locust/event.py
@@ -207,6 +207,11 @@ class Events:
     Fired when the Reset Stats button is clicked in the web UI.
     """
 
+    cpu_warning: EventHook
+    """
+    Fired when the CPU usage exceeds runners.CPU_WARNING_THRESHOLD (90% by default)
+    """
+
     def __init__(self):
         # For backwarde compatibility use also values of class attributes
         for name, value in vars(type(self)).items():

diff --git a/locust/runners.py b/locust/runners.py
@@ -56,6 +56,7 @@
 ]
 WORKER_REPORT_INTERVAL = 3.0
 CPU_MONITOR_INTERVAL = 5.0
+CPU_WARNING_THRESHOLD = 90
 HEARTBEAT_INTERVAL = 1
 HEARTBEAT_LIVENESS = 3
 HEARTBEAT_DEAD_INTERNAL = -60
@@ -288,11 +289,13 @@ def monitor_cpu_and_memory(self):
         while True:
             self.current_cpu_usage = process.cpu_percent()
             self.current_memory_usage = process.memory_info().rss
-            if self.current_cpu_usage > 90 and not self.cpu_warning_emitted:
-                logging.warning(
-                    "CPU usage above 90%! This may constrain your throughput and may even give inconsistent response time measurements! See https://docs.locust.io/en/stable/running-locust-distributed.html for how to distribute the load over multiple CPU cores or machines"
-                )
-                self.cpu_warning_emitted = True
+            if self.current_cpu_usage > CPU_WARNING_THRESHOLD:
+                self.environment.events.cpu_warning.fire(environment=self.environment, cpu_usage=self.current_cpu_usage)
+                if not self.cpu_warning_emitted:
+                    logging.warning(
+                        f"CPU usage above {CPU_WARNING_THRESHOLD}%! This may constrain your throughput and may even give inconsistent response time measurements! See https://docs.locust.io/en/stable/running-locust-distributed.html for how to distribute the load over multiple CPU cores or machines"
+                    )
+                    self.cpu_warning_emitted = True
             gevent.sleep(CPU_MONITOR_INTERVAL)
 
     def start(self, user_count: int, spawn_rate: float, wait: bool = False):

diff --git a/locust/test/test_runners.py b/locust/test/test_runners.py
@@ -144,10 +144,19 @@ def cpu_task(self):
                         _ = 3 / 2
 
             environment = Environment(user_classes=[CpuUser])
+            environment._cpu_warning_event_triggered = False
+
+            def cpu_warning(environment, cpu_usage, **kwargs):
+                environment._cpu_warning_event_triggered = True
+                environment._cpu_usage = cpu_usage
+
+            environment.events.cpu_warning.add_listener(cpu_warning)
             runner = LocalRunner(environment)
             self.assertFalse(runner.cpu_warning_emitted)
             runner.spawn_users({CpuUser.__name__: 1}, wait=False)
             sleep(2.5)
+            self.assertTrue(environment._cpu_warning_event_triggered)
+            self.assertGreater(environment._cpu_usage, 90)
             runner.quit()
             self.assertTrue(runner.cpu_warning_emitted)
         finally: