Add Python object reference tracking

This feature adds the capability to track Python objects throughout
their lifecycle in memray. The implementation introduces a reference
tracking system that records PyObject creation and destruction events,
storing them as TrackedObject records in the memray capture file with
associated stack traces. The Tracker class now accepts a
reference_tracking flag to enable this functionality and provides a
get_surviving_objects() method to retrieve live objects at the end of
tracking.

Object tracking works across threads and supports native, Python, and
hybrid stack traces, providing insights into Python object
lifetimes. This enhancement extends memray beyond memory allocation
tracking to give developers a powerful tool for understanding object
lifecycles and identifying potential reference leaks.

With this addition, memray can now correlate memory allocations with the
Python objects that trigger them, enabling more precise debugging of
memory usage patterns. The feature is particularly useful for
identifying objects that persist longer than expected or detecting
unexpected object creation patterns. By tracking both memory allocations
and object references, memray provides a more complete picture of a
program's memory behavior, helping developers optimize memory usage and
resolve leaks more effectively.

This implementation leverages the new Python reference tracing API that
was introduced in CPython 3.13.

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

Author: pablogsal

news/752.feature.rst

@@ -0,0 +1 @@
+Add object tracking capabilities to memray

src/memray/_memray.pyi

@@ -65,6 +65,35 @@ class AllocationRecord:
     def __lt__(self, other: Any) -> Any: ...
     def __ne__(self, other: Any) -> Any: ...
 
+class TrackedObjectRecord:
+    @property
+    def tid(self) -> int: ...
+    @property
+    def address(self) -> int: ...
+    @property
+    def is_created(self) -> bool: ...
+    @property
+    def frame_index(self) -> int: ...
+    @property
+    def native_frame_id(self) -> int: ...
+    @property
+    def native_segment_generation(self) -> int: ...
+    def toPythonObject(self) -> Any: ...
+    def hybrid_stack_trace(
+        self,
+        max_stacks: Optional[int] = None,
+    ) -> List[Union[PythonStackElement, NativeStackElement]]: ...
+    def native_stack_trace(
+        self, max_stacks: Optional[int] = None
+    ) -> List[NativeStackElement]: ...
+    def stack_trace(
+        self, max_stacks: Optional[int] = None
+    ) -> List[PythonStackElement]: ...
+    def __eq__(self, other: Any) -> Any: ...
+    def __hash__(self) -> Any: ...
+    def __repr__(self) -> str: ...
+    def __str__(self) -> str: ...
+
 class Interval:
     def __init__(
         self,
@@ -169,6 +198,9 @@ class FileReader:
     @property
     def closed(self) -> bool: ...
     def close(self) -> None: ...
+    def get_tracked_objects(
+        self, filter_objs: Optional[Iterable[Any]] = None
+    ) -> Iterable[TrackedObjectRecord]: ...
 
 def compute_statistics(
     file_name: Union[str, Path],
@@ -212,6 +244,7 @@ class Tracker:
         follow_fork: bool = ...,
         trace_python_allocators: bool = ...,
         file_format: FileFormat = ...,
+        reference_tracking: bool = ...,
     ) -> None: ...
     @overload
     def __init__(
@@ -223,6 +256,7 @@ class Tracker:
         follow_fork: bool = ...,
         trace_python_allocators: bool = ...,
         file_format: FileFormat = ...,
+        reference_tracking: bool = ...,
     ) -> None: ...
     def __enter__(self) -> Any: ...
     def __exit__(

src/memray/_memray.pyx

@@ -32,6 +32,7 @@ from _memray.records cimport Allocation as _Allocation
 from _memray.records cimport FileFormat as _FileFormat
 from _memray.records cimport MemoryRecord
 from _memray.records cimport MemorySnapshot as _MemorySnapshot
+from _memray.records cimport TrackedObject
 from _memray.sink cimport FileSink
 from _memray.sink cimport NullSink
 from _memray.sink cimport Sink
@@ -56,7 +57,9 @@ from _memray.source cimport SocketSource
 from _memray.tracking_api cimport Tracker as NativeTracker
 from _memray.tracking_api cimport install_trace_function
 from _memray.tracking_api cimport set_up_pthread_fork_handlers
+from cpython cimport Py_DECREF
 from cpython cimport PyErr_CheckSignals
+from cpython cimport PyObject
 from libc.math cimport ceil
 from libc.stdint cimport uint64_t
 from libcpp cimport bool
@@ -67,6 +70,7 @@ from libcpp.memory cimport shared_ptr
 from libcpp.memory cimport unique_ptr
 from libcpp.string cimport string as cppstring
 from libcpp.unordered_map cimport unordered_map
+from libcpp.unordered_set cimport unordered_set
 from libcpp.utility cimport move
 from libcpp.vector cimport vector
 
@@ -362,6 +366,100 @@ cdef class AllocationRecord:
                 f"allocations={self.n_allocations}>")
 
 
+@cython.freelist(1024)
+cdef class TrackedObjectRecord:
+    cdef object _tuple
+    cdef dict _stack_trace_cache
+    cdef shared_ptr[RecordReader] _reader
+
+    def __init__(self, record):
+        self._tuple = record
+        self._stack_trace_cache = {}
+
+    def __eq__(self, other):
+        cdef TrackedObjectRecord _other
+        if isinstance(other, TrackedObjectRecord):
+            _other = other
+            return self._tuple == _other._tuple
+        return NotImplemented
+
+    def __hash__(self):
+        return hash(self._tuple)
+
+    @property
+    def tid(self):
+        return self._tuple[0]
+
+    @property
+    def address(self):
+        return self._tuple[1]
+
+    @property
+    def is_created(self):
+        return self._tuple[2]
+
+    @property
+    def stack_id(self):
+        return self._tuple[3]
+
+    @property
+    def native_stack_id(self):
+        return self._tuple[4]
+
+    @property
+    def native_segment_generation(self):
+        return self._tuple[5]
+
+    @property
+    def thread_name(self):
+        if self.tid == -1:
+            return "merged thread"
+        assert self._reader.get() != NULL, "Cannot get thread name without reader."
+        return self._reader.get().getThreadName(self.tid)
+
+    def stack_trace(self, max_stacks=None):
+        cache_key = ("python", max_stacks)
+        if cache_key not in self._stack_trace_cache:
+            self._stack_trace_cache[cache_key] = stack_trace(
+                self._reader.get(),
+                self.tid,
+                None,  # No allocator for objects
+                self.stack_id,
+                max_stacks,
+            )
+        return self._stack_trace_cache[cache_key]
+
+    def native_stack_trace(self, max_stacks=None):
+        cache_key = ("native", max_stacks)
+        if cache_key not in self._stack_trace_cache:
+            self._stack_trace_cache[cache_key] = native_stack_trace(
+                self._reader.get(),
+                None,  # No allocator for objects
+                self.native_stack_id,
+                self.native_segment_generation,
+                max_stacks,
+            )
+        return self._stack_trace_cache[cache_key]
+
+    def hybrid_stack_trace(self, max_stacks=None):
+        cache_key = ("hybrid", max_stacks)
+        if cache_key not in self._stack_trace_cache:
+            self._stack_trace_cache[cache_key] = hybrid_stack_trace(
+                self._reader.get(),
+                self.tid,
+                None,  # No allocator for objects
+                self.stack_id,
+                self.native_stack_id,
+                self.native_segment_generation,
+                max_stacks,
+            )
+        return self._stack_trace_cache[cache_key]
+
+    def __repr__(self):
+        return (f"TrackedObjectRecord<tid={hex(self.tid)}, address={hex(self.address)}, "
+                f"is_created={self.is_created}, stack_id={self.stack_id}>")
+
+
 @cython.freelist(1024)
 cdef class Interval:
     cdef public size_t allocated_before_snapshot
@@ -628,12 +726,14 @@ cdef class Tracker:
             of supported file formats and their limitations.
     """
     cdef bool _native_traces
+    cdef bool _reference_tracking
     cdef unsigned int _memory_interval_ms
     cdef bool _follow_fork
     cdef bool _trace_python_allocators
     cdef object _previous_profile_func
     cdef object _previous_thread_profile_func
     cdef unique_ptr[RecordWriter] _writer
+    cdef object _surviving_objects
 
     cdef unique_ptr[Sink] _make_writer(self, destination) except*:
         # Creating a Sink can raise Python exceptions (if is interrupted by signal
@@ -658,12 +758,20 @@ 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,
-                  FileFormat file_format=FileFormat.ALL_ALLOCATIONS):
+                  reference_tracking=False, FileFormat file_format=FileFormat.ALL_ALLOCATIONS):
         if (file_name, destination).count(None) != 1:
             raise TypeError("Exactly one of 'file_name' or 'destination' argument must be specified")
 
+        # Check Python version if reference tracking is enabled
+        if reference_tracking and sys.version_info < (3, 13):
+            raise RuntimeError(
+                "Python object reference tracking requires Python 3.13 or later. "
+                f"Current version: {sys.version_info.major}.{sys.version_info.minor}"
+            )
+
         cdef cppstring command_line = " ".join(sys.argv)
         self._native_traces = native_traces
+        self._reference_tracking = reference_tracking
         self._memory_interval_ms = memory_interval_ms
         self._follow_fork = follow_fork
         self._trace_python_allocators = trace_python_allocators
@@ -714,17 +822,21 @@ cdef class Tracker:
             if "greenlet" in sys.modules:
                 NativeTracker.beginTrackingGreenlets()
 
+            self._surviving_objects = []
+
             NativeTracker.createTracker(
                 move(writer),
                 self._native_traces,
                 self._memory_interval_ms,
                 self._follow_fork,
                 self._trace_python_allocators,
+                self._reference_tracking,
             )
             return self
 
     @cython.profile(False)
     def __exit__(self, exc_type, exc_value, exc_traceback):
+        self._populate_suriving_objects()
         with tracker_creation_lock:
             NativeTracker.destroyTracker()
             sys.setprofile(self._previous_profile_func)
@@ -733,6 +845,27 @@ cdef class Tracker:
             for attr in ("_name", "_ident"):
                 delattr(threading.Thread, attr)
 
+    cdef void _populate_suriving_objects(self):
+        assert NativeTracker.getTracker() != NULL
+        cdef unordered_set[PyObject*] objects = NativeTracker.getTracker().getSurvivingObjects()
+        for obj in objects:
+            pass
+            self._surviving_objects.append(<object>obj)
+            Py_DECREF(<object>obj)
+
+    def get_surviving_objects(self):
+        """Get a list of objects that were alive at the end of the tracking period.
+
+        Returns:
+            list: A list of objects that were alive at the end of the tracking period.
+        """
+        if sys.version_info < (3, 13):
+            raise RuntimeError(
+                "Python object reference tracking requires Python 3.13 or later. "
+                f"Current version: {sys.version_info.major}.{sys.version_info.minor}"
+            )
+        return self._surviving_objects
+
 
 def start_thread_trace(frame, event, arg):
     if event in {"call", "c_call"}:
@@ -946,6 +1079,8 @@ cdef class FileReader:
                     )
                 elif ret == RecordResult.RecordResultMemorySnapshot:
                     self._memory_snapshots.push_back(reader.getLatestMemorySnapshot())
+                elif ret == RecordResult.RecordResultObjectRecord:
+                    pass
                 else:
                     break
 
@@ -1239,6 +1374,58 @@ cdef class FileReader:
         for record in self._memory_snapshots:
             yield MemorySnapshot(record.ms_since_epoch, record.rss, record.heap)
 
+    def get_tracked_objects(self, filter_objs=None):
+        """Return all tracked objects from the file.
+
+        This method yields TrackedObjectRecord instances representing Python objects
+        that were tracked during the recording session.
+
+        Args:
+            filter_objs (list, optional): A list of objects to filter the results by.
+
+        Returns:
+            An iterator of TrackedObjectRecord instances.
+        """
+        self._ensure_not_closed()
+
+        # Check Python version
+        if sys.version_info < (3, 13):
+            raise RuntimeError(
+                "Python object reference tracking requires Python 3.13 or later. "
+                f"Current version: {sys.version_info.major}.{sys.version_info.minor}"
+            )
+
+        object_ids = set()
+        if filter_objs is not None:
+            object_ids = {id(obj) for obj in filter_objs}
+
+        cdef shared_ptr[RecordReader] reader_sp = make_shared[RecordReader](
+            unique_ptr[FileSource](new FileSource(self._path))
+        )
+        cdef RecordReader* reader = reader_sp.get()
+
+        cdef TrackedObject tracked_object
+        while True:
+            PyErr_CheckSignals()
+            ret = reader.nextRecord()
+            if ret == RecordResult.RecordResultObjectRecord:
+                tracked_object = reader.getLatestObject()
+                if filter_objs is not None and tracked_object.address not in object_ids:
+                    continue
+                object_record = TrackedObjectRecord(tracked_object.toPythonObject())
+                object_record._reader = reader_sp
+                yield object_record
+            elif ret == RecordResult.RecordResultAllocationRecord:
+                pass
+            elif ret == RecordResult.RecordResultMemoryRecord:
+                pass
+            elif ret == RecordResult.RecordResultMemorySnapshot:
+                pass
+            else:
+                break
+
+        reader.close()
+
     @property
     def metadata(self):
         return _create_metadata(self._header, self._high_watermark.peak_memory)
@@ -1283,6 +1470,8 @@ def compute_statistics(
                 pass
             elif ret == RecordResult.RecordResultMemorySnapshot:
                 pass
+            elif ret == RecordResult.RecordResultObjectRecord:
+                pass
             else:
                 assert ret != RecordResult.RecordResultMemorySnapshot
                 assert ret != RecordResult.RecordResultAggregatedAllocationRecord

src/memray/_memray/compat.h

@@ -90,4 +90,16 @@ threadStateGetInterpreter(PyThreadState* tstate)
 void
 setprofileAllThreads(Py_tracefunc func, PyObject* arg);
 
+typedef int (*refTracer)(PyObject*, int event, void* data);
+
+inline int
+refTracerSetTracer(refTracer tracer, void* data)
+{
+#if PY_VERSION_HEX >= 0x030D0000
+    return PyRefTracer_SetTracer(reinterpret_cast<PyRefTracer>(tracer), data);
+#else
+    return 0;
+#endif
+}
+
 }  // namespace memray::compat

src/memray/_memray/hooks.cpp

@@ -486,4 +486,11 @@ PyGILState_Ensure() noexcept
     return ret;
 }
 
+int
+pyreftracer(PyObject* obj, int event, void* data) noexcept
+{
+    tracking_api::Tracker::trackObject(obj, event);
+    return 0;
+}
+
 }  // namespace memray::intercept