"""Thin Python wrapper around C binary writer for profiling data.""" import sys import time import _remote_debugging from .collector import Collector # Compression type constants (must match binary_io.h) COMPRESSION_NONE = 0 COMPRESSION_ZSTD = 1 CAPTURE_FEATURES = { "all_threads": 1 << 0, "native": 1 << 1, "gc": 1 << 2, "opcodes": 1 << 3, "blocking": 1 << 4, } def encode_capture_config(capture_config): if capture_config is None: return -1 return sum( bit for name, bit in CAPTURE_FEATURES.items() if capture_config.get(name, False) ) def _resolve_compression(compression): """Resolve compression type from string or int. Args: compression: 'auto', 'zstd', 'none', or int (0/1) Returns: int: Compression type constant """ if isinstance(compression, int): return compression compression = compression.lower() if compression == 'none': return COMPRESSION_NONE elif compression == 'zstd': return COMPRESSION_ZSTD elif compression == 'auto': # Auto: use zstd if available, otherwise none if _remote_debugging.zstd_available(): return COMPRESSION_ZSTD return COMPRESSION_NONE else: raise ValueError(f"Unknown compression type: {compression}") class BinaryCollector(Collector): """High-performance binary collector using C implementation. This collector writes profiling data directly to a binary file format with optional zstd compression. All I/O is performed in C for maximum throughput. The binary format uses string/frame deduplication and varint encoding for efficient storage. """ def __init__(self, filename, sample_interval_usec, *, skip_idle=False, compression='auto', mode=None, capture_config=None): """Create a new binary collector. Args: filename: Path to output binary file sample_interval_usec: Sampling interval in microseconds skip_idle: If True, skip idle threads (not used in binary format) compression: 'auto', 'zstd', 'none', or int (0=none, 1=zstd) mode: Profiling mode, or None if unknown capture_config: Mapping of capture feature names to booleans, or None if the capture configuration is unknown """ self.filename = filename self.sample_interval_usec = sample_interval_usec self.skip_idle = skip_idle self.running = True compression_type = _resolve_compression(compression) start_time_us = int(time.monotonic() * 1_000_000) self._writer = _remote_debugging.BinaryWriter( filename, sample_interval_usec, start_time_us, compression=compression_type, mode=-1 if mode is None else mode, capture_features=encode_capture_config(capture_config), ) def collect(self, stack_frames, timestamp_us=None): """Collect profiling data from stack frames. This passes stack_frames directly to the C writer which handles all encoding and buffering. Args: stack_frames: List of InterpreterInfo objects from _remote_debugging timestamp_us: Optional timestamp in microseconds. If not provided, uses time.monotonic() to generate one. """ if not self.running: return if timestamp_us is None: timestamp_us = int(time.monotonic() * 1_000_000) try: self._writer.write_sample(stack_frames, timestamp_us) except OverflowError as e: if not self._writer.limit_reached: raise self.running = False print(f"Warning: {e}; stopping early and keeping the data " "collected so far.", file=sys.stderr) def collect_failed_sample(self): """Record a failed sample attempt (no-op for binary format).""" pass def set_stats(self, sample_interval_usec, duration_sec, sample_rate, error_rate=None, missed_samples=None, **kwargs): """Persist measured statistics for later replay.""" self._writer.set_stats( duration_sec, sample_rate, error_rate, missed_samples ) def export(self, filename=None): """Finalize and close the binary file. Args: filename: Ignored (binary files are written incrementally) """ self._writer.finalize() return True @property def total_samples(self): return self._writer.total_samples def get_stats(self): """Get encoding statistics. Returns: Dict with encoding statistics including repeat/full/suffix/pop-push record counts, frames written/saved, and compression ratio. """ return self._writer.get_stats() def __enter__(self): return self def __exit__(self, exc_type, exc_val, exc_tb): """Finalize if the writer can still produce a valid file.""" return self._writer.__exit__(exc_type, exc_val, exc_tb)