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*
+      sequence, *filters* is a list of :class:`DomainFilter` and
-      is an empty list, return a new :class:`Snapshot` instance with a copy of
+      :class:`Filter` instances.  If *filters* is an empty list, return a new
-      the traces.
+      :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))),
+        (0, 10, (('a.py', 2), ('b.py', 4))),
-        (10, (('a.py', 2), ('b.py', 4))),
+        (0, 10, (('a.py', 2), ('b.py', 4))),
-        (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))),
+        (0, 10, (('a.py', 2), ('b.py', 4))),
-        (10, (('a.py', 2), ('b.py', 4))),
+        (0, 10, (('a.py', 2), ('b.py', 4))),
-        (10, (('a.py', 2), ('b.py', 4))),
+        (0, 10, (('a.py', 2), ('b.py', 4))),
 
-        (2, (('a.py', 5), ('b.py', 4))),
+        (2, 2, (('a.py', 5), ('b.py', 4))),
-        (5000, (('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
+        domain1, size1, traceback1 = trace1
-        size2, traceback2 = trace2
+        domain2, size2, traceback2 = trace2
-        self.assertEqual(traceback2, traceback1)
         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))),
+            (0, 10, (('a.py', 2), ('b.py', 4))),
-            (10, (('a.py', 2), ('b.py', 4))),
+            (0, 10, (('a.py', 2), ('b.py', 4))),
-            (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))),
-            (7, (('<unknown>', 0),)),
+            (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))),
+            (0, 10, (('a.py', 2), ('b.py', 4))),
-            (10, (('a.py', 2), ('b.py', 4))),
+            (0, 10, (('a.py', 2), ('b.py', 4))),
-            (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))),
         ])
 
         # 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
+        # trace is a tuple: (domain: int, size: int, traceback: tuple).
-        # for the format of the 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>"
+        return ("<Trace domain=%s size=%s, traceback=%r>"
-                % (_format_size(self.size, False), self.traceback))
+                % (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
+        is a list of Filter or DomainFilter instances.  If filters is an empty
-        new Snapshot instance with a copy of the traces.
+        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.
 

Modules/_tracemalloc.c

@@ -39,7 +39,11 @@ static struct {
     /* limit of the number of frames in a traceback, 1 by default.
        Variable protected by the GIL. */
     int max_nframe;
-} tracemalloc_config = {TRACEMALLOC_NOT_INITIALIZED, 0, 1};
+
+    /* use domain in trace key?
+       Variable protected by the GIL. */
+    int use_domain;
+} tracemalloc_config = {TRACEMALLOC_NOT_INITIALIZED, 0, 1, 1};
 
 #if defined(TRACE_RAW_MALLOC) && defined(WITH_THREAD)
 /* This lock is needed because tracemalloc_free() is called without
@@ -54,10 +58,23 @@ static PyThread_type_lock tables_lock;
 #  define TABLES_UNLOCK()
 #endif
 
+
+#define DEFAULT_DOMAIN 0
+
+typedef unsigned int domain_t;
+
+/* Pack the frame_t structure to reduce the memory footprint. */
+typedef struct
+#ifdef __GNUC__
+__attribute__((packed))
+#endif
+{
+    Py_uintptr_t ptr;
+    domain_t domain;
+} pointer_t;
+
 /* Pack the frame_t structure to reduce the memory footprint on 64-bit
-   architectures: 12 bytes instead of 16. This optimization might produce
+   architectures: 12 bytes instead of 16. */
-   SIGBUS on architectures not supporting unaligned memory accesses (64-bit
-   MIPS CPU?): on such architecture, the structure must not be packed. */
 typedef struct
 #ifdef __GNUC__
 __attribute__((packed))
@@ -71,6 +88,7 @@ _declspec(align(4))
     unsigned int lineno;
 } frame_t;
 
+
 typedef struct {
     Py_uhash_t hash;
     int nframe;
@@ -83,6 +101,7 @@ typedef struct {
 #define MAX_NFRAME \
         ((INT_MAX - (int)sizeof(traceback_t)) / (int)sizeof(frame_t) + 1)
 
+
 static PyObject *unknown_filename = NULL;
 static traceback_t tracemalloc_empty_traceback;
 
@@ -95,6 +114,7 @@ typedef struct {
     traceback_t *traceback;
 } trace_t;
 
+
 /* Size in bytes of currently traced memory.
    Protected by TABLES_LOCK(). */
 static size_t tracemalloc_traced_memory = 0;
@@ -121,6 +141,7 @@ static _Py_hashtable_t *tracemalloc_tracebacks = NULL;
    Protected by TABLES_LOCK(). */
 static _Py_hashtable_t *tracemalloc_traces = NULL;
 
+
 #ifdef TRACE_DEBUG
 static void
 tracemalloc_error(const char *format, ...)
@@ -135,6 +156,7 @@ tracemalloc_error(const char *format, ...)
 }
 #endif
 
+
 #if defined(WITH_THREAD) && defined(TRACE_RAW_MALLOC)
 #define REENTRANT_THREADLOCAL
 
@@ -196,6 +218,7 @@ set_reentrant(int reentrant)
 }
 #endif
 
+
 static Py_uhash_t
 hashtable_hash_pyobject(size_t key_size, const void *pkey)
 {
@@ -205,21 +228,53 @@ hashtable_hash_pyobject(size_t key_size, const void *pkey)
     return PyObject_Hash(obj);
 }
 
+
 static int
 hashtable_compare_unicode(size_t key_size, const void *pkey,
                           const _Py_hashtable_entry_t *entry)
 {
-    PyObject *key, *entry_key;
+    PyObject *key1, *key2;
 
-    _Py_HASHTABLE_READ_KEY(key_size, pkey, key);
+    _Py_HASHTABLE_READ_KEY(key_size, pkey, key1);
-    _Py_HASHTABLE_ENTRY_READ_KEY(key_size, entry, entry_key);
+    _Py_HASHTABLE_ENTRY_READ_KEY(key_size, entry, key2);
 
-    if (key != NULL && entry_key != NULL)
+    if (key1 != NULL && key2 != NULL)
-        return (PyUnicode_Compare(key, entry_key) == 0);
+        return (PyUnicode_Compare(key1, key2) == 0);
     else
-        return key == entry_key;
+        return key1 == key2;
 }
 
+
+static Py_uhash_t
+hashtable_hash_pointer_t(size_t key_size, const void *pkey)
+{
+    pointer_t ptr;
+    Py_uhash_t hash;
+
+    _Py_HASHTABLE_READ_KEY(key_size, pkey, ptr);
+
+    hash = (Py_uhash_t)_Py_HashPointer((void*)ptr.ptr);
+    hash ^= ptr.domain;
+    return hash;
+}
+
+
+int
+hashtable_compare_pointer_t(size_t key_size, const void *pkey,
+                            const _Py_hashtable_entry_t *entry)
+{
+    pointer_t ptr1, ptr2;
+
+    _Py_HASHTABLE_READ_KEY(key_size, pkey, ptr1);
+    _Py_HASHTABLE_ENTRY_READ_KEY(key_size, entry, ptr2);
+
+    /* compare pointer before domain, because pointer is more likely to be
+       different */
+    return (ptr1.ptr == ptr2.ptr && ptr1.domain == ptr2.domain);
+
+}
+
+
 static _Py_hashtable_t *
 hashtable_new(size_t key_size, size_t data_size,
               _Py_hashtable_hash_func hash_func,
@@ -231,6 +286,7 @@ hashtable_new(size_t key_size, size_t data_size,
                                   &hashtable_alloc);
 }
 
+
 static void*
 raw_malloc(size_t size)
 {
@@ -243,6 +299,7 @@ raw_free(void *ptr)
     allocators.raw.free(allocators.raw.ctx, ptr);
 }
 
+
 static Py_uhash_t
 hashtable_hash_traceback(size_t key_size, const void *pkey)
 {
@@ -252,6 +309,7 @@ hashtable_hash_traceback(size_t key_size, const void *pkey)
     return traceback->hash;
 }
 
+
 static int
 hashtable_compare_traceback(size_t key_size, const void *pkey,
                             const _Py_hashtable_entry_t *he)
@@ -281,6 +339,7 @@ hashtable_compare_traceback(size_t key_size, const void *pkey,
     return 1;
 }
 
+
 static void
 tracemalloc_get_frame(PyFrameObject *pyframe, frame_t *frame)
 {
@@ -353,6 +412,7 @@ tracemalloc_get_frame(PyFrameObject *pyframe, frame_t *frame)
     frame->filename = filename;
 }
 
+
 static Py_uhash_t
 traceback_hash(traceback_t *traceback)
 {
@@ -377,6 +437,7 @@ traceback_hash(traceback_t *traceback)
     return x;
 }
 
+
 static void
 traceback_get_frames(traceback_t *traceback)
 {
@@ -404,6 +465,7 @@ traceback_get_frames(traceback_t *traceback)
     }
 }
 
+
 static traceback_t *
 traceback_new(void)
 {
@@ -455,42 +517,73 @@ traceback_new(void)
     return traceback;
 }
 
+
+static void
+tracemalloc_remove_trace(domain_t domain, Py_uintptr_t ptr)
+{
+    trace_t trace;
+    int removed;
+
+    if (tracemalloc_config.use_domain) {
+        pointer_t key = {ptr, domain};
+        removed = _Py_HASHTABLE_POP(tracemalloc_traces, key, trace);
+    }
+    else {
+        removed = _Py_HASHTABLE_POP(tracemalloc_traces, ptr, trace);
+    }
+    if (!removed) {
+        return;
+    }
+
+    assert(tracemalloc_traced_memory >= trace.size);
+    tracemalloc_traced_memory -= trace.size;
+}
+
+#define REMOVE_TRACE(ptr) \
+            tracemalloc_remove_trace(DEFAULT_DOMAIN, (Py_uintptr_t)(ptr))
+
+
 static int
-tracemalloc_add_trace(void *ptr, size_t size)
+tracemalloc_add_trace(domain_t domain, Py_uintptr_t ptr, size_t size)
 {
     traceback_t *traceback;
     trace_t trace;
     int res;
 
+    /* first, remove the previous trace (if any) */
+    tracemalloc_remove_trace(domain, ptr);
+
     traceback = traceback_new();
-    if (traceback == NULL)
+    if (traceback == NULL) {
         return -1;
+    }
 
     trace.size = size;
     trace.traceback = traceback;
 
-    res = _Py_HASHTABLE_SET(tracemalloc_traces, ptr, trace);
+    if (tracemalloc_config.use_domain) {
-    if (res == 0) {
+        pointer_t key = {ptr, domain};
-        assert(tracemalloc_traced_memory <= PY_SIZE_MAX - size);
+        res = _Py_HASHTABLE_SET(tracemalloc_traces, key, trace);
-        tracemalloc_traced_memory += size;
+    }
-        if (tracemalloc_traced_memory > tracemalloc_peak_traced_memory)
+    else {
-            tracemalloc_peak_traced_memory = tracemalloc_traced_memory;
+        res = _Py_HASHTABLE_SET(tracemalloc_traces, ptr, trace);
     }
 
-    return res;
+    if (res != 0) {
-}
+        return res;
-
-static void
-tracemalloc_remove_trace(void *ptr)
-{
-    trace_t trace;
-
-    if (_Py_HASHTABLE_POP(tracemalloc_traces, ptr, trace)) {
-        assert(tracemalloc_traced_memory >= trace.size);
-        tracemalloc_traced_memory -= trace.size;
     }
+
+    assert(tracemalloc_traced_memory <= PY_SIZE_MAX - size);
+    tracemalloc_traced_memory += size;
+    if (tracemalloc_traced_memory > tracemalloc_peak_traced_memory)
+        tracemalloc_peak_traced_memory = tracemalloc_traced_memory;
+    return 0;
 }
 
+#define ADD_TRACE(ptr, size) \
+            tracemalloc_add_trace(DEFAULT_DOMAIN, (Py_uintptr_t)(ptr), size)
+
+
 static void*
 tracemalloc_alloc(int use_calloc, void *ctx, size_t nelem, size_t elsize)
 {
@@ -507,7 +600,7 @@ tracemalloc_alloc(int use_calloc, void *ctx, size_t nelem, size_t elsize)
         return NULL;
 
     TABLES_LOCK();
-    if (tracemalloc_add_trace(ptr, nelem * elsize) < 0) {
+    if (ADD_TRACE(ptr, nelem * elsize) < 0) {
         /* Failed to allocate a trace for the new memory block */
         TABLES_UNLOCK();
         alloc->free(alloc->ctx, ptr);
@@ -517,6 +610,7 @@ tracemalloc_alloc(int use_calloc, void *ctx, size_t nelem, size_t elsize)
     return ptr;
 }
 
+
 static void*
 tracemalloc_realloc(void *ctx, void *ptr, size_t new_size)
 {
@@ -531,9 +625,9 @@ tracemalloc_realloc(void *ctx, void *ptr, size_t new_size)
         /* an existing memory block has been resized */
 
         TABLES_LOCK();
-        tracemalloc_remove_trace(ptr);
+        REMOVE_TRACE(ptr);
 
-        if (tracemalloc_add_trace(ptr2, new_size) < 0) {
+        if (ADD_TRACE(ptr2, new_size) < 0) {
             /* Memory allocation failed. The error cannot be reported to
                the caller, because realloc() may already have shrinked the
                memory block and so removed bytes.
@@ -551,7 +645,7 @@ tracemalloc_realloc(void *ctx, void *ptr, size_t new_size)
         /* new allocation */
 
         TABLES_LOCK();
-        if (tracemalloc_add_trace(ptr2, new_size) < 0) {
+        if (ADD_TRACE(ptr2, new_size) < 0) {
             /* Failed to allocate a trace for the new memory block */
             TABLES_UNLOCK();
             alloc->free(alloc->ctx, ptr2);
@@ -562,6 +656,7 @@ tracemalloc_realloc(void *ctx, void *ptr, size_t new_size)
     return ptr2;
 }
 
+
 static void
 tracemalloc_free(void *ctx, void *ptr)
 {
@@ -576,10 +671,11 @@ tracemalloc_free(void *ctx, void *ptr)
     alloc->free(alloc->ctx, ptr);
 
     TABLES_LOCK();
-    tracemalloc_remove_trace(ptr);
+    REMOVE_TRACE(ptr);
     TABLES_UNLOCK();
 }
 
+
 static void*
 tracemalloc_alloc_gil(int use_calloc, void *ctx, size_t nelem, size_t elsize)
 {
@@ -604,18 +700,21 @@ tracemalloc_alloc_gil(int use_calloc, void *ctx, size_t nelem, size_t elsize)
     return ptr;
 }
 
+
 static void*
 tracemalloc_malloc_gil(void *ctx, size_t size)
 {
     return tracemalloc_alloc_gil(0, ctx, 1, size);
 }
 
+
 static void*
 tracemalloc_calloc_gil(void *ctx, size_t nelem, size_t elsize)
 {
     return tracemalloc_alloc_gil(1, ctx, nelem, elsize);
 }
 
+
 static void*
 tracemalloc_realloc_gil(void *ctx, void *ptr, size_t new_size)
 {
@@ -631,7 +730,7 @@ tracemalloc_realloc_gil(void *ctx, void *ptr, size_t new_size)
         ptr2 = alloc->realloc(alloc->ctx, ptr, new_size);
         if (ptr2 != NULL && ptr != NULL) {
             TABLES_LOCK();
-            tracemalloc_remove_trace(ptr);
+            REMOVE_TRACE(ptr);
             TABLES_UNLOCK();
         }
         return ptr2;
@@ -648,6 +747,7 @@ tracemalloc_realloc_gil(void *ctx, void *ptr, size_t new_size)
     return ptr2;
 }
 
+
 #ifdef TRACE_RAW_MALLOC
 static void*
 tracemalloc_raw_alloc(int use_calloc, void *ctx, size_t nelem, size_t elsize)
@@ -682,18 +782,21 @@ tracemalloc_raw_alloc(int use_calloc, void *ctx, size_t nelem, size_t elsize)
     return ptr;
 }
 
+
 static void*
 tracemalloc_raw_malloc(void *ctx, size_t size)
 {
     return tracemalloc_raw_alloc(0, ctx, 1, size);
 }
 
+
 static void*
 tracemalloc_raw_calloc(void *ctx, size_t nelem, size_t elsize)
 {
     return tracemalloc_raw_alloc(1, ctx, nelem, elsize);
 }
 
+
 static void*
 tracemalloc_raw_realloc(void *ctx, void *ptr, size_t new_size)
 {
@@ -710,7 +813,7 @@ tracemalloc_raw_realloc(void *ctx, void *ptr, size_t new_size)
 
         if (ptr2 != NULL && ptr != NULL) {
             TABLES_LOCK();
-            tracemalloc_remove_trace(ptr);
+            REMOVE_TRACE(ptr);
             TABLES_UNLOCK();
         }
         return ptr2;
@@ -734,6 +837,7 @@ tracemalloc_raw_realloc(void *ctx, void *ptr, size_t new_size)
 }
 #endif   /* TRACE_RAW_MALLOC */
 
+
 static int
 tracemalloc_clear_filename(_Py_hashtable_t *ht, _Py_hashtable_entry_t *entry,
                            void *user_data)
@@ -745,6 +849,7 @@ tracemalloc_clear_filename(_Py_hashtable_t *ht, _Py_hashtable_entry_t *entry,
     return 0;
 }
 
+
 static int
 traceback_free_traceback(_Py_hashtable_t *ht, _Py_hashtable_entry_t *entry,
                          void *user_data)
@@ -756,6 +861,7 @@ traceback_free_traceback(_Py_hashtable_t *ht, _Py_hashtable_entry_t *entry,
     return 0;
 }
 
+
 /* reentrant flag must be set to call this function and GIL must be held */
 static void
 tracemalloc_clear_traces(void)
@@ -782,6 +888,7 @@ tracemalloc_clear_traces(void)
     _Py_hashtable_clear(tracemalloc_filenames);
 }
 
+
 static int
 tracemalloc_init(void)
 {
@@ -826,9 +933,18 @@ tracemalloc_init(void)
                                            hashtable_hash_traceback,
                                            hashtable_compare_traceback);
 
-    tracemalloc_traces = hashtable_new(sizeof(void*), sizeof(trace_t),
+    if (tracemalloc_config.use_domain) {
-                                       _Py_hashtable_hash_ptr,
+        tracemalloc_traces = hashtable_new(sizeof(pointer_t),
-                                       _Py_hashtable_compare_direct);
+                                           sizeof(trace_t),
+                                           hashtable_hash_pointer_t,
+                                           hashtable_compare_pointer_t);
+    }
+    else {
+        tracemalloc_traces = hashtable_new(sizeof(Py_uintptr_t),
+                                           sizeof(trace_t),
+                                           _Py_hashtable_hash_ptr,
+                                           _Py_hashtable_compare_direct);
+    }
 
     if (tracemalloc_filenames == NULL || tracemalloc_tracebacks == NULL
        || tracemalloc_traces == NULL) {
@@ -856,6 +972,7 @@ tracemalloc_init(void)
     return 0;
 }
 
+
 static void
 tracemalloc_deinit(void)
 {
@@ -884,6 +1001,7 @@ tracemalloc_deinit(void)
     Py_XDECREF(unknown_filename);
 }
 
+
 static int
 tracemalloc_start(int max_nframe)
 {
@@ -941,6 +1059,7 @@ tracemalloc_start(int max_nframe)
     return 0;
 }
 
+
 static void
 tracemalloc_stop(void)
 {
@@ -974,6 +1093,7 @@ PyDoc_STRVAR(tracemalloc_is_tracing_doc,
     "True if the tracemalloc module is tracing Python memory allocations,\n"
     "False otherwise.");
 
+
 static PyObject*
 py_tracemalloc_is_tracing(PyObject *self)
 {
@@ -985,6 +1105,7 @@ PyDoc_STRVAR(tracemalloc_clear_traces_doc,
     "\n"
     "Clear traces of memory blocks allocated by Python.");
 
+
 static PyObject*
 py_tracemalloc_clear_traces(PyObject *self)
 {
@@ -998,6 +1119,7 @@ py_tracemalloc_clear_traces(PyObject *self)
     Py_RETURN_NONE;
 }
 
+
 static PyObject*
 frame_to_pyobject(frame_t *frame)
 {
@@ -1020,6 +1142,7 @@ frame_to_pyobject(frame_t *frame)
     return frame_obj;
 }
 
+
 static PyObject*
 traceback_to_pyobject(traceback_t *traceback, _Py_hashtable_t *intern_table)
 {
@@ -1058,33 +1181,43 @@ traceback_to_pyobject(traceback_t *traceback, _Py_hashtable_t *intern_table)
     return frames;
 }
 
+
 static PyObject*
-trace_to_pyobject(trace_t *trace, _Py_hashtable_t *intern_tracebacks)
+trace_to_pyobject(domain_t domain, trace_t *trace,
+                  _Py_hashtable_t *intern_tracebacks)
 {
     PyObject *trace_obj = NULL;
-    PyObject *size, *traceback;
+    PyObject *obj;
 
-    trace_obj = PyTuple_New(2);
+    trace_obj = PyTuple_New(3);
     if (trace_obj == NULL)
         return NULL;
 
-    size = PyLong_FromSize_t(trace->size);
+    obj = PyLong_FromSize_t(domain);
-    if (size == NULL) {
+    if (obj == NULL) {
         Py_DECREF(trace_obj);
         return NULL;
     }
-    PyTuple_SET_ITEM(trace_obj, 0, size);
+    PyTuple_SET_ITEM(trace_obj, 0, obj);
 
-    traceback = traceback_to_pyobject(trace->traceback, intern_tracebacks);
+    obj = PyLong_FromSize_t(trace->size);
-    if (traceback == NULL) {
+    if (obj == NULL) {
         Py_DECREF(trace_obj);
         return NULL;
     }
-    PyTuple_SET_ITEM(trace_obj, 1, traceback);
+    PyTuple_SET_ITEM(trace_obj, 1, obj);
+
+    obj = traceback_to_pyobject(trace->traceback, intern_tracebacks);
+    if (obj == NULL) {
+        Py_DECREF(trace_obj);
+        return NULL;
+    }
+    PyTuple_SET_ITEM(trace_obj, 2, obj);
 
     return trace_obj;
 }
 
+
 typedef struct {
     _Py_hashtable_t *traces;
     _Py_hashtable_t *tracebacks;
@@ -1096,13 +1229,22 @@ tracemalloc_get_traces_fill(_Py_hashtable_t *traces, _Py_hashtable_entry_t *entr
                             void *user_data)
 {
     get_traces_t *get_traces = user_data;
+    domain_t domain;
     trace_t *trace;
     PyObject *tracemalloc_obj;
     int res;
 
+    if (tracemalloc_config.use_domain) {
+        pointer_t key;
+        _Py_HASHTABLE_ENTRY_READ_KEY(traces->key_size, entry, key);
+        domain = key.domain;
+    }
+    else {
+        domain = DEFAULT_DOMAIN;
+    }
     trace = (trace_t *)_Py_HASHTABLE_ENTRY_DATA(traces, entry);
 
-    tracemalloc_obj = trace_to_pyobject(trace, get_traces->tracebacks);
+    tracemalloc_obj = trace_to_pyobject(domain, trace, get_traces->tracebacks);
     if (tracemalloc_obj == NULL)
         return 1;
 
@@ -1114,6 +1256,7 @@ tracemalloc_get_traces_fill(_Py_hashtable_t *traces, _Py_hashtable_entry_t *entr
     return 0;
 }
 
+
 static int
 tracemalloc_pyobject_decref_cb(_Py_hashtable_t *tracebacks,
                                _Py_hashtable_entry_t *entry,
@@ -1125,6 +1268,7 @@ tracemalloc_pyobject_decref_cb(_Py_hashtable_t *tracebacks,
     return 0;
 }
 
+
 PyDoc_STRVAR(tracemalloc_get_traces_doc,
     "_get_traces() -> list\n"
     "\n"
@@ -1194,8 +1338,9 @@ finally:
     return get_traces.list;
 }
 
+
 static traceback_t*
-tracemalloc_get_traceback(const void *ptr)
+tracemalloc_get_traceback(domain_t domain, const void *ptr)
 {
     trace_t trace;
     int found;
@@ -1204,7 +1349,13 @@ tracemalloc_get_traceback(const void *ptr)
         return NULL;
 
     TABLES_LOCK();
-    found = _Py_HASHTABLE_GET(tracemalloc_traces, ptr, trace);
+    if (tracemalloc_config.use_domain) {
+        pointer_t key = {(Py_uintptr_t)ptr, domain};
+        found = _Py_HASHTABLE_GET(tracemalloc_traces, key, trace);
+    }
+    else {
+        found = _Py_HASHTABLE_GET(tracemalloc_traces, ptr, trace);
+    }
     TABLES_UNLOCK();
 
     if (!found)
@@ -1213,6 +1364,7 @@ tracemalloc_get_traceback(const void *ptr)
     return trace.traceback;
 }
 
+
 PyDoc_STRVAR(tracemalloc_get_object_traceback_doc,
     "_get_object_traceback(obj)\n"
     "\n"
@@ -1235,13 +1387,14 @@ py_tracemalloc_get_object_traceback(PyObject *self, PyObject *obj)
     else
         ptr = (void *)obj;
 
-    traceback = tracemalloc_get_traceback(ptr);
+    traceback = tracemalloc_get_traceback(DEFAULT_DOMAIN, ptr);
     if (traceback == NULL)
         Py_RETURN_NONE;
 
     return traceback_to_pyobject(traceback, NULL);
 }
 
+
 #define PUTS(fd, str) _Py_write_noraise(fd, str, (int)strlen(str))
 
 static void
@@ -1262,7 +1415,7 @@ _PyMem_DumpTraceback(int fd, const void *ptr)
     traceback_t *traceback;
     int i;
 
-    traceback = tracemalloc_get_traceback(ptr);
+    traceback = tracemalloc_get_traceback(DEFAULT_DOMAIN, ptr);
     if (traceback == NULL)
         return;
 
@@ -1275,6 +1428,7 @@ _PyMem_DumpTraceback(int fd, const void *ptr)
 
 #undef PUTS
 
+
 PyDoc_STRVAR(tracemalloc_start_doc,
     "start(nframe: int=1)\n"
     "\n"
@@ -1310,6 +1464,7 @@ PyDoc_STRVAR(tracemalloc_stop_doc,
     "Stop tracing Python memory allocations and clear traces\n"
     "of memory blocks allocated by Python.");
 
+
 static PyObject*
 py_tracemalloc_stop(PyObject *self)
 {
@@ -1317,6 +1472,7 @@ py_tracemalloc_stop(PyObject *self)
     Py_RETURN_NONE;
 }
 
+
 PyDoc_STRVAR(tracemalloc_get_traceback_limit_doc,
     "get_traceback_limit() -> int\n"
     "\n"
@@ -1332,6 +1488,7 @@ py_tracemalloc_get_traceback_limit(PyObject *self)
     return PyLong_FromLong(tracemalloc_config.max_nframe);
 }
 
+
 PyDoc_STRVAR(tracemalloc_get_tracemalloc_memory_doc,
     "get_tracemalloc_memory() -> int\n"
     "\n"
@@ -1355,6 +1512,7 @@ tracemalloc_get_tracemalloc_memory(PyObject *self)
     return Py_BuildValue("N", size_obj);
 }
 
+
 PyDoc_STRVAR(tracemalloc_get_traced_memory_doc,
     "get_traced_memory() -> (int, int)\n"
     "\n"
@@ -1380,6 +1538,7 @@ tracemalloc_get_traced_memory(PyObject *self)
     return Py_BuildValue("NN", size_obj, peak_size_obj);
 }
 
+
 static PyMethodDef module_methods[] = {
     {"is_tracing", (PyCFunction)py_tracemalloc_is_tracing,
      METH_NOARGS, tracemalloc_is_tracing_doc},
@@ -1430,6 +1589,7 @@ PyInit__tracemalloc(void)
     return m;
 }
 
+
 static int
 parse_sys_xoptions(PyObject *value)
 {
@@ -1458,6 +1618,7 @@ parse_sys_xoptions(PyObject *value)
     return Py_SAFE_DOWNCAST(nframe, long, int);
 }
 
+
 int
 _PyTraceMalloc_Init(void)
 {
@@ -1516,6 +1677,7 @@ _PyTraceMalloc_Init(void)
     return tracemalloc_start(nframe);
 }
 
+
 void
 _PyTraceMalloc_Fini(void)
 {
@@ -1524,4 +1686,3 @@ _PyTraceMalloc_Fini(void)
 #endif
     tracemalloc_deinit();
 }
-