Merge pull request #1354 from stratosphereips/alya/fix_performance_profilers

Fix CPU profiler "dev" mode

Author: AlyaGomaa

README.md

@@ -23,6 +23,7 @@ Slips v1.1.7
 [![License](https://img.shields.io/badge/Blog-Stratosphere-cyan)](https://www.stratosphereips.org/blog/tag/slips)
 [![Discord](https://img.shields.io/discord/761894295376494603?label=&logo=discord&logoColor=ffffff&color=7389D8&labelColor=6A7EC2)](https://discord.gg/zu5HwMFy5C)
 ![Twitter Follow](https://img.shields.io/twitter/follow/StratosphereIPS?style=social)
+
 <hr>
 
 

config/slips.yaml

@@ -455,7 +455,8 @@ Docker:
 #############################
 Profiling:
   # CPU profiling
-  # enable cpu profiling [yes,no]
+  # enable cpu profiling [true/false]
+  # NOTE: the cpu profiler uses port 9001 to show the results.
   cpu_profiler_enable: false
 
   # Available options are [dev,live]
@@ -466,25 +467,31 @@ Profiling:
   # time. it is accessible from web interface
   cpu_profiler_mode: dev
 
-  # profile all subprocesses in dev mode [yes,no].
-  cpu_profiler_multiprocess: true
+  # decides whether the profiler tracks all processes or only one.
+  # only used  in dev mode [true,false].
+  cpu_profiler_multiprocess: false
 
   # set number of tracer entries (dev mode only)
-  cpu_profiler_dev_mode_entries: 1000000
+  # VizTracer uses a circular buffer to store the entries.
+  # When there are too many entries, it will only store the latest ones
+  # so you know what happened recently.
+  # the more the entries, the more RAM viztracer is going to use.
+  # https://viztracer.readthedocs.io/en/latest/basic_usage.html#circular-buffer-size
+  cpu_profiler_dev_mode_entries: 500000
 
   # set maximum output lines (live mode only)
   cpu_profiler_output_limit: 20
 
   # set the wait time between sampling sequences in seconds (live mode only)
   cpu_profiler_sampling_interval: 20
 
-  # enable memory profiling [yes,no]
+  # enable memory profiling [true,false]
   memory_profiler_enable: false
 
   # set profiling mode [dev,live]
   memory_profiler_mode: live
 
-  # profile all subprocesses [yes,no]
+  # profile all subprocesses [true,false]
   memory_profiler_multiprocess: true
 
 #############################

docker/Dockerfile

@@ -30,7 +30,6 @@ ENV NVM_DIR=/root/.nvm
 SHELL ["/bin/bash", "-c"]
 
 
-# Install wget and add Zeek and redis repositories to our sources.
 RUN apt update && apt install -y --no-install-recommends \
     wget \
     ca-certificates \

docs/profiling_slips.md

@@ -85,7 +85,7 @@ Use WASD to zoom in and move left/right.
 ### CPU Profiler Live Mode
 #### Step 1
 Go to slips.yaml and make sure the settings are set correctly.
-```cpu_profiler_enable = yes```
+```cpu_profiler_enable = True```
 ```cpu_profiler_mode = live```
 
 You can also set maximum output lines (live mode only)  to adjust profiler behavior.
@@ -112,15 +112,15 @@ If we print out the data getting sent to the “cpu_profile” redis channel, it
 
 The memory profiler settings are much simpler.
 in slips.yaml, first, enable memory profiling
-```memory_profiler_enable = yes```
+```memory_profiler_enable = True```
 
 
 and set profiling mode:
 
 ```memory_profiler_mode = dev```
 
 now, profile all subprocesses
-```memory_profiler_multiprocess = yes```
+```memory_profiler_multiprocess = True```
 
 #### Step 2
 
@@ -158,9 +158,9 @@ Under the table directory, the files are much simpler. They just show a table of
 Go to ```slips.yaml``` and use the following settings
 
 ```
-memory_profiler_enable = yes
-memory_profiler_mode = yes
-memory_profiler_multiprocess = yes
+memory_profiler_enable = True
+memory_profiler_mode = True
+memory_profiler_multiprocess = True
 ```
 
 ### Step 2

managers/process_manager.py

@@ -721,9 +721,6 @@ def shutdown_gracefully(self):
                 format_ = self.main.conf.export_labeled_flows_to().lower()
                 self.main.db.export_labeled_flows(format_)
 
-            self.main.profilers_manager.cpu_profiler_release()
-            self.main.profilers_manager.memory_profiler_release()
-
             # if store_a_copy_of_zeek_files is set to yes in slips.yaml
             # copy the whole zeek_files dir to the output dir
             self.main.store_zeek_dir_copy()
@@ -740,6 +737,9 @@ def shutdown_gracefully(self):
                 f"finished in {analysis_time:.2f} minutes"
             )
 
+            self.main.profilers_manager.cpu_profiler_release()
+            self.main.profilers_manager.memory_profiler_release()
+
             self.main.db.close()
             if graceful_shutdown:
                 print(

managers/profilers_manager.py

@@ -3,11 +3,13 @@
 import os
 import subprocess
 import sys
+from slips_files.common.style import green
 
 
 class ProfilersManager:
     def __init__(self, main):
         self.main = main
+        self.args = self.main.args
         self.read_configurations()
 
     def read_configurations(self):
@@ -50,28 +52,38 @@ def cpu_profiler_init(self):
                 args = sys.argv
                 if args[-1] != "--no-recurse":
                     tracer_entries = str(self.cpu_profiler_dev_mode_entries)
+                    output_file = str(
+                        os.path.join(
+                            self.args.output,
+                            "cpu_profiling_result.json",
+                        )
+                    )
                     viz_args = [
                         "viztracer",
                         "--tracer_entries",
                         tracer_entries,
                         "--max_stack_depth",
-                        "10",
+                        "5",
                         "-o",
-                        str(
-                            os.path.join(
-                                self.args.output,
-                                "cpu_profiling_result.json",
-                            )
-                        ),
+                        output_file,
+                        # viztracer takes -- as a separator between arguments
+                        # to viztracer and positional arguments to your script.
+                        "--",
                     ]
+                    # add slips args
                     viz_args.extend(args)
+                    # add --no-recurse to avoid infinite recursion
                     viz_args.append("--no-recurse")
                     print(
-                        "Starting multiprocess profiling recursive subprocess"
+                        f"Starting multiprocess profiling recursive "
+                        f"subprocess using command: "
+                        f"{green(' '.join(viz_args))}"
                     )
                     subprocess.run(viz_args)
                     exit(0)
             else:
+                # reaching here means slips is now running using the vistracer
+                # command
                 self.cpu_profiler = CPUProfiler(
                     db=self.main.db,
                     output=self.args.output,

slips/main.py

@@ -45,7 +45,6 @@ def __init__(self, testing=False):
         self.conf = ConfigParser()
         self.metadata_man = MetadataManager(self)
         self.ui_man = UIManager(self)
-        self.profilers_manager = ProfilersManager(self)
         self.version = utils.get_slips_version()
         # will be filled later
         self.commit = "None"
@@ -57,6 +56,7 @@ def __init__(self, testing=False):
         # TODO use mocks instead of this testing param
         if not testing:
             self.args = self.conf.get_args()
+            self.profilers_manager = ProfilersManager(self)
             self.pid = os.getpid()
             self.checker.check_given_flags()
 

slips_files/common/performance_profilers/cpu_profiler.py

@@ -15,6 +15,9 @@
 
 class CPUProfiler(IPerformanceProfiler):
     def __init__(self, db, output, mode="dev", limit=20, interval=20):
+        """
+        :param output: the currently used slips output directory
+        """
         valid_modes = ["dev", "live"]
         if mode not in valid_modes:
             print(
@@ -46,6 +49,9 @@ def print(self):
 
 class DevProfiler(IPerformanceProfiler):
     def __init__(self, output):
+        """
+        :param output: the currently used slips output directory
+        """
         self.profiler = self._create_profiler()
         self.output = output
 
@@ -59,8 +65,12 @@ def stop(self):
         self.profiler.stop()
 
     def print(self):
+        """
+        Prints the profiling result to cpu_profiling_result in slips output
+        directory
+        """
         result_path = os.path.join(self.output, "cpu_profiling_result.json")
-        self.profiler.save(result_path)
+        self.profiler.save(output_file=result_path)
 
 
 class LiveProfiler(IPerformanceProfiler):

slips_files/core/database/sqlite_db/database.py

@@ -434,12 +434,3 @@ def execute(self, query: str, params=None) -> None:
 
                 elif "database is locked" in str(err):
                     sleep(5)
-                else:
-                    self.print(
-                        f"Re-trying to execute query "
-                        f"({query} - {params}) - {err}. "
-                        f"Trial number: {trial}.",
-                        0,
-                        1,
-                        log_to_logfiles_only=True,
-                    )

tests/module_factory.py

@@ -12,6 +12,7 @@
 
 from managers.host_ip_manager import HostIPManager
 from managers.metadata_manager import MetadataManager
+from managers.profilers_manager import ProfilersManager
 from modules.flowalerts.conn import Conn
 from modules.threat_intelligence.circl_lu import Circllu
 from modules.threat_intelligence.spamhaus import Spamhaus
@@ -364,6 +365,13 @@ def create_redis_manager_obj(self, mock_db):
         main.args = Mock()
         return RedisManager(main)
 
+    @patch(MODULE_DB_MANAGER, name="mock_db")
+    def create_profilers_manager_obj(self, mock_db):
+        main = self.create_main_obj()
+        main.db = mock_db
+        main.args = Mock()
+        return ProfilersManager(main)
+
     @patch(MODULE_DB_MANAGER, name="mock_db")
     def create_host_ip_manager_obj(self, mock_db):
         main = self.create_main_obj()

tests/test_main.py

@@ -60,60 +60,6 @@ def test_is_total_flows_unknown(args, input_type, expected_result):
     assert main.is_total_flows_unknown() == expected_result
 
 
-@pytest.mark.parametrize(
-    "cpu_profiler_multiprocess, expected_stop_calls, expected_print_calls",
-    [  # Testcase1: CPU profiler enabled, not multiprocess
-        (False, 1, 1),
-        # Testcase2: CPU profiler enabled, multiprocess
-        (True, 0, 0),
-    ],
-)
-def test_cpu_profiler_release_enabled(
-    cpu_profiler_multiprocess,
-    expected_stop_calls,
-    expected_print_calls,
-):
-    main = ModuleFactory().create_main_obj()
-    main.profilers_manager.cpu_profiler_enabled = True
-    main.profilers_manager.cpu_profiler_multiprocess = (
-        cpu_profiler_multiprocess
-    )
-    main.profilers_manager.cpu_profiler = MagicMock()
-    main.profilers_manager.cpu_profiler_release()
-
-    assert (
-        main.profilers_manager.cpu_profiler.stop.call_count
-        == expected_stop_calls
-    )
-    assert (
-        main.profilers_manager.cpu_profiler.print.call_count
-        == expected_print_calls
-    )
-
-
-def test_cpu_profiler_release_disabled():
-    main = ModuleFactory().create_main_obj()
-    main.profilers_manager.cpu_profiler_enabled = False
-    main.profilers_manager.cpu_profiler_release()
-    assert not hasattr(main.profilers_manager, "memory_profiler")
-
-
-def test_memory_profiler_release_enabled():
-    main = ModuleFactory().create_main_obj()
-    main.profilers_manager.memory_profiler_enabled = True
-    main.profilers_manager.memory_profiler = MagicMock()
-    main.profilers_manager.memory_profiler_release()
-    main.profilers_manager.memory_profiler.stop.assert_called_once()
-
-
-def test_memory_profiler_release_disabled():
-    main = ModuleFactory().create_main_obj()
-    main.profilers_manager.memory_profiler_enabled = False
-    main.profilers_manager.memory_profiler_release()
-
-    assert not hasattr(main.profilers_manager, "memory_profiler")
-
-
 @pytest.mark.parametrize(
     "mode, time_diff, expected_calls",
     [  # Testcase1: Should update stats

tests/test_profilers_manager.py

@@ -0,0 +1,51 @@
+# SPDX-FileCopyrightText: 2021 Sebastian Garcia <[email protected]>
+# SPDX-License-Identifier: GPL-2.0-only
+from unittest.mock import MagicMock
+import pytest
+from tests.module_factory import ModuleFactory
+
+
+@pytest.mark.parametrize(
+    "cpu_profiler_multiprocess, expected_stop_calls, expected_print_calls",
+    [  # Testcase1: CPU profiler enabled, not multiprocess
+        (False, 1, 1),
+        # Testcase2: CPU profiler enabled, multiprocess
+        (True, 0, 0),
+    ],
+)
+def test_cpu_profiler_release_enabled(
+    cpu_profiler_multiprocess,
+    expected_stop_calls,
+    expected_print_calls,
+):
+    handler = ModuleFactory().create_profilers_manager_obj()
+    handler.cpu_profiler_enabled = True
+    handler.cpu_profiler_multiprocess = cpu_profiler_multiprocess
+    handler.cpu_profiler = MagicMock()
+    handler.cpu_profiler_release()
+
+    assert handler.cpu_profiler.stop.call_count == expected_stop_calls
+    assert handler.cpu_profiler.print.call_count == expected_print_calls
+
+
+def test_cpu_profiler_release_disabled():
+    handler = ModuleFactory().create_profilers_manager_obj()
+    handler.cpu_profiler_enabled = False
+    handler.cpu_profiler_release()
+    assert not hasattr(handler, "memory_profiler")
+
+
+def test_memory_profiler_release_enabled():
+    handler = ModuleFactory().create_profilers_manager_obj()
+    handler.memory_profiler_enabled = True
+    handler.memory_profiler = MagicMock()
+    handler.memory_profiler_release()
+    handler.memory_profiler.stop.assert_called_once()
+
+
+def test_memory_profiler_release_disabled():
+    handler = ModuleFactory().create_profilers_manager_obj()
+    handler.memory_profiler_enabled = False
+    handler.memory_profiler_release()
+
+    assert not hasattr(handler, "memory_profiler")