Add speedscope export timestamp fallback

news/899.feature.rst

@@ -0,0 +1,4 @@
+Add ``allocation_timestamps`` support to ``Tracker`` and a ``speedscope``
+output format for ``memray transform``. Speedscope exports now fall back to
+temporal allocation records to preserve chronological ordering when a capture
+does not include per-allocation timestamps.

src/memray/_ipython/flamegraph.py

@@ -20,6 +20,11 @@
 from memray.commands.common import warn_if_not_enough_symbols
 from memray.reporters.flamegraph import FlameGraphReporter
 
+_typed_cell_magic = cast(
+    Callable[[Callable[..., Any]], Callable[..., Any]],
+    cell_magic,
+)
+
 TEMPLATE = """\
 from memray import Tracker, FileFormat
 with Tracker(
@@ -133,7 +138,7 @@ def argument_parser() -> argparse.ArgumentParser:
 
 @magics_class
 class FlamegraphMagics(Magics):
-    @cell_magic  # type: ignore
+    @_typed_cell_magic
     def memray_flamegraph(self, line: str, cell: str) -> None:
         """Memory profile the code in the cell and display a flame graph."""
         if self.shell is None:

src/memray/_memray.pyi

@@ -48,6 +48,8 @@ class AllocationRecord:
     @property
     def native_segment_generation(self) -> int: ...
     @property
+    def timestamp_us(self) -> int: ...
+    @property
     def thread_name(self) -> str: ...
     def hybrid_stack_trace(
         self,
@@ -245,6 +247,7 @@ class Tracker:
         memory_interval_ms: int = ...,
         follow_fork: bool = ...,
         trace_python_allocators: bool = ...,
+        allocation_timestamps: bool = ...,
         file_format: FileFormat = ...,
         reference_tracking: bool = ...,
         track_object_lifetimes: bool = ...,
@@ -258,6 +261,7 @@ class Tracker:
         memory_interval_ms: int = ...,
         follow_fork: bool = ...,
         trace_python_allocators: bool = ...,
+        allocation_timestamps: bool = ...,
         file_format: FileFormat = ...,
         reference_tracking: bool = ...,
     ) -> None: ...
@@ -335,6 +339,7 @@ class RecordWriterTestHarness:
         file_path: str,
         native_traces: bool = False,
         trace_python_allocators: bool = False,
+        allocation_timestamps: bool = False,
         file_format: FileFormat = FileFormat.ALL_ALLOCATIONS,
         main_tid: int = 1,
         skipped_frames: int = 0,
@@ -358,6 +363,7 @@ class RecordWriterTestHarness:
         size: int,
         allocator: int,
         native_frame_id: int = 0,
+        timestamp_us: int = 0,
     ) -> bool: ...
     def write_frame_push(
         self,

src/memray/_memray.pyx

@@ -337,6 +337,10 @@ cdef class AllocationRecord:
     def native_segment_generation(self):
         return self._tuple[7]
 
+    @property
+    def timestamp_us(self):
+        return self._tuple[8]
+
     @property
     def thread_name(self):
         if self.tid == -1:
@@ -385,7 +389,7 @@ cdef class AllocationRecord:
     def __repr__(self):
         return (f"AllocationRecord<tid={hex(self.tid)}, address={hex(self.address)}, "
                 f"size={'N/A' if not self.size else size_fmt(self.size)}, allocator={self.allocator!r}, "
-                f"allocations={self.n_allocations}>")
+                f"allocations={self.n_allocations}, timestamp_us={self.timestamp_us}>")
 
 
 @cython.freelist(1024)
@@ -732,6 +736,8 @@ cdef class Tracker:
             created during the tracking session and still alive at the end (or
             in other words, what objects are leaked by the code being tracked).
             Defaults to False.
+        allocation_timestamps (bool): Whether or not to record a timestamp for
+            every allocation and deallocation event. Defaults to False.
         follow_fork (bool): Whether or not to continue tracking in a subprocess
             that is forked from the tracked process (see :ref:`Tracking across
             Forks`). Defaults to False.
@@ -752,6 +758,7 @@ cdef class Tracker:
     cdef unsigned int _memory_interval_ms
     cdef bool _follow_fork
     cdef bool _trace_python_allocators
+    cdef bool _allocation_timestamps
     cdef object _previous_profile_func
     cdef object _previous_thread_profile_func
     cdef unique_ptr[RecordWriter] _writer
@@ -780,6 +787,7 @@ cdef class Tracker:
     def __cinit__(self, object file_name=None, *, object destination=None,
                   bool native_traces=False, unsigned int memory_interval_ms = 10,
                   bool follow_fork=False, bool trace_python_allocators=False,
+                  bool allocation_timestamps=False,
                   bool track_object_lifetimes=False,
                   FileFormat file_format=FileFormat.ALL_ALLOCATIONS):
         if (file_name, destination).count(None) != 1:
@@ -798,6 +806,7 @@ cdef class Tracker:
         self._memory_interval_ms = memory_interval_ms
         self._follow_fork = follow_fork
         self._trace_python_allocators = trace_python_allocators
+        self._allocation_timestamps = allocation_timestamps
 
         if file_name is not None:
             destination = FileDestination(path=file_name)
@@ -809,6 +818,11 @@ cdef class Tracker:
             if file_format != FileFormat.ALL_ALLOCATIONS:
                 raise RuntimeError("AGGREGATED_ALLOCATIONS requires an output file")
 
+        if allocation_timestamps and file_format != FileFormat.ALL_ALLOCATIONS:
+            raise RuntimeError(
+                "allocation_timestamps requires FileFormat.ALL_ALLOCATIONS"
+            )
+
         self._writer = move(
             createRecordWriter(
                 move(self._make_writer(destination)),
@@ -817,6 +831,7 @@ cdef class Tracker:
                 file_format,
                 trace_python_allocators,
                 track_object_lifetimes,
+                allocation_timestamps,
             )
         )
 
@@ -861,6 +876,7 @@ cdef class Tracker:
                 self._follow_fork,
                 self._trace_python_allocators,
                 self._track_object_lifetimes,
+                self._allocation_timestamps,
             )
             return self
 
@@ -956,6 +972,7 @@ cdef _create_metadata(header, peak_memory):
         has_native_traces=header["native_traces"],
         trace_python_allocators=header["trace_python_allocators"],
         file_format=FileFormat(header["file_format"]),
+        has_allocation_timestamps=header["has_allocation_timestamps"],
     )
 
 
@@ -1806,6 +1823,7 @@ cdef class RecordWriterTestHarness:
         str file_path,
         bool native_traces=False,
         bool trace_python_allocators=False,
+        bool allocation_timestamps=False,
         track_object_lifetimes=False,
         records.FileFormat file_format=records.FileFormat.ALL_ALLOCATIONS,
         records.thread_id_t main_tid=1,
@@ -1828,6 +1846,7 @@ cdef class RecordWriterTestHarness:
             file_format,
             trace_python_allocators,
             track_object_lifetimes,
+            allocation_timestamps,
         )
         self._writer.get().setMainTidAndSkippedFrames(main_tid, skipped_frames)
         self.write_header(False)
@@ -1892,13 +1911,15 @@ cdef class RecordWriterTestHarness:
 
     def write_allocation_record(self, records.thread_id_t tid, uintptr_t address,
                                      size_t size, unsigned char allocator,
-                                     size_t native_frame_id=0) -> bool:
+                                     size_t native_frame_id=0,
+                                     uint64_t timestamp_us=0) -> bool:
         """Write a native allocation record to the file."""
         cdef records.AllocationRecord record
         record.address = address
         record.size = size
         record.allocator = <records.Allocator>allocator
         record.native_frame_id = native_frame_id
+        record.timestamp_us = timestamp_us
         return self._writer.get().writeThreadSpecificRecord(tid, record)
 
     def write_frame_push(

src/memray/_memray/record_reader.cpp

@@ -98,7 +98,10 @@ RecordReader::readHeader(HeaderRecord& header)
                 sizeof(header.trace_python_allocators))
         || !d_input->read(
                 reinterpret_cast<char*>(&header.track_object_lifetimes),
-                sizeof(header.track_object_lifetimes)))
+                sizeof(header.track_object_lifetimes))
+        || !d_input->read(
+                reinterpret_cast<char*>(&header.has_allocation_timestamps),
+                sizeof(header.has_allocation_timestamps)))
     {
         throw std::ios_base::failure("Failed to read input file header.");
     }
@@ -318,6 +321,17 @@ RecordReader::parseAllocationRecord(AllocationRecord* record, unsigned int flags
         return false;
     }
 
+    if (d_header.has_allocation_timestamps) {
+        uint64_t delta_us = 0;
+        if (!readVarint(&delta_us)) {
+            return false;
+        }
+        d_last.allocation_timestamp_us += delta_us;
+        record->timestamp_us = d_last.allocation_timestamp_us;
+    } else {
+        record->timestamp_us = 0;
+    }
+
     return true;
 }
 
@@ -342,6 +356,7 @@ RecordReader::processAllocationRecord(const AllocationRecord& record)
         d_latest_allocation.native_segment_generation = 0;
     }
     d_latest_allocation.n_allocations = 1;
+    d_latest_allocation.timestamp_us = record.timestamp_us;
     return true;
 }
 
@@ -1151,7 +1166,7 @@ RecordReader::dumpAllRecords()
            " n_allocations=%zd n_frames=%zd start_time=%lld end_time=%lld"
            " pid=%d main_tid=%lu skipped_frames_on_main_tid=%zd"
            " command_line=%s python_allocator=%s trace_python_allocators=%s"
-           " track_object_lifetimes=%s\n",
+           " track_object_lifetimes=%s has_allocation_timestamps=%s\n",
            (int)sizeof(d_header.magic),
            d_header.magic,
            d_header.version,
@@ -1168,7 +1183,8 @@ RecordReader::dumpAllRecords()
            d_header.command_line.c_str(),
            python_allocator.c_str(),
            d_header.trace_python_allocators ? "true" : "false",
-           d_header.track_object_lifetimes ? "true" : "false");
+           d_header.track_object_lifetimes ? "true" : "false",
+           d_header.has_allocation_timestamps ? "true" : "false");
 
     switch (d_header.file_format) {
         case FileFormat::ALL_ALLOCATIONS:
@@ -1222,11 +1238,21 @@ RecordReader::dumpAllRecordsFromAllAllocationsFile()
                             "<unknown allocator " + std::to_string((int)record.allocator) + ">";
                     allocator = unknownAllocator.c_str();
                 }
-                printf("address=%p size=%zd allocator=%s native_frame_id=%zd\n",
-                       (void*)record.address,
-                       record.size,
-                       allocator,
-                       record.native_frame_id);
+                if (d_header.has_allocation_timestamps) {
+                    printf("address=%p size=%zd allocator=%s native_frame_id=%zd timestamp_us=%" PRIu64
+                           "\n",
+                           (void*)record.address,
+                           record.size,
+                           allocator,
+                           record.native_frame_id,
+                           record.timestamp_us);
+                } else {
+                    printf("address=%p size=%zd allocator=%s native_frame_id=%zd\n",
+                           (void*)record.address,
+                           record.size,
+                           allocator,
+                           record.native_frame_id);
+                }
             } break;
             case RecordType::FRAME_PUSH: {
                 printf("FRAME_PUSH ");