diff --git a/bbot/test/benchmarks/_memory_helpers.py b/bbot/test/benchmarks/_memory_helpers.py new file mode 100644 index 0000000000..36d87b09c0 --- /dev/null +++ b/bbot/test/benchmarks/_memory_helpers.py @@ -0,0 +1,345 @@ +""" +Shared memory measurement helpers for scan benchmarks. + +Each helper is independent — pick the ones you need. Designed to be used +from subprocess scan scripts (so tracemalloc + psutil polling are not +contaminated by pytest's own allocations). + +Three measurement angles, each answering a different question: + + - ``RSSSampler``: actual OS-level RAM (catches Rust / lxml / yara that + tracemalloc misses). Headline metrics: peak, end, retention (median + over the last 25% of samples — the metric most sensitive to + "stuck for the rest of the scan" pathology). + + - ``event_census``: who is alive right now, grouped by event type, with + HTTP_RESPONSE body bytes broken out separately. Answers + "where is the memory going?". + + - ``lineage_census``: walk every live event's parent chain back to the + seed and bucket pinned events by seed. Answers "is the chain + holding things alive?". +""" + +import gc +import json +import threading +import time +import weakref + +import psutil + + +class RSSSampler: + """ + Background thread that polls process RSS at a fixed interval. + + Usage:: + + sampler = RSSSampler(interval_s=0.2) + sampler.start() + # ... do work ... + sampler.stop() + m = sampler.metrics() # peak_rss_mb, end_rss_mb, retention_rss_mb + """ + + def __init__(self, interval_s=0.2): + self.interval_s = interval_s + self.samples = [] # list of (t, rss_mb) + self._stop = threading.Event() + self._thread = None + self._proc = psutil.Process() + + def start(self): + self._thread = threading.Thread(target=self._loop, daemon=True) + self._thread.start() + + def _loop(self): + t0 = time.monotonic() + while not self._stop.is_set(): + rss_mb = self._proc.memory_info().rss / 1024 / 1024 + self.samples.append((time.monotonic() - t0, rss_mb)) + self._stop.wait(self.interval_s) + + def stop(self): + self._stop.set() + if self._thread is not None: + self._thread.join(timeout=1.0) + + def metrics(self): + if not self.samples: + return { + "peak_rss_mb": 0.0, + "end_rss_mb": 0.0, + "retention_rss_mb": 0.0, + "samples": 0, + "duration_s": 0.0, + } + rss_values = [r for _, r in self.samples] + peak_rss_mb = max(rss_values) + end_rss_mb = rss_values[-1] + # median over the last 25% of samples — robust to a single + # transient spike at the tail and sensitive to baseline drift. + last_quartile_start = max(1, int(len(rss_values) * 0.75)) + last_quartile = sorted(rss_values[last_quartile_start:]) + retention_rss_mb = last_quartile[len(last_quartile) // 2] + duration_s = self.samples[-1][0] + return { + "peak_rss_mb": round(peak_rss_mb, 2), + "end_rss_mb": round(end_rss_mb, 2), + "retention_rss_mb": round(retention_rss_mb, 2), + "samples": len(self.samples), + "duration_s": round(duration_s, 2), + } + + +class LiveEventTracker: + """ + Mid-scan event census via weakref, without scanning the full Python + object graph. + + ``event_census()`` calls ``gc.get_objects()`` and walks every Python + object — fine at scan-end (one-shot), too expensive to call every + few hundred events. ``LiveEventTracker`` patches + ``BaseEvent.__init__`` to register newcomers into a ``WeakSet``, so + ``census()`` is O(live events) instead of O(every Python object). + + Usage:: + + tracker = LiveEventTracker() + tracker.install() + # ... run scan, calling tracker.census() periodically ... + tracker.uninstall() # optional; not needed if the process exits + """ + + def __init__(self): + self._events = weakref.WeakSet() + self._original_init = None + self._patched = False + + def install(self): + if self._patched: + return + from bbot.core.event.base import BaseEvent + + # Seed with any events that already exist (e.g. SCAN root_event + # created before install). + for obj in gc.get_objects(): + if isinstance(obj, BaseEvent): + self._events.add(obj) + + original_init = BaseEvent.__init__ + events_ref = self._events + + def patched_init(self_evt, *a, **kw): + original_init(self_evt, *a, **kw) + try: + events_ref.add(self_evt) + except TypeError: + # Not weakref-able for some reason; skip silently. + pass + + self._original_init = original_init + BaseEvent.__init__ = patched_init + self._patched = True + + def uninstall(self): + if not self._patched: + return + from bbot.core.event.base import BaseEvent + + BaseEvent.__init__ = self._original_init + self._patched = False + + def census(self): + """O(live events) census. Same shape as ``event_census()``.""" + by_type = {} + body_bytes = 0 + body_count = 0 + for obj in self._events: + by_type[obj.type] = by_type.get(obj.type, 0) + 1 + if obj.type == "HTTP_RESPONSE": + data = getattr(obj, "data", None) + body = data.get("body") if isinstance(data, dict) else None + if body: + body_count += 1 + body_bytes += len(body) + return { + "live_events": sum(by_type.values()), + "by_type": by_type, + "http_response_body_mb": round(body_bytes / 1024 / 1024, 2), + "http_response_with_body": body_count, + } + + +def event_census(): + """ + Walk live BaseEvent instances and classify by type + HTTP_RESPONSE body bytes. + + Returns a dict:: + + { + "live_events": int, + "by_type": {"DNS_NAME": int, "HTTP_RESPONSE": int, ...}, + "http_response_body_mb": float, + "http_response_with_body": int, + } + """ + from bbot.core.event.base import BaseEvent + + gc.collect() + by_type = {} + body_bytes = 0 + body_count = 0 + + for obj in gc.get_objects(): + if not isinstance(obj, BaseEvent): + continue + by_type[obj.type] = by_type.get(obj.type, 0) + 1 + if obj.type == "HTTP_RESPONSE": + data = getattr(obj, "data", None) + body = data.get("body") if isinstance(data, dict) else None + if body: + body_count += 1 + body_bytes += len(body) + + return { + "live_events": sum(by_type.values()), + "by_type": by_type, + "http_response_body_mb": round(body_bytes / 1024 / 1024, 2), + "http_response_with_body": body_count, + } + + +def lineage_census(top_n=20): + """ + Walk every live BaseEvent's parent chain back to the seed and bucket + pinned events by seed. + + A long-lived seed pinning hundreds of events is the signature of the + chain-retention pathology — this is the metric that proves a fix. + + Returns:: + + { + "seeds": [{"seed": str, "pinned_events": int, + "max_chain_depth": int, + "types_pinned": {type: count, ...}}, ...], # top N + "total_pinned_events": int, + "max_chain_depth": int, + "live_events_walked": int, + } + """ + from bbot.core.event.base import BaseEvent + + gc.collect() + seeds = {} + max_chain_depth = 0 + walked = 0 + + for obj in gc.get_objects(): + if not isinstance(obj, BaseEvent): + continue + walked += 1 + depth = 0 + node = obj + # Walk up via .parent. Root events are self-parented (node.parent is node). + while node is not None and node is not getattr(node, "parent", None): + depth += 1 + parent = getattr(node, "parent", None) + if parent is None or parent is node: + break + node = parent + if node is None: + continue + max_chain_depth = max(max_chain_depth, depth) + seed_data = getattr(node, "data", None) + seed_data_str = str(seed_data)[:80] if seed_data is not None else "" + seed_key = f"{node.type}:{seed_data_str}" + s = seeds.setdefault( + seed_key, + {"pinned_events": 0, "max_chain_depth": 0, "types_pinned": {}}, + ) + s["pinned_events"] += 1 + s["max_chain_depth"] = max(s["max_chain_depth"], depth) + s["types_pinned"][obj.type] = s["types_pinned"].get(obj.type, 0) + 1 + + sorted_seeds = sorted( + ({"seed": k, **v} for k, v in seeds.items()), + key=lambda x: -x["pinned_events"], + ) + return { + "seeds": sorted_seeds[:top_n], + "total_pinned_events": sum(v["pinned_events"] for v in seeds.values()), + "max_chain_depth": max_chain_depth, + "live_events_walked": walked, + } + + +def queue_residence(scan, tracker): + """ + Classify live events by where they currently live in the scan + pipeline. + + Two angles, since they answer different questions: + + - Per-module queue depth (``by_queue``): which module is the + current bottleneck — events are piling up where queue depth is + large. + - In-pipeline vs. chain-only (``in_pipeline`` / ``chain_only``): + of the live events, how many are still being processed by some + module (``event._module_consumers > 0``) vs. how many are + held alive *only* by the parent chain. The chain_only count is + the chain-retention pathology made directly visible. + + Returns:: + + { + "in_pipeline": int, # _module_consumers > 0 + "chain_only": int, # _module_consumers == 0 + "queue_total": int, # sum of all module queue depths + "by_queue": {module_name: {"incoming": int, "outgoing": int}, ...}, + } + """ + by_queue = {} + for module in scan.modules.values(): + m_name = getattr(module, "name", str(module)) + in_q = getattr(module, "_incoming_event_queue", None) + out_q = getattr(module, "_outgoing_event_queue", None) + in_count = 0 + out_count = 0 + # asyncio.Queue / ShuffleQueue both expose the underlying ``_queue`` deque. + if in_q and in_q is not False and hasattr(in_q, "_queue"): + in_count = len(in_q._queue) + if out_q and hasattr(out_q, "_queue"): + out_count = len(out_q._queue) + if in_count or out_count: + by_queue[m_name] = {"incoming": in_count, "outgoing": out_count} + + in_pipeline = 0 + chain_only = 0 + # Iterate the WeakSet directly — O(live events). + for ev in tracker._events: + if getattr(ev, "_module_consumers", 0) > 0: + in_pipeline += 1 + else: + chain_only += 1 + + return { + "in_pipeline": in_pipeline, + "chain_only": chain_only, + "queue_total": sum(d["incoming"] + d["outgoing"] for d in by_queue.values()), + "by_queue": by_queue, + } + + +def emit_metrics_json(**metrics): + """ + Print a single ``METRICS_JSON:`` line that ``test_scan_memory.py`` parses. + + Backward-compat: also prints the legacy ``PEAK_MB:`` line when + ``peak_tracemalloc_mb`` is supplied, so any external readers still work. + """ + print(f"METRICS_JSON:{json.dumps(metrics)}") + if "peak_tracemalloc_mb" in metrics: + print(f"PEAK_MB:{metrics['peak_tracemalloc_mb']}") diff --git a/bbot/test/benchmarks/_scan_memory_deep_chain.py b/bbot/test/benchmarks/_scan_memory_deep_chain.py new file mode 100644 index 0000000000..f13663e003 --- /dev/null +++ b/bbot/test/benchmarks/_scan_memory_deep_chain.py @@ -0,0 +1,150 @@ +""" +Subprocess script for the deep-chain memory benchmark. + +Spawns a local HTTP server that serves a strict linear chain — page N +links only to page N+1, no siblings — and runs a BBOT scan that follows +the entire chain. The discovery pipeline produces a deep parent lineage +(each hop adds URL → HTTP_RESPONSE → URL_UNVERIFIED → URL → … to the +chain), exposing chain-retention pathology that the wide-and-shallow +``_scan_memory_web_crawl.py`` workload masks. + +Invoked by test_scan_memory.py — not meant to be run directly. +""" + +import gc +import sys +import asyncio +import threading +import tracemalloc +import importlib.util +from http.server import HTTPServer, BaseHTTPRequestHandler + +import psutil + +from bbot.scanner import Scanner +from bbot.test.benchmarks._memory_helpers import ( + LiveEventTracker, + RSSSampler, + event_census, + lineage_census, + queue_residence, + emit_metrics_json, +) + +CHAIN_LENGTH = int(sys.argv[1]) +BODY_SIZE = int(sys.argv[2]) +CHECKPOINT_EVERY = 25 # mid-scan census cadence (events seen) + +HTTP_MODULE = "httpx" if importlib.util.find_spec("bbot.modules.httpx") else "http" + + +class H(BaseHTTPRequestHandler): + def do_GET(self): + path = self.path.rstrip("/") + if path in ("", "/"): + i = 0 + elif path.startswith("/page"): + try: + i = int(path[len("/page") :]) + except ValueError: + self.send_response(404) + self.end_headers() + return + else: + self.send_response(404) + self.end_headers() + return + + # Strict chain: page i links only to page i+1; the last page has no link. + if i + 1 < CHAIN_LENGTH: + link = f'next' + else: + link = "" + body = f"

Page {i}

{link}{'A' * BODY_SIZE}" + + self.send_response(200) + self.send_header("Content-Type", "text/html") + self.end_headers() + self.wfile.write(body.encode()) + + def log_message(self, *a): + pass + + +server = HTTPServer(("127.0.0.1", 0), H) +port = server.server_address[1] +threading.Thread(target=server.serve_forever, daemon=True).start() + +# spider_distance/depth must exceed chain length so the spider follows +# the full chain. spider_links_per_page=1 matches the server (one link). +scan = Scanner( + f"http://127.0.0.1:{port}/", + modules=[HTTP_MODULE], + output_modules=["python"], + config={ + "dns": {"disable": True}, + "scope": {"search_distance": 0}, + "web": { + "spider_distance": CHAIN_LENGTH + 5, + "spider_depth": CHAIN_LENGTH + 5, + "spider_links_per_page": 2, + }, + "speculate": True, + "excavate": True, + "aggregate": False, + "cloudcheck": False, + }, + force_start=True, +) + + +async def run(): + await scan._prep() + gc.collect() + if tracemalloc.is_tracing(): + tracemalloc.stop() + + tracker = LiveEventTracker() + tracker.install() + sampler = RSSSampler(interval_s=0.2) + sampler.start() + tracemalloc.start() + + # Count emitted events without holding strong refs. Holding them in a + # list would inflate live-event counts artificially — production + # callers iterate and discard. + events_seen = 0 + checkpoints = [] + proc = psutil.Process() + async for event in scan.async_start(): + del event + events_seen += 1 + if events_seen % CHECKPOINT_EVERY == 0: + checkpoints.append( + { + "events_seen": events_seen, + "rss_mb": round(proc.memory_info().rss / 1024 / 1024, 2), + **tracker.census(), + "residence": queue_residence(scan, tracker), + } + ) + + sampler.stop() + _, peak = tracemalloc.get_traced_memory() + tracemalloc.stop() + + emit_metrics_json( + peak_tracemalloc_mb=round(peak / 1024 / 1024, 2), + chain_length=CHAIN_LENGTH, + body_size=BODY_SIZE, + events_collected=events_seen, + rss=sampler.metrics(), + census=event_census(), + lineage=lineage_census(), + end_residence=queue_residence(scan, tracker), + checkpoints=checkpoints, + ) + + +asyncio.run(run()) +server.shutdown() diff --git a/bbot/test/benchmarks/_scan_memory_parallel_chains.py b/bbot/test/benchmarks/_scan_memory_parallel_chains.py new file mode 100644 index 0000000000..9bd1cf634c --- /dev/null +++ b/bbot/test/benchmarks/_scan_memory_parallel_chains.py @@ -0,0 +1,158 @@ +""" +Subprocess script for the parallel-chains memory benchmark. + +Mirrors the real-scale pattern: many independent targets being scanned +concurrently, each producing its own deep chain. Even when each chain +is naturally serial in its fetch cadence, the union across hundreds +of chains is where bodies pile up — the pathology a single-seed +``_scan_memory_deep_chain.py`` cannot reproduce. + +The HTTP server serves a path scheme of ``/seed{S}/page{N}``. Page N +of seed S links to page N+1 of the same seed (no cross-seed links, +strict per-seed chain). NUM_SEEDS targets are passed to the scanner +so the spider runs all chains concurrently. + +Invoked by test_scan_memory.py — not meant to be run directly. +""" + +import gc +import sys +import asyncio +import threading +import tracemalloc +import importlib.util +from http.server import HTTPServer, BaseHTTPRequestHandler + +import psutil + +from bbot.scanner import Scanner +from bbot.test.benchmarks._memory_helpers import ( + LiveEventTracker, + RSSSampler, + event_census, + lineage_census, + queue_residence, + emit_metrics_json, +) + +NUM_SEEDS = int(sys.argv[1]) +CHAIN_LENGTH = int(sys.argv[2]) +BODY_SIZE = int(sys.argv[3]) +CHECKPOINT_EVERY = 500 # mid-scan census cadence (events seen) + +HTTP_MODULE = "httpx" if importlib.util.find_spec("bbot.modules.httpx") else "http" + + +class H(BaseHTTPRequestHandler): + def do_GET(self): + # Path format: /seed{S}/page{N} + parts = self.path.strip("/").split("/") + if len(parts) != 2 or not parts[0].startswith("seed") or not parts[1].startswith("page"): + self.send_response(404) + self.end_headers() + return + try: + seed_idx = int(parts[0][len("seed") :]) + page_idx = int(parts[1][len("page") :]) + except ValueError: + self.send_response(404) + self.end_headers() + return + if seed_idx >= NUM_SEEDS or page_idx >= CHAIN_LENGTH: + self.send_response(404) + self.end_headers() + return + + # Page N links only to page N+1 within the same seed (no cross-seed links). + if page_idx + 1 < CHAIN_LENGTH: + link = f'next' + else: + link = "" + body = f"

seed{seed_idx} page{page_idx}

{link}{'A' * BODY_SIZE}" + self.send_response(200) + self.send_header("Content-Type", "text/html") + self.end_headers() + self.wfile.write(body.encode()) + + def log_message(self, *a): + pass + + +server = HTTPServer(("127.0.0.1", 0), H) +port = server.server_address[1] +threading.Thread(target=server.serve_forever, daemon=True).start() + +# NUM_SEEDS independent targets — the spider runs all chains concurrently. +targets = [f"http://127.0.0.1:{port}/seed{i}/page0" for i in range(NUM_SEEDS)] +scan = Scanner( + *targets, + modules=[HTTP_MODULE], + output_modules=["python"], + config={ + "dns": {"disable": True}, + "scope": {"search_distance": 0}, + "web": { + "spider_distance": CHAIN_LENGTH + 5, + "spider_depth": CHAIN_LENGTH + 5, + "spider_links_per_page": 2, + }, + "speculate": True, + "excavate": True, + "aggregate": False, + "cloudcheck": False, + }, + force_start=True, +) + + +async def run(): + await scan._prep() + gc.collect() + if tracemalloc.is_tracing(): + tracemalloc.stop() + + tracker = LiveEventTracker() + tracker.install() + sampler = RSSSampler(interval_s=0.2) + sampler.start() + tracemalloc.start() + + # Count emitted events without holding strong refs. Holding them in a + # list would inflate live-event counts artificially — production + # callers iterate and discard. + events_seen = 0 + checkpoints = [] + proc = psutil.Process() + async for event in scan.async_start(): + del event + events_seen += 1 + if events_seen % CHECKPOINT_EVERY == 0: + checkpoints.append( + { + "events_seen": events_seen, + "rss_mb": round(proc.memory_info().rss / 1024 / 1024, 2), + **tracker.census(), + "residence": queue_residence(scan, tracker), + } + ) + + sampler.stop() + _, peak = tracemalloc.get_traced_memory() + tracemalloc.stop() + + emit_metrics_json( + peak_tracemalloc_mb=round(peak / 1024 / 1024, 2), + num_seeds=NUM_SEEDS, + chain_length=CHAIN_LENGTH, + body_size=BODY_SIZE, + events_collected=events_seen, + rss=sampler.metrics(), + census=event_census(), + lineage=lineage_census(), + end_residence=queue_residence(scan, tracker), + checkpoints=checkpoints, + ) + + +asyncio.run(run()) +server.shutdown() diff --git a/bbot/test/benchmarks/_scan_memory_subdomain_enum.py b/bbot/test/benchmarks/_scan_memory_subdomain_enum.py index 33b4b997f6..55b456e044 100644 --- a/bbot/test/benchmarks/_scan_memory_subdomain_enum.py +++ b/bbot/test/benchmarks/_scan_memory_subdomain_enum.py @@ -12,12 +12,23 @@ import asyncio import tracemalloc +import psutil + from bbot.scanner import Scanner +from bbot.test.benchmarks._memory_helpers import ( + LiveEventTracker, + RSSSampler, + event_census, + lineage_census, + queue_residence, + emit_metrics_json, +) SUBDOMAIN_ENUM_COUNT = int(sys.argv[1]) +CHECKPOINT_EVERY = 1000 # mid-scan census cadence (events seen) scan = Scanner( - "blacklanternsecurity.com", + "example.com", modules=[], output_modules=["python"], config={ @@ -38,25 +49,62 @@ async def run(): gc.collect() if tracemalloc.is_tracing(): tracemalloc.stop() + + tracker = LiveEventTracker() + tracker.install() + sampler = RSSSampler(interval_s=0.2) + sampler.start() tracemalloc.start() - events = [] + + # Count emitted events without holding strong refs. Holding them in a + # list would inflate live-event counts artificially — production + # callers iterate and discard. + events_seen = 0 + checkpoints = [] injected = False + proc = psutil.Process() async for event in scan.async_start(): - events.append(event) - if event.type == "SCAN" and not injected: + events_seen += 1 + is_scan = event.type == "SCAN" + # Read what we need; don't hold the event reference. + if is_scan and not injected: injected = True root_event = scan.root_event for i in range(SUBDOMAIN_ENUM_COUNT): dns_event = scan.make_event( - f"sub{i}.blacklanternsecurity.com", + f"sub{i}.example.com", "DNS_NAME", parent=root_event, context=f"benchmark DNS_NAME {i}", ) await scan.ingress_module.queue_event(dns_event, {}) + del dns_event + del root_event + del event + if events_seen % CHECKPOINT_EVERY == 0: + checkpoints.append( + { + "events_seen": events_seen, + "rss_mb": round(proc.memory_info().rss / 1024 / 1024, 2), + **tracker.census(), + "residence": queue_residence(scan, tracker), + } + ) + + sampler.stop() + _, peak = tracemalloc.get_traced_memory() + tracemalloc.stop() + + emit_metrics_json( + peak_tracemalloc_mb=round(peak / 1024 / 1024, 2), + num_subdomains=SUBDOMAIN_ENUM_COUNT, + events_collected=events_seen, + rss=sampler.metrics(), + census=event_census(), + lineage=lineage_census(), + end_residence=queue_residence(scan, tracker), + checkpoints=checkpoints, + ) asyncio.run(run()) -_, peak = tracemalloc.get_traced_memory() -tracemalloc.stop() -print(f"PEAK_MB:{round(peak / 1024 / 1024, 2)}") diff --git a/bbot/test/benchmarks/_scan_memory_web_crawl.py b/bbot/test/benchmarks/_scan_memory_web_crawl.py index e609220c74..c5fc94245e 100644 --- a/bbot/test/benchmarks/_scan_memory_web_crawl.py +++ b/bbot/test/benchmarks/_scan_memory_web_crawl.py @@ -15,10 +15,21 @@ import importlib.util from http.server import HTTPServer, BaseHTTPRequestHandler +import psutil + from bbot.scanner import Scanner +from bbot.test.benchmarks._memory_helpers import ( + LiveEventTracker, + RSSSampler, + event_census, + lineage_census, + queue_residence, + emit_metrics_json, +) NUM_PAGES = int(sys.argv[1]) BODY_SIZE = int(sys.argv[2]) +CHECKPOINT_EVERY = 200 # mid-scan census cadence (events seen) HTTP_MODULE = "httpx" if importlib.util.find_spec("bbot.modules.httpx") else "http" @@ -73,14 +84,48 @@ async def run(): gc.collect() if tracemalloc.is_tracing(): tracemalloc.stop() + + tracker = LiveEventTracker() + tracker.install() + sampler = RSSSampler(interval_s=0.2) + sampler.start() tracemalloc.start() - events = [] + + # Count emitted events without holding strong refs. Holding them in a + # list would inflate live-event counts artificially — production + # callers iterate and discard. + events_seen = 0 + checkpoints = [] + proc = psutil.Process() async for event in scan.async_start(): - events.append(event) + del event + events_seen += 1 + if events_seen % CHECKPOINT_EVERY == 0: + checkpoints.append( + { + "events_seen": events_seen, + "rss_mb": round(proc.memory_info().rss / 1024 / 1024, 2), + **tracker.census(), + "residence": queue_residence(scan, tracker), + } + ) + + sampler.stop() + _, peak = tracemalloc.get_traced_memory() + tracemalloc.stop() + + emit_metrics_json( + peak_tracemalloc_mb=round(peak / 1024 / 1024, 2), + num_pages=NUM_PAGES, + body_size=BODY_SIZE, + events_collected=events_seen, + rss=sampler.metrics(), + census=event_census(), + lineage=lineage_census(), + end_residence=queue_residence(scan, tracker), + checkpoints=checkpoints, + ) asyncio.run(run()) -_, peak = tracemalloc.get_traced_memory() -tracemalloc.stop() server.shutdown() -print(f"PEAK_MB:{round(peak / 1024 / 1024, 2)}") diff --git a/bbot/test/benchmarks/test_scan_memory.py b/bbot/test/benchmarks/test_scan_memory.py index af4f76dfa6..760be8568c 100644 --- a/bbot/test/benchmarks/test_scan_memory.py +++ b/bbot/test/benchmarks/test_scan_memory.py @@ -1,12 +1,29 @@ """ Memory benchmarks for BBOT scan patterns. -Each benchmark launches a scan as a subprocess so tracemalloc measurements -are not contaminated by pytest's own allocations. The subprocess writes -peak memory (MB) to stdout, which the test reads and stores in -benchmark extra_info["total_memory_mb"]. +Each benchmark launches a scan as a subprocess so tracemalloc / RSS +measurements are not contaminated by pytest's own allocations. The +subprocess emits a single ``METRICS_JSON:`` line containing four +measurement angles, all of which are surfaced to pytest-benchmark via +``benchmark.extra_info``: + + - ``peak_tracemalloc_mb``: Python-side peak (legacy headline metric; + misses Rust / lxml / yara allocations). + - ``rss.peak_rss_mb`` / ``rss.end_rss_mb`` / ``rss.retention_rss_mb``: + OS-level RSS sampled every ~200 ms. ``retention_rss_mb`` is the + median of the last 25% of samples — the metric most sensitive to + "stuck for the rest of the scan" pathology. + - ``census``: live BaseEvent count by type + HTTP_RESPONSE body MB + at scan-end. Answers "where is the memory going?". + - ``lineage``: every live event's parent chain walked back to the + seed; reports max chain depth and which seeds pin the most events. + Answers "is the chain holding things alive?". + +A backward-compatible ``PEAK_MB:`` line is still emitted by +``_memory_helpers.emit_metrics_json`` for any external readers. """ +import json import subprocess import sys from pathlib import Path @@ -14,15 +31,24 @@ import pytest -NUM_PAGES = 500 -BODY_SIZE = 500_000 # 500 KB per page -SUBDOMAIN_ENUM_COUNT = 5000 +NUM_PAGES = 1000 +BODY_SIZE = 1_000_000 # 1 MB per page → 1 GB served +SUBDOMAIN_ENUM_COUNT = 20_000 +DEEP_CHAIN_LENGTH = 200 +DEEP_CHAIN_BODY_SIZE = 1_000_000 # 1 MB per page + +# Parallel-chains workload: simulates real-scale scans with many concurrent +# targets. Even when each chain is naturally serial, the union across chains +# is where bodies pile up — the pathology a single deep chain cannot reproduce. +PARALLEL_NUM_SEEDS = 50 +PARALLEL_CHAIN_LENGTH = 30 +PARALLEL_BODY_SIZE = 500_000 _BENCHMARKS_DIR = Path(__file__).parent -def _run_scan_subprocess(script_path: Path, *args: str) -> float: - """Run a scan script in a clean subprocess, return peak memory in MB.""" +def _run_scan_subprocess(script_path: Path, *args: str) -> dict: + """Run a scan script in a clean subprocess; return parsed metrics dict.""" result = subprocess.run( [sys.executable, str(script_path), *args], capture_output=True, @@ -31,36 +57,141 @@ def _run_scan_subprocess(script_path: Path, *args: str) -> float: ) if result.returncode != 0: raise RuntimeError(f"Scan subprocess failed:\n{result.stderr[-2000:]}") + metrics = None for line in result.stdout.strip().splitlines(): - if line.startswith("PEAK_MB:"): - return float(line.split(":", 1)[1]) - raise RuntimeError(f"No PEAK_MB in subprocess output:\n{result.stdout[-2000:]}") + if line.startswith("METRICS_JSON:"): + metrics = json.loads(line[len("METRICS_JSON:") :]) + if metrics is None: + raise RuntimeError(f"No METRICS_JSON in subprocess output:\n{result.stdout[-2000:]}") + return metrics + + +def _record(benchmark, metrics: dict, **extra) -> None: + """Flatten the metrics dict into ``benchmark.extra_info``.""" + benchmark.extra_info["total_memory_mb"] = metrics["peak_tracemalloc_mb"] + benchmark.extra_info["events_collected"] = metrics["events_collected"] + + rss = metrics["rss"] + benchmark.extra_info["peak_rss_mb"] = rss["peak_rss_mb"] + benchmark.extra_info["end_rss_mb"] = rss["end_rss_mb"] + benchmark.extra_info["retention_rss_mb"] = rss["retention_rss_mb"] + benchmark.extra_info["rss_samples"] = rss["samples"] + benchmark.extra_info["scan_duration_s"] = rss["duration_s"] + + census = metrics["census"] + benchmark.extra_info["live_events"] = census["live_events"] + benchmark.extra_info["live_by_type"] = census["by_type"] + benchmark.extra_info["http_response_body_mb"] = census["http_response_body_mb"] + benchmark.extra_info["http_response_with_body"] = census["http_response_with_body"] + + lineage = metrics["lineage"] + benchmark.extra_info["max_chain_depth"] = lineage["max_chain_depth"] + benchmark.extra_info["total_pinned_events"] = lineage["total_pinned_events"] + benchmark.extra_info["top_pinning_seeds"] = lineage["seeds"][:5] + + # Mid-scan checkpoints expose body-byte peaks before minimize() runs + # and the live-event growth curve that end-of-scan census misses. + checkpoints = metrics.get("checkpoints", []) + benchmark.extra_info["checkpoints"] = checkpoints + if checkpoints: + benchmark.extra_info["peak_live_events_midscan"] = max(c["live_events"] for c in checkpoints) + benchmark.extra_info["peak_body_mb_midscan"] = max(c["http_response_body_mb"] for c in checkpoints) + # Residence peaks: max in-pipeline (queued or being handled) vs + # max chain-only (alive solely because the parent chain pins them). + residences = [c["residence"] for c in checkpoints if "residence" in c] + if residences: + benchmark.extra_info["peak_in_pipeline_midscan"] = max(r["in_pipeline"] for r in residences) + benchmark.extra_info["peak_chain_only_midscan"] = max(r["chain_only"] for r in residences) + benchmark.extra_info["peak_queue_total_midscan"] = max(r["queue_total"] for r in residences) + + end_residence = metrics.get("end_residence") + if end_residence is not None: + benchmark.extra_info["end_residence"] = end_residence + + for k, v in extra.items(): + benchmark.extra_info[k] = v class TestWebCrawlMemory: - """Measures peak memory during a realistic web crawl with large pages.""" + """Wide-and-shallow web crawl: every page hangs off ``/``.""" @pytest.mark.benchmark(group="memory_scan_patterns") def test_memory_use_web_crawl(self, benchmark): - peak_mb = _run_scan_subprocess( + metrics = _run_scan_subprocess( _BENCHMARKS_DIR / "_scan_memory_web_crawl.py", str(NUM_PAGES), str(BODY_SIZE), ) - benchmark.extra_info["total_memory_mb"] = peak_mb - benchmark.extra_info["num_pages"] = NUM_PAGES + _record(benchmark, metrics, num_pages=NUM_PAGES, body_size=BODY_SIZE) benchmark.pedantic(lambda: None, iterations=1, rounds=1, warmup_rounds=0) class TestSubdomainEnumMemory: - """Measures peak memory during a large subdomain enumeration.""" + """Synthetic breadth-only workload: ``SUBDOMAIN_ENUM_COUNT`` injected DNS_NAME events.""" @pytest.mark.benchmark(group="memory_scan_patterns") def test_memory_use_subdomain_enum(self, benchmark): - peak_mb = _run_scan_subprocess( + metrics = _run_scan_subprocess( _BENCHMARKS_DIR / "_scan_memory_subdomain_enum.py", str(SUBDOMAIN_ENUM_COUNT), ) - benchmark.extra_info["total_memory_mb"] = peak_mb - benchmark.extra_info["num_subdomains"] = SUBDOMAIN_ENUM_COUNT + _record(benchmark, metrics, num_subdomains=SUBDOMAIN_ENUM_COUNT) + benchmark.pedantic(lambda: None, iterations=1, rounds=1, warmup_rounds=0) + + +class TestDeepChainMemory: + """ + Deep-and-narrow spider: page N links only to page N+1. + + Designed to expose chain-retention pathology — long parent lineage + with HTTP_RESPONSE bodies pinned via ``parent`` references. Compare + ``retention_rss_mb`` and ``max_chain_depth`` against the wide + workload to gauge how much memory is held by lineage vs. burst. + """ + + @pytest.mark.benchmark(group="memory_scan_patterns") + def test_memory_use_deep_chain(self, benchmark): + metrics = _run_scan_subprocess( + _BENCHMARKS_DIR / "_scan_memory_deep_chain.py", + str(DEEP_CHAIN_LENGTH), + str(DEEP_CHAIN_BODY_SIZE), + ) + _record( + benchmark, + metrics, + chain_length=DEEP_CHAIN_LENGTH, + body_size=DEEP_CHAIN_BODY_SIZE, + ) + benchmark.pedantic(lambda: None, iterations=1, rounds=1, warmup_rounds=0) + + +class TestParallelChainsMemory: + """ + Many independent chains running concurrently — closest analogue to + a real-scale scan over many domains. + + A single deep chain doesn't expose body retention because the + spider is naturally serial within a chain (one fetch at a time). + Across N independent chains, body windows overlap and bodies pile + up even though each chain is individually well-behaved. This is + the workload most relevant to the "20-hour scan over thousands of + domains" pattern where both body retention and chain retention + compound. + """ + + @pytest.mark.benchmark(group="memory_scan_patterns") + def test_memory_use_parallel_chains(self, benchmark): + metrics = _run_scan_subprocess( + _BENCHMARKS_DIR / "_scan_memory_parallel_chains.py", + str(PARALLEL_NUM_SEEDS), + str(PARALLEL_CHAIN_LENGTH), + str(PARALLEL_BODY_SIZE), + ) + _record( + benchmark, + metrics, + num_seeds=PARALLEL_NUM_SEEDS, + chain_length=PARALLEL_CHAIN_LENGTH, + body_size=PARALLEL_BODY_SIZE, + ) benchmark.pedantic(lambda: None, iterations=1, rounds=1, warmup_rounds=0)