tracemalloc now supports domains

Issue #26588:

* The _tracemalloc now supports tracing memory allocations of multiple address
  spaces (domains).
* Add domain parameter to tracemalloc_add_trace() and
  tracemalloc_remove_trace().
* tracemalloc_add_trace() now starts by removing the previous trace, if any.
* _tracemalloc._get_traces() now returns a list of (domain, size,
  traceback_frames): the domain is new.
* Add tracemalloc.DomainFilter
* tracemalloc.Filter: add an optional domain parameter to the constructor and a
  domain attribute
* Sublte change: use Py_uintptr_t rather than void* in the traces key.
* Add tracemalloc_config.use_domain, currently hardcoded to 1

Doc/library/tracemalloc.rst

@@ -355,10 +355,32 @@ Functions
    See also the :func:`get_object_traceback` function.
 
 
+DomainFilter
+^^^^^^^^^^^^
+
+.. class:: DomainFilter(inclusive: bool, domain: int)
+
+   Filter traces of memory blocks by their address space (domain).
+
+   .. versionadded:: 3.6
+
+   .. attribute:: inclusive
+
+      If *inclusive* is ``True`` (include), match memory blocks allocated
+      in the address space :attr:`domain`.
+
+      If *inclusive* is ``False`` (exclude), match memory blocks not allocated
+      in the address space :attr:`domain`.
+
+   .. attribute:: domain
+
+      Address space of a memory block (``int``). Read-only property.
+
+
 Filter
 ^^^^^^
 
-.. class:: Filter(inclusive: bool, filename_pattern: str, lineno: int=None, all_frames: bool=False)
+.. class:: Filter(inclusive: bool, filename_pattern: str, lineno: int=None, all_frames: bool=False, domain: int=None)
 
    Filter on traces of memory blocks.
 
@@ -378,9 +400,17 @@ Filter
    .. versionchanged:: 3.5
       The ``'.pyo'`` file extension is no longer replaced with ``'.py'``.
 
+   .. versionchanged:: 3.6
+      Added the :attr:`domain` attribute.
+
+
+   .. attribute:: domain
+
+      Address space of a memory block (``int`` or ``None``).
+
    .. attribute:: inclusive
 
-      If *inclusive* is ``True`` (include), only trace memory blocks allocated
+      If *inclusive* is ``True`` (include), only match memory blocks allocated
       in a file with a name matching :attr:`filename_pattern` at line number
       :attr:`lineno`.
 
@@ -395,7 +425,7 @@ Filter
 
    .. attribute:: filename_pattern
 
-      Filename pattern of the filter (``str``).
+      Filename pattern of the filter (``str``). Read-only property.
 
    .. attribute:: all_frames
 
@@ -458,14 +488,17 @@ Snapshot
    .. method:: filter_traces(filters)
 
       Create a new :class:`Snapshot` instance with a filtered :attr:`traces`
-      sequence, *filters* is a list of :class:`Filter` instances.  If *filters*
-      is an empty list, return a new :class:`Snapshot` instance with a copy of
-      the traces.
+      sequence, *filters* is a list of :class:`DomainFilter` and
+      :class:`Filter` instances.  If *filters* is an empty list, return a new
+      :class:`Snapshot` instance with a copy of the traces.
 
       All inclusive filters are applied at once, a trace is ignored if no
       inclusive filters match it. A trace is ignored if at least one exclusive
       filter matches it.
 
+      .. versionchanged:: 3.6
+         :class:`DomainFilter` instances are now also accepted in *filters*.
+
 
    .. classmethod:: load(filename)
 

Lib/test/test_tracemalloc.py

@@ -37,28 +37,31 @@ def allocate_bytes(size):
 def create_snapshots():
     traceback_limit = 2
 
+    # _tracemalloc._get_traces() returns a list of (domain, size,
+    # traceback_frames) tuples. traceback_frames is a tuple of (filename,
+    # line_number) tuples.
     raw_traces = [
-        (10, (('a.py', 2), ('b.py', 4))),
-        (10, (('a.py', 2), ('b.py', 4))),
-        (10, (('a.py', 2), ('b.py', 4))),
+        (0, 10, (('a.py', 2), ('b.py', 4))),
+        (0, 10, (('a.py', 2), ('b.py', 4))),
+        (0, 10, (('a.py', 2), ('b.py', 4))),
 
-        (2, (('a.py', 5), ('b.py', 4))),
+        (1, 2, (('a.py', 5), ('b.py', 4))),
 
-        (66, (('b.py', 1),)),
+        (2, 66, (('b.py', 1),)),
 
-        (7, (('<unknown>', 0),)),
+        (3, 7, (('<unknown>', 0),)),
     ]
     snapshot = tracemalloc.Snapshot(raw_traces, traceback_limit)
 
     raw_traces2 = [
-        (10, (('a.py', 2), ('b.py', 4))),
-        (10, (('a.py', 2), ('b.py', 4))),
-        (10, (('a.py', 2), ('b.py', 4))),
+        (0, 10, (('a.py', 2), ('b.py', 4))),
+        (0, 10, (('a.py', 2), ('b.py', 4))),
+        (0, 10, (('a.py', 2), ('b.py', 4))),
 
-        (2, (('a.py', 5), ('b.py', 4))),
-        (5000, (('a.py', 5), ('b.py', 4))),
+        (2, 2, (('a.py', 5), ('b.py', 4))),
+        (2, 5000, (('a.py', 5), ('b.py', 4))),
 
-        (400, (('c.py', 578),)),
+        (4, 400, (('c.py', 578),)),
     ]
     snapshot2 = tracemalloc.Snapshot(raw_traces2, traceback_limit)
 
@@ -126,7 +129,7 @@ class TestTracemallocEnabled(unittest.TestCase):
 
     def find_trace(self, traces, traceback):
         for trace in traces:
-            if trace[1] == traceback._frames:
+            if trace[2] == traceback._frames:
                 return trace
 
         self.fail("trace not found")
@@ -140,7 +143,7 @@ class TestTracemallocEnabled(unittest.TestCase):
         trace = self.find_trace(traces, obj_traceback)
 
         self.assertIsInstance(trace, tuple)
-        size, traceback = trace
+        domain, size, traceback = trace
         self.assertEqual(size, obj_size)
         self.assertEqual(traceback, obj_traceback._frames)
 
@@ -167,9 +170,8 @@ class TestTracemallocEnabled(unittest.TestCase):
 
         trace1 = self.find_trace(traces, obj1_traceback)
         trace2 = self.find_trace(traces, obj2_traceback)
-        size1, traceback1 = trace1
-        size2, traceback2 = trace2
-        self.assertEqual(traceback2, traceback1)
+        domain1, size1, traceback1 = trace1
+        domain2, size2, traceback2 = trace2
         self.assertIs(traceback2, traceback1)
 
     def test_get_traced_memory(self):
@@ -292,7 +294,7 @@ class TestSnapshot(unittest.TestCase):
     maxDiff = 4000
 
     def test_create_snapshot(self):
-        raw_traces = [(5, (('a.py', 2),))]
+        raw_traces = [(0, 5, (('a.py', 2),))]
 
         with contextlib.ExitStack() as stack:
             stack.enter_context(patch.object(tracemalloc, 'is_tracing',
@@ -322,11 +324,11 @@ class TestSnapshot(unittest.TestCase):
         # exclude b.py
         snapshot3 = snapshot.filter_traces((filter1,))
         self.assertEqual(snapshot3.traces._traces, [
-            (10, (('a.py', 2), ('b.py', 4))),
-            (10, (('a.py', 2), ('b.py', 4))),
-            (10, (('a.py', 2), ('b.py', 4))),
-            (2, (('a.py', 5), ('b.py', 4))),
-            (7, (('<unknown>', 0),)),
+            (0, 10, (('a.py', 2), ('b.py', 4))),
+            (0, 10, (('a.py', 2), ('b.py', 4))),
+            (0, 10, (('a.py', 2), ('b.py', 4))),
+            (1, 2, (('a.py', 5), ('b.py', 4))),
+            (3, 7, (('<unknown>', 0),)),
         ])
 
         # filter_traces() must not touch the original snapshot
@@ -335,10 +337,10 @@ class TestSnapshot(unittest.TestCase):
         # only include two lines of a.py
         snapshot4 = snapshot3.filter_traces((filter2, filter3))
         self.assertEqual(snapshot4.traces._traces, [
-            (10, (('a.py', 2), ('b.py', 4))),
-            (10, (('a.py', 2), ('b.py', 4))),
-            (10, (('a.py', 2), ('b.py', 4))),
-            (2, (('a.py', 5), ('b.py', 4))),
+            (0, 10, (('a.py', 2), ('b.py', 4))),
+            (0, 10, (('a.py', 2), ('b.py', 4))),
+            (0, 10, (('a.py', 2), ('b.py', 4))),
+            (1, 2, (('a.py', 5), ('b.py', 4))),
         ])
 
         # No filter: just duplicate the snapshot
@@ -349,6 +351,54 @@ class TestSnapshot(unittest.TestCase):
 
         self.assertRaises(TypeError, snapshot.filter_traces, filter1)
 
+    def test_filter_traces_domain(self):
+        snapshot, snapshot2 = create_snapshots()
+        filter1 = tracemalloc.Filter(False, "a.py", domain=1)
+        filter2 = tracemalloc.Filter(True, "a.py", domain=1)
+
+        original_traces = list(snapshot.traces._traces)
+
+        # exclude a.py of domain 1
+        snapshot3 = snapshot.filter_traces((filter1,))
+        self.assertEqual(snapshot3.traces._traces, [
+            (0, 10, (('a.py', 2), ('b.py', 4))),
+            (0, 10, (('a.py', 2), ('b.py', 4))),
+            (0, 10, (('a.py', 2), ('b.py', 4))),
+            (2, 66, (('b.py', 1),)),
+            (3, 7, (('<unknown>', 0),)),
+        ])
+
+        # include domain 1
+        snapshot3 = snapshot.filter_traces((filter1,))
+        self.assertEqual(snapshot3.traces._traces, [
+            (0, 10, (('a.py', 2), ('b.py', 4))),
+            (0, 10, (('a.py', 2), ('b.py', 4))),
+            (0, 10, (('a.py', 2), ('b.py', 4))),
+            (2, 66, (('b.py', 1),)),
+            (3, 7, (('<unknown>', 0),)),
+        ])
+
+    def test_filter_traces_domain_filter(self):
+        snapshot, snapshot2 = create_snapshots()
+        filter1 = tracemalloc.DomainFilter(False, domain=3)
+        filter2 = tracemalloc.DomainFilter(True, domain=3)
+
+        # exclude domain 2
+        snapshot3 = snapshot.filter_traces((filter1,))
+        self.assertEqual(snapshot3.traces._traces, [
+            (0, 10, (('a.py', 2), ('b.py', 4))),
+            (0, 10, (('a.py', 2), ('b.py', 4))),
+            (0, 10, (('a.py', 2), ('b.py', 4))),
+            (1, 2, (('a.py', 5), ('b.py', 4))),
+            (2, 66, (('b.py', 1),)),
+        ])
+
+        # include domain 2
+        snapshot3 = snapshot.filter_traces((filter2,))
+        self.assertEqual(snapshot3.traces._traces, [
+            (3, 7, (('<unknown>', 0),)),
+        ])
+
     def test_snapshot_group_by_line(self):
         snapshot, snapshot2 = create_snapshots()
         tb_0 = traceback_lineno('<unknown>', 0)

Lib/tracemalloc.py

@@ -244,17 +244,21 @@ class Trace:
     __slots__ = ("_trace",)
 
     def __init__(self, trace):
-        # trace is a tuple: (size, traceback), see Traceback constructor
-        # for the format of the traceback tuple
+        # trace is a tuple: (domain: int, size: int, traceback: tuple).
+        # See Traceback constructor for the format of the traceback tuple.
         self._trace = trace
 
     @property
-    def size(self):
+    def domain(self):
         return self._trace[0]
 
+    @property
+    def size(self):
+        return self._trace[1]
+
     @property
     def traceback(self):
-        return Traceback(self._trace[1])
+        return Traceback(self._trace[2])
 
     def __eq__(self, other):
         return (self._trace == other._trace)
@@ -266,8 +270,8 @@ class Trace:
         return "%s: %s" % (self.traceback, _format_size(self.size, False))
 
     def __repr__(self):
-        return ("<Trace size=%s, traceback=%r>"
-                % (_format_size(self.size, False), self.traceback))
+        return ("<Trace domain=%s size=%s, traceback=%r>"
+                % (self.domain, _format_size(self.size, False), self.traceback))
 
 
 class _Traces(Sequence):
@@ -302,19 +306,29 @@ def _normalize_filename(filename):
     return filename
 
 
-class Filter:
+class BaseFilter:
+    def __init__(self, inclusive):
+        self.inclusive = inclusive
+
+    def _match(self, trace):
+        raise NotImplementedError
+
+
+class Filter(BaseFilter):
     def __init__(self, inclusive, filename_pattern,
-                 lineno=None, all_frames=False):
+                 lineno=None, all_frames=False, domain=None):
+        super().__init__(inclusive)
         self.inclusive = inclusive
         self._filename_pattern = _normalize_filename(filename_pattern)
         self.lineno = lineno
         self.all_frames = all_frames
+        self.domain = domain
 
     @property
     def filename_pattern(self):
         return self._filename_pattern
 
-    def __match_frame(self, filename, lineno):
+    def _match_frame_impl(self, filename, lineno):
         filename = _normalize_filename(filename)
         if not fnmatch.fnmatch(filename, self._filename_pattern):
             return False
@@ -324,11 +338,11 @@ class Filter:
             return (lineno == self.lineno)
 
     def _match_frame(self, filename, lineno):
-        return self.__match_frame(filename, lineno) ^ (not self.inclusive)
+        return self._match_frame_impl(filename, lineno) ^ (not self.inclusive)
 
     def _match_traceback(self, traceback):
         if self.all_frames:
-            if any(self.__match_frame(filename, lineno)
+            if any(self._match_frame_impl(filename, lineno)
                    for filename, lineno in traceback):
                 return self.inclusive
             else:
@@ -337,6 +351,30 @@ class Filter:
             filename, lineno = traceback[0]
             return self._match_frame(filename, lineno)
 
+    def _match(self, trace):
+        domain, size, traceback = trace
+        res = self._match_traceback(traceback)
+        if self.domain is not None:
+            if self.inclusive:
+                return res and (domain == self.domain)
+            else:
+                return res or (domain != self.domain)
+        return res
+
+
+class DomainFilter(BaseFilter):
+    def __init__(self, inclusive, domain):
+        super().__init__(inclusive)
+        self._domain = domain
+
+    @property
+    def domain(self):
+        return self._domain
+
+    def _match(self, trace):
+        domain, size, traceback = trace
+        return (domain == self.domain) ^ (not self.inclusive)
+
 
 class Snapshot:
     """
@@ -365,13 +403,12 @@ class Snapshot:
             return pickle.load(fp)
 
     def _filter_trace(self, include_filters, exclude_filters, trace):
-        traceback = trace[1]
         if include_filters:
-            if not any(trace_filter._match_traceback(traceback)
+            if not any(trace_filter._match(trace)
                        for trace_filter in include_filters):
                 return False
         if exclude_filters:
-            if any(not trace_filter._match_traceback(traceback)
+            if any(not trace_filter._match(trace)
                    for trace_filter in exclude_filters):
                 return False
         return True
@@ -379,8 +416,8 @@ class Snapshot:
     def filter_traces(self, filters):
         """
         Create a new Snapshot instance with a filtered traces sequence, filters
-        is a list of Filter instances.  If filters is an empty list, return a
-        new Snapshot instance with a copy of the traces.
+        is a list of Filter or DomainFilter instances.  If filters is an empty
+        list, return a new Snapshot instance with a copy of the traces.
         """
         if not isinstance(filters, Iterable):
             raise TypeError("filters must be a list of filters, not %s"
@@ -412,7 +449,7 @@ class Snapshot:
         tracebacks = {}
         if not cumulative:
             for trace in self.traces._traces:
-                size, trace_traceback = trace
+                domain, size, trace_traceback = trace
                 try:
                     traceback = tracebacks[trace_traceback]
                 except KeyError:
@@ -433,7 +470,7 @@ class Snapshot:
         else:
             # cumulative statistics
             for trace in self.traces._traces:
-                size, trace_traceback = trace
+                domain, size, trace_traceback = trace
                 for frame in trace_traceback:
                     try:
                         traceback = tracebacks[frame]

Misc/NEWS

@@ -232,6 +232,9 @@ Core and Builtins
 Library
 -------
 
+- Issue #26588: The _tracemalloc now supports tracing memory allocations of
+  multiple address spaces (domains).
+
 - Issue #24266: Ctrl+C during Readline history search now cancels the search
   mode when compiled with Readline 7.