Expose the record writer for validating arbitrary files

Adds TestRecordWriter class to enable testing and validation of memray
record files by allowing direct writing of test records. This provides a
way to create controlled test files with known contents for validating
the file format and reader functionality. The implementation supports
writing all record types including memory snapshots, allocations,
frames, and thread information, making it possible to test edge cases
and complex scenarios that would be difficult to trigger in normal
operation.

Signed-off-by: Pablo Galindo <[email protected]>

Author: pablogsal

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: ...