Fixes and more tests

Author: pablogsal

Lib/profiling/sampling/collector.py

@@ -44,14 +44,16 @@ def extract_lineno(location):
 
 class Collector(ABC):
     @abstractmethod
-    def collect(self, stack_frames, timestamp_us=None):
+    def collect(self, stack_frames, timestamps_us=None):
         """Collect profiling data from stack frames.
 
         Args:
             stack_frames: List of InterpreterInfo objects
-            timestamp_us: Optional timestamp in microseconds. If provided (from
-                binary replay), use this instead of current time. If None,
-                collectors should use time.monotonic() or similar.
+            timestamps_us: Optional list of timestamps in microseconds. If provided
+                (from binary replay with RLE batching), use these instead of current
+                time. If None, collectors should use time.monotonic() or similar.
+                The list may contain multiple timestamps when samples are batched
+                together (same stack, different times).
         """
 
     def collect_failed_sample(self):

Lib/test/test_profiling/test_sampling_profiler/test_binary_format.py

@@ -1012,5 +1012,70 @@ def test_same_thread_id_multiple_interpreters_stress(self):
             )
 
 
+class TimestampCollector:
+    """Collector that captures timestamps for verification."""
+
+    def __init__(self):
+        self.all_timestamps = []
+
+    def collect(self, stack_frames, timestamps_us=None):
+        if timestamps_us is not None:
+            self.all_timestamps.extend(timestamps_us)
+
+    def export(self, filename):
+        pass
+
+
+class TestTimestampPreservation(BinaryFormatTestBase):
+    """Tests for timestamp preservation during binary round-trip."""
+
+    def test_timestamp_preservation(self):
+        """Timestamps are preserved during round-trip."""
+        frame = make_frame("test.py", 10, "func")
+        timestamps = [1000000, 2000000, 3000000]
+
+        with tempfile.NamedTemporaryFile(suffix=".bin", delete=False) as f:
+            filename = f.name
+        self.temp_files.append(filename)
+
+        collector = BinaryCollector(filename, 1000, compression="none")
+        for ts in timestamps:
+            sample = [make_interpreter(0, [make_thread(1, [frame])])]
+            collector.collect(sample, timestamp_us=ts)
+        collector.export(None)
+
+        ts_collector = TimestampCollector()
+        with BinaryReader(filename) as reader:
+            count = reader.replay_samples(ts_collector)
+
+        self.assertEqual(count, 3)
+        self.assertEqual(ts_collector.all_timestamps, timestamps)
+
+    def test_timestamp_preservation_with_rle(self):
+        """RLE-batched samples preserve individual timestamps."""
+        frame = make_frame("rle.py", 42, "rle_func")
+
+        with tempfile.NamedTemporaryFile(suffix=".bin", delete=False) as f:
+            filename = f.name
+        self.temp_files.append(filename)
+
+        # Identical samples (triggers RLE) with different timestamps
+        collector = BinaryCollector(filename, 1000, compression="none")
+        expected_timestamps = []
+        for i in range(50):
+            ts = 1000000 + i * 100
+            expected_timestamps.append(ts)
+            sample = [make_interpreter(0, [make_thread(1, [frame])])]
+            collector.collect(sample, timestamp_us=ts)
+        collector.export(None)
+
+        ts_collector = TimestampCollector()
+        with BinaryReader(filename) as reader:
+            count = reader.replay_samples(ts_collector)
+
+        self.assertEqual(count, 50)
+        self.assertEqual(ts_collector.all_timestamps, expected_timestamps)
+
+
 if __name__ == "__main__":
     unittest.main()

Modules/_remote_debugging/binary_io_reader.c

@@ -586,6 +586,14 @@ decode_stack_suffix(ReaderThreadState *ts, const uint8_t *data,
 
     /* Move shared frames (from bottom of stack) to make room for new frames at the top */
     if (new_count > 0 && shared > 0) {
+        /* Defensive check: ensure subtraction won't underflow.
+         * This should already be guaranteed by the check above, but we add
+         * this assertion as defense-in-depth against stack corruption. */
+        if (ts->current_stack_depth < shared) {
+            PyErr_SetString(PyExc_ValueError,
+                "Internal error: stack corruption detected in suffix decoding");
+            return -1;
+        }
         size_t prev_shared_start = ts->current_stack_depth - shared;
         memmove(&ts->current_stack[new_count],
                 &ts->current_stack[prev_shared_start],

Modules/_remote_debugging/binary_io_writer.c

@@ -447,12 +447,13 @@ writer_get_or_create_thread_entry(BinaryWriter *writer, uint64_t thread_id,
     }
 
     if (writer->thread_count >= writer->thread_capacity) {
-        writer->thread_entries = grow_array(writer->thread_entries,
-                                            &writer->thread_capacity,
-                                            sizeof(ThreadEntry));
-        if (!writer->thread_entries) {
+        ThreadEntry *new_entries = grow_array(writer->thread_entries,
+                                              &writer->thread_capacity,
+                                              sizeof(ThreadEntry));
+        if (!new_entries) {
             return NULL;
         }
+        writer->thread_entries = new_entries;
     }
 
     ThreadEntry *entry = &writer->thread_entries[writer->thread_count];

Modules/_remote_debugging/module.c

@@ -1307,11 +1307,13 @@ _remote_debugging_BinaryWriter___exit___impl(BinaryWriterObject *self,
 /*[clinic end generated code: output=61831f47c72a53c6 input=12334ce1009af37f]*/
 {
     if (self->writer) {
-        /* Finalize on normal exit */
-        if (binary_writer_finalize(self->writer) < 0) {
-            binary_writer_destroy(self->writer);
-            self->writer = NULL;
-            return NULL;
+        /* Only finalize on normal exit (no exception) */
+        if (exc_type == Py_None) {
+            if (binary_writer_finalize(self->writer) < 0) {
+                binary_writer_destroy(self->writer);
+                self->writer = NULL;
+                return NULL;
+            }
         }
         binary_writer_destroy(self->writer);
         self->writer = NULL;