gh-151292: _remote_debugging: Do not corrupt the binary file when hitting OverflowError (#152892)

* the kolektor

* test

* better test

* news

* =Q, move const to the base, not self.running

* gh-151292: Track binary writer finalization state

---------

Co-authored-by: Pablo Galindo Salgado <Pablogsal@gmail.com>

Author: maurycy

Lib/profiling/sampling/binary_collector.py

@@ -1,5 +1,6 @@
 """Thin Python wrapper around C binary writer for profiling data."""
 
+import sys
 import time
 
 import _remote_debugging
@@ -81,6 +82,7 @@ def __init__(self, filename, sample_interval_usec, *, skip_idle=False,
         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)
@@ -102,9 +104,19 @@ def collect(self, stack_frames, timestamp_us=None):
             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)
-        self._writer.write_sample(stack_frames, timestamp_us)
+        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)."""
@@ -143,9 +155,5 @@ def __enter__(self):
         return self
 
     def __exit__(self, exc_type, exc_val, exc_tb):
-        """Context manager exit - finalize unless there was an error."""
-        if exc_type is None:
-            self._writer.finalize()
-        else:
-            self._writer.close()
-        return False
+        """Finalize if the writer can still produce a valid file."""
+        return self._writer.__exit__(exc_type, exc_val, exc_tb)

Lib/test/test_profiling/test_sampling_profiler/test_binary_format.py

@@ -9,6 +9,8 @@
 import unittest
 from collections import defaultdict
 
+from test.support import captured_stderr
+
 try:
     import _remote_debugging
     from _remote_debugging import (
@@ -1051,6 +1053,142 @@ def test_writer_total_samples_after_close_returns_zero(self):
         w.close()
         self.assertEqual(w.total_samples, 0)
 
+    def test_binary_collector_stops_gracefully_on_overflow(self):
+        """OverflowError from the writer stops collection via the running
+        protocol instead of propagating and corrupting the file.
+        See gh-151292."""
+        with tempfile.NamedTemporaryFile(suffix=".bin", delete=False) as f:
+            filename = f.name
+        self.temp_files.append(filename)
+
+        collector = BinaryCollector(filename, 1000, compression="none")
+        self.assertTrue(collector.running)
+
+        sample = [
+            make_interpreter(0, [make_thread(1, [make_frame("a.py", 1, "f")])])
+        ]
+
+        # Collect real samples first, then hit the limit.
+        for i in range(3):
+            collector.collect(sample, timestamp_us=(i + 1) * 1000)
+        self.assertTrue(collector.running)
+
+        bad = [make_interpreter(2**32, sample[0].threads)]
+        with captured_stderr() as stderr:
+            collector.collect(bad, timestamp_us=4000)
+            collector.collect(sample, timestamp_us=5000)
+
+        self.assertFalse(collector.running)
+        self.assertTrue(collector._writer.limit_reached)
+        self.assertEqual(stderr.getvalue().count("Warning:"), 1)
+        self.assertIn("interpreter_id", stderr.getvalue())
+
+        collector.export(None)
+
+        self.assertEqual(collector.total_samples, 3)
+
+        reader_collector = RawCollector()
+        with BinaryReader(filename) as reader:
+            self.assertEqual(reader.replay_samples(reader_collector), 3)
+
+    def test_interpreter_id_overflow_rejected(self):
+        """An interpreter_id wider than u32 raises OverflowError before any
+        writer state is mutated: subsequent valid samples are still accepted
+        and finalize produces a readable file."""
+        with tempfile.NamedTemporaryFile(suffix=".bin", delete=False) as f:
+            filename = f.name
+        self.temp_files.append(filename)
+
+        good = [
+            make_interpreter(0, [make_thread(1, [make_frame("a.py", 1, "f")])])
+        ]
+        bad = [
+            make_interpreter(2**32, [make_thread(1, [make_frame("a.py", 1, "f")])])
+        ]
+
+        writer = _remote_debugging.BinaryWriter(filename, 1000, 0, compression=0)
+        writer.write_sample(good, 1000)
+        with self.assertRaises(OverflowError):
+            writer.write_sample(bad, 2000)
+        writer.write_sample(good, 3000)
+        writer.finalize()
+        self.assertEqual(writer.total_samples, 2)
+
+        reader_collector = RawCollector()
+        with BinaryReader(filename) as reader:
+            self.assertEqual(reader.replay_samples(reader_collector), 2)
+
+    def test_writer_finalizes_after_format_limit(self):
+        for compression in (0, 1) if ZSTD_AVAILABLE else (0,):
+            with self.subTest(compression=compression):
+                with tempfile.NamedTemporaryFile(suffix=".bin", delete=False) as f:
+                    filename = f.name
+                self.temp_files.append(filename)
+                good = [make_interpreter(0, [
+                    make_thread(1, [make_frame("a.py", 1, "f")])
+                ])]
+                bad = [make_interpreter(2**32, good[0].threads)]
+                writer = _remote_debugging.BinaryWriter(
+                    filename, 1000, 0, compression=compression
+                )
+                with self.assertRaises(OverflowError):
+                    with writer:
+                        writer.write_sample(good, 1000)
+                        writer.write_sample(good, 2000)
+                        # The first interpreter is committed before the limit.
+                        writer.write_sample(good + bad, 3000)
+                self.assertEqual(writer.total_samples, 3)
+                with BinaryReader(filename) as reader:
+                    self.assertEqual(reader.replay_samples(RawCollector()), 3)
+
+    def test_collector_does_not_swallow_unrelated_overflow(self):
+        class BadStatus:
+            def __index__(self):
+                raise OverflowError("status conversion failed")
+
+        with tempfile.NamedTemporaryFile(suffix=".bin", delete=False) as f:
+            filename = f.name
+        self.temp_files.append(filename)
+        collector = BinaryCollector(filename, 1000, compression="none")
+        self.addCleanup(collector._writer.close)
+        sample = [make_interpreter(0, [make_thread(1, [], BadStatus())])]
+        with captured_stderr() as stderr:
+            with self.assertRaisesRegex(OverflowError, "status conversion failed"):
+                collector.collect(sample, timestamp_us=1000)
+        self.assertEqual(stderr.getvalue(), "")
+        self.assertFalse(collector._writer.limit_reached)
+        with self.assertRaisesRegex(ValueError, "broken"):
+            collector.export()
+        with self.assertRaisesRegex(ValueError, "broken"):
+            collector._writer.write_sample([], 2000)
+        # Closing a broken writer must not attempt to finalize it.
+        collector.__exit__(None, None, None)
+
+    def test_collector_finalizes_after_external_exception(self):
+        with tempfile.NamedTemporaryFile(suffix=".bin", delete=False) as f:
+            filename = f.name
+        self.temp_files.append(filename)
+        with self.assertRaisesRegex(RuntimeError, "sampling failed"):
+            with BinaryCollector(filename, 1000, compression="none") as collector:
+                collector.collect([make_interpreter(0, [make_thread(1, [])])])
+                raise RuntimeError("sampling failed")
+        self.assertEqual(collector.total_samples, 1)
+        with BinaryReader(filename) as reader:
+            self.assertEqual(reader.replay_samples(RawCollector()), 1)
+
+    @unittest.skipUnless(os.path.exists("/dev/full"), "requires /dev/full")
+    def test_finalize_failure_breaks_writer(self):
+        writer = _remote_debugging.BinaryWriter("/dev/full", 1000, 0)
+        self.addCleanup(writer.close)
+        writer.write_sample([make_interpreter(0, [make_thread(1, [])])], 1000)
+        with self.assertRaises(OSError):
+            writer.finalize()
+        self.assertFalse(writer.limit_reached)
+        with self.assertRaisesRegex(ValueError, "broken"):
+            writer.finalize()
+        with self.assertRaisesRegex(ValueError, "broken"):
+            writer.write_sample([], 2000)
+
 
 class TestBinaryFormatValidation(BinaryFormatTestBase):
     """Tests for malformed binary files."""

Misc/NEWS.d/next/Library/2026-07-02-15-26-51.gh-issue-151292.nmnQlp.rst

@@ -0,0 +1,3 @@
+Fix ``profiling.sampling --binary`` leaving unreadable profile files when
+the binary format reaches a size limit. Preserve collected samples when the
+writer can still finalize safely. Patch by Maurycy Pawłowski-Wieroński.

Modules/_remote_debugging/binary_io.h

@@ -290,9 +290,18 @@ typedef struct {
     size_t pending_rle_samples;
 } ThreadEntry;
 
+/* Limit errors occur before emitting an incomplete sample. Other write
+ * failures may leave partial records and must prevent finalization. */
+typedef enum {
+    BINARY_WRITER_OPEN,
+    BINARY_WRITER_LIMIT_REACHED,
+    BINARY_WRITER_BROKEN,
+} BinaryWriterState;
+
 /* Main binary writer structure */
 typedef struct {
     FILE *fp;
+    BinaryWriterState state;
 
     /* Write buffer for batched I/O */
     uint8_t *write_buffer;

Modules/_remote_debugging/binary_io_writer.c

@@ -371,6 +371,7 @@ writer_intern_string(BinaryWriter *writer, PyObject *string, uint32_t *index)
     }
 
     if (writer->string_count >= UINT32_MAX) {
+        writer->state = BINARY_WRITER_LIMIT_REACHED;
         PyErr_SetString(PyExc_OverflowError,
             "too many strings for binary format");
         return -1;
@@ -380,6 +381,9 @@ writer_intern_string(BinaryWriter *writer, PyObject *string, uint32_t *index)
                                   (void **)&writer->string_lengths,
                                   &writer->string_capacity,
                                   sizeof(char *), sizeof(size_t)) < 0) {
+            if (PyErr_ExceptionMatches(PyExc_OverflowError)) {
+                writer->state = BINARY_WRITER_LIMIT_REACHED;
+            }
             return -1;
         }
     }
@@ -390,6 +394,7 @@ writer_intern_string(BinaryWriter *writer, PyObject *string, uint32_t *index)
         return -1;
     }
     if ((uintmax_t)str_len > UINT32_MAX) {
+        writer->state = BINARY_WRITER_LIMIT_REACHED;
         PyErr_Format(PyExc_OverflowError,
             "string length %zd exceeds binary format maximum %u",
             str_len, UINT32_MAX);
@@ -438,12 +443,16 @@ writer_intern_frame(BinaryWriter *writer, const FrameEntry *entry, uint32_t *ind
     }
 
     if (writer->frame_count >= UINT32_MAX) {
+        writer->state = BINARY_WRITER_LIMIT_REACHED;
         PyErr_SetString(PyExc_OverflowError,
             "too many frames for binary format");
         return -1;
     }
     if (GROW_ARRAY(writer->frame_entries, writer->frame_count,
                    writer->frame_capacity, FrameEntry) < 0) {
+        if (PyErr_ExceptionMatches(PyExc_OverflowError)) {
+            writer->state = BINARY_WRITER_LIMIT_REACHED;
+        }
         return -1;
     }
 
@@ -487,6 +496,7 @@ writer_get_or_create_thread_entry(BinaryWriter *writer, uint64_t thread_id,
     }
 
     if (writer->thread_count >= UINT32_MAX) {
+        writer->state = BINARY_WRITER_LIMIT_REACHED;
         PyErr_SetString(PyExc_OverflowError,
             "too many threads for binary format");
         return NULL;
@@ -496,6 +506,9 @@ writer_get_or_create_thread_entry(BinaryWriter *writer, uint64_t thread_id,
                                               &writer->thread_capacity,
                                               sizeof(ThreadEntry));
         if (!new_entries) {
+            if (PyErr_ExceptionMatches(PyExc_OverflowError)) {
+                writer->state = BINARY_WRITER_LIMIT_REACHED;
+            }
             return NULL;
         }
         writer->thread_entries = new_entries;
@@ -928,6 +941,12 @@ static int
 process_thread_sample(BinaryWriter *writer, PyObject *thread_info,
                       uint32_t interpreter_id, uint64_t timestamp_us)
 {
+    if (writer->total_samples == UINT64_MAX) {
+        writer->state = BINARY_WRITER_LIMIT_REACHED;
+        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);
@@ -950,7 +969,6 @@ process_thread_sample(BinaryWriter *writer, PyObject *thread_info,
 
     /* Calculate timestamp delta */
     uint64_t delta = timestamp_us - entry->prev_timestamp;
-    entry->prev_timestamp = timestamp_us;
 
     /* Process frames and build current stack */
     uint32_t curr_stack[MAX_STACK_DEPTH];
@@ -1006,6 +1024,7 @@ process_thread_sample(BinaryWriter *writer, PyObject *thread_info,
         entry->prev_stack_depth = curr_depth;
     }
 
+    entry->prev_timestamp = timestamp_us;
     writer->total_samples++;
     return 0;
 }
@@ -1025,15 +1044,16 @@ binary_writer_write_sample(BinaryWriter *writer, PyObject *stack_frames, uint64_
         PyObject *interp_id_obj = PyStructSequence_GET_ITEM(interp_info, 0);
         PyObject *threads = PyStructSequence_GET_ITEM(interp_info, 1);
 
-        unsigned long interp_id_long = PyLong_AsUnsignedLong(interp_id_obj);
-        if (interp_id_long == (unsigned long)-1 && PyErr_Occurred()) {
+        unsigned long long interp_id_long = PyLong_AsUnsignedLongLong(interp_id_obj);
+        if (interp_id_long == (unsigned long long)-1 && PyErr_Occurred()) {
             return -1;
         }
         /* Bounds check: interpreter_id is stored as uint32_t in binary format */
         if (interp_id_long > UINT32_MAX) {
+            writer->state = BINARY_WRITER_LIMIT_REACHED;
             PyErr_Format(PyExc_OverflowError,
-                "interpreter_id %lu exceeds maximum value %lu",
-                interp_id_long, (unsigned long)UINT32_MAX);
+                "interpreter_id %llu exceeds maximum value %u",
+                interp_id_long, UINT32_MAX);
             return -1;
         }
         uint32_t interpreter_id = (uint32_t)interp_id_long;

Modules/_remote_debugging/module.c

@@ -1789,7 +1789,15 @@ _remote_debugging_BinaryWriter_write_sample_impl(BinaryWriterObject *self,
         return NULL;
     }
 
+    if (self->writer->state == BINARY_WRITER_BROKEN) {
+        PyErr_SetString(PyExc_ValueError, "Writer is broken");
+        return NULL;
+    }
+    self->writer->state = BINARY_WRITER_OPEN;
     if (binary_writer_write_sample(self->writer, stack_frames, timestamp_us) < 0) {
+        if (self->writer->state != BINARY_WRITER_LIMIT_REACHED) {
+            self->writer->state = BINARY_WRITER_BROKEN;
+        }
         return NULL;
     }
 
@@ -1852,7 +1860,12 @@ _remote_debugging_BinaryWriter_set_stats_impl(BinaryWriterObject *self,
 static int
 binary_writer_finalize_and_cache(BinaryWriterObject *self)
 {
+    if (self->writer->state == BINARY_WRITER_BROKEN) {
+        PyErr_SetString(PyExc_ValueError, "Writer is broken");
+        return -1;
+    }
     if (binary_writer_finalize(self->writer) < 0) {
+        self->writer->state = BINARY_WRITER_BROKEN;
         return -1;
     }
     self->cached_total_samples = self->writer->total_samples;
@@ -1933,8 +1946,7 @@ _remote_debugging_BinaryWriter___exit___impl(BinaryWriterObject *self,
 /*[clinic end generated code: output=61831f47c72a53c6 input=12334ce1009af37f]*/
 {
     if (self->writer) {
-        /* Only finalize on normal exit (no exception) */
-        if (exc_type == Py_None) {
+        if (self->writer->state != BINARY_WRITER_BROKEN) {
             if (binary_writer_finalize_and_cache(self) < 0) {
                 if (self->writer) {
                     binary_writer_destroy(self->writer);
@@ -1983,8 +1995,17 @@ BinaryWriter_get_total_samples(PyObject *op, void *closure)
     return PyLong_FromUnsignedLongLong(self->writer->total_samples);
 }
 
+static PyObject *
+BinaryWriter_get_limit_reached(PyObject *op, void *closure)
+{
+    BinaryWriter *writer = BinaryWriter_CAST(op)->writer;
+    return PyBool_FromLong(writer && writer->state == BINARY_WRITER_LIMIT_REACHED);
+}
+
 static PyGetSetDef BinaryWriter_getset[] = {
     {"total_samples", BinaryWriter_get_total_samples, NULL, "Total samples written", NULL},
+    {"limit_reached", BinaryWriter_get_limit_reached, NULL,
+     "A format limit was reached; the collected samples can still be finalized", NULL},
     {NULL}
 };