mirror of
https://github.com/mvt-project/mvt.git
synced 2026-10-19 05:57:20 +02:00
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:
1 parent
eceafdc785
commit
06cea62a40
2 files changed
+46
No files matched your search
@@ -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)
|
||||
|
||||
@@ -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]
|
||||
|
||||
Reference in new issue
Block a user