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