Address Laszlo's review

Author: pablogsal

Doc/library/profiling.sampling.rst

@@ -1084,7 +1084,7 @@ The :option:`--compression` option controls data compression:
 
 - ``auto`` (default): Use zstd compression if available, otherwise no
   compression
-- ``zstd``: Force zstd compression (requires zstd support)
+- ``zstd``: Force zstd compression (requires :mod:`compression.zstd` support)
 - ``none``: Disable compression
 
 ::

Lib/profiling/sampling/binary_collector.py

@@ -2,6 +2,8 @@
 
 import time
 
+import _remote_debugging
+
 from .collector import Collector
 
 # Compression type constants (must match binary_io.h)
@@ -28,7 +30,6 @@ def _resolve_compression(compression):
         return COMPRESSION_ZSTD
     elif compression == 'auto':
         # Auto: use zstd if available, otherwise none
-        import _remote_debugging
         if _remote_debugging.zstd_available():
             return COMPRESSION_ZSTD
         return COMPRESSION_NONE
@@ -57,8 +58,6 @@ def __init__(self, filename, sample_interval_usec, *, skip_idle=False,
             skip_idle: If True, skip idle threads (not used in binary format)
             compression: 'auto', 'zstd', 'none', or int (0=none, 1=zstd)
         """
-        import _remote_debugging
-
         self.filename = filename
         self.sample_interval_usec = sample_interval_usec
         self.skip_idle = skip_idle

Lib/profiling/sampling/cli.py

@@ -17,7 +17,7 @@
 from .heatmap_collector import HeatmapCollector
 from .gecko_collector import GeckoCollector
 from .binary_collector import BinaryCollector
-from .binary_reader import BinaryReader, convert_binary_to_format
+from .binary_reader import BinaryReader
 from .constants import (
     PROFILING_MODE_ALL,
     PROFILING_MODE_WALL,

Modules/_remote_debugging/binary_io.h

@@ -22,7 +22,37 @@ extern "C" {
  * ============================================================================ */
 
 #define BINARY_FORMAT_MAGIC     0x54414348  /* "TACH" (Tachyon) */
-#define BINARY_FORMAT_VERSION   2
+#define BINARY_FORMAT_VERSION   3
+
+/* Header field offsets and sizes */
+#define HDR_OFF_MAGIC        0
+#define HDR_SIZE_MAGIC       4
+#define HDR_OFF_VERSION      (HDR_OFF_MAGIC + HDR_SIZE_MAGIC)
+#define HDR_SIZE_VERSION     4
+#define HDR_OFF_PY_VERSION   (HDR_OFF_VERSION + HDR_SIZE_VERSION)
+#define HDR_SIZE_PY_VERSION  4   /* 3 bytes: major, minor, micro + 1 reserved */
+#define HDR_OFF_PY_MAJOR     HDR_OFF_PY_VERSION
+#define HDR_OFF_PY_MINOR     (HDR_OFF_PY_VERSION + 1)
+#define HDR_OFF_PY_MICRO     (HDR_OFF_PY_VERSION + 2)
+#define HDR_OFF_START_TIME   (HDR_OFF_PY_VERSION + HDR_SIZE_PY_VERSION)
+#define HDR_SIZE_START_TIME  8
+#define HDR_OFF_INTERVAL     (HDR_OFF_START_TIME + HDR_SIZE_START_TIME)
+#define HDR_SIZE_INTERVAL    8
+#define HDR_OFF_SAMPLES      (HDR_OFF_INTERVAL + HDR_SIZE_INTERVAL)
+#define HDR_SIZE_SAMPLES     4
+#define HDR_OFF_THREADS      (HDR_OFF_SAMPLES + HDR_SIZE_SAMPLES)
+#define HDR_SIZE_THREADS     4
+#define HDR_OFF_STR_TABLE    (HDR_OFF_THREADS + HDR_SIZE_THREADS)
+#define HDR_SIZE_STR_TABLE   8
+#define HDR_OFF_FRAME_TABLE  (HDR_OFF_STR_TABLE + HDR_SIZE_STR_TABLE)
+#define HDR_SIZE_FRAME_TABLE 8
+#define HDR_OFF_COMPRESSION  (HDR_OFF_FRAME_TABLE + HDR_SIZE_FRAME_TABLE)
+#define HDR_SIZE_COMPRESSION 4
+#define FILE_HEADER_SIZE     (HDR_OFF_COMPRESSION + HDR_SIZE_COMPRESSION)
+#define FILE_HEADER_PLACEHOLDER_SIZE 64
+
+static_assert(FILE_HEADER_SIZE <= FILE_HEADER_PLACEHOLDER_SIZE,
+              "FILE_HEADER_SIZE exceeds FILE_HEADER_PLACEHOLDER_SIZE");
 
 /* Buffer sizes: 512KB balances syscall amortization against memory use,
  * and aligns well with filesystem block sizes and zstd dictionary windows */
@@ -253,6 +283,9 @@ typedef struct {
     size_t decompressed_size;
 
     /* Header metadata */
+    uint8_t py_major;
+    uint8_t py_minor;
+    uint8_t py_micro;
     uint64_t start_time_us;
     uint64_t sample_interval_us;
     uint32_t sample_count;
@@ -415,10 +448,22 @@ grow_array(void *ptr, size_t *capacity, size_t elem_size)
     return new_ptr;
 }
 
-/* Macro wrapper for type safety with grow_array */
-#define GROW_ARRAY(ptr, count, capacity, type) \
-    ((count) < (capacity) ? 0 : \
-     ((ptr) = grow_array((ptr), &(capacity), sizeof(type))) ? 0 : -1)
+static inline int
+grow_array_inplace(void **ptr_addr, size_t count, size_t *capacity, size_t elem_size)
+{
+    if (count < *capacity) {
+        return 0;
+    }
+    void *tmp = grow_array(*ptr_addr, capacity, elem_size);
+    if (tmp == NULL) {
+        return -1;
+    }
+    *ptr_addr = tmp;
+    return 0;
+}
+
+#define GROW_ARRAY(ptr, count, cap, type) \
+    grow_array_inplace((void**)&(ptr), (count), &(cap), sizeof(type))
 
 /* ============================================================================
  * BINARY WRITER API

Modules/_remote_debugging/binary_io_reader.c

@@ -22,9 +22,7 @@
  * ============================================================================ */
 
 /* File structure sizes */
-#define FILE_HEADER_PLACEHOLDER_SIZE 64  /* Placeholder written at file start */
-#define FILE_HEADER_SIZE 52              /* Actual header content size */
-#define FILE_FOOTER_SIZE 32              /* Footer size */
+#define FILE_FOOTER_SIZE 32
 #define MIN_DECOMPRESS_BUFFER_SIZE (64 * 1024)  /* Minimum decompression buffer */
 
 /* Progress callback frequency */
@@ -57,17 +55,35 @@ reader_parse_header(BinaryReader *reader, const uint8_t *data, size_t file_size)
     }
 
     if (version != BINARY_FORMAT_VERSION) {
-        PyErr_Format(PyExc_ValueError, "Unsupported version: %u", version);
+        if (version > BINARY_FORMAT_VERSION && file_size >= HDR_OFF_PY_MICRO + 1) {
+            /* Newer format - try to read Python version for better error */
+            uint8_t py_major = data[HDR_OFF_PY_MAJOR];
+            uint8_t py_minor = data[HDR_OFF_PY_MINOR];
+            uint8_t py_micro = data[HDR_OFF_PY_MICRO];
+            PyErr_Format(PyExc_ValueError,
+                "Binary file was created with Python %u.%u.%u (format version %u), "
+                "but this is Python %d.%d.%d (format version %d)",
+                py_major, py_minor, py_micro, version,
+                PY_MAJOR_VERSION, PY_MINOR_VERSION, PY_MICRO_VERSION,
+                BINARY_FORMAT_VERSION);
+        } else {
+            PyErr_Format(PyExc_ValueError,
+                "Unsupported format version %u (this reader supports version %d)",
+                version, BINARY_FORMAT_VERSION);
+        }
         return -1;
     }
 
-    memcpy(&reader->start_time_us, &data[8], sizeof(reader->start_time_us));
-    memcpy(&reader->sample_interval_us, &data[16], sizeof(reader->sample_interval_us));
-    memcpy(&reader->sample_count, &data[24], sizeof(reader->sample_count));
-    memcpy(&reader->thread_count, &data[28], sizeof(reader->thread_count));
-    memcpy(&reader->string_table_offset, &data[32], sizeof(reader->string_table_offset));
-    memcpy(&reader->frame_table_offset, &data[40], sizeof(reader->frame_table_offset));
-    memcpy(&reader->compression_type, &data[48], sizeof(reader->compression_type));
+    reader->py_major = data[HDR_OFF_PY_MAJOR];
+    reader->py_minor = data[HDR_OFF_PY_MINOR];
+    reader->py_micro = data[HDR_OFF_PY_MICRO];
+    memcpy(&reader->start_time_us, &data[HDR_OFF_START_TIME], HDR_SIZE_START_TIME);
+    memcpy(&reader->sample_interval_us, &data[HDR_OFF_INTERVAL], HDR_SIZE_INTERVAL);
+    memcpy(&reader->sample_count, &data[HDR_OFF_SAMPLES], HDR_SIZE_SAMPLES);
+    memcpy(&reader->thread_count, &data[HDR_OFF_THREADS], HDR_SIZE_THREADS);
+    memcpy(&reader->string_table_offset, &data[HDR_OFF_STR_TABLE], HDR_SIZE_STR_TABLE);
+    memcpy(&reader->frame_table_offset, &data[HDR_OFF_FRAME_TABLE], HDR_SIZE_FRAME_TABLE);
+    memcpy(&reader->compression_type, &data[HDR_OFF_COMPRESSION], HDR_SIZE_COMPRESSION);
 
     return 0;
 }
@@ -1023,9 +1039,15 @@ binary_reader_replay(BinaryReader *reader, PyObject *collector, PyObject *progre
 PyObject *
 binary_reader_get_info(BinaryReader *reader)
 {
+    PyObject *py_version = Py_BuildValue("(B,B,B)",
+        reader->py_major, reader->py_minor, reader->py_micro);
+    if (py_version == NULL) {
+        return NULL;
+    }
     return Py_BuildValue(
-        "{s:I, s:K, s:K, s:I, s:I, s:I, s:I, s:i}",
+        "{s:I, s:N, s:K, s:K, s:I, s:I, s:I, s:I, s:i}",
         "version", BINARY_FORMAT_VERSION,
+        "python_version", py_version,
         "start_time_us", reader->start_time_us,
         "sample_interval_us", reader->sample_interval_us,
         "sample_count", reader->sample_count,

Modules/_remote_debugging/binary_io_writer.c

@@ -30,9 +30,7 @@
 #define MAX_FRAME_BUFFER_SIZE ((MAX_STACK_DEPTH * MAX_VARINT_SIZE_U32) + MAX_VARINT_SIZE_U32 + 16)
 
 /* File structure sizes */
-#define FILE_HEADER_PLACEHOLDER_SIZE 64  /* Placeholder written at file start */
-#define FILE_HEADER_SIZE 52              /* Actual header content size */
-#define FILE_FOOTER_SIZE 32              /* Footer size */
+#define FILE_FOOTER_SIZE 32
 
 /* ============================================================================
  * WRITER-SPECIFIC UTILITY HELPERS
@@ -1073,19 +1071,22 @@ binary_writer_finalize(BinaryWriter *writer)
     uint64_t string_table_offset_u64 = (uint64_t)string_table_offset;
     uint64_t frame_table_offset_u64 = (uint64_t)frame_table_offset;
 
-    uint8_t header[52] = {0};
+    uint8_t header[FILE_HEADER_SIZE] = {0};
     uint32_t magic = BINARY_FORMAT_MAGIC;
     uint32_t version = BINARY_FORMAT_VERSION;
-    memcpy(header + 0, &magic, 4);
-    memcpy(header + 4, &version, 4);
-    memcpy(header + 8, &writer->start_time_us, 8);
-    memcpy(header + 16, &writer->sample_interval_us, 8);
-    memcpy(header + 24, &writer->total_samples, 4);
-    memcpy(header + 28, &writer->thread_count, 4);
-    memcpy(header + 32, &string_table_offset_u64, 8);
-    memcpy(header + 40, &frame_table_offset_u64, 8);
-    memcpy(header + 48, &writer->compression_type, 4);
-    if (fwrite_checked_allow_threads(header, 52, writer->fp) < 0) {
+    memcpy(header + HDR_OFF_MAGIC, &magic, HDR_SIZE_MAGIC);
+    memcpy(header + HDR_OFF_VERSION, &version, HDR_SIZE_VERSION);
+    header[HDR_OFF_PY_MAJOR] = PY_MAJOR_VERSION;
+    header[HDR_OFF_PY_MINOR] = PY_MINOR_VERSION;
+    header[HDR_OFF_PY_MICRO] = PY_MICRO_VERSION;
+    memcpy(header + HDR_OFF_START_TIME, &writer->start_time_us, HDR_SIZE_START_TIME);
+    memcpy(header + HDR_OFF_INTERVAL, &writer->sample_interval_us, HDR_SIZE_INTERVAL);
+    memcpy(header + HDR_OFF_SAMPLES, &writer->total_samples, HDR_SIZE_SAMPLES);
+    memcpy(header + HDR_OFF_THREADS, &writer->thread_count, HDR_SIZE_THREADS);
+    memcpy(header + HDR_OFF_STR_TABLE, &string_table_offset_u64, HDR_SIZE_STR_TABLE);
+    memcpy(header + HDR_OFF_FRAME_TABLE, &frame_table_offset_u64, HDR_SIZE_FRAME_TABLE);
+    memcpy(header + HDR_OFF_COMPRESSION, &writer->compression_type, HDR_SIZE_COMPRESSION);
+    if (fwrite_checked_allow_threads(header, FILE_HEADER_SIZE, writer->fp) < 0) {
         return -1;
     }