test: Enable profiling in integration tests (#5130)

tests/integration_tests/assets/enable_profile.py

@@ -0,0 +1,25 @@
+from pathlib import Path
+
+services = [
+    "cloud-init-local.service",
+    "cloud-init.service",
+    "cloud-config.service",
+    "cloud-final.service",
+]
+service_dir = Path("/lib/systemd/system/")
+
+# Check for the existence of the service files
+for service in services:
+    if not (service_dir / service).is_file():
+        print(f"Error: {service} does not exist in {service_dir}")
+        exit(1)
+
+# Prepend the ExecStart= line with 'python3 -m coverage run'
+for service in services:
+    file_path = service_dir / service
+    content = file_path.read_text()
+    content = content.replace(
+        "ExecStart=/usr",
+        (f"ExecStart=python3 -m cProfile -o /var/log/{service}.stats /usr"),
+    )
+    file_path.write_text(content)

tests/integration_tests/conftest.py

@@ -118,17 +118,30 @@ def setup_image(session_cloud: IntegrationCloud, request):
     """
     source = get_validated_source(session_cloud)
     if not (
-        source.installs_new_version() or integration_settings.INCLUDE_COVERAGE
+        source.installs_new_version()
+        or integration_settings.INCLUDE_COVERAGE
+        or integration_settings.INCLUDE_PROFILE
     ):
         return
     log.info("Setting up source image")
     client = session_cloud.launch()
     if source.installs_new_version():
         log.info("Installing cloud-init from %s", source.name)
         client.install_new_cloud_init(source)
+    if (
+        integration_settings.INCLUDE_PROFILE
+        and integration_settings.INCLUDE_COVERAGE
+    ):
+        log.error(
+            "Invalid configuration, cannot enable both profile and coverage."
+        )
+        raise ValueError()
     if integration_settings.INCLUDE_COVERAGE:
         log.info("Installing coverage")
         client.install_coverage()
+    elif integration_settings.INCLUDE_PROFILE:
+        log.info("Installing profiler")
+        client.install_profile()
     # All done customizing the image, so snapshot it and make it global
     snapshot_id = client.snapshot()
     client.cloud.snapshot_id = snapshot_id
@@ -169,6 +182,30 @@ def _collect_coverage(instance: IntegrationInstance, log_dir: Path):
         log.error("Failed to pull coverage for: %s", e)
 
 
+def _collect_profile(instance: IntegrationInstance, log_dir: Path):
+    log.info("Writing profile to %s", log_dir)
+    try:
+        (log_dir / "profile").mkdir(parents=True)
+        instance.pull_file(
+            "/var/log/cloud-init-local.service.stats",
+            log_dir / "profile" / "local.stats",
+        )
+        instance.pull_file(
+            "/var/log/cloud-init.service.stats",
+            log_dir / "profile" / "network.stats",
+        )
+        instance.pull_file(
+            "/var/log/cloud-config.service.stats",
+            log_dir / "profile" / "config.stats",
+        )
+        instance.pull_file(
+            "/var/log/cloud-final.service.stats",
+            log_dir / "profile" / "final.stats",
+        )
+    except Exception as e:
+        log.error("Failed to pull profile for: %s", e)
+
+
 def _setup_artifact_paths(node_id: str):
     parent_dir = Path(integration_settings.LOCAL_LOG_PATH, session_start_time)
 
@@ -209,7 +246,12 @@ def _collect_artifacts(
         integration_settings.COLLECT_LOGS == "ON_ERROR" and test_failed
     )
     should_collect_coverage = integration_settings.INCLUDE_COVERAGE
-    if not (should_collect_logs or should_collect_coverage):
+    should_collect_profile = integration_settings.INCLUDE_PROFILE
+    if not (
+        should_collect_logs
+        or should_collect_coverage
+        or should_collect_profile
+    ):
         return
 
     log_dir = _setup_artifact_paths(node_id)
@@ -220,6 +262,9 @@ def _collect_artifacts(
     if should_collect_coverage:
         _collect_coverage(instance, log_dir)
 
+    elif should_collect_profile:
+        _collect_profile(instance, log_dir)
+
 
 @contextmanager
 def _client(
@@ -385,7 +430,20 @@ def _generate_coverage_report() -> None:
     log.info("Coverage report generated")
 
 
+def _generate_profile_report() -> None:
+    log.info("Profile reports generated, run the following to view:")
+    command = (
+        "python3 -m snakeviz /tmp/cloud_init_test_logs/"
+        "last/tests/integration_tests/*/*/*/profile/%s"
+    )
+    log.info(command, "local.stats")
+    log.info(command, "network.stats")
+    log.info(command, "config.stats")
+    log.info(command, "final.stats")
+
+
 def pytest_sessionfinish(session, exitstatus) -> None:
-    if not integration_settings.INCLUDE_COVERAGE:
-        return
-    _generate_coverage_report()
+    if integration_settings.INCLUDE_COVERAGE:
+        _generate_coverage_report()
+    elif integration_settings.INCLUDE_PROFILE:
+        _generate_profile_report()

tests/integration_tests/instances.py

@@ -166,6 +166,13 @@ def install_coverage(self):
         )
         assert self.execute("python3 /var/tmp/enable_coverage.py").ok
 
+    def install_profile(self):
+        self.push_file(
+            local_path=ASSETS_DIR / "enable_profile.py",
+            remote_path="/var/tmp/enable_profile.py",
+        )
+        assert self.execute("python3 /var/tmp/enable_profile.py").ok
+
     def install_new_cloud_init(
         self,
         source: CloudInitSource,

tests/integration_tests/integration_settings.py

@@ -88,6 +88,12 @@
 # `html` directory with the coverage report.
 INCLUDE_COVERAGE = False
 
+# We default our profile to False because it involves modifying the
+# cloud-init systemd services, which is too intrusive of a change to
+# enable by default. If changed to true, the test directory corresponding
+# to the test run under LOCAL_LOG_PATH defined above will contain a report
+INCLUDE_PROFILE = False
+
 ##################################################################
 # USER SETTINGS OVERRIDES
 ##################################################################