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