Fallback speedscope exports to temporal records

Keep the normal transform flow for existing formats, but route
speedscope exports through temporal allocation records when a capture
lacks exact per-allocation timestamps.

The reporter uses snapshot times to order surviving intervals for leak
reports and the intervals contributing to the peak snapshot for
high-water-mark reports.

Signed-off-by: Pablo Galindo Salgado <pablogsal@gmail.com>

Author: pablogsal

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/commands/transform.py

@@ -1,14 +1,19 @@
 import argparse
 import importlib.util
+import os
 import shutil
 import sys
 
 from rich import print as pprint
 
+from memray import FileReader
 from memray._errors import MemrayCommandError
+from memray._memray import FileFormat
 
 from ..reporters.transform import TransformReporter
 from .common import HighWatermarkCommand
+from .common import warn_if_file_is_not_aggregated_and_is_too_big
+from .common import warn_if_not_enough_symbols
 
 
 class TransformCommand(HighWatermarkCommand):
@@ -72,3 +77,103 @@ def post_run_gprof2dot(self) -> None:
         print()
         print("To generate a graph from the transform file, run for example:")
         print(f"{command} -f json {self.output_file} | dot -Tpng -o output.png")
+
+    def write_report(
+        self,
+        result_path,
+        output_file,
+        show_memory_leaks,
+        temporary_allocation_threshold,
+        merge_threads=None,
+        inverted=None,
+        temporal=False,
+        max_memory_records=None,
+        no_web=False,
+    ) -> None:
+        if self.reporter_name != "speedscope":
+            return super().write_report(
+                result_path=result_path,
+                output_file=output_file,
+                show_memory_leaks=show_memory_leaks,
+                temporary_allocation_threshold=temporary_allocation_threshold,
+                merge_threads=merge_threads,
+                inverted=inverted,
+                temporal=temporal,
+                max_memory_records=max_memory_records,
+                no_web=no_web,
+            )
+
+        try:
+            kwargs = {}
+            if max_memory_records is not None:
+                kwargs["max_memory_records"] = max_memory_records
+            reader = FileReader(os.fspath(result_path), report_progress=True, **kwargs)
+            merge_threads = True if merge_threads is None else merge_threads
+            inverted = False if inverted is None else inverted
+
+            if reader.metadata.has_native_traces:
+                warn_if_not_enough_symbols()
+
+            if not temporal and temporary_allocation_threshold < 0:
+                warn_if_file_is_not_aggregated_and_is_too_big(reader, result_path)
+
+            memory_records = tuple(reader.get_memory_snapshots())
+            reporter_kwargs = {
+                "memory_records": memory_records,
+                "native_traces": reader.metadata.has_native_traces,
+            }
+
+            use_temporal_fallback = (
+                reader.metadata.file_format == FileFormat.ALL_ALLOCATIONS
+                and not reader.metadata.has_allocation_timestamps
+                and temporary_allocation_threshold < 0
+            )
+
+            if use_temporal_fallback:
+                if show_memory_leaks:
+                    allocations = reader.get_temporal_allocation_records(
+                        merge_threads=merge_threads
+                    )
+                    reporter = self.reporter_factory(allocations, **reporter_kwargs)
+                else:
+                    (
+                        allocations,
+                        high_water_mark_by_snapshot,
+                    ) = reader.get_temporal_high_water_mark_allocation_records(
+                        merge_threads=merge_threads
+                    )
+                    reporter = self.reporter_factory(
+                        allocations,
+                        high_water_mark_by_snapshot=high_water_mark_by_snapshot,
+                        **reporter_kwargs,
+                    )
+            else:
+                if show_memory_leaks:
+                    allocations = reader.get_leaked_allocation_records(
+                        merge_threads=merge_threads
+                    )
+                elif temporary_allocation_threshold >= 0:
+                    allocations = reader.get_temporary_allocation_records(
+                        threshold=temporary_allocation_threshold,
+                        merge_threads=merge_threads,
+                    )
+                else:
+                    allocations = reader.get_high_watermark_allocation_records(
+                        merge_threads=merge_threads
+                    )
+                reporter = self.reporter_factory(allocations, **reporter_kwargs)
+        except OSError as e:
+            raise MemrayCommandError(
+                f"Failed to parse allocation records in {result_path}\nReason: {e}",
+                exit_code=1,
+            )
+
+        with open(os.fspath(output_file.expanduser()), "w") as f:
+            reporter.render(
+                outfile=f,
+                metadata=reader.metadata,
+                show_memory_leaks=show_memory_leaks,
+                merge_threads=merge_threads,
+                inverted=inverted,
+                no_web=no_web,
+            )

src/memray/reporters/transform.py

@@ -7,12 +7,14 @@
 from typing import Optional
 from typing import TextIO
 from typing import Tuple
+from typing import Union
 
 from memray import AllocationRecord
 from memray import AllocatorType
 from memray import MemorySnapshot
 from memray import Metadata
 from memray import __version__
+from memray._memray import TemporalAllocationRecord
 from memray.reporters.common import format_thread_name
 
 Location = Tuple[str, str]
@@ -34,13 +36,15 @@ def __init__(
         format: str,
         native_traces: bool,
         memory_records: Iterable[MemorySnapshot],
+        high_water_mark_by_snapshot: Optional[List[int]] = None,
         **kwargs: Any,
     ) -> None:
         super().__init__()
         self.allocations = allocations
         self.format = format
         self.native_traces = native_traces
-        self.memory_records = memory_records
+        self.memory_records = tuple(memory_records)
+        self.high_water_mark_by_snapshot = high_water_mark_by_snapshot
 
     def render_as_gprof2dot(
         self,
@@ -74,7 +78,9 @@ def render_as_gprof2dot(
         }
         json.dump(result, outfile)
 
-    def _stack_trace_for_record(self, record: AllocationRecord) -> Tuple[Tuple[str, str, int], ...]:
+    def _stack_trace_for_record(
+        self, record: Union[AllocationRecord, TemporalAllocationRecord]
+    ) -> Tuple[Tuple[str, str, int], ...]:
         return (
             tuple(record.hybrid_stack_trace())
             if self.native_traces
@@ -83,7 +89,7 @@ def _stack_trace_for_record(self, record: AllocationRecord) -> Tuple[Tuple[str,
 
     def _speedscope_sample_for_record(
         self,
-        record: AllocationRecord,
+        record: Union[AllocationRecord, TemporalAllocationRecord],
         *,
         location_to_index: Dict[FrameLocation, int],
         frames: List[Dict[str, Any]],
@@ -140,7 +146,9 @@ def _aggregate_snapshot_speedscope_samples(
         frames: List[Dict[str, Any]] = []
         sample_weights: Dict[FrameSample, List[int]] = {}
         sample_order: Dict[FrameSample, int] = {}
-        has_exact_timestamps = metadata is not None and metadata.has_allocation_timestamps
+        has_exact_timestamps = (
+            metadata is not None and metadata.has_allocation_timestamps
+        )
 
         for sequence, record in enumerate(allocations):
             sample = self._speedscope_sample_for_record(
@@ -164,6 +172,111 @@ def _aggregate_snapshot_speedscope_samples(
         )
         return frames, ordered_samples
 
+    def _snapshot_order_key(self, snapshot_index: int) -> int:
+        if 0 <= snapshot_index < len(self.memory_records):
+            # Convert ms-since-epoch to µs for comparison with timestamp_us
+            return self.memory_records[snapshot_index].time * 1000
+        if self.memory_records:
+            return self.memory_records[-1].time * 1000
+        return snapshot_index
+
+    def _peak_snapshot_index(self) -> int:
+        high_water_mark_by_snapshot = self.high_water_mark_by_snapshot or [0]
+        return max(
+            range(len(high_water_mark_by_snapshot)),
+            key=high_water_mark_by_snapshot.__getitem__,
+        )
+
+    def _contribution_for_temporal_record(
+        self,
+        record: TemporalAllocationRecord,
+        *,
+        show_memory_leaks: bool,
+        peak_snapshot: Optional[int] = None,
+    ) -> Tuple[int, int, Optional[int]]:
+        size = 0
+        n_allocations = 0
+        first_snapshot = None
+
+        if show_memory_leaks:
+            for interval in record.intervals:
+                if interval.deallocated_before_snapshot is not None:
+                    continue
+                size += interval.n_bytes
+                n_allocations += interval.n_allocations
+                snapshot = interval.allocated_before_snapshot
+                if first_snapshot is None or snapshot < first_snapshot:
+                    first_snapshot = snapshot
+            return size, n_allocations, first_snapshot
+
+        if peak_snapshot is None:
+            peak_snapshot = self._peak_snapshot_index()
+        for interval in record.intervals:
+            if interval.allocated_before_snapshot > peak_snapshot:
+                continue
+            if (
+                interval.deallocated_before_snapshot is not None
+                and peak_snapshot >= interval.deallocated_before_snapshot
+            ):
+                continue
+            size += interval.n_bytes
+            n_allocations += interval.n_allocations
+            snapshot = interval.allocated_before_snapshot
+            if first_snapshot is None or snapshot < first_snapshot:
+                first_snapshot = snapshot
+        return size, n_allocations, first_snapshot
+
+    def _aggregate_temporal_speedscope_samples(
+        self,
+        allocations: Iterable[TemporalAllocationRecord],
+        *,
+        show_memory_leaks: bool,
+    ) -> Tuple[List[Dict[str, Any]], List[Tuple[FrameSample, List[int]]]]:
+        location_to_index: Dict[FrameLocation, int] = {}
+        frames: List[Dict[str, Any]] = []
+        sample_weights: Dict[FrameSample, List[int]] = {}
+        sample_order: Dict[FrameSample, int] = {}
+
+        peak_snapshot = None if show_memory_leaks else self._peak_snapshot_index()
+
+        for sequence, record in enumerate(allocations):
+            sample = self._speedscope_sample_for_record(
+                record,
+                location_to_index=location_to_index,
+                frames=frames,
+            )
+            (
+                size,
+                n_allocations,
+                first_snapshot,
+            ) = self._contribution_for_temporal_record(
+                record,
+                show_memory_leaks=show_memory_leaks,
+                peak_snapshot=peak_snapshot,
+            )
+            if size <= 0 and n_allocations <= 0:
+                continue
+
+            order_key = (
+                self._snapshot_order_key(first_snapshot)
+                if first_snapshot is not None
+                else sequence
+            )
+            self._add_speedscope_sample(
+                sample=sample,
+                size=size,
+                n_allocations=n_allocations,
+                order_key=order_key,
+                sample_weights=sample_weights,
+                sample_order=sample_order,
+            )
+
+        ordered_samples = sorted(
+            sample_weights.items(),
+            key=lambda item: (sample_order[item[0]], item[0]),
+        )
+        return frames, ordered_samples
+
     def _create_speedscope_profile(
         self,
         *,
@@ -198,10 +311,19 @@ def render_as_speedscope(
         **kwargs: Any,
     ) -> None:
         metadata = kwargs.get("metadata")
-        frames, sample_weights = self._aggregate_snapshot_speedscope_samples(
-            self.allocations,
-            metadata=metadata,
-        )
+        show_memory_leaks = kwargs.get("show_memory_leaks", False)
+        allocations = list(self.allocations)
+
+        if allocations and self._is_temporal_record(allocations[0]):
+            frames, sample_weights = self._aggregate_temporal_speedscope_samples(
+                allocations,
+                show_memory_leaks=show_memory_leaks,
+            )
+        else:
+            frames, sample_weights = self._aggregate_snapshot_speedscope_samples(
+                allocations,
+                metadata=metadata,
+            )
 
         result = {
             "$schema": "https://www.speedscope.app/file-format-schema.json",
@@ -272,3 +394,9 @@ def render_as_csv(
                     "|".join(f"{func};{mod};{line}" for func, mod, line in stack_trace),
                 ]
             )
+
+    @staticmethod
+    def _is_temporal_record(
+        record: Union[AllocationRecord, TemporalAllocationRecord]
+    ) -> bool:
+        return hasattr(record, "intervals")