import collections
import datetime
import mmap
import signal
import subprocess
import sys
import textwrap
import threading
import time
from pathlib import Path

import pytest

from memray import AllocatorType
from memray import FileFormat
from memray import FileReader
from memray import Tracker
from memray._memray import compute_statistics
from memray._test import MemoryAllocator
from memray._test import MmapAllocator
from memray._test import PrimeCaches
from memray._test import PymallocDomain
from memray._test import PymallocMemoryAllocator
from memray._test import _cython_allocate_in_two_places
from memray._test import allocate_cpp_vector
from memray._test import fill_cpp_vector
from tests.utils import filter_relevant_allocations
from tests.utils import run_without_tracer

ALLOCATORS = [
    ("malloc", AllocatorType.MALLOC),
    ("valloc", AllocatorType.VALLOC),
    ("pvalloc", AllocatorType.PVALLOC),
    ("calloc", AllocatorType.CALLOC),
    ("memalign", AllocatorType.MEMALIGN),
    ("posix_memalign", AllocatorType.POSIX_MEMALIGN),
    ("aligned_alloc", AllocatorType.ALIGNED_ALLOC),
    ("realloc", AllocatorType.REALLOC),
]

PYMALLOC_ALLOCATORS = [
    ("malloc", AllocatorType.PYMALLOC_MALLOC),
    ("calloc", AllocatorType.PYMALLOC_CALLOC),
    ("realloc", AllocatorType.PYMALLOC_REALLOC),
]

PYMALLOC_DOMAINS = [
    PymallocDomain.PYMALLOC_RAW,
    PymallocDomain.PYMALLOC_MEM,
    PymallocDomain.PYMALLOC_OBJECT,
]


PAGE_SIZE = mmap.PAGESIZE
ALLOC_SIZE = 123 * 8


@pytest.mark.skipif(
    sys.platform == "darwin", reason="Test triggers some extra allocations in macOS"
)
def test_no_allocations_while_tracking(tmp_path):
    output = tmp_path / "test.bin"
    with Tracker(output):
        pass

    records = list(FileReader(output).get_allocation_records())
    assert not records


@pytest.mark.parametrize(["allocator_func", "allocator_type"], ALLOCATORS)
def test_simple_allocation_tracking(allocator_func, allocator_type, tmp_path):
    # GIVEN
    allocator = MemoryAllocator()
    output = tmp_path / "test.bin"

    # WHEN
    with Tracker(output):
        res = getattr(allocator, allocator_func)(ALLOC_SIZE)
        if res:
            allocator.free()

    if not res:
        pytest.skip(f"Allocator {allocator_func} not supported in this platform")

    # THEN
    allocations = list(FileReader(output).get_allocation_records())
    allocs = [
        event
        for event in allocations
        if event.size == ALLOC_SIZE and event.allocator == allocator_type
    ]
    assert len(allocs) == 1
    (alloc,) = allocs

    frees = [
        event
        for event in allocations
        if event.address == alloc.address and event.allocator == AllocatorType.FREE
    ]
    assert len(frees) >= 1


def test_simple_cpp_allocation_tracking(tmp_path):
    # GIVEN
    output = tmp_path / "test.bin"

    # WHEN
    with Tracker(output):
        allocate_cpp_vector(ALLOC_SIZE)

    # THEN
    allocations = list(FileReader(output).get_allocation_records())
    allocs = [event for event in allocations if event.size == ALLOC_SIZE]
    assert len(allocs) == 1
    (alloc,) = allocs

    frees = [
        event
        for event in allocations
        if event.address == alloc.address and event.allocator == AllocatorType.FREE
    ]
    assert len(frees) >= 1


@pytest.mark.parametrize("domain", PYMALLOC_DOMAINS)
@pytest.mark.parametrize(["allocator_func", "allocator_type"], PYMALLOC_ALLOCATORS)
def test_simple_pymalloc_allocation_tracking(
    allocator_func, allocator_type, domain, tmp_path
):
    # GIVEN
    allocator = PymallocMemoryAllocator(domain)
    output = tmp_path / "test.bin"

    # WHEN
    the_allocator = getattr(allocator, allocator_func)
    with Tracker(output, trace_python_allocators=True):
        res = the_allocator(ALLOC_SIZE)
        if res:
            allocator.free()

    if not res:
        pytest.skip(f"Allocator {allocator_func} not supported in this platform")

    # THEN
    allocations = list(FileReader(output).get_allocation_records())
    allocs = [
        event
        for event in allocations
        if event.size == ALLOC_SIZE and event.allocator == allocator_type
    ]
    assert len(allocs) == 1
    (alloc,) = allocs

    frees = [
        event
        for event in allocations
        if event.address == alloc.address
        and event.allocator == AllocatorType.PYMALLOC_FREE
    ]
    assert len(frees) >= 1


@pytest.mark.parametrize("domain", PYMALLOC_DOMAINS)
@pytest.mark.parametrize(["allocator_func", "allocator_type"], PYMALLOC_ALLOCATORS)
def test_pymalloc_allocation_tracking_deactivated(
    allocator_func, allocator_type, domain, tmp_path
):
    # GIVEN
    allocator = PymallocMemoryAllocator(domain)
    output = tmp_path / "test.bin"

    # WHEN
    the_allocator = getattr(allocator, allocator_func)
    with Tracker(output, trace_python_allocators=False):
        res = the_allocator(ALLOC_SIZE)
        if res:
            allocator.free()

    if not res:
        pytest.skip(f"Allocator {allocator_func} not supported in this platform")

    # THEN
    allocations = list(FileReader(output).get_allocation_records())
    allocs = [event for event in allocations if event.allocator == allocator_type]
    assert not allocs
    frees = [
        event for event in allocations if event.allocator == AllocatorType.PYMALLOC_FREE
    ]
    assert not frees


def test_mmap_tracking(tmp_path):
    # GIVEN / WHEN
    output = tmp_path / "test.bin"
    with Tracker(output):
        with mmap.mmap(-1, length=2048, access=mmap.ACCESS_WRITE) as mmap_obj:
            mmap_obj[0:100] = b"a" * 100

    # THEN
    records = list(FileReader(output).get_allocation_records())
    assert len(records) >= 2

    mmap_records = [
        record
        for record in records
        if AllocatorType.MMAP == record.allocator and record.size == 2048
    ]
    assert len(mmap_records) == 1
    mmunmap_record = [
        record for record in records if AllocatorType.MUNMAP == record.allocator
    ]
    assert len(mmunmap_record) == 1


def test_pthread_tracking(tmp_path):
    # GIVEN
    allocator = MemoryAllocator()

    def tracking_function():  # pragma: no cover
        allocator.valloc(ALLOC_SIZE)
        allocator.free()

    # WHEN
    output = tmp_path / "test.bin"
    with Tracker(output):
        allocator.run_in_pthread(tracking_function)

    # THEN
    allocations = list(FileReader(output).get_allocation_records())
    allocs = [
        event
        for event in allocations
        if event.size == ALLOC_SIZE and event.allocator == AllocatorType.VALLOC
    ]
    assert len(allocs) == 1
    (alloc,) = allocs

    frees = [
        event
        for event in allocations
        if event.address == alloc.address and event.allocator == AllocatorType.FREE
    ]
    assert len(frees) >= 1


def test_tracking_with_SIGKILL(tmpdir):
    """Verify that we can successfully retrieve the allocations after SIGKILL."""
    # GIVEN
    output = Path(tmpdir) / "test.bin"
    subprocess_code = textwrap.dedent(
        f"""
        import os
        import signal
        from memray import Tracker
        from memray._test import MemoryAllocator

        allocator = MemoryAllocator()
        output = "{output}"

        with Tracker(output) as tracker:
            num_flushes = 0
            last_size = os.stat(output).st_size

            # Loop until two flushes occur, since the first flush might include
            # a frame record but not an allocation record.
            while num_flushes < 2:
                allocator.valloc(1024)
                new_size = os.stat(output).st_size
                if new_size != last_size:
                    last_size = new_size
                    num_flushes += 1

            # Kill ourselves without letting the tracker clean itself up.
            os.kill(os.getpid(), signal.SIGKILL)
    """
    )

    # WHEN
    process = subprocess.run([sys.executable, "-c", subprocess_code], timeout=5)

    # THEN
    assert process.returncode == -signal.SIGKILL

    records = list(FileReader(output).get_allocation_records())
    vallocs = [
        record
        for record in filter_relevant_allocations(records)
        if record.allocator == AllocatorType.VALLOC
    ]
    (allocation, *rest) = vallocs
    assert allocation.size == 1024


def test_no_allocations(tmpdir):
    """Verify that we can successfully read a file that has no allocations."""
    # GIVEN
    output = Path(tmpdir) / "test.bin"
    subprocess_code = textwrap.dedent(
        f"""
        import os
        from memray import Tracker
        output = "{output}"
        tracker = Tracker(output)
        with tracker:
            os._exit(0)
    """
    )

    # WHEN
    process = subprocess.run([sys.executable, "-c", subprocess_code], timeout=5)

    # THEN
    assert process.returncode == 0

    records = list(FileReader(output).get_allocation_records())
    assert not records


def test_unsupported_operations_on_aggregated_capture(tmpdir):
    """Verify that we can successfully read a file that has no allocations."""
    # GIVEN
    output = Path(tmpdir) / "test.bin"
    subprocess_code = textwrap.dedent(
        f"""
        import os
        from memray import Tracker, FileFormat
        output = "{output}"
        tracker = Tracker(output, file_format=FileFormat.AGGREGATED_ALLOCATIONS)
        with tracker:
            pass
        """
    )
    process = subprocess.run([sys.executable, "-c", subprocess_code], timeout=5)
    assert process.returncode == 0
    reader = FileReader(output)

    # WHEN / THEN
    with pytest.raises(
        NotImplementedError,
        match="Can't find temporary allocations using a pre-aggregated capture file",
    ):
        list(reader.get_temporary_allocation_records())

    with pytest.raises(
        NotImplementedError,
        match="Can't get all allocations from a pre-aggregated capture file",
    ):
        list(reader.get_allocation_records())

    with pytest.raises(
        NotImplementedError,
        match="Can't compute statistics using a pre-aggregated capture file",
    ):
        compute_statistics(str(output))


@pytest.mark.parametrize(
    "file_format",
    [
        pytest.param(FileFormat.ALL_ALLOCATIONS, id="ALL_ALLOCATIONS"),
        pytest.param(FileFormat.AGGREGATED_ALLOCATIONS, id="AGGREGATED_ALLOCATIONS"),
    ],
)
class TestHighWatermark:
    def test_no_allocations_while_tracking(self, tmp_path, file_format):
        # GIVEN
        output = tmp_path / "test.bin"

        def tracking_function():
            pass

        with PrimeCaches():
            tracking_function()

        # WHEN
        with Tracker(output, native_traces=True, file_format=file_format):
            tracking_function()

        # THEN
        assert list(FileReader(output).get_high_watermark_allocation_records()) == []

    @pytest.mark.parametrize(["allocator_func", "allocator_type"], ALLOCATORS)
    def test_simple_allocation_tracking(
        self, tmp_path, allocator_func, allocator_type, file_format
    ):
        # GIVEN
        allocator = MemoryAllocator()
        output = tmp_path / "test.bin"

        # WHEN
        with Tracker(output, file_format=file_format):
            res = getattr(allocator, allocator_func)(ALLOC_SIZE)
            if res:
                allocator.free()

        if not res:
            pytest.skip(f"Allocator {allocator_func} not supported in this platform")

        # THEN
        peak_allocations_unfiltered = FileReader(
            output
        ).get_high_watermark_allocation_records()
        peak_allocations = [
            record
            for record in peak_allocations_unfiltered
            if record.size == ALLOC_SIZE
        ]
        assert len(peak_allocations) == 1, peak_allocations

        record = peak_allocations[0]
        assert record.allocator == allocator_type
        assert record.n_allocations == 1

    def test_multiple_high_watermark(self, tmp_path, file_format):
        # GIVEN
        allocator = MemoryAllocator()
        output = tmp_path / "test.bin"

        def test_function():
            for _ in range(2):
                allocator.valloc(1024)
                allocator.free()

        with PrimeCaches():
            test_function()

        # WHEN
        with Tracker(output, file_format=file_format):
            test_function()

        # THEN
        reader = FileReader(output)

        if file_format == FileFormat.ALL_ALLOCATIONS:
            all_allocations = list(
                filter_relevant_allocations(reader.get_allocation_records())
            )
            assert len(all_allocations) == 4

        peak_allocations = list(
            filter_relevant_allocations(reader.get_high_watermark_allocation_records())
        )
        assert len(peak_allocations) == 1
        record = peak_allocations[0]

        assert record.allocator == AllocatorType.VALLOC
        assert record.size == 1024
        assert record.n_allocations == 1

    def test_freed_before_high_watermark_do_not_appear(self, tmp_path, file_format):
        # GIVEN
        allocator1 = MemoryAllocator()
        allocator2 = MemoryAllocator()
        output = tmp_path / "test.bin"

        def test_function():
            allocator1.valloc(1024)
            allocator1.free()
            allocator2.valloc(2048)
            allocator2.free()

        with PrimeCaches():
            test_function()

        # WHEN
        with Tracker(output, file_format=file_format):
            test_function()

        # THEN
        reader = FileReader(output)

        if file_format == FileFormat.ALL_ALLOCATIONS:
            all_allocations = list(
                filter_relevant_allocations(reader.get_allocation_records())
            )
            assert len(all_allocations) == 4

        peak_allocations = list(
            filter_relevant_allocations(reader.get_high_watermark_allocation_records())
        )
        assert len(peak_allocations) == 1

        record = peak_allocations[0]
        assert record.allocator == AllocatorType.VALLOC
        assert record.size == 2048
        assert record.n_allocations == 1

    def test_freed_after_high_watermark_do_not_appear(self, tmp_path, file_format):
        # GIVEN
        allocator1 = MemoryAllocator()
        allocator2 = MemoryAllocator()
        output = tmp_path / "test.bin"

        def test_function():
            allocator2.valloc(2048)
            allocator2.free()
            allocator1.valloc(1024)
            allocator1.free()

        with PrimeCaches():
            test_function()

        # WHEN
        with Tracker(output, file_format=file_format):
            test_function()

        # THEN
        reader = FileReader(output)

        if file_format == FileFormat.ALL_ALLOCATIONS:
            all_allocations = list(
                filter_relevant_allocations(reader.get_allocation_records())
            )
            assert len(all_allocations) == 4

        peak_allocations = list(
            filter_relevant_allocations(reader.get_high_watermark_allocation_records())
        )
        assert len(peak_allocations) == 1

        record = peak_allocations[0]
        assert record.allocator == AllocatorType.VALLOC
        assert record.size == 2048
        assert record.n_allocations == 1

    def test_allocations_aggregation_on_same_line(self, tmp_path, file_format):
        # GIVEN
        allocators = []
        output = tmp_path / "test.bin"

        # WHEN
        with Tracker(output, file_format=file_format):
            for _ in range(2):
                allocator = MemoryAllocator()
                allocators.append(allocator)

                allocator.valloc(1024)

            for allocator in allocators:
                allocator.free()

        # THEN
        reader = FileReader(output)

        if file_format == FileFormat.ALL_ALLOCATIONS:
            all_allocations = list(
                filter_relevant_allocations(reader.get_allocation_records())
            )
            assert len(all_allocations) == 4

        peak_allocations = list(
            filter_relevant_allocations(reader.get_high_watermark_allocation_records())
        )
        assert len(peak_allocations) == 1

        record = peak_allocations[0]
        assert record.allocator == AllocatorType.VALLOC
        assert record.size == 2048
        assert record.n_allocations == 2

    def test_aggregation_same_python_stack_and_same_native_stack(
        self, tmp_path, file_format
    ):
        # GIVEN
        allocators = []
        output = tmp_path / "test.bin"

        # WHEN
        with Tracker(output, native_traces=True, file_format=file_format):
            for _ in range(2):
                allocator = MemoryAllocator()
                allocators.append(allocator)

                allocator.valloc(1024)

            for allocator in allocators:
                allocator.free()

        # THEN
        reader = FileReader(output)

        if file_format == FileFormat.ALL_ALLOCATIONS:
            all_allocations = list(
                filter_relevant_allocations(reader.get_allocation_records())
            )
            assert len(all_allocations) == 4

        peak_allocations = list(
            filter_relevant_allocations(reader.get_high_watermark_allocation_records())
        )
        assert len(peak_allocations) == 1

        record = peak_allocations[0]
        assert record.allocator == AllocatorType.VALLOC
        assert record.size == 2048
        assert record.n_allocations == 2

    def test_allocations_aggregation_on_different_lines(self, tmp_path, file_format):
        # GIVEN
        allocator1 = MemoryAllocator()
        allocator2 = MemoryAllocator()
        output = tmp_path / "test.bin"

        # WHEN
        with Tracker(output, file_format=file_format):
            allocator1.valloc(1024)
            allocator2.valloc(2048)
            allocator1.free()
            allocator2.free()

        # THEN
        reader = FileReader(output)

        if file_format == FileFormat.ALL_ALLOCATIONS:
            all_allocations = list(
                filter_relevant_allocations(reader.get_allocation_records())
            )
            assert len(all_allocations) == 4

        peak_allocations = list(
            filter_relevant_allocations(reader.get_high_watermark_allocation_records())
        )
        assert len(peak_allocations) == 2
        assert sum(record.size for record in peak_allocations) == 1024 + 2048
        assert all(record.n_allocations == 1 for record in peak_allocations)

    def test_aggregation_same_python_stack_but_different_native_stack(
        self, tmp_path, file_format
    ):
        # GIVEN
        output = tmp_path / "test.bin"

        # WHEN
        with Tracker(output, native_traces=True, file_format=file_format):
            _cython_allocate_in_two_places(ALLOC_SIZE)

        # THEN
        reader = FileReader(output)

        if file_format == FileFormat.ALL_ALLOCATIONS:
            all_allocations = list(
                filter_relevant_allocations(reader.get_allocation_records())
            )
            assert len(all_allocations) == 4

        peak_allocations = list(
            filter_relevant_allocations(reader.get_high_watermark_allocation_records())
        )
        assert len(peak_allocations) == 2
        assert (
            sum(record.size for record in peak_allocations) == ALLOC_SIZE + ALLOC_SIZE
        )
        assert all(record.n_allocations == 1 for record in peak_allocations)

    def test_non_freed_allocations_are_accounted_for(self, tmp_path, file_format):
        # GIVEN
        allocator1 = MemoryAllocator()
        allocator2 = MemoryAllocator()
        output = tmp_path / "test.bin"

        # WHEN
        with Tracker(output, file_format=file_format):
            allocator1.valloc(1024)
            allocator2.valloc(2048)
            allocator1.free()
            allocator2.free()

        # THEN
        reader = FileReader(output)

        if file_format == FileFormat.ALL_ALLOCATIONS:
            all_allocations = list(
                filter_relevant_allocations(reader.get_allocation_records())
            )
            assert len(all_allocations) == 4

        peak_allocations = list(
            filter_relevant_allocations(reader.get_high_watermark_allocation_records())
        )
        assert len(peak_allocations) == 2
        assert sum(record.size for record in peak_allocations) == 1024 + 2048
        assert all(record.n_allocations == 1 for record in peak_allocations)

    def test_final_allocation_is_peak(self, tmp_path, file_format):
        # GIVEN
        allocator1 = MemoryAllocator()
        allocator2 = MemoryAllocator()
        output = tmp_path / "test.bin"

        # WHEN
        with Tracker(output, file_format=file_format):
            allocator1.valloc(1024)
            allocator1.free()
            allocator2.valloc(2048)
        allocator2.free()

        # THEN
        reader = FileReader(output)

        if file_format == FileFormat.ALL_ALLOCATIONS:
            all_allocations = list(
                filter_relevant_allocations(reader.get_allocation_records())
            )
            assert len(all_allocations) == 3

        peak_allocations = list(
            filter_relevant_allocations(reader.get_high_watermark_allocation_records())
        )
        assert len(peak_allocations) == 1

        record = peak_allocations[0]
        assert record.n_allocations == 1
        assert record.allocator == AllocatorType.VALLOC
        assert record.size == 2048

    def test_spiky_generally_increasing_to_final_peak(self, tmp_path, file_format):
        """Checks multiple aspects with an interesting toy function."""

        # GIVEN
        def recursive(n, chunk_size):
            """Mimics generally-increasing but spiky usage"""
            if not n:
                return

            allocator = MemoryAllocator()
            print(f"+{n:>2} kB")
            allocator.valloc(n * chunk_size)

            # Don't keep allocated memory when recursing, ~50% of the calls.
            if n % 2:
                allocator.free()
                print(f"-{n:>2} kB")
                recursive(n - 1, chunk_size)
            else:
                recursive(n - 1, chunk_size)
                allocator.free()
                print(f"-{n:>2} kB")

        # WHEN
        output = tmp_path / "test.bin"
        with Tracker(output, file_format=file_format):
            recursive(10, 1024)

        # THEN
        reader = FileReader(output)

        if file_format == FileFormat.ALL_ALLOCATIONS:
            all_allocations = list(
                filter_relevant_allocations(reader.get_allocation_records())
            )
            assert len(all_allocations) == 20
            assert sum(record.size for record in all_allocations) == 56320

        peak_allocations = list(
            filter_relevant_allocations(reader.get_high_watermark_allocation_records())
        )
        assert all(record.n_allocations == 1 for record in peak_allocations)

        expected = {10, 8, 6, 4, 2, 1}
        assert len(peak_allocations) == len(expected)
        assert {record.size / 1024 for record in peak_allocations} == expected

    def test_allocations_after_high_watermark_is_freed_do_not_appear(
        self, tmp_path, file_format
    ):
        # GIVEN
        allocator = MemoryAllocator()
        output = tmp_path / "test.bin"

        def test_function():
            allocator.valloc(2048)
            allocator.free()
            allocator.valloc(1024)

        with PrimeCaches():
            test_function()
            allocator.free()

        # WHEN
        with Tracker(output, file_format=file_format):
            test_function()
        allocator.free()

        # THEN
        reader = FileReader(output)

        if file_format == FileFormat.ALL_ALLOCATIONS:
            all_allocations = list(
                filter_relevant_allocations(reader.get_allocation_records())
            )
            assert len(all_allocations) == 3

        peak_allocations = list(
            filter_relevant_allocations(reader.get_high_watermark_allocation_records())
        )
        assert len(peak_allocations) == 1

        record = peak_allocations[0]
        assert record.n_allocations == 1
        assert record.allocator == AllocatorType.VALLOC
        assert record.size == 2048

    def test_partial_munmap(self, tmp_path, file_format):
        """Partial munmap operations should be accurately tracked: we should
        only account for the removal of the actually munmap'd chunk and not
        the entire mmap'd region when a partial munmap is performed."""

        # GIVEN
        output = tmp_path / "test.bin"

        def test_function():
            # Mmap some memory and free just the first page. This should not register
            # the deallocation of the entire mmap'd region, only one page.
            alloc = MmapAllocator(2 * PAGE_SIZE)
            alloc.munmap(PAGE_SIZE)

            # Now perform the peak allocation. This should be detected as the peak index.
            MmapAllocator(10 * PAGE_SIZE)
            # At this point we should have 11 * PAGE_SIZE allocated

        with PrimeCaches():
            test_function()

        # WHEN
        with Tracker(output, file_format=file_format):
            test_function()

        # THEN
        reader = FileReader(output)
        peak_allocations = list(reader.get_high_watermark_allocation_records())
        assert len(peak_allocations) == 2
        peak_memory = sum(x.size for x in peak_allocations)
        assert peak_memory == 11 * PAGE_SIZE

    def test_partial_munmap_gap(self, tmp_path, file_format):
        """Validate that removing chunks from a mmap'd region correctly accounts
        for the parts removed. This test allocates 4 pages and removes the first
        and last pages of the mmap'd region."""

        # GIVEN
        output = tmp_path / "test.bin"

        def test_function():
            # Mmap some memory and free two pages: one at the beginning and one at the
            # end of the region.
            alloc = MmapAllocator(4 * PAGE_SIZE)
            alloc.munmap(PAGE_SIZE)
            alloc.munmap(PAGE_SIZE, 3 * PAGE_SIZE)

            # Now perform the peak allocation. This should be detected as the peak index.
            MmapAllocator(10 * PAGE_SIZE)
            # At this point we should have 12 * PAGE_SIZE allocated

        with PrimeCaches():
            test_function()

        # WHEN
        with Tracker(output, file_format=file_format):
            test_function()

        # THEN
        reader = FileReader(output)
        peak_allocations = list(reader.get_high_watermark_allocation_records())
        assert len(peak_allocations) == 2
        peak_memory = sum(x.size for x in peak_allocations)
        assert peak_memory == 12 * PAGE_SIZE

    def test_munmap_multiple_mmaps(self, tmp_path, file_format):
        """Allocate multiple contiguous mmap'd regions and then deallocate all of them
        with munmap in one go."""

        # GIVEN
        output = tmp_path / "test.bin"
        with Tracker(output, file_format=file_format):
            # Ensure we have a long enough free buffer for the contiguous
            # mmap's. We also need to make sure the allocation addresses are
            # page-aligned later, and mmap does that for us. We can then use
            # the pointer from this allocation for the actual test.
            buf = MmapAllocator(8 * PAGE_SIZE)
            buf.munmap(8 * PAGE_SIZE)

            # WHEN
            # Allocate 2 contiguous chunks of 4 pages (8 pages in total) and free them
            # with a single munmap
            alloc1 = MmapAllocator(4 * PAGE_SIZE, buf.address)
            MmapAllocator(4 * PAGE_SIZE, alloc1.address + (4 * PAGE_SIZE))
            alloc1.munmap(8 * PAGE_SIZE)

            # Now perform the peak allocation. This should be detected as the peak index.
            MmapAllocator(10 * PAGE_SIZE)

        # THEN
        reader = FileReader(output)
        peak_allocations = list(
            filter_relevant_allocations(
                reader.get_high_watermark_allocation_records(), ranged=True
            )
        )
        peak_memory = sum(x.size for x in peak_allocations)
        assert peak_memory == 10 * PAGE_SIZE

    def test_munmap_multiple_mmaps_multiple_munmaps(self, tmp_path, file_format):
        """Allocate multiple contiguous mmap'd regions and then with multiple munmap's, each
        deallocating several mmap'd areas in one go."""
        # GIVEN
        output = tmp_path / "test.bin"
        with Tracker(output, file_format=file_format):
            # Ensure we have a long enough free buffer for the contiguous
            # mmap's. We also need to make sure the allocation addresses are
            # page-aligned later, and mmap does that for us. We can then use
            # the pointer from this allocation for the actual test.
            buf = MmapAllocator(8 * PAGE_SIZE)
            buf.munmap(8 * PAGE_SIZE)

            # WHEN
            alloc1 = MmapAllocator(2 * PAGE_SIZE, buf.address)
            MmapAllocator(2 * PAGE_SIZE, buf.address + (2 * PAGE_SIZE))
            MmapAllocator(2 * PAGE_SIZE, buf.address + (4 * PAGE_SIZE))
            MmapAllocator(2 * PAGE_SIZE, buf.address + (6 * PAGE_SIZE))
            alloc1.munmap(4 * PAGE_SIZE)
            alloc1.munmap(4 * PAGE_SIZE, 4 * PAGE_SIZE)

            # Now perform the peak allocation. This should be detected as the peak index.
            MmapAllocator(10 * PAGE_SIZE)

        # THEN
        reader = FileReader(output)
        peak_allocations = list(
            filter_relevant_allocations(
                reader.get_high_watermark_allocation_records(), ranged=True
            )
        )
        peak_memory = sum(x.size for x in peak_allocations)
        assert peak_memory == 10 * PAGE_SIZE

    def test_partial_munmap_multiple_split_in_middle(self, tmp_path, file_format):
        """Deallocate pages in of a larger mmap'd area, splitting it into 3 areas."""
        # GIVEN
        output = tmp_path / "test.bin"

        def test_function():
            alloc = MmapAllocator(5 * PAGE_SIZE)
            alloc.munmap(PAGE_SIZE, 1 * PAGE_SIZE)
            alloc.munmap(PAGE_SIZE, 3 * PAGE_SIZE)
            MmapAllocator(10 * PAGE_SIZE)

        with PrimeCaches():
            test_function()

        # WHEN
        with Tracker(output, file_format=file_format):
            test_function()

        # THEN
        reader = FileReader(output)
        peak_allocations = list(reader.get_high_watermark_allocation_records())
        assert len(peak_allocations) == 2
        peak_memory = sum(x.size for x in peak_allocations)
        assert peak_memory == 13 * PAGE_SIZE

    def test_partial_munmap_split_in_middle(self, tmp_path, file_format):
        """Deallocate a single page in the middle of a larger mmap'd area."""
        # GIVEN
        output = tmp_path / "test.bin"

        def test_function():
            alloc = MmapAllocator(8 * PAGE_SIZE)
            alloc.munmap(PAGE_SIZE, 4 * PAGE_SIZE)

            MmapAllocator(10 * PAGE_SIZE)

        with PrimeCaches():
            test_function()

        # WHEN
        with Tracker(output, file_format=file_format):
            test_function()

        # THEN
        reader = FileReader(output)
        peak_allocations = list(reader.get_high_watermark_allocation_records())
        assert len(peak_allocations) == 2
        peak_memory = sum(x.size for x in peak_allocations)
        assert peak_memory == 17 * PAGE_SIZE


@pytest.mark.parametrize(
    "file_format",
    [
        pytest.param(FileFormat.ALL_ALLOCATIONS, id="ALL_ALLOCATIONS"),
        pytest.param(FileFormat.AGGREGATED_ALLOCATIONS, id="AGGREGATED_ALLOCATIONS"),
    ],
)
class TestLeaks:
    def test_leaks_allocations_are_detected(self, tmp_path, file_format):
        # GIVEN
        allocator = MemoryAllocator()
        output = tmp_path / "test.bin"

        # WHEN
        with Tracker(output, file_format=file_format):
            allocator.valloc(1024)
        allocator.free()

        # THEN
        reader = FileReader(output)

        if file_format == FileFormat.ALL_ALLOCATIONS:
            all_allocations = list(
                filter_relevant_allocations(reader.get_allocation_records())
            )
            assert len(all_allocations) == 1

        leaked_allocations = list(
            filter_relevant_allocations(reader.get_leaked_allocation_records())
        )
        assert len(leaked_allocations) == 1

        record = leaked_allocations[0]
        assert record.n_allocations == 1
        assert record.allocator == AllocatorType.VALLOC
        assert record.size == 1024

    def test_allocations_that_are_freed_do_not_appear_as_leaks(
        self, tmp_path, file_format
    ):
        # GIVEN
        allocator = MemoryAllocator()
        output = tmp_path / "test.bin"

        # WHEN
        with Tracker(output, file_format=file_format):
            allocator.valloc(1024)
            allocator.free()
            allocator.valloc(1024)
            allocator.free()
            allocator.valloc(1024)
            allocator.free()

        # THEN
        reader = FileReader(output)

        if file_format == FileFormat.ALL_ALLOCATIONS:
            all_allocations = list(
                filter_relevant_allocations(reader.get_allocation_records())
            )
            assert len(all_allocations) == 6

        leaked_allocations = list(
            filter_relevant_allocations(reader.get_leaked_allocation_records())
        )
        assert not leaked_allocations

    def test_leak_that_happens_in_the_middle_is_detected(self, tmp_path, file_format):
        # GIVEN
        allocator = MemoryAllocator()
        leak_allocator = MemoryAllocator()
        output = tmp_path / "test.bin"

        # WHEN
        with Tracker(output, file_format=file_format):
            allocator.valloc(1024)
            allocator.free()
            allocator.valloc(1024)
            leak_allocator.valloc(2048)
            allocator.free()
            allocator.valloc(1024)
            allocator.free()
        leak_allocator.free()

        # THEN
        reader = FileReader(output)

        if file_format == FileFormat.ALL_ALLOCATIONS:
            all_allocations = list(
                filter_relevant_allocations(reader.get_allocation_records())
            )
            assert len(all_allocations) == 7

        leaked_allocations = list(
            filter_relevant_allocations(reader.get_leaked_allocation_records())
        )

        assert len(leaked_allocations) == 1

        record = leaked_allocations[0]
        assert record.n_allocations == 1
        assert record.allocator == AllocatorType.VALLOC
        assert record.size == 2048

    def test_leaks_that_happen_in_different_lines(self, tmp_path, file_format):
        # GIVEN
        allocator1 = MemoryAllocator()
        allocator2 = MemoryAllocator()
        output = tmp_path / "test.bin"

        # WHEN
        with Tracker(output, file_format=file_format):
            allocator1.valloc(1024)
            allocator2.valloc(2048)

        allocator1.free()
        allocator2.free()

        # THEN
        leaked_allocations = list(
            filter_relevant_allocations(
                FileReader(output).get_leaked_allocation_records()
            )
        )
        assert len(leaked_allocations) == 2
        assert sum(record.size for record in leaked_allocations) == 1024 + 2048
        assert all(record.n_allocations == 1 for record in leaked_allocations)

    def test_leaks_that_happen_in_the_same_function_are_aggregated(
        self, tmp_path, file_format
    ):
        # GIVEN
        allocators = []
        output = tmp_path / "test.bin"

        def foo():
            allocator = MemoryAllocator()
            allocator.valloc(1024)
            allocators.append(allocator)

        # WHEN
        with Tracker(output, file_format=file_format):
            for _ in range(10):
                foo()

        for allocator in allocators:
            allocator.free()

        # THEN
        reader = FileReader(output)

        if file_format == FileFormat.ALL_ALLOCATIONS:
            all_allocations = list(
                filter_relevant_allocations(reader.get_allocation_records())
            )
            assert len(all_allocations) == 10

        leaked_allocations = list(
            filter_relevant_allocations(reader.get_leaked_allocation_records())
        )
        assert len(leaked_allocations) == 1
        (allocation,) = leaked_allocations
        assert allocation.size == 1024 * 10
        assert allocation.n_allocations == 10

    def test_unmatched_deallocations_are_not_reported(self, tmp_path, file_format):
        # GIVEN
        allocator = MemoryAllocator()
        output = tmp_path / "test.bin"

        # WHEN
        allocator.valloc(ALLOC_SIZE)
        with Tracker(output, file_format=file_format):
            allocator.free()

        # THEN
        reader = FileReader(output)

        if file_format == FileFormat.ALL_ALLOCATIONS:
            all_allocations = list(reader.get_allocation_records())
            assert len(all_allocations) >= 1

        assert not list(
            filter_relevant_allocations(reader.get_leaked_allocation_records())
        )

    def test_thread_allocations_multiple_threads(self, tmpdir, file_format):
        # GIVEN
        def allocating_function(allocator, amount, stop_flag):
            allocator.posix_memalign(amount)
            allocator.posix_memalign(amount)
            # We need a barrier as pthread might reuse the same thread ID if the first thread
            # finishes before the other one starts
            stop_flag.wait()

        # WHEN
        alloc1 = MemoryAllocator()
        stop_flag1 = threading.Event()
        alloc2 = MemoryAllocator()
        stop_flag2 = threading.Event()
        output = Path(tmpdir) / "test.bin"
        with Tracker(output, file_format=file_format):
            t1 = threading.Thread(
                target=allocating_function, args=(alloc1, 2048, stop_flag1)
            )
            t1.start()
            t2 = threading.Thread(
                target=allocating_function, args=(alloc2, 2048, stop_flag2)
            )
            t2.start()

            stop_flag1.set()
            t1.join()
            stop_flag2.set()
            t2.join()

        # THEN
        reader = FileReader(output)

        if file_format == FileFormat.ALL_ALLOCATIONS:
            all_allocations = [
                record
                for record in reader.get_allocation_records()
                if record.allocator == AllocatorType.POSIX_MEMALIGN
            ]
            assert len(all_allocations) == 4

        high_watermark_records = (
            record
            for record in reader.get_high_watermark_allocation_records(
                merge_threads=False
            )
            if record.allocator == AllocatorType.POSIX_MEMALIGN
        )
        # Group the allocations per thread
        records = collections.defaultdict(list)
        for record in high_watermark_records:
            records[record.tid].append(record)
        assert len(records.keys()) == 2
        # Each thread should have 4096 bytes total allocations
        for allocations in records.values():
            assert sum(allocation.size for allocation in allocations) == 4096


class TestTemporaryAllocations:
    def test_temporary_allocations_are_detected(self, tmp_path):
        # GIVEN
        allocator = MemoryAllocator()
        output = tmp_path / "test.bin"

        # WHEN
        with Tracker(output):
            allocator.valloc(1024)
            allocator.free()

        # THEN
        reader = FileReader(output)
        all_allocations = list(
            filter_relevant_allocations(reader.get_allocation_records())
        )
        assert len(all_allocations) == 2

        temporary_allocations = list(
            filter_relevant_allocations(reader.get_temporary_allocation_records())
        )
        assert len(temporary_allocations) == 1

        record = temporary_allocations[0]
        assert record.n_allocations == 1
        assert record.allocator == AllocatorType.VALLOC
        assert record.size == 1024

    def test_temporary_allocations_with_two_allocators_are_detected(self, tmp_path):
        # GIVEN
        allocator1 = MemoryAllocator()
        allocator2 = MemoryAllocator()
        output = tmp_path / "test.bin"

        # WHEN
        with Tracker(output):
            allocator1.valloc(1024)
            allocator1.free()
            allocator2.valloc(1024)
            allocator2.free()
        # THEN
        reader = FileReader(output)
        all_allocations = list(
            filter_relevant_allocations(reader.get_allocation_records())
        )
        assert len(all_allocations) == 4

        temporary_allocations = list(
            filter_relevant_allocations(reader.get_temporary_allocation_records())
        )
        assert len(temporary_allocations) == 2

        for record in temporary_allocations:
            assert record.n_allocations == 1
            assert record.allocator == AllocatorType.VALLOC
            assert record.size == 1024

    @pytest.mark.parametrize(
        "buffer_size",
        [
            1,
            2,
            5,
            10,
        ],
    )
    def test_temporary_allocations_outside_buffer_are_not_detected(
        self, tmp_path, buffer_size
    ):
        # GIVEN
        allocators = [MemoryAllocator() for _ in range(buffer_size + 1)]
        output = tmp_path / "test.bin"

        # WHEN
        with Tracker(output):
            for allocator in allocators:
                allocator.valloc(1024)
            for allocator in reversed(allocators):
                allocator.free()

        # THEN
        reader = FileReader(output)
        all_allocations = list(
            filter_relevant_allocations(reader.get_allocation_records())
        )
        assert len(all_allocations) == 2 * (buffer_size + 1)

        temporary_allocations = list(
            filter_relevant_allocations(
                reader.get_temporary_allocation_records(threshold=buffer_size - 1)
            )
        )
        assert len(temporary_allocations) == 1
        (allocation,) = temporary_allocations
        assert allocation.n_allocations == buffer_size

    def test_temporary_allocations_that_happen_in_different_lines(self, tmp_path):
        # GIVEN
        allocator1 = MemoryAllocator()
        allocator2 = MemoryAllocator()
        output = tmp_path / "test.bin"

        # WHEN
        with Tracker(output):
            allocator1.valloc(1024)
            allocator1.free()
            allocator2.valloc(2048)
            allocator2.free()

        # THEN
        temporary_allocations = list(
            filter_relevant_allocations(
                FileReader(output).get_temporary_allocation_records()
            )
        )
        assert len(temporary_allocations) == 2
        assert sum(record.size for record in temporary_allocations) == 1024 + 2048
        assert all(record.n_allocations == 1 for record in temporary_allocations)

    def test_temporary_allocations_that_happen_in_the_same_function_are_aggregated(
        self, tmp_path
    ):
        # GIVEN
        output = tmp_path / "test.bin"

        def foo():
            allocator = MemoryAllocator()
            allocator.valloc(1024)
            allocator.free()

        # WHEN
        with Tracker(output):
            for _ in range(10):
                foo()

        # THEN
        reader = FileReader(output)
        all_allocations = list(
            filter_relevant_allocations(reader.get_allocation_records())
        )
        assert len(all_allocations) == 10 + 10  # 10 x valloc + 10 x free

        temporary_allocations = list(
            filter_relevant_allocations(reader.get_temporary_allocation_records())
        )
        assert len(temporary_allocations) == 1
        (allocation,) = temporary_allocations
        assert allocation.size == 1024 * 10
        assert allocation.n_allocations == 10

    def test_unmatched_allocations_are_not_reported(self, tmp_path):
        # GIVEN
        allocator = MemoryAllocator()
        output = tmp_path / "test.bin"

        # WHEN
        with Tracker(output):
            allocator.valloc(ALLOC_SIZE)
        allocator.free()

        # THEN
        reader = FileReader(output)
        all_allocations = list(reader.get_allocation_records())
        assert len(all_allocations) >= 1
        assert not list(
            filter_relevant_allocations(reader.get_temporary_allocation_records())
        )

    def test_thread_allocations_multiple_threads(self, tmpdir):
        # GIVEN
        def allocating_function(allocator, amount, stop_flag):
            allocator.posix_memalign(amount)
            allocator.free()
            allocator.posix_memalign(amount)
            allocator.free()
            # We need a barrier as pthread might reuse the same thread ID if the first thread
            # finishes before the other one starts
            stop_flag.wait()

        # WHEN
        alloc1 = MemoryAllocator()
        stop_flag1 = threading.Event()
        alloc2 = MemoryAllocator()
        stop_flag2 = threading.Event()
        output = Path(tmpdir) / "test.bin"
        with Tracker(output):
            t1 = threading.Thread(
                target=allocating_function, args=(alloc1, 2048, stop_flag1)
            )
            t1.start()
            t2 = threading.Thread(
                target=allocating_function, args=(alloc2, 2048, stop_flag2)
            )
            t2.start()

            stop_flag1.set()
            t1.join()
            stop_flag2.set()
            t2.join()

        # THEN
        reader = FileReader(output)
        all_allocations = [
            record
            for record in reader.get_allocation_records()
            if record.allocator == AllocatorType.POSIX_MEMALIGN
        ]
        assert len(all_allocations) == 4

        temporary_allocation_records = (
            record
            for record in reader.get_temporary_allocation_records(merge_threads=False)
            if record.allocator == AllocatorType.POSIX_MEMALIGN
        )
        # Group the allocations per thread
        records = collections.defaultdict(list)
        for record in temporary_allocation_records:
            records[record.tid].append(record)
        assert len(records.keys()) == 2
        # Each thread should have 4096 bytes total allocations
        for allocations in records.values():
            assert sum(allocation.size for allocation in allocations) == 4096

    def test_intertwined_temporary_allocations_in_threads(self, tmpdir):
        """This test checks that temporary allocations are correctly detected
        when they happen in different threads with interleaved calls to malloc
        and free.

        Note that there is not an easy way to synchronize this test because
        using condition variables may allocate, disrupting the test as there
        will be other spurious allocations between our call to posix_memalign
        and the call to free.

        The best effort is that both threads are synchronized by spinlocking
        around a python boolean that uses the GIL as means to ensure the correct
        ordering of allocations.
        """

        # GIVEN
        thread1_allocated = False
        thread1_should_free = False

        def thread1_body(allocator):
            nonlocal thread1_allocated
            allocator.posix_memalign(1234)
            thread1_allocated = True
            while not thread1_should_free:
                pass
            allocator.free()

        def thread2_body(allocator):
            nonlocal thread1_should_free
            while not thread1_allocated:
                pass
            allocator.posix_memalign(1234)
            allocator.free()
            thread1_should_free = True

        # WHEN
        alloc1 = MemoryAllocator()
        alloc2 = MemoryAllocator()
        output = Path(tmpdir) / "test.bin"
        with Tracker(output):
            t1 = threading.Thread(target=thread1_body, args=[alloc1])
            t1.start()
            t2 = threading.Thread(
                target=thread2_body,
                args=[alloc2],
            )
            t2.start()

            t1.join()
            t2.join()

        # THEN
        reader = FileReader(output)
        all_allocations = [
            record
            for record in reader.get_allocation_records()
            if record.allocator == AllocatorType.POSIX_MEMALIGN
        ]
        assert len(all_allocations) == 2

        temporary_allocation_records = (
            record
            for record in reader.get_temporary_allocation_records(merge_threads=False)
            if record.allocator == AllocatorType.POSIX_MEMALIGN
        )
        # Group the allocations per thread
        records = collections.defaultdict(list)
        for record in temporary_allocation_records:
            records[record.tid].append(record)
        assert len(records.keys()) == 2
        # Each thread should have 1234 bytes total allocations
        for allocations in records.values():
            assert sum(allocation.size for allocation in allocations) == 1234


class TestHeader:
    def test_get_header(self, monkeypatch, tmpdir):
        # GIVEN
        allocator = MemoryAllocator()
        output = Path(tmpdir) / "test.bin"

        with PrimeCaches():
            pass

        # WHEN

        monkeypatch.setattr(sys, "argv", ["python", "-m", "pytest"])
        with run_without_tracer():
            start_time = datetime.datetime.now().astimezone()
            with Tracker(output):
                for _ in range(100):
                    allocator.valloc(1024)
            end_time = datetime.datetime.now().astimezone()

        reader = FileReader(output)
        n_records = len(list(reader.get_allocation_records()))
        metadata = reader.metadata

        # THEN
        assert metadata.end_time > metadata.start_time
        assert abs(metadata.start_time - start_time).seconds < 1
        assert abs(metadata.end_time - end_time).seconds < 1
        assert metadata.total_allocations == n_records
        assert metadata.command_line == "python -m pytest"
        assert metadata.peak_memory == 1024 * 100

    def test_get_header_after_snapshot(self, monkeypatch, tmpdir):
        """Verify that we can successfully retrieve the metadata after querying
        the high watermark snapshot."""
        # GIVEN
        allocator = MemoryAllocator()
        output = Path(tmpdir) / "test.bin"

        with PrimeCaches():
            pass

        # WHEN

        monkeypatch.setattr(sys, "argv", ["python", "-m", "pytest"])
        start_time = datetime.datetime.now().astimezone()
        with Tracker(output):
            for _ in range(100):
                allocator.valloc(1024)
        end_time = datetime.datetime.now().astimezone()

        reader = FileReader(output)
        peak, *_ = list(reader.get_high_watermark_allocation_records())
        metadata = reader.metadata

        # THEN
        assert metadata.end_time > metadata.start_time
        assert abs(metadata.start_time - start_time).seconds < 1
        assert abs(metadata.end_time - end_time).seconds < 1
        assert metadata.total_allocations == peak.n_allocations
        assert metadata.command_line == "python -m pytest"
        assert metadata.peak_memory == 1024 * 100

    @pytest.mark.parametrize(
        "allocator, allocator_name",
        [
            ("malloc", "malloc"),
            ("pymalloc", "pymalloc"),
            ("pymalloc_debug", "pymalloc debug"),
        ],
    )
    def test_header_allocator(self, allocator, allocator_name, tmpdir):
        # GIVEN
        output = Path(tmpdir) / "test.bin"

        # WHEN

        subprocess_code = textwrap.dedent(
            f"""
        from memray import Tracker
        from memray._test import MemoryAllocator
        allocator = MemoryAllocator()

        with Tracker('{output}'):
            allocator.valloc(1024)
            allocator.free()
        """
        )

        subprocess.run(
            [sys.executable, "-c", subprocess_code],
            timeout=5,
            env={"PYTHONMALLOC": allocator},
        )

        reader = FileReader(output)
        metadata = reader.metadata

        # THEN
        assert metadata.python_allocator == allocator_name


class TestMemorySnapshots:
    @pytest.mark.valgrind
    def test_memory_snapshots_are_written(self, tmp_path):
        # GIVEN
        allocator = MemoryAllocator()
        output = tmp_path / "test.bin"

        # WHEN
        with Tracker(output):
            allocator.valloc(ALLOC_SIZE)
            time.sleep(0.11)
            allocator.free()

        memory_snapshots = list(FileReader(output).get_memory_snapshots())

        assert memory_snapshots
        assert all(record.rss > 0 for record in memory_snapshots)
        assert any(record.heap >= ALLOC_SIZE for record in memory_snapshots)
        assert sorted(memory_snapshots, key=lambda r: r.time) == memory_snapshots
        assert all(
            _next.time - prev.time >= 10
            for prev, _next in zip(memory_snapshots, memory_snapshots[1:])
        )

    def test_memory_snapshots_tick_interval(self, tmp_path):
        # GIVEN
        allocator = MemoryAllocator()
        output = tmp_path / "test.bin"

        # WHEN
        with Tracker(output, memory_interval_ms=20):
            allocator.valloc(ALLOC_SIZE)
            time.sleep(1)

        memory_snapshots = list(FileReader(output).get_memory_snapshots())

        assert len(memory_snapshots)
        assert all(record.rss > 0 for record in memory_snapshots)
        assert any(record.heap >= ALLOC_SIZE for record in memory_snapshots)
        assert sorted(memory_snapshots, key=lambda r: r.time) == memory_snapshots
        assert all(
            _next.time - prev.time >= 20
            for prev, _next in zip(memory_snapshots, memory_snapshots[1:])
        )

    def test_memory_snapshots_limit_when_reading(self, tmp_path):
        # GIVEN
        allocator = MemoryAllocator()
        output = tmp_path / "test.bin"

        # WHEN
        with Tracker(output):
            for _ in range(2):
                allocator.valloc(ALLOC_SIZE)
                time.sleep(0.11)
                allocator.free()

        reader = FileReader(output)
        memory_snapshots = list(reader.get_memory_snapshots())
        temporal_records = list(reader.get_temporal_allocation_records())

        assert memory_snapshots
        n_snapshots = len(memory_snapshots)
        n_temporal_records = len(temporal_records)

        reader = FileReader(output, max_memory_records=n_snapshots // 2)
        memory_snapshots = list(reader.get_memory_snapshots())
        temporal_records = list(reader.get_temporal_allocation_records())

        assert memory_snapshots
        assert len(memory_snapshots) <= n_snapshots // 2 + 1
        assert len(temporal_records) <= n_temporal_records // 2 + 1

    def test_temporary_allocations_when_filling_vector_without_preallocating(
        self, tmp_path
    ):
        # GIVEN
        output = tmp_path / "test.bin"

        # WHEN
        with Tracker(output):
            elements = fill_cpp_vector(2 << 10)

        # THEN
        reader = FileReader(output)
        temporary_allocations = [
            alloc
            for alloc in reader.get_temporary_allocation_records(threshold=1)
            if __file__ in alloc.stack_trace()[0][1]
        ]
        assert elements == 512
        assert len(temporary_allocations) == 1
        (record,) = temporary_allocations
        assert record.n_allocations == 10
        assert record.allocator == AllocatorType.MALLOC
        assert record.size >= (2 << 10)

    def test_temporary_allocations_when_filling_vector_without_preallocating_small_buffer(
        self, tmp_path
    ):
        # GIVEN
        output = tmp_path / "test.bin"

        # WHEN
        with Tracker(output):
            elements = fill_cpp_vector(2 << 10)

        # THEN
        reader = FileReader(output)
        temporary_allocations = [
            alloc
            for alloc in reader.get_temporary_allocation_records(threshold=0)
            if __file__ in alloc.stack_trace()[0][1]
        ]
        assert elements == 512
        assert len(temporary_allocations) == 1
        (record,) = temporary_allocations
        # With a buffer of 1 we only see the last allocation
        assert record.n_allocations == 1
        assert record.allocator == AllocatorType.MALLOC
        assert record.size == 2 << 10