Lib/profiling/sampling/binary_collector.py

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