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

1 files changed, 77 insertions, 27 deletions

diff --git a/Lib/test/test_tracemalloc.py b/Lib/test/test_tracemalloc.py
index f65e361..7b92b87 100644
--- a/Lib/test/test_tracemalloc.py
+++ b/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)