From 06cea62a40b1a7c0b28c4dafc17cb83617b287f0 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Donncha=20=C3=93=20Cearbhaill?=
 <donncha.ocearbhaill@amnesty.org>
Date: Wed, 16 Sep 2026 00:45:47 +0200
Subject: [PATCH] Record wall time and peak memory growth per module

Command.run() times every module and measures how much the process's
peak resident set size grew while it ran. Both are logged at INFO in
one line per module and written to info.json under module_stats, in
run order, so any run says which module cost what and CI can compare a
branch against main.

ru_maxrss is the process peak, so a module's growth is how much it
raised the high-water mark, which is what matters on a host without
swap; a module running after a bigger one reports 0. Where the resource
module is unavailable the growth is reported as 0.
---
 src/mvt/common/command.py    | 31 +++++++++++++++++++++++++++++++
 tests/common/test_command.py | 15 +++++++++++++++
 2 files changed, 46 insertions(+)

diff --git a/src/mvt/common/command.py b/src/mvt/common/command.py
index 6848e8c0..cb0e0f37 100644
--- a/src/mvt/common/command.py
+++ b/src/mvt/common/command.py
@@ -7,6 +7,7 @@ import json
 import logging
 import os
 import sys
+import time
 from datetime import datetime
 from functools import cached_property
 from heapq import heappop, heappush
@@ -36,6 +37,17 @@ from .utils import (
 from .version import MVT_VERSION
 
 
+def _peak_rss_mb() -> float:
+    """The process's peak resident set size so far, in MB (0 where unknown)."""
+    try:
+        import resource
+    except ImportError:
+        return 0.0
+    peak = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
+    # ru_maxrss is in bytes on macOS and in kilobytes elsewhere.
+    return peak / (1024 * 1024 if sys.platform == "darwin" else 1024)
+
+
 class Command:
     def __init__(
         self,
@@ -77,6 +89,9 @@ class Command:
         # declaring `replaces` to the module which took their place.
         self.module_replacements: dict[type[MVTModule], type[MVTModule]] = {}
 
+        # Wall time and peak-RSS growth of every module run, in run order.
+        self.module_stats: list[dict[str, Any]] = []
+
         # This list will contain all executed modules.
         # We can use this to reference e.g. self.executed[0].results.
         self.executed: list[MVTModule] = []
@@ -190,6 +205,7 @@ class Command:
             "date": convert_datetime_to_iso(datetime.now()),
             "ioc_files": [],
             "hashes": [],
+            "module_stats": self.module_stats,
         }
 
         for coll in self.iocs.ioc_collections:
@@ -766,6 +782,8 @@ class Command:
             except NotImplementedError:
                 pass
 
+            started = time.perf_counter()
+            peak_before = _peak_rss_mb()
             try:
                 run_module(m)
             except EncryptedBackupError:
@@ -775,6 +793,19 @@ class Command:
                 )
                 return
 
+            stats = {
+                "module": module.__name__,
+                "seconds": round(time.perf_counter() - started, 2),
+                "peak_rss_growth_mb": round(_peak_rss_mb() - peak_before),
+            }
+            self.module_stats.append(stats)
+            self.log.info(
+                "Module %s ran in %.1fs and raised the peak memory by %d MB",
+                stats["module"],
+                stats["seconds"],
+                stats["peak_rss_growth_mb"],
+            )
+
             self.executed.append(m)
             executed_by_type[module] = m
             self.timeline.extend(m.timeline)
diff --git a/tests/common/test_command.py b/tests/common/test_command.py
index 357d3b7d..2c96d4a3 100644
--- a/tests/common/test_command.py
+++ b/tests/common/test_command.py
@@ -298,6 +298,21 @@ class TestCommand:
         alerts = json.loads((tmp_path / "alerts.json").read_text())
         assert alerts[0]["event"]["payload"] == "\\xa8\\xa9"
 
+    def test_records_time_and_memory_per_module(self, tmp_path):
+        cmd = RecordingCommand(results_path=str(tmp_path))
+        cmd.modules = [FirstModule, SecondModule]
+
+        cmd.run()
+
+        info = json.loads((tmp_path / "info.json").read_text())
+        assert [stat["module"] for stat in info["module_stats"]] == [
+            "FirstModule",
+            "SecondModule",
+        ]
+        for stat in info["module_stats"]:
+            assert stat["seconds"] >= 0
+            assert stat["peak_rss_growth_mb"] >= 0
+
     def test_stores_collected_urls(self, tmp_path):
         cmd = RecordingCommand(results_path=str(tmp_path))
         cmd.modules = [URLRecordingModule]