2025-09-22 09:06:20 +02:00

132 lines
4.7 KiB
Python

#!/usr/bin/env python3
"""
memory_leak.py
Starts with sensible defaults if no CLI args are provided.
Usage:
python3 memory_leak.py # runs with defaults
python3 memory_leak.py -c 20M # override chunk size, etc.
"""
import argparse
import time
import platform
try:
import psutil
except Exception:
psutil = None
_leak_store = []
def parse_size(s: str) -> int:
"""Parse sizes like '10M', '512K', '1G' into bytes."""
s = str(s).strip().upper()
if s.endswith("G"):
return int(float(s[:-1]) * 1024**3)
if s.endswith("M"):
return int(float(s[:-1]) * 1024**2)
if s.endswith("K"):
return int(float(s[:-1]) * 1024)
return int(s)
def rss_bytes() -> int:
"""Return current process RSS in bytes (best-effort)."""
if psutil:
return psutil.Process().memory_info().rss
if platform.system() == "Linux":
try:
with open("/proc/self/statm", "r") as f:
parts = f.read().split()
if len(parts) >= 2:
pages = int(parts[1])
import os
return pages * os.sysconf("SC_PAGE_SIZE")
except Exception:
pass
return 0
def human_readable(n: int) -> str:
for unit in ["B", "K", "M", "G", "T"]:
if n < 1024.0:
return f"{n:.1f}{unit}"
n /= 1024.0
return f"{n:.1f}P"
def run_leak(chunk_size: int, interval: float, max_bytes: int, max_iters: int, verbose: bool, report_interval: float):
total_allocated = 0
iters = 0
last_report = time.time()
try:
while True:
# Stop conditions
if max_bytes and total_allocated >= max_bytes:
if verbose:
print(f"[info] reached max-bytes {human_readable(max_bytes)}; stopping allocation.")
break
if max_iters and iters >= max_iters:
if verbose:
print(f"[info] reached max-iters {max_iters}; stopping allocation.")
break
try:
chunk = bytearray(b"\x41") * chunk_size
except MemoryError:
print("[error] MemoryError during allocation — exiting allocation loop.")
break
_leak_store.append(chunk)
total_allocated += chunk_size
iters += 1
now = time.time()
if verbose and (now - last_report >= report_interval):
rss = rss_bytes()
print(f"[alloc #{iters}] chunk={human_readable(chunk_size)}, total_alloc={human_readable(total_allocated)}, RSS={human_readable(rss)}")
last_report = now
if interval:
time.sleep(interval)
except KeyboardInterrupt:
print("\n[stopped] KeyboardInterrupt received — stopping allocations.")
rss = rss_bytes()
print("\n=== Summary ===")
print(f"iterations: {iters}")
print(f"total allocated (kept references): {human_readable(total_allocated)}")
print(f"final RSS: {human_readable(rss)}")
print(f"stored objects in leak store: {len(_leak_store)}")
def main():
parser = argparse.ArgumentParser(description="Simple memory leak generator for testing (runs with defaults if no args).")
parser.add_argument("--chunk-size", "-c", default="10M", help="Size per allocation chunk (e.g. 10M). Default 10M.")
parser.add_argument("--interval", "-i", type=float, default=0.1, help="Seconds between allocations. Default 0.1s.")
parser.add_argument("--max-bytes", "-m", default=None, help="Stop after allocating this many bytes (e.g. 200M). Optional.")
parser.add_argument("--max-iters", type=int, default=0, help="Stop after N allocations (0 = unlimited).")
parser.add_argument("--report-interval", type=float, default=1.0, help="How often (s) to print allocation status when verbose.")
parser.add_argument("--quiet", action="store_true", help="Do not print allocation messages.")
args = parser.parse_args()
chunk_size = parse_size(args.chunk_size)
max_bytes = parse_size(args.max_bytes) if args.max_bytes else 0
max_iters = args.max_iters if args.max_iters > 0 else 0
interval = float(args.interval)
verbose = not args.quiet
print("memory_leak starting (defaults used when no args provided):")
print(f" chunk_size = {human_readable(chunk_size)}")
print(f" interval = {interval} s")
if max_bytes:
print(f" max_bytes = {human_readable(max_bytes)}")
if max_iters:
print(f" max_iters = {max_iters}")
print(" verbose =", verbose)
print("Press Ctrl-C to stop manually.\n")
run_leak(chunk_size=chunk_size, interval=interval, max_bytes=max_bytes, max_iters=max_iters, verbose=verbose, report_interval=args.report_interval)
if __name__ == "__main__":
main()