| #!/usr/bin/env python3 |
| # |
| # Author: Claude Sonnet 4.5 as driven by gpshead |
| # |
| """ |
| Microbenchmark for pickle module chunked reading performance (GH PR #119204). |
| |
| This script generates Python data structures that act as antagonistic load |
| tests for the chunked reading code introduced to prevent memory exhaustion when |
| unpickling large objects. |
| |
| The PR adds chunked reading (1MB chunks) for: |
| - BINBYTES8 (large bytes) |
| - BINUNICODE8 (large strings) |
| - BYTEARRAY8 (large bytearrays) |
| - FRAME (large frames) |
| - LONG4 (large integers) |
| |
| Including an antagonistic mode that exercies memory denial of service pickles. |
| |
| Usage: |
| python memory_dos_impact.py --help |
| """ |
| |
| import argparse |
| import gc |
| import io |
| import json |
| import os |
| import pickle |
| import statistics |
| import struct |
| import subprocess |
| import sys |
| import tempfile |
| import tracemalloc |
| from pathlib import Path |
| from time import perf_counter |
| from typing import Any, Dict, List, Tuple, Optional |
| |
| |
| # Configuration |
| MIN_READ_BUF_SIZE = 1 << 20 # 1MB - matches pickle.py _MIN_READ_BUF_SIZE |
| |
| # Test sizes in MiB |
| DEFAULT_SIZES_MIB = [1, 2, 5, 10, 20, 50, 100, 200] |
| |
| # Convert to bytes, plus threshold boundary tests |
| DEFAULT_SIZES = ( |
| [999_000] # Below 1MiB (no chunking) |
| + [size * (1 << 20) for size in DEFAULT_SIZES_MIB] # MiB to bytes |
| + [1_048_577] # Just above 1MiB (minimal chunking overhead) |
| ) |
| DEFAULT_SIZES.sort() |
| |
| # Baseline benchmark configuration |
| BASELINE_BENCHMARK_TIMEOUT_SECONDS = 600 # 10 minutes |
| |
| # Sparse memo attack test configuration |
| # Format: test_name -> (memo_index, baseline_memory_note) |
| SPARSE_MEMO_TESTS = { |
| "sparse_memo_1M": (1_000_000, "~8 MB array"), |
| "sparse_memo_100M": (100_000_000, "~800 MB array"), |
| "sparse_memo_1B": (1_000_000_000, "~8 GB array"), |
| } |
| |
| |
| # Utility functions |
| |
| def _extract_size_mb(size_key: str) -> float: |
| """Extract numeric MiB value from size_key like '10.00MB' or '1.00MiB'. |
| |
| Returns 0.0 for non-numeric keys (they'll be sorted last). |
| """ |
| try: |
| return float(size_key.replace('MB', '').replace('MiB', '')) |
| except ValueError: |
| return 999999.0 # Put non-numeric keys last |
| |
| |
| def _format_output(results: Dict[str, Dict[str, Any]], format_type: str, is_antagonistic: bool) -> str: |
| """Format benchmark results according to requested format. |
| |
| Args: |
| results: Benchmark results dictionary |
| format_type: Output format ('text', 'markdown', or 'json') |
| is_antagonistic: Whether these are antagonistic (DoS) test results |
| |
| Returns: |
| Formatted output string |
| """ |
| if format_type == 'json': |
| return Reporter.format_json(results) |
| elif is_antagonistic: |
| # Antagonistic mode uses specialized formatter for text/markdown |
| return Reporter.format_antagonistic(results) |
| elif format_type == 'text': |
| return Reporter.format_text(results) |
| elif format_type == 'markdown': |
| return Reporter.format_markdown(results) |
| else: |
| # Default to text format |
| return Reporter.format_text(results) |
| |
| |
| class AntagonisticGenerator: |
| """Generate malicious/truncated pickles for DoS protection testing. |
| |
| These pickles claim large sizes but provide minimal data, causing them to fail |
| during unpickling. They demonstrate the memory protection of chunked reading. |
| """ |
| |
| @staticmethod |
| def truncated_binbytes8(claimed_size: int, actual_size: int = 1024) -> bytes: |
| """BINBYTES8 claiming `claimed_size` but providing only `actual_size` bytes. |
| |
| This will fail with UnpicklingError but demonstrates peak memory usage. |
| Before PR: Allocates full claimed_size |
| After PR: Allocates in 1MB chunks, fails fast |
| """ |
| return b'\x8e' + struct.pack('<Q', claimed_size) + b'x' * actual_size |
| |
| @staticmethod |
| def truncated_binunicode8(claimed_size: int, actual_size: int = 1024) -> bytes: |
| """BINUNICODE8 claiming `claimed_size` but providing only `actual_size` bytes.""" |
| return b'\x8d' + struct.pack('<Q', claimed_size) + b'x' * actual_size |
| |
| @staticmethod |
| def truncated_bytearray8(claimed_size: int, actual_size: int = 1024) -> bytes: |
| """BYTEARRAY8 claiming `claimed_size` but providing only `actual_size` bytes.""" |
| return b'\x96' + struct.pack('<Q', claimed_size) + b'x' * actual_size |
| |
| @staticmethod |
| def truncated_frame(claimed_size: int) -> bytes: |
| """FRAME claiming `claimed_size` but providing minimal data.""" |
| return b'\x95' + struct.pack('<Q', claimed_size) + b'N.' |
| |
| @staticmethod |
| def sparse_memo_attack(index: int) -> bytes: |
| """LONG_BINPUT with huge sparse index. |
| |
| Before PR: Tries to allocate array with `index` slots (OOM) |
| After PR: Uses dict-based memo for sparse indices |
| """ |
| return (b'(]r' + struct.pack('<I', index & 0xFFFFFFFF) + |
| b'j' + struct.pack('<I', index & 0xFFFFFFFF) + b't.') |
| |
| @staticmethod |
| def multi_claim_attack(count: int, size_each: int) -> bytes: |
| """Multiple BINBYTES8 claims in sequence. |
| |
| Tests that multiple large claims don't accumulate memory. |
| """ |
| data = b'(' # MARK |
| for _ in range(count): |
| data += b'\x8e' + struct.pack('<Q', size_each) + b'x' * 1024 |
| data += b't.' # TUPLE + STOP |
| return data |
| |
| |
| class DataGenerator: |
| """Generate various types of large data structures for pickle testing.""" |
| |
| @staticmethod |
| def large_bytes(size: int) -> bytes: |
| """Generate random bytes of specified size.""" |
| return os.urandom(size) |
| |
| @staticmethod |
| def large_string_ascii(size: int) -> str: |
| """Generate ASCII string of specified size.""" |
| return 'x' * size |
| |
| @staticmethod |
| def large_string_multibyte(size: int) -> str: |
| """Generate multibyte UTF-8 string (3 bytes per char for €).""" |
| # Each € is 3 bytes in UTF-8 |
| return '€' * (size // 3) |
| |
| @staticmethod |
| def large_bytearray(size: int) -> bytearray: |
| """Generate bytearray of specified size.""" |
| return bytearray(os.urandom(size)) |
| |
| @staticmethod |
| def list_of_large_bytes(item_size: int, count: int) -> List[bytes]: |
| """Generate list containing multiple large bytes objects.""" |
| return [os.urandom(item_size) for _ in range(count)] |
| |
| @staticmethod |
| def dict_with_large_values(value_size: int, count: int) -> Dict[str, bytes]: |
| """Generate dict with large bytes values.""" |
| return { |
| f'key_{i}': os.urandom(value_size) |
| for i in range(count) |
| } |
| |
| @staticmethod |
| def nested_structure(size: int) -> Dict[str, Any]: |
| """Generate nested structure with various large objects.""" |
| chunk_size = size // 4 |
| return { |
| 'name': 'test_object', |
| 'data': { |
| 'bytes': os.urandom(chunk_size), |
| 'string': 's' * chunk_size, |
| 'bytearray': bytearray(b'b' * chunk_size), |
| }, |
| 'items': [os.urandom(chunk_size // 4) for _ in range(4)], |
| 'metadata': { |
| 'size': size, |
| 'type': 'nested', |
| }, |
| } |
| |
| @staticmethod |
| def tuple_of_large_objects(size: int) -> Tuple[bytes, str, bytearray]: |
| """Generate tuple with large objects (immutable, different pickle path).""" |
| chunk_size = size // 3 |
| return ( |
| os.urandom(chunk_size), |
| 'x' * chunk_size, |
| bytearray(b'y' * chunk_size), |
| ) |
| |
| |
| class PickleBenchmark: |
| """Benchmark pickle unpickling performance and memory usage.""" |
| |
| def __init__(self, obj: Any, protocol: int = 5, iterations: int = 3): |
| self.obj = obj |
| self.protocol = protocol |
| self.iterations = iterations |
| self.pickle_data = pickle.dumps(obj, protocol=protocol) |
| self.pickle_size = len(self.pickle_data) |
| |
| def benchmark_time(self) -> Dict[str, float]: |
| """Measure unpickling time over multiple iterations.""" |
| times = [] |
| |
| for _ in range(self.iterations): |
| start = perf_counter() |
| result = pickle.loads(self.pickle_data) |
| elapsed = perf_counter() - start |
| times.append(elapsed) |
| |
| # Verify correctness (first iteration only) |
| if len(times) == 1: |
| if result != self.obj: |
| raise ValueError("Unpickled object doesn't match original!") |
| |
| return { |
| 'mean': statistics.mean(times), |
| 'median': statistics.median(times), |
| 'stdev': statistics.stdev(times) if len(times) > 1 else 0.0, |
| 'min': min(times), |
| 'max': max(times), |
| } |
| |
| def benchmark_memory(self) -> int: |
| """Measure peak memory usage during unpickling.""" |
| tracemalloc.start() |
| |
| # Warmup |
| pickle.loads(self.pickle_data) |
| |
| # Actual measurement |
| gc.collect() |
| tracemalloc.reset_peak() |
| result = pickle.loads(self.pickle_data) |
| current, peak = tracemalloc.get_traced_memory() |
| |
| tracemalloc.stop() |
| |
| # Verify correctness |
| if result != self.obj: |
| raise ValueError("Unpickled object doesn't match original!") |
| |
| return peak |
| |
| def run_all(self) -> Dict[str, Any]: |
| """Run all benchmarks and return comprehensive results.""" |
| time_stats = self.benchmark_time() |
| peak_memory = self.benchmark_memory() |
| |
| return { |
| 'pickle_size_bytes': self.pickle_size, |
| 'pickle_size_mb': self.pickle_size / (1 << 20), |
| 'protocol': self.protocol, |
| 'time': time_stats, |
| 'memory_peak_bytes': peak_memory, |
| 'memory_peak_mb': peak_memory / (1 << 20), |
| 'iterations': self.iterations, |
| } |
| |
| |
| class AntagonisticBenchmark: |
| """Benchmark antagonistic/malicious pickles that demonstrate DoS protection. |
| |
| These pickles are designed to FAIL unpickling, but we measure peak memory |
| usage before the failure to demonstrate the memory protection. |
| """ |
| |
| def __init__(self, pickle_data: bytes, name: str): |
| self.pickle_data = pickle_data |
| self.name = name |
| |
| def measure_peak_memory(self, expect_success: bool = False) -> Dict[str, Any]: |
| """Measure peak memory when attempting to unpickle antagonistic data. |
| |
| Args: |
| expect_success: If True, test expects successful unpickling (e.g., sparse memo). |
| If False, test expects failure (e.g., truncated data). |
| """ |
| tracemalloc.start() |
| gc.collect() |
| tracemalloc.reset_peak() |
| |
| error_type = None |
| error_msg = None |
| succeeded = False |
| |
| try: |
| result = pickle.loads(self.pickle_data) |
| succeeded = True |
| if expect_success: |
| error_type = "Success (expected)" |
| else: |
| error_type = "WARNING: Expected failure but succeeded" |
| except (pickle.UnpicklingError, EOFError, ValueError, OverflowError) as e: |
| if expect_success: |
| error_type = f"UNEXPECTED FAILURE: {type(e).__name__}" |
| error_msg = str(e)[:100] |
| else: |
| # Expected failure for truncated data tests |
| error_type = type(e).__name__ |
| error_msg = str(e)[:100] |
| |
| current, peak = tracemalloc.get_traced_memory() |
| tracemalloc.stop() |
| |
| return { |
| 'test_name': self.name, |
| 'peak_memory_bytes': peak, |
| 'peak_memory_mb': peak / (1 << 20), |
| 'error_type': error_type, |
| 'error_msg': error_msg, |
| 'pickle_size_bytes': len(self.pickle_data), |
| 'expected_outcome': 'success' if expect_success else 'failure', |
| 'succeeded': succeeded, |
| } |
| |
| |
| class AntagonisticTestSuite: |
| """Manage a suite of antagonistic (DoS protection) tests.""" |
| |
| # Default sizes in MB to claim (will provide only 1KB actual data) |
| DEFAULT_ANTAGONISTIC_SIZES_MB = [10, 50, 100, 500, 1000, 5000] |
| |
| def __init__(self, claimed_sizes_mb: List[int]): |
| self.claimed_sizes_mb = claimed_sizes_mb |
| |
| def _run_truncated_test( |
| self, |
| test_type: str, |
| generator_func, |
| claimed_bytes: int, |
| claimed_mb: int, |
| size_key: str, |
| all_results: Dict[str, Dict[str, Any]] |
| ) -> None: |
| """Run a single truncated data test and store results. |
| |
| Args: |
| test_type: Type identifier (e.g., 'binbytes8', 'binunicode8') |
| generator_func: Function to generate malicious pickle data |
| claimed_bytes: Size claimed in the pickle (bytes) |
| claimed_mb: Size claimed in the pickle (MB) |
| size_key: Result key for this size (e.g., '10MB') |
| all_results: Dictionary to store results in |
| """ |
| test_name = f"{test_type}_{size_key}_claim" |
| data = generator_func(claimed_bytes) |
| bench = AntagonisticBenchmark(data, test_name) |
| result = bench.measure_peak_memory(expect_success=False) |
| result['claimed_mb'] = claimed_mb |
| all_results[size_key][test_name] = result |
| |
| def run_all_tests(self) -> Dict[str, Dict[str, Any]]: |
| """Run comprehensive antagonistic test suite.""" |
| all_results = {} |
| |
| for claimed_mb in self.claimed_sizes_mb: |
| claimed_bytes = claimed_mb << 20 |
| size_key = f"{claimed_mb}MB" |
| all_results[size_key] = {} |
| |
| # Run truncated data tests (expect failure) |
| self._run_truncated_test('binbytes8', AntagonisticGenerator.truncated_binbytes8, |
| claimed_bytes, claimed_mb, size_key, all_results) |
| self._run_truncated_test('binunicode8', AntagonisticGenerator.truncated_binunicode8, |
| claimed_bytes, claimed_mb, size_key, all_results) |
| self._run_truncated_test('bytearray8', AntagonisticGenerator.truncated_bytearray8, |
| claimed_bytes, claimed_mb, size_key, all_results) |
| self._run_truncated_test('frame', AntagonisticGenerator.truncated_frame, |
| claimed_bytes, claimed_mb, size_key, all_results) |
| |
| # Test 5: Sparse memo (expect success - dict-based memo works!) |
| all_results["Sparse Memo (Success Expected)"] = {} |
| for test_name, (index, baseline_note) in SPARSE_MEMO_TESTS.items(): |
| data = AntagonisticGenerator.sparse_memo_attack(index) |
| bench = AntagonisticBenchmark(data, test_name) |
| result = bench.measure_peak_memory(expect_success=True) |
| result['claimed_mb'] = "N/A" |
| result['baseline_note'] = f"Without PR: {baseline_note}" |
| all_results["Sparse Memo (Success Expected)"][test_name] = result |
| |
| # Test 6: Multi-claim attack (expect failure) |
| test_name = "multi_claim_10x100MB" |
| data = AntagonisticGenerator.multi_claim_attack(10, 100 << 20) |
| bench = AntagonisticBenchmark(data, test_name) |
| result = bench.measure_peak_memory(expect_success=False) |
| result['claimed_mb'] = 1000 # 10 * 100MB |
| all_results["Multi-Claim (Failure Expected)"] = {test_name: result} |
| |
| return all_results |
| |
| |
| class TestSuite: |
| """Manage a suite of benchmark tests.""" |
| |
| def __init__(self, sizes: List[int], protocol: int = 5, iterations: int = 3): |
| self.sizes = sizes |
| self.protocol = protocol |
| self.iterations = iterations |
| self.results = {} |
| |
| def run_test(self, name: str, obj: Any) -> Dict[str, Any]: |
| """Run benchmark for a single test object.""" |
| bench = PickleBenchmark(obj, self.protocol, self.iterations) |
| results = bench.run_all() |
| results['test_name'] = name |
| results['object_type'] = type(obj).__name__ |
| return results |
| |
| def run_all_tests(self) -> Dict[str, Dict[str, Any]]: |
| """Run comprehensive test suite across all sizes and types.""" |
| all_results = {} |
| |
| for size in self.sizes: |
| size_key = f"{size / (1 << 20):.2f}MB" |
| all_results[size_key] = {} |
| |
| # Test 1: Large bytes object (BINBYTES8) |
| test_name = f"bytes_{size_key}" |
| obj = DataGenerator.large_bytes(size) |
| all_results[size_key][test_name] = self.run_test(test_name, obj) |
| |
| # Test 2: Large ASCII string (BINUNICODE8) |
| test_name = f"string_ascii_{size_key}" |
| obj = DataGenerator.large_string_ascii(size) |
| all_results[size_key][test_name] = self.run_test(test_name, obj) |
| |
| # Test 3: Large multibyte UTF-8 string |
| if size >= 3: |
| test_name = f"string_utf8_{size_key}" |
| obj = DataGenerator.large_string_multibyte(size) |
| all_results[size_key][test_name] = self.run_test(test_name, obj) |
| |
| # Test 4: Large bytearray (BYTEARRAY8, protocol 5) |
| if self.protocol >= 5: |
| test_name = f"bytearray_{size_key}" |
| obj = DataGenerator.large_bytearray(size) |
| all_results[size_key][test_name] = self.run_test(test_name, obj) |
| |
| # Test 5: List of large objects (repeated chunking) |
| if size >= MIN_READ_BUF_SIZE * 2: |
| test_name = f"list_large_items_{size_key}" |
| item_size = size // 5 |
| obj = DataGenerator.list_of_large_bytes(item_size, 5) |
| all_results[size_key][test_name] = self.run_test(test_name, obj) |
| |
| # Test 6: Dict with large values |
| if size >= MIN_READ_BUF_SIZE * 2: |
| test_name = f"dict_large_values_{size_key}" |
| value_size = size // 3 |
| obj = DataGenerator.dict_with_large_values(value_size, 3) |
| all_results[size_key][test_name] = self.run_test(test_name, obj) |
| |
| # Test 7: Nested structure |
| if size >= MIN_READ_BUF_SIZE: |
| test_name = f"nested_{size_key}" |
| obj = DataGenerator.nested_structure(size) |
| all_results[size_key][test_name] = self.run_test(test_name, obj) |
| |
| # Test 8: Tuple (immutable) |
| if size >= 3: |
| test_name = f"tuple_{size_key}" |
| obj = DataGenerator.tuple_of_large_objects(size) |
| all_results[size_key][test_name] = self.run_test(test_name, obj) |
| |
| return all_results |
| |
| |
| class Comparator: |
| """Compare benchmark results between current and baseline interpreters.""" |
| |
| @staticmethod |
| def _extract_json_from_output(output: str) -> Dict[str, Dict[str, Any]]: |
| """Extract JSON data from subprocess output. |
| |
| Skips any print statements before the JSON output and parses the JSON. |
| |
| Args: |
| output: Raw stdout from subprocess |
| |
| Returns: |
| Parsed JSON as dictionary |
| |
| Raises: |
| SystemExit: If JSON cannot be found or parsed |
| """ |
| output_lines = output.strip().split('\n') |
| json_start = -1 |
| for i, line in enumerate(output_lines): |
| if line.strip().startswith('{'): |
| json_start = i |
| break |
| |
| if json_start == -1: |
| print("Error: Could not find JSON output from baseline", file=sys.stderr) |
| sys.exit(1) |
| |
| json_output = '\n'.join(output_lines[json_start:]) |
| try: |
| return json.loads(json_output) |
| except json.JSONDecodeError as e: |
| print(f"Error: Could not parse baseline JSON output: {e}", file=sys.stderr) |
| sys.exit(1) |
| |
| @staticmethod |
| def run_baseline_benchmark(baseline_python: str, args: argparse.Namespace) -> Dict[str, Dict[str, Any]]: |
| """Run the benchmark using the baseline Python interpreter.""" |
| # Build command to run this script with baseline Python |
| cmd = [ |
| baseline_python, |
| __file__, |
| '--format', 'json', |
| '--protocol', str(args.protocol), |
| '--iterations', str(args.iterations), |
| ] |
| |
| if args.sizes is not None: |
| cmd.extend(['--sizes'] + [str(s) for s in args.sizes]) |
| |
| if args.antagonistic: |
| cmd.append('--antagonistic') |
| |
| print(f"\nRunning baseline benchmark with: {baseline_python}") |
| print(f"Command: {' '.join(cmd)}\n") |
| |
| try: |
| result = subprocess.run( |
| cmd, |
| capture_output=True, |
| text=True, |
| timeout=BASELINE_BENCHMARK_TIMEOUT_SECONDS, |
| ) |
| |
| if result.returncode != 0: |
| print(f"Error running baseline benchmark:", file=sys.stderr) |
| print(result.stderr, file=sys.stderr) |
| sys.exit(1) |
| |
| # Extract and parse JSON from output |
| return Comparator._extract_json_from_output(result.stdout) |
| |
| except subprocess.TimeoutExpired: |
| print("Error: Baseline benchmark timed out", file=sys.stderr) |
| sys.exit(1) |
| |
| @staticmethod |
| def calculate_change(baseline_value: float, current_value: float) -> float: |
| """Calculate percentage change from baseline to current.""" |
| if baseline_value == 0: |
| return 0.0 |
| return ((current_value - baseline_value) / baseline_value) * 100 |
| |
| @staticmethod |
| def format_comparison( |
| current_results: Dict[str, Dict[str, Any]], |
| baseline_results: Dict[str, Dict[str, Any]] |
| ) -> str: |
| """Format comparison results as readable text.""" |
| lines = [] |
| lines.append("=" * 100) |
| lines.append("Pickle Unpickling Benchmark Comparison") |
| lines.append("=" * 100) |
| lines.append("") |
| lines.append("Legend: Current vs Baseline | % Change (+ is slower/more memory, - is faster/less memory)") |
| lines.append("") |
| |
| # Sort size keys numerically |
| for size_key in sorted(current_results.keys(), key=_extract_size_mb): |
| if size_key not in baseline_results: |
| continue |
| |
| lines.append(f"\n{size_key} Comparison") |
| lines.append("-" * 100) |
| |
| current_tests = current_results[size_key] |
| baseline_tests = baseline_results[size_key] |
| |
| for test_name in sorted(current_tests.keys()): |
| if test_name not in baseline_tests: |
| continue |
| |
| curr = current_tests[test_name] |
| base = baseline_tests[test_name] |
| |
| time_change = Comparator.calculate_change( |
| base['time']['mean'], curr['time']['mean'] |
| ) |
| mem_change = Comparator.calculate_change( |
| base['memory_peak_mb'], curr['memory_peak_mb'] |
| ) |
| |
| lines.append(f"\n {curr['test_name']}") |
| lines.append(f" Time: {curr['time']['mean']*1000:6.2f}ms vs {base['time']['mean']*1000:6.2f}ms | " |
| f"{time_change:+6.1f}%") |
| lines.append(f" Memory: {curr['memory_peak_mb']:6.2f}MB vs {base['memory_peak_mb']:6.2f}MB | " |
| f"{mem_change:+6.1f}%") |
| |
| lines.append("\n" + "=" * 100) |
| lines.append("\nSummary:") |
| |
| # Calculate overall statistics |
| time_changes = [] |
| mem_changes = [] |
| |
| for size_key in current_results.keys(): |
| if size_key not in baseline_results: |
| continue |
| for test_name in current_results[size_key].keys(): |
| if test_name not in baseline_results[size_key]: |
| continue |
| curr = current_results[size_key][test_name] |
| base = baseline_results[size_key][test_name] |
| |
| time_changes.append(Comparator.calculate_change( |
| base['time']['mean'], curr['time']['mean'] |
| )) |
| mem_changes.append(Comparator.calculate_change( |
| base['memory_peak_mb'], curr['memory_peak_mb'] |
| )) |
| |
| if time_changes: |
| lines.append(f" Time change: mean={statistics.mean(time_changes):+.1f}%, " |
| f"median={statistics.median(time_changes):+.1f}%") |
| if mem_changes: |
| lines.append(f" Memory change: mean={statistics.mean(mem_changes):+.1f}%, " |
| f"median={statistics.median(mem_changes):+.1f}%") |
| |
| lines.append("=" * 100) |
| return "\n".join(lines) |
| |
| @staticmethod |
| def format_antagonistic_comparison( |
| current_results: Dict[str, Dict[str, Any]], |
| baseline_results: Dict[str, Dict[str, Any]] |
| ) -> str: |
| """Format antagonistic benchmark comparison results.""" |
| lines = [] |
| lines.append("=" * 100) |
| lines.append("Antagonistic Pickle Benchmark Comparison (Memory DoS Protection)") |
| lines.append("=" * 100) |
| lines.append("") |
| lines.append("Legend: Current vs Baseline | Memory Change (- is better, shows memory saved)") |
| lines.append("") |
| lines.append("This compares TWO types of DoS protection:") |
| lines.append(" 1. Truncated data → Baseline allocates full claimed size, Current uses chunked reading") |
| lines.append(" 2. Sparse memo → Baseline uses huge arrays, Current uses dict-based memo") |
| lines.append("") |
| |
| # Track statistics |
| truncated_memory_changes = [] |
| sparse_memory_changes = [] |
| |
| # Sort size keys numerically |
| for size_key in sorted(current_results.keys(), key=_extract_size_mb): |
| if size_key not in baseline_results: |
| continue |
| |
| lines.append(f"\n{size_key} Comparison") |
| lines.append("-" * 100) |
| |
| current_tests = current_results[size_key] |
| baseline_tests = baseline_results[size_key] |
| |
| for test_name in sorted(current_tests.keys()): |
| if test_name not in baseline_tests: |
| continue |
| |
| curr = current_tests[test_name] |
| base = baseline_tests[test_name] |
| |
| curr_peak_mb = curr['peak_memory_mb'] |
| base_peak_mb = base['peak_memory_mb'] |
| expected_outcome = curr.get('expected_outcome', 'failure') |
| |
| mem_change = Comparator.calculate_change(base_peak_mb, curr_peak_mb) |
| mem_saved_mb = base_peak_mb - curr_peak_mb |
| |
| lines.append(f"\n {curr['test_name']}") |
| lines.append(f" Memory: {curr_peak_mb:6.2f}MB vs {base_peak_mb:6.2f}MB | " |
| f"{mem_change:+6.1f}% ({mem_saved_mb:+.2f}MB saved)") |
| |
| # Track based on test type |
| if expected_outcome == 'success': |
| sparse_memory_changes.append(mem_change) |
| if curr.get('baseline_note'): |
| lines.append(f" Note: {curr['baseline_note']}") |
| else: |
| truncated_memory_changes.append(mem_change) |
| claimed_mb = curr.get('claimed_mb', 'N/A') |
| if claimed_mb != 'N/A': |
| lines.append(f" Claimed: {claimed_mb:,}MB") |
| |
| # Show status |
| curr_status = curr.get('error_type', 'Unknown') |
| base_status = base.get('error_type', 'Unknown') |
| if curr_status != base_status: |
| lines.append(f" Status: {curr_status} (baseline: {base_status})") |
| else: |
| lines.append(f" Status: {curr_status}") |
| |
| lines.append("\n" + "=" * 100) |
| lines.append("\nSummary:") |
| lines.append("") |
| |
| if truncated_memory_changes: |
| lines.append(" Truncated Data Protection (chunked reading):") |
| lines.append(f" Mean memory change: {statistics.mean(truncated_memory_changes):+.1f}%") |
| lines.append(f" Median memory change: {statistics.median(truncated_memory_changes):+.1f}%") |
| avg_change = statistics.mean(truncated_memory_changes) |
| if avg_change < -50: |
| lines.append(f" Result: ✓ Dramatic memory reduction ({avg_change:.1f}%) - DoS protection working!") |
| elif avg_change < 0: |
| lines.append(f" Result: ✓ Memory reduced ({avg_change:.1f}%)") |
| else: |
| lines.append(f" Result: ⚠ Memory increased ({avg_change:.1f}%) - unexpected!") |
| lines.append("") |
| |
| if sparse_memory_changes: |
| lines.append(" Sparse Memo Protection (dict-based memo):") |
| lines.append(f" Mean memory change: {statistics.mean(sparse_memory_changes):+.1f}%") |
| lines.append(f" Median memory change: {statistics.median(sparse_memory_changes):+.1f}%") |
| avg_change = statistics.mean(sparse_memory_changes) |
| if avg_change < -50: |
| lines.append(f" Result: ✓ Dramatic memory reduction ({avg_change:.1f}%) - Dict optimization working!") |
| elif avg_change < 0: |
| lines.append(f" Result: ✓ Memory reduced ({avg_change:.1f}%)") |
| else: |
| lines.append(f" Result: ⚠ Memory increased ({avg_change:.1f}%) - unexpected!") |
| |
| lines.append("") |
| lines.append("=" * 100) |
| return "\n".join(lines) |
| |
| |
| class Reporter: |
| """Format and display benchmark results.""" |
| |
| @staticmethod |
| def format_text(results: Dict[str, Dict[str, Any]]) -> str: |
| """Format results as readable text.""" |
| lines = [] |
| lines.append("=" * 80) |
| lines.append("Pickle Unpickling Benchmark Results") |
| lines.append("=" * 80) |
| lines.append("") |
| |
| for size_key, tests in results.items(): |
| lines.append(f"\n{size_key} Test Results") |
| lines.append("-" * 80) |
| |
| for test_name, data in tests.items(): |
| lines.append(f"\n Test: {data['test_name']}") |
| lines.append(f" Type: {data['object_type']}") |
| lines.append(f" Pickle size: {data['pickle_size_mb']:.2f} MB") |
| lines.append(f" Time (mean): {data['time']['mean']*1000:.2f} ms") |
| lines.append(f" Time (stdev): {data['time']['stdev']*1000:.2f} ms") |
| lines.append(f" Peak memory: {data['memory_peak_mb']:.2f} MB") |
| lines.append(f" Protocol: {data['protocol']}") |
| |
| lines.append("\n" + "=" * 80) |
| return "\n".join(lines) |
| |
| @staticmethod |
| def format_markdown(results: Dict[str, Dict[str, Any]]) -> str: |
| """Format results as markdown table.""" |
| lines = [] |
| lines.append("# Pickle Unpickling Benchmark Results\n") |
| |
| for size_key, tests in results.items(): |
| lines.append(f"## {size_key}\n") |
| lines.append("| Test | Type | Pickle Size (MB) | Time (ms) | Stdev (ms) | Peak Memory (MB) |") |
| lines.append("|------|------|------------------|-----------|------------|------------------|") |
| |
| for test_name, data in tests.items(): |
| lines.append( |
| f"| {data['test_name']} | " |
| f"{data['object_type']} | " |
| f"{data['pickle_size_mb']:.2f} | " |
| f"{data['time']['mean']*1000:.2f} | " |
| f"{data['time']['stdev']*1000:.2f} | " |
| f"{data['memory_peak_mb']:.2f} |" |
| ) |
| lines.append("") |
| |
| return "\n".join(lines) |
| |
| @staticmethod |
| def format_json(results: Dict[str, Dict[str, Any]]) -> str: |
| """Format results as JSON.""" |
| import json |
| return json.dumps(results, indent=2) |
| |
| @staticmethod |
| def format_antagonistic(results: Dict[str, Dict[str, Any]]) -> str: |
| """Format antagonistic benchmark results.""" |
| lines = [] |
| lines.append("=" * 100) |
| lines.append("Antagonistic Pickle Benchmark (Memory DoS Protection Test)") |
| lines.append("=" * 100) |
| lines.append("") |
| lines.append("This benchmark tests TWO types of DoS protection:") |
| lines.append(" 1. Truncated data attacks → Expect FAILURE with minimal memory before failure") |
| lines.append(" 2. Sparse memo attacks → Expect SUCCESS with dict-based memo (vs huge array)") |
| lines.append("") |
| |
| # Sort size keys numerically |
| for size_key in sorted(results.keys(), key=_extract_size_mb): |
| tests = results[size_key] |
| |
| # Determine test type from first test |
| if tests: |
| first_test = next(iter(tests.values())) |
| expected_outcome = first_test.get('expected_outcome', 'failure') |
| claimed_mb = first_test.get('claimed_mb', 'N/A') |
| |
| # Header varies by test type |
| if "Sparse Memo" in size_key: |
| lines.append(f"\n{size_key}") |
| lines.append("-" * 100) |
| elif "Multi-Claim" in size_key: |
| lines.append(f"\n{size_key}") |
| lines.append("-" * 100) |
| elif claimed_mb != 'N/A': |
| lines.append(f"\n{size_key} Claimed (actual: 1KB) - Expect Failure") |
| lines.append("-" * 100) |
| else: |
| lines.append(f"\n{size_key}") |
| lines.append("-" * 100) |
| |
| for test_name, data in tests.items(): |
| peak_mb = data['peak_memory_mb'] |
| claimed = data.get('claimed_mb', 'N/A') |
| expected_outcome = data.get('expected_outcome', 'failure') |
| succeeded = data.get('succeeded', False) |
| baseline_note = data.get('baseline_note', '') |
| |
| lines.append(f" {data['test_name']}") |
| |
| # Format output based on test type |
| if expected_outcome == 'success': |
| # Sparse memo test - show success with dict |
| status_icon = "✓" if succeeded else "✗" |
| lines.append(f" Peak memory: {peak_mb:8.2f} MB {status_icon}") |
| lines.append(f" Status: {data['error_type']}") |
| if baseline_note: |
| lines.append(f" {baseline_note}") |
| else: |
| # Truncated data test - show savings before failure |
| if claimed != 'N/A': |
| saved_mb = claimed - peak_mb |
| savings_pct = (saved_mb / claimed * 100) if claimed > 0 else 0 |
| lines.append(f" Peak memory: {peak_mb:8.2f} MB (claimed: {claimed:,} MB, saved: {saved_mb:.2f} MB, {savings_pct:.1f}%)") |
| else: |
| lines.append(f" Peak memory: {peak_mb:8.2f} MB") |
| lines.append(f" Status: {data['error_type']}") |
| |
| lines.append("\n" + "=" * 100) |
| |
| # Calculate statistics by test type |
| truncated_claimed = 0 |
| truncated_peak = 0 |
| truncated_count = 0 |
| |
| sparse_peak_total = 0 |
| sparse_count = 0 |
| |
| for size_key, tests in results.items(): |
| for test_name, data in tests.items(): |
| expected_outcome = data.get('expected_outcome', 'failure') |
| |
| if expected_outcome == 'failure': |
| # Truncated data test |
| claimed = data.get('claimed_mb', 0) |
| if claimed != 'N/A' and claimed > 0: |
| truncated_claimed += claimed |
| truncated_peak += data['peak_memory_mb'] |
| truncated_count += 1 |
| else: |
| # Sparse memo test |
| sparse_peak_total += data['peak_memory_mb'] |
| sparse_count += 1 |
| |
| lines.append("\nSummary:") |
| lines.append("") |
| |
| if truncated_count > 0: |
| avg_claimed = truncated_claimed / truncated_count |
| avg_peak = truncated_peak / truncated_count |
| avg_saved = avg_claimed - avg_peak |
| avg_savings_pct = (avg_saved / avg_claimed * 100) if avg_claimed > 0 else 0 |
| |
| lines.append(" Truncated Data Protection (chunked reading):") |
| lines.append(f" Average claimed: {avg_claimed:,.1f} MB") |
| lines.append(f" Average peak: {avg_peak:,.2f} MB") |
| lines.append(f" Average saved: {avg_saved:,.2f} MB ({avg_savings_pct:.1f}% reduction)") |
| lines.append(f" Status: ✓ Fails fast with minimal memory") |
| lines.append("") |
| |
| if sparse_count > 0: |
| avg_sparse_peak = sparse_peak_total / sparse_count |
| lines.append(" Sparse Memo Protection (dict-based memo):") |
| lines.append(f" Average peak: {avg_sparse_peak:,.2f} MB") |
| lines.append(f" Status: ✓ Succeeds with dict (vs GB-sized arrays without PR)") |
| lines.append(f" Note: Compare with --baseline to see actual memory savings") |
| |
| lines.append("") |
| lines.append("=" * 100) |
| return "\n".join(lines) |
| |
| |
| def main(): |
| parser = argparse.ArgumentParser( |
| description="Benchmark pickle unpickling performance for large objects" |
| ) |
| parser.add_argument( |
| '--sizes', |
| type=int, |
| nargs='+', |
| default=None, |
| metavar='MiB', |
| help=f'Object sizes to test in MiB (default: {DEFAULT_SIZES_MIB})' |
| ) |
| parser.add_argument( |
| '--protocol', |
| type=int, |
| default=5, |
| choices=[0, 1, 2, 3, 4, 5], |
| help='Pickle protocol version (default: 5)' |
| ) |
| parser.add_argument( |
| '--iterations', |
| type=int, |
| default=3, |
| help='Number of benchmark iterations (default: 3)' |
| ) |
| parser.add_argument( |
| '--format', |
| choices=['text', 'markdown', 'json'], |
| default='text', |
| help='Output format (default: text)' |
| ) |
| parser.add_argument( |
| '--baseline', |
| type=str, |
| metavar='PYTHON', |
| help='Path to baseline Python interpreter for comparison (e.g., ../main-build/python)' |
| ) |
| parser.add_argument( |
| '--antagonistic', |
| action='store_true', |
| help='Run antagonistic/malicious pickle tests (DoS protection benchmark)' |
| ) |
| |
| args = parser.parse_args() |
| |
| # Handle antagonistic mode |
| if args.antagonistic: |
| # Antagonistic mode uses claimed sizes in MB, not actual data sizes |
| if args.sizes is None: |
| claimed_sizes_mb = AntagonisticTestSuite.DEFAULT_ANTAGONISTIC_SIZES_MB |
| else: |
| claimed_sizes_mb = args.sizes |
| |
| print(f"Running ANTAGONISTIC pickle benchmark (DoS protection test)...") |
| print(f"Claimed sizes: {claimed_sizes_mb} MiB (actual data: 1KB each)") |
| print(f"NOTE: These pickles will FAIL to unpickle (expected)") |
| print() |
| |
| # Run antagonistic benchmark suite |
| suite = AntagonisticTestSuite(claimed_sizes_mb) |
| results = suite.run_all_tests() |
| |
| # Format and display results |
| if args.baseline: |
| # Verify baseline Python exists |
| baseline_path = Path(args.baseline) |
| if not baseline_path.exists(): |
| print(f"Error: Baseline Python not found: {args.baseline}", file=sys.stderr) |
| return 1 |
| |
| # Run baseline benchmark |
| baseline_results = Comparator.run_baseline_benchmark(args.baseline, args) |
| |
| # Show comparison |
| comparison_output = Comparator.format_antagonistic_comparison(results, baseline_results) |
| print(comparison_output) |
| else: |
| # Format and display results |
| output = _format_output(results, args.format, is_antagonistic=True) |
| print(output) |
| |
| else: |
| # Normal mode: legitimate pickle benchmarks |
| # Convert sizes from MiB to bytes |
| if args.sizes is None: |
| sizes_bytes = DEFAULT_SIZES |
| else: |
| sizes_bytes = [size * (1 << 20) for size in args.sizes] |
| |
| print(f"Running pickle benchmark with protocol {args.protocol}...") |
| print(f"Test sizes: {[f'{s/(1<<20):.2f}MiB' for s in sizes_bytes]}") |
| print(f"Iterations per test: {args.iterations}") |
| print() |
| |
| # Run benchmark suite |
| suite = TestSuite(sizes_bytes, args.protocol, args.iterations) |
| results = suite.run_all_tests() |
| |
| # If baseline comparison requested, run baseline and compare |
| if args.baseline: |
| # Verify baseline Python exists |
| baseline_path = Path(args.baseline) |
| if not baseline_path.exists(): |
| print(f"Error: Baseline Python not found: {args.baseline}", file=sys.stderr) |
| return 1 |
| |
| # Run baseline benchmark |
| baseline_results = Comparator.run_baseline_benchmark(args.baseline, args) |
| |
| # Show comparison |
| comparison_output = Comparator.format_comparison(results, baseline_results) |
| print(comparison_output) |
| |
| else: |
| # Format and display results |
| output = _format_output(results, args.format, is_antagonistic=False) |
| print(output) |
| |
| return 0 |
| |
| |
| if __name__ == '__main__': |
| sys.exit(main()) |
