gh-151292: Change total_samples:u32 in the binary format to total_samples:u64 (GH-153425)

(cherry picked from commit 1d19650)

Co-authored-by: Maurycy Pawłowski-Wieroński <maurycy@maurycy.com>

Author: maurycy

Lib/test/test_profiling/test_sampling_profiler/test_binary_format.py

@@ -999,9 +999,9 @@ class TestBinaryFormatValidation(BinaryFormatTestBase):
     """Tests for malformed binary files."""
 
     HDR_OFF_SAMPLES = 28
-    HDR_OFF_THREADS = 32
-    HDR_OFF_STR_TABLE = 36
-    HDR_OFF_FRAME_TABLE = 44
+    HDR_OFF_THREADS = 36
+    HDR_OFF_STR_TABLE = 40
+    HDR_OFF_FRAME_TABLE = 48
     FILE_HEADER_PLACEHOLDER_SIZE = 64
     FILE_FOOTER_SIZE = 32
     FTR_OFF_STRINGS = 0
@@ -1039,7 +1039,7 @@ def test_replay_rejects_sample_count_mismatch(self):
 
         with open(filename, "r+b") as raw:
             raw.seek(self.HDR_OFF_SAMPLES)
-            raw.write(struct.pack("=I", 2))
+            raw.write(struct.pack("=Q", 2))
 
         with BinaryReader(filename) as reader:
             self.assertEqual(reader.get_info()["sample_count"], 2)
@@ -1149,6 +1149,31 @@ def test_open_accepts_frame_count_at_capacity_boundary(self):
             f"possible {max_frames}",
         )
 
+    def test_sample_count_reads_full_64_bits(self):
+        """sample_count values requiring the upper 32 bits decode correctly."""
+        filename = self.create_binary_file([], compression="none")
+        big_count = 0x1_0002_0003
+
+        with open(filename, "r+b") as raw:
+            raw.seek(self.HDR_OFF_SAMPLES)
+            raw.write(struct.pack("=Q", big_count))
+
+        with BinaryReader(filename) as reader:
+            self.assertEqual(reader.get_info()["sample_count"], big_count)
+
+    def test_sample_count_boundary_values(self):
+        """Values above the old u32 ceiling decode fine."""
+        filename = self.create_binary_file([], compression="none")
+
+        for value in (0xFFFFFFFF - 1, 0xFFFFFFFF, 0xFFFFFFFF + 1):
+            with self.subTest(value=value):
+                with open(filename, "r+b") as raw:
+                    raw.seek(self.HDR_OFF_SAMPLES)
+                    raw.write(struct.pack("=Q", value))
+
+                with BinaryReader(filename) as reader:
+                    self.assertEqual(reader.get_info()["sample_count"], value)
+
 
 class TestBinaryEncodings(BinaryFormatTestBase):
     """Tests specifically targeting different stack encodings."""

Misc/NEWS.d/next/Library/2026-07-09-15-18-23.gh-issue-151292.f4NQSW.rst

@@ -0,0 +1,5 @@
+Store the sample count in :mod:`profiling.sampling` binary format, as a
+64-bit integer, widening from previously used 32-bit integer. This breaks
+the existing recordings. Long or thread-heavy profiling sessions will no
+longer fail because of :exc:`OverflowError`. Patch by Maurycy
+Pawłowski-Wieroński.

Modules/_remote_debugging/binary_io.h

@@ -50,7 +50,7 @@ extern "C" {
 #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_SIZE_SAMPLES     8
 #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)
@@ -260,7 +260,7 @@ typedef struct {
     /* Metadata */
     uint64_t start_time_us;
     uint64_t sample_interval_us;
-    uint32_t total_samples;
+    uint64_t total_samples;
 
     /* String hash table: PyObject* -> uint32_t index */
     _Py_hashtable_t *string_hash;
@@ -326,7 +326,7 @@ typedef struct {
     int needs_swap;  /* Non-zero if file was written on different-endian system */
     uint64_t start_time_us;
     uint64_t sample_interval_us;
-    uint32_t sample_count;
+    uint64_t sample_count;
     uint32_t thread_count;
     uint64_t string_table_offset;
     uint64_t frame_table_offset;

Modules/_remote_debugging/binary_io_reader.c

@@ -83,7 +83,8 @@ reader_parse_header(BinaryReader *reader, const uint8_t *data, size_t file_size)
 
     /* Read header fields with byte-swapping if needed */
     uint64_t start_time_us, sample_interval_us, string_table_offset, frame_table_offset;
-    uint32_t sample_count, thread_count, compression_type;
+    uint64_t sample_count;
+    uint32_t thread_count, compression_type;
 
     memcpy(&start_time_us, &data[HDR_OFF_START_TIME], HDR_SIZE_START_TIME);
     memcpy(&sample_interval_us, &data[HDR_OFF_INTERVAL], HDR_SIZE_INTERVAL);
@@ -95,7 +96,7 @@ reader_parse_header(BinaryReader *reader, const uint8_t *data, size_t file_size)
 
     reader->start_time_us = SWAP64_IF(reader->needs_swap, start_time_us);
     reader->sample_interval_us = SWAP64_IF(reader->needs_swap, sample_interval_us);
-    reader->sample_count = SWAP32_IF(reader->needs_swap, sample_count);
+    reader->sample_count = SWAP64_IF(reader->needs_swap, sample_count);
     reader->thread_count = SWAP32_IF(reader->needs_swap, thread_count);
     reader->string_table_offset = SWAP64_IF(reader->needs_swap, string_table_offset);
     reader->frame_table_offset = SWAP64_IF(reader->needs_swap, frame_table_offset);
@@ -993,10 +994,10 @@ emit_batch(RemoteDebuggingState *state, PyObject *collector,
 
 /* Helper to invoke progress callback, returns -1 on error */
 static inline int
-invoke_progress_callback(PyObject *callback, Py_ssize_t current, uint32_t total)
+invoke_progress_callback(PyObject *callback, Py_ssize_t current, uint64_t total)
 {
     if (callback && callback != Py_None) {
-        PyObject *result = PyObject_CallFunction(callback, "nI", current, total);
+        PyObject *result = PyObject_CallFunction(callback, "nK", current, (unsigned long long)total);
         if (result) {
             Py_DECREF(result);
         } else {
@@ -1248,8 +1249,8 @@ binary_reader_replay(BinaryReader *reader, PyObject *collector, PyObject *progre
 
     if ((uint64_t)replayed != reader->sample_count) {
         PyErr_Format(PyExc_ValueError,
-            "Sample count mismatch: header declares %u samples but replay decoded %zd",
-            reader->sample_count, replayed);
+            "Sample count mismatch: header declares %llu samples but replay decoded %zd",
+            (unsigned long long)reader->sample_count, replayed);
         return -1;
     }
 
@@ -1270,7 +1271,7 @@ binary_reader_get_info(BinaryReader *reader)
         return NULL;
     }
     return Py_BuildValue(
-        "{s:I, s:N, s:K, s:K, s:I, s:I, s:I, s:I, s:i}",
+        "{s:I, s:N, s:K, s:K, s:K, s:I, s:I, s:I, s:i}",
         "version", BINARY_FORMAT_VERSION,
         "python_version", py_version,
         "start_time_us", reader->start_time_us,

Modules/_remote_debugging/binary_io_writer.c

@@ -924,12 +924,6 @@ static int
 process_thread_sample(BinaryWriter *writer, PyObject *thread_info,
                       uint32_t interpreter_id, uint64_t timestamp_us)
 {
-    if (writer->total_samples >= UINT32_MAX) {
-        PyErr_SetString(PyExc_OverflowError,
-            "too many samples for binary format");
-        return -1;
-    }
-
     PyObject *thread_id_obj = PyStructSequence_GET_ITEM(thread_info, 0);
     PyObject *status_obj = PyStructSequence_GET_ITEM(thread_info, 1);
     PyObject *frame_list = PyStructSequence_GET_ITEM(thread_info, 2);

Modules/_remote_debugging/module.c

@@ -14,7 +14,7 @@
 typedef struct {
     PyObject_HEAD
     BinaryWriter *writer;
-    uint32_t cached_total_samples;  /* Preserved after finalize */
+    uint64_t cached_total_samples;  /* Preserved after finalize */
 } BinaryWriterObject;
 
 typedef struct {
@@ -1916,9 +1916,9 @@ BinaryWriter_get_total_samples(PyObject *op, void *closure)
     BinaryWriterObject *self = BinaryWriter_CAST(op);
     if (!self->writer) {
         /* Use cached value after finalize/close */
-        return PyLong_FromUnsignedLong(self->cached_total_samples);
+        return PyLong_FromUnsignedLongLong(self->cached_total_samples);
     }
-    return PyLong_FromUnsignedLong(self->writer->total_samples);
+    return PyLong_FromUnsignedLongLong(self->writer->total_samples);
 }
 
 static PyGetSetDef BinaryWriter_getset[] = {
@@ -2141,7 +2141,7 @@ BinaryReader_get_sample_count(BinaryReaderObject *self, void *closure)
     if (!self->reader) {
         return PyLong_FromLong(0);
     }
-    return PyLong_FromUnsignedLong(self->reader->sample_count);
+    return PyLong_FromUnsignedLongLong(self->reader->sample_count);
 }
 
 static PyObject *