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.
This commit is contained in:
Donncha Ó Cearbhaill committed 2026-09-16 00:45:47 +02:00
1 parent eceafdc785
commit 06cea62a40
2 files changed
+46

No files matched your search

+31
View File
@@ -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)
+15
View File
@@ -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]