From bdc8737f420468f10825043e14758188f85d4d8a Mon Sep 17 00:00:00 2001
From: Pablo Galindo <pablogsal@gmail.com>
Date: Wed, 6 Mar 2024 14:03:46 +0000
Subject: [PATCH] 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 <pablogsal@gmail.com>
---
 news/560.feature.rst              |  1 +
 src/memray/_memray.pyx            |  1 +
 src/memray/_metadata.py           |  1 +
 tests/unit/test_reader.py         | 26 ++++++++++++++++++++++++++
 tests/unit/test_stats_reporter.py |  2 ++
 5 files changed, 31 insertions(+)
 create mode 100644 news/560.feature.rst

diff --git a/news/560.feature.rst b/news/560.feature.rst
new file mode 100644
index 0000000000..7ce4615145
--- /dev/null
+++ b/news/560.feature.rst
@@ -0,0 +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 0423215e9b..9db72791b6 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 a44caa9718..5d5913b498 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
diff --git a/tests/unit/test_reader.py b/tests/unit/test_reader.py
index f94bba82a8..4ac0a12e0f 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 e6883601b6..193f5e4b2b 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())