Merge pull request #1917 from solowalker27/track-memory

Track worker memory

locust/runners.py

@@ -87,7 +87,8 @@ def __init__(self, environment):
         self.current_cpu_usage = 0
         self.cpu_warning_emitted = False
         self.worker_cpu_warning_emitted = False
-        self.greenlet.spawn(self.monitor_cpu).link_exception(greenlet_exception_handler)
+        self.current_memory_usage = 0
+        self.greenlet.spawn(self.monitor_cpu_and_memory).link_exception(greenlet_exception_handler)
         self.exceptions = {}
         # Because of the way the ramp-up/ramp-down is implemented, target_user_classes_count
         # is only updated at the end of the ramp-up/ramp-down.
@@ -282,10 +283,11 @@ def stop_users(self, user_classes_stop_count: Dict[str, int]):
             "%g users have been stopped, %g still running", sum(user_classes_stop_count.values()), self.user_count
         )
 
-    def monitor_cpu(self):
+    def monitor_cpu_and_memory(self):
         process = psutil.Process()
         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"
@@ -534,6 +536,7 @@ def __init__(self, id: str, state=STATE_INIT, heartbeat_liveness=HEARTBEAT_LIVEN
         self.heartbeat = heartbeat_liveness
         self.cpu_usage = 0
         self.cpu_warning_emitted = False
+        self.memory_usage = 0
         # The reported users running on the worker
         self.user_classes_count: Dict[str, int] = {}
 
@@ -928,6 +931,8 @@ def client_listener(self):
                         logger.warning(
                             "Worker %s exceeded cpu threshold (will only log this once per worker)" % (msg.node_id)
                         )
+                    if "current_memory_usage" in msg.data:
+                        c.memory_usage = msg.data["current_memory_usage"]
             elif msg.type == "stats":
                 self.environment.events.worker_report.fire(client_id=msg.node_id, data=msg.data)
             elif msg.type == "spawning":
@@ -1103,6 +1108,7 @@ def heartbeat(self):
                         {
                             "state": self.worker_state,
                             "current_cpu_usage": self.current_cpu_usage,
+                            "current_memory_usage": self.current_memory_usage,
                         },
                         self.client_id,
                     )

locust/static/locust.js

@@ -45,6 +45,19 @@ $(".close_link").click(function(event) {
     $(this).parent().parent().hide();
 });
 
+function formatBytes(bytes, decimals = 2) {
+    if (bytes === 0) return '0 Bytes';
+    if (bytes === 0) return 'N/A';
+
+    const k = 1024;
+    const dm = decimals < 0 ? 0 : decimals;
+    const sizes = ['Bytes', 'KB', 'MB', 'GB', 'TB', 'PB', 'EB', 'ZB', 'YB'];
+
+    const i = Math.floor(Math.log(bytes) / Math.log(k));
+
+    return parseFloat((bytes / Math.pow(k, i)).toFixed(dm)) + ' ' + sizes[i];
+}
+
 $("ul.tabs").tabs("div.panes > div").on("onClick", function (event) {
     // trigger resizing of charts
     resizeCharts();

locust/templates/index.html

@@ -219,6 +219,7 @@ <h2>Edit running load test</h2>
                                 <th class="stats_label" href="#" data-sortkey="state">State</th>
                                 <th class="stats_label numeric" href="#" data-sortkey="user_count" title="Number of users on this worker"># users</th>
                                 <th class="stats_label numeric" href="#" data-sortkey="cpu_usage" title="CPU usage of process">CPU usage</th>
+                                <th class="stats_label numeric" href="#" data-sortkey="memory_usage" title="Memory usage of process">Memory usage</th>
                             </tr>
                         </thead>
                         <tbody>
@@ -321,6 +322,7 @@ <h2>Version <a href="https://github.com/locustio/locust/releases/tag/{{version}}
             <td><%= this.state %></td>
             <td class="numeric"><%= this.user_count %></td>
             <td class="numeric"><%= this.cpu_usage %>%</td>
+            <td class="numeric"><%= formatBytes(this.memory_usage) %></td>
         </tr>
         <% alternate = !alternate; %>
         ]]>

locust/test/test_runners.py

@@ -1781,13 +1781,25 @@ def my_task(self):
 
             sleep(0.2)
             server.mocked_send(
-                Message("heartbeat", {"state": STATE_RUNNING, "current_cpu_usage": 50, "count": 1}, "fake_client1")
+                Message(
+                    "heartbeat",
+                    {"state": STATE_RUNNING, "current_cpu_usage": 50, "current_memory_usage": 200, "count": 1},
+                    "fake_client1",
+                )
             )
             server.mocked_send(
-                Message("heartbeat", {"state": STATE_RUNNING, "current_cpu_usage": 50, "count": 1}, "fake_client2")
+                Message(
+                    "heartbeat",
+                    {"state": STATE_RUNNING, "current_cpu_usage": 50, "current_memory_usage": 200, "count": 1},
+                    "fake_client2",
+                )
             )
             server.mocked_send(
-                Message("heartbeat", {"state": STATE_RUNNING, "current_cpu_usage": 50, "count": 1}, "fake_client3")
+                Message(
+                    "heartbeat",
+                    {"state": STATE_RUNNING, "current_cpu_usage": 50, "current_memory_usage": 200, "count": 1},
+                    "fake_client3",
+                )
             )
 
             sleep(0.2)
@@ -1798,7 +1810,11 @@ def my_task(self):
             )
 
             server.mocked_send(
-                Message("heartbeat", {"state": STATE_RUNNING, "current_cpu_usage": 50, "count": 1}, "fake_client1")
+                Message(
+                    "heartbeat",
+                    {"state": STATE_RUNNING, "current_cpu_usage": 50, "current_memory_usage": 200, "count": 1},
+                    "fake_client1",
+                )
             )
 
             sleep(0.4)
@@ -2762,9 +2778,10 @@ def my_task(self):
                 sleep(0.1)
 
             message = next((m for m in reversed(client.outbox) if m.type == "heartbeat"))
-            self.assertEqual(len(message.data), 2)
+            self.assertEqual(len(message.data), 3)
             self.assertIn("state", message.data)
             self.assertIn("current_cpu_usage", message.data)
+            self.assertIn("current_memory_usage", message.data)
 
             worker.quit()
 

locust/web.py

@@ -308,6 +308,7 @@ def request_stats():
                             "state": worker.state,
                             "user_count": worker.user_count,
                             "cpu_usage": worker.cpu_usage,
+                            "memory_usage": worker.memory_usage,
                         }
                     )