From 138b3497012dc43e2a7af822ba77a092c10e6653 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>
---
 src/memray/_memray.pyx            |  1 +
 src/memray/_metadata.py           |  1 +
 tests/unit/test_reader.py         | 26 ++++++++++++++++++++++++++
 tests/unit/test_stats_reporter.py |  2 ++
 4 files changed, 30 insertions(+)

diff --git a/src/memray/_memray.pyx b/src/memray/_memray.pyx
index 32053383cc..6c6bc9cfd4 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_tid=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 64cb2d7b55..14946986f7 100644
--- a/src/memray/_metadata.py
+++ b/src/memray/_metadata.py
@@ -11,6 +11,7 @@ class Metadata:
     peak_memory: int
     command_line: str
     pid: int
+    main_thread_tid: 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..f1943cfffd 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_tid
+    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 be25e2b51c..0674d74a7f 100644
--- a/tests/unit/test_stats_reporter.py
+++ b/tests/unit/test_stats_reporter.py
@@ -97,6 +97,7 @@ def fake_stats():
             python_allocator="pymalloc",
             has_native_traces=False,
             trace_python_allocators=True,
+            main_thread_tid=0x1,
         ),
         total_num_allocations=20,
         total_memory_allocated=sum(mem_allocation_list),
@@ -435,6 +436,7 @@ def test_stats_output_json(fake_stats, tmp_path):
             "python_allocator": "pymalloc",
             "has_native_traces": False,
             "trace_python_allocators": True,
+            "main_thread_tid": 0x1,
         },
     }
     actual = json.loads(output_file.read_text())