Expose the record writer for validating arbitrary files

setup.py

@@ -267,12 +267,19 @@ def build_js_files(self):
     name="memray._test_utils",
     sources=[
         "src/memray/_memray_test_utils.pyx",
+        "src/memray/_memray/sink.cpp",
+        "src/memray/_memray/records.cpp",
+        "src/memray/_memray/snapshot.cpp",
+        "src/memray/_memray/record_writer.cpp",
+        "src/memray/_memray/hooks.cpp",
+        "src/memray/_memray/logging.cpp",
     ],
     language="c++",
     extra_compile_args=["-std=c++17", "-Wall", *EXTRA_COMPILE_ARGS],
     extra_link_args=["-std=c++17", *EXTRA_LINK_ARGS],
     define_macros=DEFINE_MACROS,
     undef_macros=UNDEF_MACROS,
+    **library_flags,
 )
 
 MEMRAY_INJECT_EXTENSION = Extension(

src/memray/_memray/record_writer.pxd

@@ -1,13 +1,36 @@
+from _memray.records cimport AllocationRecord
 from _memray.records cimport FileFormat
+from _memray.records cimport FramePop
+from _memray.records cimport FramePush
+from _memray.records cimport HeaderRecord
+from _memray.records cimport ImageSegments
+from _memray.records cimport MemoryRecord
+from _memray.records cimport NativeAllocationRecord
+from _memray.records cimport ThreadRecord
+from _memray.records cimport UnresolvedNativeFrame
+from _memray.records cimport thread_id_t
 from _memray.sink cimport Sink
 from libcpp cimport bool
 from libcpp.memory cimport unique_ptr
 from libcpp.string cimport string
+from libcpp.vector cimport vector
 
 
-cdef extern from "record_writer.h" namespace "memray::api":
+cdef extern from "record_writer.h" namespace "memray::tracking_api":
     cdef cppclass RecordWriter:
-        pass
+        bool writeRecord(const MemoryRecord& record) except+
+        bool writeRecord(const UnresolvedNativeFrame& record) except+
+        bool writeMappings(const vector[ImageSegments]& mappings) except+
+        bool writeThreadSpecificRecord(thread_id_t tid, const FramePop& record) except+
+        bool writeThreadSpecificRecord(thread_id_t tid, const FramePush& record) except+
+        bool writeThreadSpecificRecord(thread_id_t tid, const AllocationRecord& record) except+
+        bool writeThreadSpecificRecord(thread_id_t tid, const NativeAllocationRecord& record) except+
+        bool writeThreadSpecificRecord(thread_id_t tid, const ThreadRecord& record) except+
+        bool writeHeader(bool seek_to_start) except+
+        bool writeTrailer() except+
+        void setMainTidAndSkippedFrames(thread_id_t main_tid, size_t skipped_frames_on_main_tid) except+
+        unique_ptr[RecordWriter] cloneInChildProcess() except+
+
     cdef unique_ptr[RecordWriter] createRecordWriter(
         unique_ptr[Sink],
         string command_line,

src/memray/_memray/records.pxd

@@ -5,15 +5,69 @@ from libcpp.string cimport string
 from libcpp.vector cimport vector
 
 
+cdef extern from "hooks.h" namespace "memray::hooks":
+    cdef enum Allocator:
+        MALLOC "memray::hooks::Allocator::MALLOC"
+        CALLOC "memray::hooks::Allocator::CALLOC"
+        REALLOC "memray::hooks::Allocator::REALLOC"
+        VALLOC "memray::hooks::Allocator::VALLOC"
+        ALIGNED_ALLOC "memray::hooks::Allocator::ALIGNED_ALLOC"
+        POSIX_MEMALIGN "memray::hooks::Allocator::POSIX_MEMALIGN"
+        MEMALIGN "memray::hooks::Allocator::MEMALIGN"
+        PVALLOC "memray::hooks::Allocator::PVALLOC"
+        FREE "memray::hooks::Allocator::FREE"
+        PYMALLOC_MALLOC "memray::hooks::Allocator::PYMALLOC_MALLOC"
+        PYMALLOC_CALLOC "memray::hooks::Allocator::PYMALLOC_CALLOC"
+        PYMALLOC_REALLOC "memray::hooks::Allocator::PYMALLOC_REALLOC"
+        PYMALLOC_FREE "memray::hooks::Allocator::PYMALLOC_FREE"
+
 cdef extern from "records.h" namespace "memray::tracking_api":
    ctypedef unsigned long thread_id_t
    ctypedef size_t frame_id_t
+   ctypedef long long millis_t
+
+   struct MemoryRecord:
+       unsigned long int ms_since_epoch
+       size_t rss
+
+   struct AllocationRecord:
+       uintptr_t address
+       size_t size
+       Allocator allocator
+
+   struct NativeAllocationRecord:
+       uintptr_t address
+       size_t size
+       Allocator allocator
+       frame_id_t native_frame_id
 
    struct Frame:
        string function_name
        string filename
        int lineno
 
+   struct FramePush:
+       frame_id_t frame_id
+
+   struct FramePop:
+       size_t count
+
+   struct ThreadRecord:
+       const char* name
+
+   struct Segment:
+       uintptr_t vaddr
+       size_t memsz
+
+   struct ImageSegments:
+       string filename
+       uintptr_t addr
+       vector[Segment] segments
+
+   struct UnresolvedNativeFrame:
+       uintptr_t ip
+       frame_id_t index
+
    struct TrackerStats:
        size_t n_allocations
        size_t n_frames

src/memray/_memray_test_utils.pyx

@@ -42,10 +42,38 @@ from cpython.pylifecycle cimport Py_FinalizeEx
 from libc.errno cimport errno
 from libc.stdint cimport uintptr_t
 from libc.stdlib cimport exit as _exit
+from libcpp.utility cimport move
 from libcpp.vector cimport vector
 
 from ._destination import Destination
 
+from _memray.record_writer cimport RecordWriter
+from _memray.record_writer cimport createRecordWriter
+from _memray.records cimport AllocationRecord
+from _memray.records cimport Allocator
+from _memray.records cimport FileFormat
+from _memray.records cimport FramePop
+from _memray.records cimport FramePush
+from _memray.records cimport HeaderRecord
+from _memray.records cimport ImageSegments
+from _memray.records cimport MemoryRecord
+from _memray.records cimport NativeAllocationRecord
+from _memray.records cimport Segment
+from _memray.records cimport ThreadRecord
+from _memray.records cimport UnresolvedNativeFrame
+from _memray.records cimport thread_id_t
+from _memray.sink cimport FileSink
+from _memray.sink cimport Sink
+from cpython.ref cimport PyObject
+from libcpp cimport bool
+from libcpp.memory cimport unique_ptr
+from libcpp.string cimport string
+
+import os
+from typing import List
+from typing import Optional
+from typing import Union
+
 
 cdef extern from *:
     """
@@ -285,3 +313,123 @@ cdef class PrimeCaches:
         return self
     def __exit__(self, *args):
         sys.setprofile(self.old_profile)
+
+
+cdef class TestRecordWriter:
+    """A Python wrapper around the C++ RecordWriter class for testing purposes."""
+
+    cdef unique_ptr[RecordWriter] _writer
+    cdef unique_ptr[Sink] _sink
+    cdef bool _native_traces
+    cdef bool _trace_python_allocators
+    cdef FileFormat _file_format
+
+    def __cinit__(self, str file_path, bool native_traces=False,
+                  bool trace_python_allocators=False,
+                  FileFormat file_format=FileFormat.ALL_ALLOCATIONS):
+        """Initialize a new TestRecordWriter.
+
+        Args:
+            file_path: Path to the output file
+            native_traces: Whether to include native traces
+            trace_python_allocators: Whether to trace Python allocators
+            file_format: The format of the output file
+        """
+        self._native_traces = native_traces
+        self._trace_python_allocators = trace_python_allocators
+        self._file_format = file_format
+
+        # Create the sink
+        cdef string cpp_path = file_path.encode('utf-8')
+        try:
+            self._sink = unique_ptr[Sink](new FileSink(cpp_path, True, False))
+        except:
+            raise IOError("Failed to create file sink")
+
+        # Create the writer
+        cdef string command_line = b" ".join(arg.encode('utf-8') for arg in sys.argv)
+        try:
+            self._writer = createRecordWriter(
+                move(self._sink),
+                command_line,
+                native_traces,
+                file_format,
+                trace_python_allocators
+            )
+        except:
+            raise IOError("Failed to create record writer")
+
+        # Write the header
+        if not self._writer.get().writeHeader(True):
+            raise RuntimeError("Failed to write header")
+
+    def write_memory_record(self, unsigned long ms_since_epoch, size_t rss) -> bool:
+        """Write a memory record to the file."""
+        cdef MemoryRecord record
+        record.ms_since_epoch = ms_since_epoch
+        record.rss = rss
+        return self._writer.get().writeRecord(record)
+
+    def write_allocation_record(self, thread_id_t tid, uintptr_t address,
+                              size_t size, unsigned char allocator) -> bool:
+        """Write an allocation record to the file."""
+        cdef AllocationRecord record
+        record.address = address
+        record.size = size
+        record.allocator = <Allocator>allocator
+        return self._writer.get().writeThreadSpecificRecord(tid, record)
+
+    def write_native_allocation_record(self, thread_id_t tid, uintptr_t address,
+                                     size_t size, unsigned char allocator,
+                                     size_t native_frame_id) -> bool:
+        """Write a native allocation record to the file."""
+        cdef NativeAllocationRecord record
+        record.address = address
+        record.size = size
+        record.allocator = <Allocator>allocator
+        record.native_frame_id = native_frame_id
+        return self._writer.get().writeThreadSpecificRecord(tid, record)
+
+    def write_frame_push(self, thread_id_t tid, size_t frame_id) -> bool:
+        """Write a frame push record to the file."""
+        cdef FramePush record
+        record.frame_id = frame_id
+        return self._writer.get().writeThreadSpecificRecord(tid, record)
+
+    def write_frame_pop(self, thread_id_t tid, size_t count) -> bool:
+        """Write a frame pop record to the file."""
+        cdef FramePop record
+        record.count = count
+        return self._writer.get().writeThreadSpecificRecord(tid, record)
+
+    def write_thread_record(self, thread_id_t tid, str name) -> bool:
+        """Write a thread record to the file."""
+        cdef ThreadRecord record
+        cdef bytes name_bytes = name.encode('utf-8')
+        record.name = name_bytes
+        return self._writer.get().writeThreadSpecificRecord(tid, record)
+
+    def write_mappings(self, list mappings) -> bool:
+        """Write memory mappings to the file."""
+        cdef vector[ImageSegments] cpp_mappings
+        cdef ImageSegments segments
+        cdef Segment segment
+        for mapping in mappings:
+            segments = ImageSegments()
+            segments.filename = mapping['filename'].encode('utf-8')
+            segments.addr = mapping['addr']
+            for seg in mapping['segments']:
+                segment.vaddr = seg['vaddr']
+                segment.memsz = seg['memsz']
+                segments.segments.push_back(segment)
+            cpp_mappings.push_back(segments)
+        return self._writer.get().writeMappings(cpp_mappings)
+
+    def write_trailer(self) -> bool:
+        """Write the trailer to the file."""
+        return self._writer.get().writeTrailer()
+
+    def set_main_tid_and_skipped_frames(self, thread_id_t main_tid,
+                                      size_t skipped_frames) -> None:
+        """Set the main thread ID and number of skipped frames."""
+        self._writer.get().setMainTidAndSkippedFrames(main_tid, skipped_frames)

src/memray/_test.py

@@ -5,6 +5,7 @@
 from ._test_utils import PrimeCaches
 from ._test_utils import PymallocDomain
 from ._test_utils import PymallocMemoryAllocator
+from ._test_utils import TestRecordWriter
 from ._test_utils import _cython_allocate_in_two_places
 from ._test_utils import _cython_nested_allocation
 from ._test_utils import allocate_cpp_vector
@@ -64,4 +65,5 @@ def run_in_pthread(self, callback: Callable[[], None]) -> None:
     "fill_cpp_vector",
     "exit",
     "PrimeCaches",
+    "TestRecordWriter",
 ]

src/memray/_test_utils.pyi

@@ -1,5 +1,6 @@
 from typing import Callable
 
+from ._memray import FileFormat
 from ._memray import PymallocDomain as PymallocDomain
 
 class MemoryAllocator:
@@ -43,3 +44,27 @@ class PrimeCaches:
     def __init__(self, size: int) -> None: ...
     def __enter__(self) -> None: ...
     def __exit__(self, *args: object) -> None: ...
+
+class TestRecordWriter:
+    def __init__(
+        self,
+        file_path: str,
+        native_traces: bool = False,
+        trace_python_allocators: bool = False,
+        file_format: FileFormat = FileFormat.ALL_ALLOCATIONS,
+    ) -> None: ...
+    def write_memory_record(self, ms_since_epoch: int, rss: int) -> bool: ...
+    def write_allocation_record(
+        self, tid: int, address: int, size: int, allocator: int
+    ) -> bool: ...
+    def write_native_allocation_record(
+        self, tid: int, address: int, size: int, allocator: int, native_frame_id: int
+    ) -> bool: ...
+    def write_frame_push(self, tid: int, frame_id: int) -> bool: ...
+    def write_frame_pop(self, tid: int, count: int) -> bool: ...
+    def write_thread_record(self, tid: int, name: str) -> bool: ...
+    def write_mappings(self, mappings: list) -> bool: ...
+    def write_trailer(self) -> bool: ...
+    def set_main_tid_and_skipped_frames(
+        self, main_tid: int, skipped_frames: int
+    ) -> None: ...