Class PickleBenchmark

Tools/picklebench/memory_dos_impact.py:227–294

Benchmark pickle unpickling performance and memory usage.

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()