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.

Signed-off-by: Pablo Galindo <[email protected]>

news/560.feature.rst

@@ -0,0 +1 @@
+Expose the main thread id in the FileReader's metadata attribute.

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"],

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

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]

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())