Repository: bloomberg/memray
Pull request #560 (2024-03-06 14:04:26)

Expose the main thread ID in the reader
To be able to filter allocations that happened only on the thread that
started the tracking, expose the main thread ID in the reader class.

*Issue number of the reported bug or feature request: #<number>*

**Describe your changes**
A clear and concise description of the changes you have made.

**Testing performed**
Describe the testing you have performed to ensure that the bug has been addressed, or that the new feature works as planned.

**Additional context**
Add any other context about your contribution here.



Patch:
diff --git a/news/560.feature.rst b/news/560.feature.rst
new file mode 100644
index 0000000..7ce4615
--- /dev/null
+++ b/news/560.feature.rst
@@ -0,0 +1,1 @@
+Expose the main thread id in the FileReader's metadata attribute.
diff --git a/src/memray/_memray.pyx b/src/memray/_memray.pyx
index 0423215..9db7279 100644
--- a/src/memray/_memray.pyx
+++ b/src/memray/_memray.pyx
@@ -754,6 +754,7 @@ cdef _create_metadata(header, peak_memory):
         peak_memory=peak_memory,
         command_line=header["command_line"],
         pid=header["pid"],
+        main_thread_id=header["main_tid"],
         python_allocator=allocator_id_to_name[header["python_allocator"]],
         has_native_traces=header["native_traces"],
         trace_python_allocators=header["trace_python_allocators"],
diff --git a/src/memray/_metadata.py b/src/memray/_metadata.py
index a44caa9..5d5913b 100644
--- a/src/memray/_metadata.py
+++ b/src/memray/_metadata.py
@@ -13,6 +13,7 @@ class Metadata:
     peak_memory: int
     command_line: str
     pid: int
+    main_thread_id: int
     python_allocator: str
     has_native_traces: bool
     trace_python_allocators: bool


Test patch:
diff --git a/tests/unit/test_reader.py b/tests/unit/test_reader.py
index f94bba8..4ac0a12 100644
--- a/tests/unit/test_reader.py
+++ b/tests/unit/test_reader.py
@@ -1,4 +1,5 @@
 import os
+import threading
 
 import pytest
 
@@ -79,3 +80,28 @@ def test_read_pid(tmp_path):
 
     # THEN
     assert FileReader(output).metadata.pid == os.getpid()
+
+
+def test_read_tid(tmp_path):
+    # GIVEN
+    output = tmp_path / "test.bin"
+    allocator = MemoryAllocator()
+
+    def func():
+        allocator.valloc(1024)
+
+    # WHEN
+    t = threading.Thread(target=func)
+    with Tracker(output):
+        func()
+        t.start()
+        t.join()
+
+    # THEN
+    reader = FileReader(output)
+    main_tid = reader.metadata.main_thread_id
+    all_allocations = reader.get_allocation_records()
+    all_tids = tuple(allocation.tid for allocation in all_allocations)
+    assert main_tid in set(all_tids)
+    # The main thread should be the first one in the list
+    assert main_tid == all_tids[0]
diff --git a/tests/unit/test_stats_reporter.py b/tests/unit/test_stats_reporter.py
index e688360..193f5e4 100644
--- a/tests/unit/test_stats_reporter.py
+++ b/tests/unit/test_stats_reporter.py
@@ -99,6 +99,7 @@ def fake_stats():
             has_native_traces=False,
             trace_python_allocators=True,
             file_format=FileFormat.ALL_ALLOCATIONS,
+            main_thread_id=0x1,
         ),
         total_num_allocations=20,
         total_memory_allocated=sum(mem_allocation_list),
@@ -438,6 +439,7 @@ def test_stats_output_json(fake_stats, tmp_path):
             "has_native_traces": False,
             "trace_python_allocators": True,
             "file_format": 0,
+            "main_thread_id": 0x1,
         },
     }
     actual = json.loads(output_file.read_text())