Add worker logs

locustio · Jun 4, 2024 · 08a9585 · 08a9585
1 parent b7caa8b
commit 08a9585
Show file tree

Hide file tree

Showing 14 changed files with 162 additions and 98 deletions.
diff --git a/locust/env.py b/locust/env.py
@@ -99,6 +99,8 @@ def __init__(
         """List of the available Tasks per User Classes to pick from in the Task Picker"""
         self.dispatcher_class = dispatcher_class
         """A user dispatcher class that decides how users are spawned, default :class:`UsersDispatcher <locust.dispatch.UsersDispatcher>`"""
+        self.worker_logs: dict[str, list[str]] = {}
+        """Captured logs from all connected workers"""
 
         self._remove_user_classes_with_weight_zero()
         self._validate_user_class_name_uniqueness()
@@ -209,6 +211,10 @@ def update_user_class(self, user_settings):
             if key == "tasks":
                 user_class.tasks = [task for task in user_tasks if task.__name__ in value]
 
+    def update_worker_logs(self, worker_log_report):
+        if worker_log_report.get("worker_id", None):
+            self.worker_logs[worker_log_report.get("worker_id")] = worker_log_report.get("logs", [])
+
     def _filter_tasks_by_tags(self) -> None:
         """
         Filter the tasks on all the user_classes recursively, according to the tags and

diff --git a/locust/log.py b/locust/log.py
@@ -75,6 +75,17 @@ def setup_logging(loglevel, logfile=None):
     logging.config.dictConfig(LOGGING_CONFIG)
 
 
+def get_logs():
+    log_reader_handler = next(
+        (handler for handler in logging.getLogger("root").handlers if handler.name == "log_reader"), None
+    )
+
+    if log_reader_handler:
+        return log_reader_handler.logs
+
+    return []
+
+
 def greenlet_exception_logger(logger, level=logging.CRITICAL):
     """
     Return a function that can be used as argument to Greenlet.link_exception() that will log the

diff --git a/locust/runners.py b/locust/runners.py
@@ -39,7 +39,7 @@
 from . import argument_parser
 from .dispatch import UsersDispatcher
 from .exception import RPCError, RPCReceiveError, RPCSendError
-from .log import greenlet_exception_logger
+from .log import get_logs, greenlet_exception_logger
 from .rpc import (
     Message,
     rpc,
@@ -66,6 +66,7 @@
     "missing",
 ]
 WORKER_REPORT_INTERVAL = 3.0
+WORKER_LOG_REPORT_INTERVAL = 10
 CPU_MONITOR_INTERVAL = 5.0
 CPU_WARNING_THRESHOLD = 90
 HEARTBEAT_INTERVAL = 1
@@ -1116,6 +1117,8 @@ def client_listener(self) -> NoReturn:
                 # a worker finished spawning (this happens multiple times during rampup)
                 self.clients[msg.node_id].state = STATE_RUNNING
                 self.clients[msg.node_id].user_classes_count = msg.data["user_classes_count"]
+            elif msg.type == "logs":
+                self.environment.update_worker_logs(msg.data)
             elif msg.type == "quit":
                 if msg.node_id in self.clients:
                     client = self.clients[msg.node_id]
@@ -1212,6 +1215,7 @@ def __init__(self, environment: Environment, master_host: str, master_port: int)
         self.client_id = socket.gethostname() + "_" + uuid4().hex
         self.master_host = master_host
         self.master_port = master_port
+        self.logs: list[str] = []
         self.worker_cpu_warning_emitted = False
         self._users_dispatcher: UsersDispatcher | None = None
         self.client = rpc.Client(master_host, master_port, self.client_id)
@@ -1220,6 +1224,7 @@ def __init__(self, environment: Environment, master_host: str, master_port: int)
         self.greenlet.spawn(self.heartbeat).link_exception(greenlet_exception_handler)
         self.greenlet.spawn(self.heartbeat_timeout_checker).link_exception(greenlet_exception_handler)
         self.greenlet.spawn(self.stats_reporter).link_exception(greenlet_exception_handler)
+        self.greenlet.spawn(self.logs_reporter).link_exception(greenlet_exception_handler)
 
         # register listener that adds the current number of spawned users to the report that is sent to the master node
         def on_report_to_master(client_id: str, data: dict[str, Any]):
@@ -1417,6 +1422,11 @@ def stats_reporter(self) -> NoReturn:
                 logger.error(f"Temporary connection lost to master server: {e}, will retry later.")
             gevent.sleep(WORKER_REPORT_INTERVAL)
 
+    def logs_reporter(self) -> NoReturn:
+        while True:
+            self._send_logs()
+            gevent.sleep(WORKER_LOG_REPORT_INTERVAL)
+
     def send_message(self, msg_type: str, data: dict[str, Any] | None = None, client_id: str | None = None) -> None:
         """
         Sends a message to master node
@@ -1433,6 +1443,13 @@ def _send_stats(self) -> None:
         self.environment.events.report_to_master.fire(client_id=self.client_id, data=data)
         self.client.send(Message("stats", data, self.client_id))
 
+    def _send_logs(self) -> None:
+        current_logs = get_logs()
+
+        if len(current_logs) > len(self.logs):
+            self.logs = current_logs
+            self.send_message("logs", {"worker_id": self.client_id, "logs": current_logs})
+
     def connect_to_master(self):
         self.retry += 1
         self.client.send(Message("client_ready", __version__, self.client_id))

diff --git a/locust/test/test_web.py b/locust/test/test_web.py
@@ -1020,7 +1020,7 @@ def test_logs(self):
 
         response = requests.get("http://127.0.0.1:%i/logs" % self.web_port)
 
-        self.assertIn(log_line, response.json().get("logs"))
+        self.assertIn(log_line, response.json().get("master"))
 
     def test_template_args(self):
         class MyUser(User):

diff --git a/locust/web.py b/locust/web.py
@@ -33,7 +33,7 @@
 from . import argument_parser
 from . import stats as stats_module
 from .html import BUILD_PATH, ROOT_PATH, STATIC_PATH, get_html_report
-from .log import greenlet_exception_logger
+from .log import get_logs, greenlet_exception_logger
 from .runners import STATE_MISSING, STATE_RUNNING, MasterRunner
 from .stats import StatsCSV, StatsCSVFileWriter, StatsErrorDict, sort_stats
 from .user.inspectuser import get_ratio
@@ -490,16 +490,7 @@ def tasks() -> dict[str, dict[str, dict[str, float]]]:
         @app.route("/logs")
         @self.auth_required_if_enabled
         def logs():
-            log_reader_handler = [
-                handler for handler in logging.getLogger("root").handlers if handler.name == "log_reader"
-            ]
-
-            if log_reader_handler:
-                logs = log_reader_handler[0].logs
-            else:
-                logs = []
-
-            return jsonify({"logs": logs})
+            return jsonify({"master": get_logs(), "workers": self.environment.worker_logs})
 
         @app.route("/login")
         def login():

diff --git a/locust/webui/dist/assets/index-cb7618a5.js → locust/webui/dist/assets/index-6b52989e.js b/locust/webui/dist/assets/index-cb7618a5.js → locust/webui/dist/assets/index-6b52989e.js
diff --git a/locust/webui/dist/auth.html b/locust/webui/dist/auth.html
@@ -7,7 +7,7 @@
     <meta name="theme-color" content="#000000" />
 
     <title>Locust</title>
-    <script type="module" crossorigin src="./assets/index-cb7618a5.js"></script>
+    <script type="module" crossorigin src="./assets/index-6b52989e.js"></script>
   </head>
   <body>
     <div id="root"></div>

diff --git a/locust/webui/dist/index.html b/locust/webui/dist/index.html
@@ -7,7 +7,7 @@
     <meta name="theme-color" content="#000000" />
 
     <title>Locust</title>
-    <script type="module" crossorigin src="./assets/index-cb7618a5.js"></script>
+    <script type="module" crossorigin src="./assets/index-6b52989e.js"></script>
   </head>
   <body>
     <div id="root"></div>

diff --git a/locust/webui/src/components/LogViewer/LogViewer.tsx b/locust/webui/src/components/LogViewer/LogViewer.tsx
@@ -1,7 +1,16 @@
-import { Box, Paper, Typography } from '@mui/material';
+import ExpandMoreIcon from '@mui/icons-material/ExpandMore';
+import {
+  Accordion,
+  AccordionDetails,
+  AccordionSummary,
+  Box,
+  Paper,
+  Typography,
+} from '@mui/material';
 import { red, orange, blue } from '@mui/material/colors';
 
 import { useSelector } from 'redux/hooks';
+import { objectLength } from 'utils/object';
 
 const getLogColor = (log: string) => {
   if (log.includes('CRITICAL')) {
@@ -19,21 +28,48 @@ const getLogColor = (log: string) => {
   return 'white';
 };
 
+function LogDisplay({ log }: { log: string }) {
+  return (
+    <Typography color={getLogColor(log)} variant='body2'>
+      {log}
+    </Typography>
+  );
+}
+
 export default function LogViewer() {
-  const logs = useSelector(({ logViewer: { logs } }) => logs);
+  const { master: masterLogs, workers: workerLogs } = useSelector(({ logViewer }) => logViewer);
 
   return (
-    <Box>
-      <Typography sx={{ mb: 2 }} variant='h5'>
-        Logs
-      </Typography>
-      <Paper elevation={3} sx={{ p: 2, fontFamily: 'monospace' }}>
-        {logs.map((log, index) => (
-          <Typography color={getLogColor(log)} key={`log-${index}`} variant='body2'>
-            {log}
+    <Box sx={{ display: 'flex', flexDirection: 'column', rowGap: 4 }}>
+      <Box>
+        <Typography sx={{ mb: 2 }} variant='h5'>
+          Master Logs
+        </Typography>
+        <Paper elevation={3} sx={{ p: 2, fontFamily: 'monospace' }}>
+          {masterLogs.map((log, index) => (
+            <LogDisplay key={`master-log-${index}`} log={log} />
+          ))}
+        </Paper>
+      </Box>
+      {!!objectLength(workerLogs) && (
+        <Box>
+          <Typography sx={{ mb: 2 }} variant='h5'>
+            Worker Logs
           </Typography>
-        ))}
-      </Paper>
+          {Object.entries(workerLogs).map(([workerId, logs], index) => (
+            <Accordion key={`worker-log-${index}`}>
+              <AccordionSummary expandIcon={<ExpandMoreIcon />}>{workerId}</AccordionSummary>
+              <AccordionDetails>
+                <Paper elevation={3} sx={{ p: 2, fontFamily: 'monospace' }}>
+                  {logs.map((log, index) => (
+                    <LogDisplay key={`worker-${workerId}-log-${index}`} log={log} />
+                  ))}
+                </Paper>
+              </AccordionDetails>
+            </Accordion>
+          ))}
+        </Box>
+      )}
     </Box>
   );
 }
diff --git a/locust/webui/src/components/LogViewer/tests/LogViewer.test.tsx b/locust/webui/src/components/LogViewer/tests/LogViewer.test.tsx
@@ -9,7 +9,7 @@ describe('LogViewer', () => {
     const { getByText } = renderWithProvider(<LogViewer />, {
       swarm: swarmStateMock,
       logViewer: {
-        logs: ['Log 1', 'Log 2', 'Log 3'],
+        master: ['Log 1', 'Log 2', 'Log 3'],
       },
     });
 

diff --git a/locust/webui/src/components/LogViewer/tests/useLogViewer.test.tsx b/locust/webui/src/components/LogViewer/tests/useLogViewer.test.tsx
@@ -11,7 +11,7 @@ import { renderWithProvider } from 'test/testUtils';
 const mockLogs = ['Log 1', 'Log 2', 'Log 3'];
 
 const server = setupServer(
-  http.get(`${TEST_BASE_API}/logs`, () => HttpResponse.json({ logs: mockLogs })),
+  http.get(`${TEST_BASE_API}/logs`, () => HttpResponse.json({ master: mockLogs })),
 );
 
 function MockHook() {
@@ -32,7 +32,7 @@ describe('useLogViewer', () => {
 
     await waitFor(() => {
       expect(getByTestId('logs').textContent).toBe(JSON.stringify(mockLogs));
-      expect(store.getState().logViewer.logs).toEqual(mockLogs);
+      expect(store.getState().logViewer.master).toEqual(mockLogs);
     });
   });
 });
diff --git a/locust/webui/src/components/LogViewer/useLogViewer.ts b/locust/webui/src/components/LogViewer/useLogViewer.ts
@@ -15,20 +15,20 @@ export default function useLogViewer() {
   const setLogs = useAction(logViewerActions.setLogs);
   const { data, refetch: refetchLogs } = useGetLogsQuery();
 
-  const logs = data ? data.logs : [];
+  const logs = data || { master: [], workers: {} };
 
   const shouldNotifyLogsUpdate = useCallback(
-    () => logs.slice(localStorage['logViewer']).some(isImportantLog),
+    () => logs.master.slice(localStorage['logViewer']).some(isImportantLog),
     [logs],
   );
 
   useInterval(refetchLogs, 5000, {
     shouldRunInterval: swarm.state === SWARM_STATE.SPAWNING || swarm.state == SWARM_STATE.RUNNING,
   });
-  useNotifications(logs, { key: 'logViewer', shouldNotify: shouldNotifyLogsUpdate });
+  useNotifications(logs.master, { key: 'logViewer', shouldNotify: shouldNotifyLogsUpdate });
 
   useEffect(() => {
-    setLogs({ logs });
+    setLogs(logs);
   }, [logs]);
 
   return logs;

diff --git a/locust/webui/src/redux/slice/logViewer.slice.ts b/locust/webui/src/redux/slice/logViewer.slice.ts
@@ -1,15 +1,15 @@
 import { PayloadAction, createSlice } from '@reduxjs/toolkit';
 
 import { updateStateWithPayload } from 'redux/utils';
+import { ILogsResponse } from 'types/ui.types';
 
-export interface ILogViewerState {
-  logs: string[];
-}
+export interface ILogViewerState extends ILogsResponse {}
 
 export type LogViewerAction = PayloadAction<ILogViewerState>;
 
 const initialState = {
-  logs: [] as string[],
+  master: [] as string[],
+  workers: {},
 };
 
 const logViewerSlice = createSlice({

diff --git a/locust/webui/src/types/ui.types.ts b/locust/webui/src/types/ui.types.ts
@@ -103,5 +103,8 @@ export interface IStatsResponse {
 }
 
 export interface ILogsResponse {
-  logs: string[];
+  master: string[];
+  workers: {
+    [key: string]: string[];
+  };
 }