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.

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,73 +5,122 @@ 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
-
-   struct Frame:
-       string function_name
-       string filename
-       int lineno
-
-   struct TrackerStats:
-       size_t n_allocations
-       size_t n_frames
-       long long start_time
-       long long end_time
-
-   cdef enum FileFormat:
-       ALL_ALLOCATIONS 'memray::tracking_api::FileFormat::ALL_ALLOCATIONS'
-       AGGREGATED_ALLOCATIONS 'memray::tracking_api::FileFormat::AGGREGATED_ALLOCATIONS'
-
-   struct HeaderRecord:
-       int version
-       bool native_traces
-       FileFormat file_format
-       TrackerStats stats
-       string command_line
-       int pid
-       size_t main_tid
-       size_t skipped_frames_on_main_tid
-       int python_allocator
-       bool trace_python_allocators
-
-   cdef cppclass Allocation:
-       thread_id_t tid
-       uintptr_t address
-       size_t size
-       Allocator allocator
-       frame_id_t native_frame_id
-       size_t frame_index
-       size_t native_segment_generation
-       size_t n_allocations
-
-       object toPythonObject()
-
-   cdef cppclass AggregatedAllocation:
-       thread_id_t tid
-       Allocator allocator
-       frame_id_t native_frame_id
-       size_t frame_index
-       size_t native_segment_generation
-
-       size_t n_allocations_in_high_water_mark
-       size_t n_allocations_leaked
-       size_t bytes_in_high_water_mark
-       size_t bytes_leaked
-
-       Allocation contributionToHighWaterMark()
-       Allocation contributionToLeaks()
-
-   struct MemoryRecord:
-       unsigned long int ms_since_epoch
-       size_t rss
-
-   struct MemorySnapshot:
-       unsigned long int ms_since_epoch
-       size_t rss
-       size_t heap
+    ctypedef unsigned long thread_id_t
+    ctypedef size_t frame_id_t
+    ctypedef long long millis_t
+
+    cdef enum FileFormat:
+        ALL_ALLOCATIONS
+        AGGREGATED_ALLOCATIONS
+
+    cdef struct MemoryRecord:
+        unsigned long int ms_since_epoch
+        size_t rss
+
+    cdef struct AllocationRecord:
+        uintptr_t address
+        size_t size
+        Allocator allocator
+
+    cdef struct NativeAllocationRecord:
+        uintptr_t address
+        size_t size
+        Allocator allocator
+        frame_id_t native_frame_id
+
+    cdef struct FramePush:
+        frame_id_t frame_id
+
+    cdef struct FramePop:
+        size_t count
+
+    cdef struct ThreadRecord:
+        const char* name
+
+    cdef struct Segment:
+        uintptr_t vaddr
+        size_t memsz
+
+    cdef struct ImageSegments:
+        string filename
+        uintptr_t addr
+        vector[Segment] segments
+
+    cdef struct UnresolvedNativeFrame:
+        uintptr_t ip
+        frame_id_t index
+
+    struct Frame:
+        string function_name
+        string filename
+        int lineno
+
+    struct TrackerStats:
+        size_t n_allocations
+        size_t n_frames
+        long long start_time
+        long long end_time
+
+    struct HeaderRecord:
+        int version
+        bool native_traces
+        FileFormat file_format
+        TrackerStats stats
+        string command_line
+        int pid
+        size_t main_tid
+        size_t skipped_frames_on_main_tid
+        int python_allocator
+        bool trace_python_allocators
+
+    cdef cppclass Allocation:
+        thread_id_t tid
+        uintptr_t address
+        size_t size
+        Allocator allocator
+        frame_id_t native_frame_id
+        size_t frame_index
+        size_t native_segment_generation
+        size_t n_allocations
+
+        object toPythonObject()
+
+    cdef cppclass AggregatedAllocation:
+        thread_id_t tid
+        Allocator allocator
+        frame_id_t native_frame_id
+        size_t frame_index
+        size_t native_segment_generation
+
+        size_t n_allocations_in_high_water_mark
+        size_t n_allocations_leaked
+        size_t bytes_in_high_water_mark
+        size_t bytes_leaked
+
+        Allocation contributionToHighWaterMark()
+        Allocation contributionToLeaks()
 
+    struct MemorySnapshot:
+        unsigned long int ms_since_epoch
+        size_t rss
+        size_t heap
 
 cdef extern from "<optional>":
    # Cython doesn't have libcpp.optional yet, so just declare this opaquely.