dial9 0.5.0-rc1

#![cfg(feature = "memory-profiling")]
#![cfg(feature = "analysis")]
#![cfg(target_os = "linux")]
//! Test that realloc emits both AllocEvent and FreeEvent when liveset is on.

mod common;

use common::{CAPTURE_BUFFER_SIZE, capture_processor, decode_all};
use dial9::analysis::analysis_events::Dial9Event;
use dial9::memory::{Dial9Allocator, MemoryProfiler, MemoryProfilingConfig};
use dial9::{MemoryBuffer, RecorderPipelineExt, RecorderTokioExt, recorder};
use std::alloc::{GlobalAlloc, Layout, System};
use std::sync::atomic::{AtomicUsize, Ordering};
use std::time::Duration;

static ALLOC_COUNT: AtomicUsize = AtomicUsize::new(0);

struct CountingAllocator;

impl CountingAllocator {
    const fn new() -> Self {
        Self
    }
}

unsafe impl GlobalAlloc for CountingAllocator {
    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
        ALLOC_COUNT.fetch_add(1, Ordering::Relaxed);
        unsafe { System.alloc(layout) }
    }
    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
        unsafe { System.dealloc(ptr, layout) }
    }
}

#[global_allocator]
static ALLOC: Dial9Allocator<CountingAllocator> = Dial9Allocator::new(CountingAllocator::new());

#[test]
fn hook_realloc_emits_alloc_and_free_when_liveset_on() {
    let (capture, batches) = capture_processor();

    let recorder = recorder(MemoryBuffer::new(CAPTURE_BUFFER_SIZE).unwrap())
        .with_custom_pipeline(|p| p.pipe(capture))
        .build();
    let (recorder, rt) = recorder
        .attach_tokio_runtime(|t| {
            t.enable_all();
            t.worker_threads(1);
        })
        .expect("attach tokio");

    let handle = recorder.handle().clone();
    let _mem_guard = MemoryProfiler::from_config(
        MemoryProfilingConfig::builder()
            .sample_rate_bytes(64)
            .track_liveset(true)
            .rng_seed(42)
            .build(),
    )
    .install(handle)
    .expect("install should succeed");

    rt.block_on(async {
        let mut v: Vec<u8> = Vec::new();
        for i in 0..1000u16 {
            v.push((i & 0xff) as u8);
        }
        std::hint::black_box(&v);
        drop(v);

        tokio::time::sleep(Duration::from_millis(200)).await;
    });

    drop(rt);
    recorder.graceful_shutdown(Duration::from_secs(1));

    assert!(
        ALLOC_COUNT.load(Ordering::Relaxed) > 0,
        "expected CountingAllocator to have been called"
    );

    let b = batches.lock().unwrap();
    let events: Vec<Dial9Event> = decode_all(&b);
    let allocs: Vec<_> = events
        .iter()
        .filter(|e| matches!(e, Dial9Event::AllocEvent(_)))
        .collect();
    let frees: Vec<_> = events
        .iter()
        .filter(|e| matches!(e, Dial9Event::FreeEvent(_)))
        .collect();

    assert!(
        !allocs.is_empty(),
        "expected at least one AllocEvent from realloc path"
    );
    assert!(
        !frees.is_empty(),
        "expected at least one FreeEvent when liveset tracking is on"
    );

    // Verify every FreeEvent matches an AllocEvent on the allocation's full
    // identity: (addr, size, alloc_timestamp_ns). The FreeEvent denormalizes
    // all three fields from the same producer-side liveset entry that the
    // matching `on_alloc` inserted, so `(addr, alloc_timestamp_ns)` uniquely
    // identifies the allocation it refers to.
    //
    // We match against *all* AllocEvents, regardless of their position in the
    // stream — NOT by replaying events in order and tracking "what is live at
    // this addr right now." The consolidator only guarantees *approximate*
    // timestamp ordering: it best-effort merge-drains two independent producer
    // queues, possibly across flush cycles (see `MemoryProfileSource::flush`).
    // Meanwhile the growing `Vec` makes the system allocator reuse freed
    // addresses, so a reuse's AllocEvent can appear in the stream *before* the
    // prior occupant's FreeEvent. An addr-keyed, in-order replay would then
    // clobber the live entry and match the free to the wrong allocation —
    // the source of this test's former flakiness. Keying on the alloc's full
    // identity makes the check order-independent.
    use std::collections::HashMap;
    let alloc_sizes: HashMap<(u64, u64), u64> = events
        .iter()
        .filter_map(|e| match e {
            Dial9Event::AllocEvent(a) => Some(((a.addr, a.timestamp_ns), a.size)),
            _ => None,
        })
        .collect();

    for free in &frees {
        let Dial9Event::FreeEvent(f) = free else {
            unreachable!("frees was filtered to FreeEvent above");
        };
        let Some(&alloc_size) = alloc_sizes.get(&(f.addr, f.alloc_timestamp_ns)) else {
            panic!(
                "FreeEvent at addr {:#x} (size={}, alloc_ts={}) has no matching \
                 AllocEvent with that (addr, timestamp)",
                f.addr, f.size, f.alloc_timestamp_ns
            );
        };
        assert_eq!(
            alloc_size, f.size,
            "FreeEvent at addr {:#x} reports size {} but its matching AllocEvent \
             (alloc_ts={}) had size {alloc_size}",
            f.addr, f.size, f.alloc_timestamp_ns
        );
    }
}