velo 0.12.0

// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0

//! Tests for the RDMA registration layer.
//!
//! Two harnesses, deliberately:
//!
//! * [`MockBackend`] — an in-process [`RdmaBackend`] that mints ids and fake
//!   keys. Everything about the pool, the budget, the guard lifecycle and the
//!   shutdown ordering is *velo* logic, and testing it against a mock makes
//!   those tests fast, deterministic, and able to inject failures UCX would
//!   never produce on demand (a `ShuttingDown` unmap, a slow map).
//! * A real [`UcxTransport`] pair over `UCX_TLS=tcp` — the same lane CI runs
//!   without RDMA hardware. These prove the wiring: that the projection of
//!   `RmaError`, the region ids, and the shutdown ordering hold against the
//!   actual progress thread, asserted through its own `live_regions` count
//!   rather than through bookkeeping this module also owns.
//!
//! A mock-only suite would test the layer against its own assumptions; a
//! UCX-only suite could not reach the failure paths. Both, then.

use std::sync::Arc;
use std::sync::atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering};
use std::time::Duration;

use bytes::Bytes;
use futures::future::BoxFuture;
use velo_ext::Transport;

use super::arena::{ArenaSet, Budget, GRANULE, RdmaPoolConfig, pool_arena_target};
use super::backend::{BackendGet, BackendRegion, RdmaBackend, RdmaError, UcxBackend};
use super::region::{Deregistered, RegionInner, RegionParts};
use super::{RdmaConfig, RdmaRegistry};

/// Generous ceiling for anything that should resolve promptly.
const T: Duration = Duration::from_secs(10);

/// Serialises the tests that assert on [`LEAKED_BUFFERS`].
///
/// The counter is process-global and `cargo test` runs these in parallel. One
/// test asserts the count *rose*; the other asserts it did **not** rise across
/// its whole body. The second is the fragile one: the first test dropping its
/// runtime mid-way through it would push the count up and fail an assertion
/// about a leak that is not its own. That is a false failure, not a false pass,
/// so it cannot hide a bug — but a flaky guard on a safety property is worth
/// less than a serialised one.
/// Async-aware, because one of the two tests holds it across awaits.
static LEAK_COUNTER: std::sync::LazyLock<tokio::sync::Mutex<()>> =
    std::sync::LazyLock::new(|| tokio::sync::Mutex::new(()));

/// An [`RdmaBackend`] that registers nothing and remembers everything.
///
/// `map` mints an id and a plausible packed key; `unmap` records the id. The
/// injectable knobs exist because the interesting registry behaviour is what it
/// does when the backend misbehaves, and UCX cannot be asked to misbehave on
/// cue.
struct MockBackend {
    next_id: AtomicU64,
    mapped: dashmap::DashMap<u64, (usize, usize)>,
    unmapped: AtomicUsize,
    /// When set, every `unmap` answers `ShuttingDown` — the case where the
    /// pages are *not* known to be released.
    refuse_unmap: AtomicBool,
    /// Artificial delay on `map`, for racing a registration against shutdown.
    map_delay: parking_lot::Mutex<Option<Duration>>,
    /// Extra bytes added to every reported `effective_len`.
    ///
    /// Lets a test produce a backend that pinned *more* than this side
    /// estimated — a larger page size, an implementation that rounds further
    /// out. Without it the top-up path in `Reservation::raise_to` is
    /// unreachable through the mock, because the mock rounds exactly the way
    /// the local estimate does.
    extra_effective: AtomicU64,
}

impl MockBackend {
    fn new() -> Arc<Self> {
        Arc::new(Self {
            next_id: AtomicU64::new(1),
            mapped: dashmap::DashMap::new(),
            unmapped: AtomicUsize::new(0),
            refuse_unmap: AtomicBool::new(false),
            map_delay: parking_lot::Mutex::new(None),
            extra_effective: AtomicU64::new(0),
        })
    }

    /// Registrations the backend still believes it holds.
    fn live(&self) -> usize {
        self.mapped.len()
    }

    fn unmap_calls(&self) -> usize {
        self.unmapped.load(Ordering::SeqCst)
    }

    /// Stand in for transport teardown force-unmapping everything it held.
    ///
    /// Not an `unmap` — it records no call and answers nobody. It is the
    /// backend simply ceasing to hold anything, which is what a torn-down
    /// progress thread looks like from outside.
    fn force_teardown(&self) {
        self.mapped.clear();
    }
}

impl RdmaBackend for MockBackend {
    fn key(&self) -> &str {
        "mock"
    }

    fn map(&self, ptr: usize, len: usize) -> BoxFuture<'_, Result<BackendRegion, RdmaError>> {
        Box::pin(async move {
            let delay = *self.map_delay.lock();
            if let Some(delay) = delay {
                tokio::time::sleep(delay).await;
            }
            let id = self.next_id.fetch_add(1, Ordering::SeqCst);
            self.mapped.insert(id, (ptr, len));
            Ok(BackendRegion {
                backend_region_id: id,
                // Rounded outward, exactly as a real registration would be, so
                // anything that mistakes the effective range for the requested
                // one is visible here rather than only on hardware.
                effective_addr: (ptr & !(GRANULE - 1)) as u64,
                effective_len: len.next_multiple_of(GRANULE) as u64
                    + self.extra_effective.load(Ordering::SeqCst),
                packed_key: Bytes::from_static(b"mock-packed-key"),
            })
        })
    }

    fn unmap(&self, backend_region_id: u64) -> BoxFuture<'_, Result<(), RdmaError>> {
        Box::pin(async move {
            self.unmapped.fetch_add(1, Ordering::SeqCst);
            if self.refuse_unmap.load(Ordering::SeqCst) {
                return Err(RdmaError::ShuttingDown);
            }
            // Idempotent, per the trait contract: an id that names nothing is
            // the state the caller asked for.
            self.mapped.remove(&backend_region_id);
            Ok(())
        })
    }

    fn live_registrations(&self) -> Option<usize> {
        Some(self.mapped.len())
    }

    fn get(&self, _req: BackendGet) -> BoxFuture<'_, Result<(), RdmaError>> {
        Box::pin(async move { Ok(()) })
    }
}

/// A pool over a mock backend with test-sized arenas.
fn mock_pool(cfg: RdmaPoolConfig) -> (Arc<MockBackend>, ArenaSet) {
    let backend = MockBackend::new();
    let budget = Arc::new(Budget::new(cfg.registered_bytes_budget, None));
    let pool = ArenaSet::new(
        Arc::clone(&backend) as Arc<dyn RdmaBackend>,
        cfg,
        Arc::clone(&budget),
        Arc::new(AtomicU64::new(1)),
        None,
    );
    (backend, pool)
}

/// Arena sizes small enough that growth and exhaustion happen in a test rather
/// than after 64 MiB of allocation.
fn small_pool_config() -> RdmaPoolConfig {
    RdmaPoolConfig {
        initial_arena_bytes: 64 * GRANULE as u64,
        max_arena_bytes: 256 * GRANULE as u64,
        dedicated_arena_min: 128 * GRANULE as u64,
        registered_bytes_budget: 1024 * GRANULE as u64,
        // Reclamation off and retention irrelevant, matching the production
        // default: the tests that exercise the sweep opt in explicitly.
        arena_reclaim_after: None,
        retain_arena_bytes: 64 * GRANULE as u64,
    }
}

/// A registry over a mock backend, for lifecycle tests that need no UCX.
fn mock_registry(cfg: RdmaConfig) -> (Arc<MockBackend>, Arc<RdmaRegistry>) {
    let backend = MockBackend::new();
    let registry = Arc::new(RdmaRegistry::new(
        Arc::clone(&backend) as Arc<dyn RdmaBackend>,
        cfg,
        tokio::runtime::Handle::current(),
        None,
    ));
    (backend, registry)
}

// ---------------------------------------------------------------------------
// Pool
// ---------------------------------------------------------------------------

/// A round trip through the pool: the buffer is writable, addresses inside the
/// arena, and describes itself consistently.
#[tokio::test]
async fn pool_alloc_roundtrip() {
    let (backend, pool) = mock_pool(small_pool_config());

    let mut buf = pool.alloc(4000).await.expect("alloc");
    assert_eq!(
        buf.len(),
        4000,
        "the exact requested length is what is handed out"
    );
    buf.fill(0xAB);
    assert!(buf.iter().all(|b| *b == 0xAB));

    let remote = buf.remote();
    assert_eq!(remote.addr, buf.addr());
    assert_eq!(remote.len, 4000);
    assert_eq!(&remote.packed_key[..], b"mock-packed-key");

    assert_eq!(
        backend.live(),
        1,
        "one arena maps once, however many buffers come out of it"
    );
    assert_eq!(pool.arena_count(), 1);
    assert_eq!(pool.live_allocations(), 1);

    drop(buf);
    assert_eq!(pool.live_allocations(), 0);
    assert_eq!(
        backend.unmap_calls(),
        0,
        "returning a suballocation must never touch the backend; the arena stays registered"
    );
}

/// Sub-granule requests still consume a whole granule, and the reported length
/// stays the requested one — the rounding must not leak into the descriptor.
#[tokio::test]
async fn pool_rounds_to_granules() {
    let (_backend, pool) = mock_pool(small_pool_config());

    let a = pool.alloc(1).await.expect("alloc a");
    let b = pool.alloc(1).await.expect("alloc b");
    assert_eq!(a.len(), 1);
    assert_eq!(b.len(), 1);
    let gap = b.addr().abs_diff(a.addr());
    assert!(
        gap >= GRANULE as u64,
        "two live buffers shared a granule: {gap} bytes apart"
    );
}

/// Zero-length allocations are refused at the boundary rather than producing a
/// buffer that names no bytes.
#[tokio::test]
async fn pool_refuses_zero_length() {
    let (_backend, pool) = mock_pool(small_pool_config());
    assert_eq!(pool.alloc(0).await.err(), Some(RdmaError::OutOfRange));
}

/// The set grows geometrically rather than mapping one arena per request, and
/// each new arena is at least as big as the last.
#[tokio::test]
async fn pool_grows_geometrically() {
    let cfg = small_pool_config();
    let (backend, pool) = mock_pool(cfg.clone());

    // Fill well past the first arena. Each buffer is a quarter of it, so the
    // fifth is the one that cannot fit.
    let quarter = cfg.initial_arena_bytes as usize / 4;
    let mut held = Vec::new();
    for _ in 0..12 {
        held.push(pool.alloc(quarter).await.expect("alloc"));
    }

    let arenas = pool.arena_count();
    assert!(arenas >= 2, "the pool never grew: {arenas} arenas");
    assert!(
        arenas < 12,
        "the pool mapped an arena per allocation ({arenas}); growth is not geometric"
    );
    assert_eq!(
        backend.live(),
        arenas,
        "every arena is one backend registration"
    );
}

/// Growth saturates at `max_arena_bytes` for every arena count, rather than
/// wrapping to zero part-way up.
///
/// `checked_shl` refuses an out-of-range *shift*, never a shifted-out *value*:
/// at the shipped defaults it answers `Some(0)` from 38 pooled arenas up to 63,
/// then `None` again at 64. A size derived from it is therefore correct at both
/// ends of the range and wrong only in the middle, which is exactly the shape
/// that survives a spot check.
#[test]
fn pool_growth_saturates_instead_of_wrapping() {
    const INITIAL: u64 = 64 << 20;
    const MAX: u64 = 1 << 30;

    for pooled in 0..=80usize {
        // Oracle in wide arithmetic, where the doubling cannot overflow.
        let want = ((INITIAL as u128) << pooled).min(MAX as u128) as u64;
        assert_eq!(
            pool_arena_target(INITIAL, MAX, pooled),
            want,
            "growth target is wrong at {pooled} pooled arenas"
        );
    }
}

/// A request at or above `dedicated_arena_min` gets an arena sized to it,
/// instead of forcing the pool up a growth step and wasting the round-up.
#[tokio::test]
async fn pool_dedicates_an_arena_to_oversize_requests() {
    let cfg = small_pool_config();
    let (backend, pool) = mock_pool(cfg.clone());

    let small = pool.alloc(GRANULE).await.expect("small alloc");
    let arenas_before = pool.arena_count();

    let big = pool
        .alloc(cfg.dedicated_arena_min as usize)
        .await
        .expect("oversize alloc");
    assert_eq!(
        pool.arena_count(),
        arenas_before + 1,
        "an oversize request must get its own arena"
    );
    assert_eq!(backend.live(), arenas_before + 1);

    // The dedicated arena is not offered to the general search, so a later
    // small request cannot land inside it and strand the big one.
    let another_small = pool.alloc(GRANULE).await.expect("second small alloc");
    let big_end = big.addr() + big.len() as u64;
    assert!(
        another_small.addr() < big.addr() || another_small.addr() >= big_end,
        "a pooled allocation landed inside the dedicated arena"
    );
    drop((small, big, another_small));
}

/// Over the registered-bytes ceiling the pool refuses with `BudgetExceeded`,
/// which is the signal Phase 3 turns into "stage chunked instead" — never a
/// panic and never a hard failure of the staging operation (D4).
#[tokio::test]
async fn pool_budget_exhaustion_is_a_refusal() {
    let cfg = RdmaPoolConfig {
        initial_arena_bytes: 16 * GRANULE as u64,
        max_arena_bytes: 16 * GRANULE as u64,
        dedicated_arena_min: 1024 * GRANULE as u64,
        // Room for exactly two arenas.
        registered_bytes_budget: 32 * GRANULE as u64,
        ..small_pool_config()
    };
    let (backend, pool) = mock_pool(cfg);

    let mut held = Vec::new();
    let mut refusal = None;
    for _ in 0..64 {
        match pool.alloc(8 * GRANULE).await {
            Ok(buf) => held.push(buf),
            Err(e) => {
                refusal = Some(e);
                break;
            }
        }
    }

    match refusal {
        Some(RdmaError::BudgetExceeded {
            registered, budget, ..
        }) => {
            assert_eq!(budget, 32 * GRANULE as u64);
            assert!(
                registered <= budget,
                "the budget was overshot before it was enforced: {registered} over {budget}"
            );
        }
        other => panic!("expected a budget refusal, got {other:?}"),
    }
    assert_eq!(
        backend.live(),
        2,
        "the pool mapped past its own ceiling before refusing"
    );
}

/// Space really comes back: fill the pool to its budget, drop everything, and
/// allocate again.
///
/// The assertion is that a subsequent allocation *succeeds*, not that some byte
/// arithmetic balances. `offset-allocator` rounds a request up to a float bin,
/// so capacity is not the sum of the requested lengths and an exact-capacity
/// assertion would be testing the allocator, not the pool.
#[tokio::test]
async fn pool_reuses_space_after_drop() {
    let cfg = RdmaPoolConfig {
        initial_arena_bytes: 16 * GRANULE as u64,
        max_arena_bytes: 16 * GRANULE as u64,
        dedicated_arena_min: 1024 * GRANULE as u64,
        registered_bytes_budget: 32 * GRANULE as u64,
        ..small_pool_config()
    };
    let (backend, pool) = mock_pool(cfg);

    let mut held = Vec::new();
    while let Ok(buf) = pool.alloc(4 * GRANULE).await {
        held.push(buf);
        assert!(held.len() < 64, "the pool never filled up");
    }
    let filled = held.len();
    assert!(filled > 0, "nothing could be allocated at all");
    let arenas_when_full = pool.arena_count();

    drop(held);
    assert_eq!(pool.live_allocations(), 0);

    let again = pool.alloc(4 * GRANULE).await;
    assert!(
        again.is_ok(),
        "space returned by a dropped PinnedBuf was not reusable: {:?}",
        again.err()
    );
    assert_eq!(
        pool.arena_count(),
        arenas_when_full,
        "reuse mapped a new arena instead of using the space that came back"
    );
    assert_eq!(backend.live(), arenas_when_full);
}

/// Concurrent allocation and release: no double-issued range, no lost space,
/// and the growth path serialises so the arena count stays sane.
///
/// Each task writes its own byte value across its whole buffer and reads it
/// back after a yield. Two buffers overlapping would show up as a mismatch,
/// which is the property that actually matters — an offset bug in the
/// suballocator is invisible to a count-based assertion.
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn pool_concurrent_alloc_and_free() {
    const TASKS: usize = 16;
    const ROUNDS: usize = 24;

    let cfg = small_pool_config();
    let (backend, pool) = mock_pool(cfg);
    let pool = Arc::new(pool);

    let mut tasks = Vec::new();
    for task in 0..TASKS {
        let pool = Arc::clone(&pool);
        tasks.push(tokio::spawn(async move {
            let tag = (task % 251) as u8;
            for round in 0..ROUNDS {
                let len = GRANULE * (1 + (round % 3));
                let mut buf = pool.alloc(len).await.expect("concurrent alloc");
                buf.fill(tag);
                tokio::task::yield_now().await;
                assert!(
                    buf.iter().all(|b| *b == tag),
                    "another allocation wrote into this range: task {task}, round {round}"
                );
            }
        }));
    }
    for task in tasks {
        task.await.expect("task panicked");
    }

    assert_eq!(pool.live_allocations(), 0, "a suballocation leaked");
    assert_eq!(
        backend.live(),
        pool.arena_count(),
        "the arena set and the backend disagree about what is mapped"
    );
    assert_eq!(
        backend.unmap_calls(),
        0,
        "no arena should have been unmapped"
    );
}

// ---------------------------------------------------------------------------
// Pool reclamation (Phase 4)
// ---------------------------------------------------------------------------

/// A pool whose empty arenas are reclaimable almost immediately.
///
/// `arena_reclaim_after` is zero rather than small: the sweep is driven
/// explicitly by these tests, not by a tick, so "how long empty" only has to
/// clear a threshold and there is nothing to wait for.
fn reclaiming_pool_config(retain_arena_bytes: u64) -> RdmaPoolConfig {
    RdmaPoolConfig {
        arena_reclaim_after: Some(Duration::ZERO),
        retain_arena_bytes,
        ..small_pool_config()
    }
}

/// The default is off, and the sweep must respect that for *pooled* arenas —
/// however long they sit empty, whatever the retention would allow.
#[tokio::test]
async fn reclaim_is_off_by_default_for_pooled_arenas() {
    let (backend, pool) = mock_pool(small_pool_config());
    assert_eq!(small_pool_config().arena_reclaim_after, None);

    // Two arenas, both emptied.
    let mut held = Vec::new();
    for _ in 0..3 {
        held.push(pool.alloc(32 * GRANULE).await.expect("alloc"));
    }
    assert!(pool.arena_count() >= 2, "the pool did not grow");
    let arenas = pool.arena_count();
    held.clear();

    assert_eq!(pool.reclaim_idle().await, 0, "the sweep ran unconfigured");
    assert_eq!(pool.arena_count(), arenas);
    assert_eq!(backend.unmap_calls(), 0);
}

/// An empty pooled arena above the retention floor is unmapped, its budget is
/// released, and the gauge behind the budget moves with it.
#[tokio::test]
async fn empty_pooled_arenas_above_the_floor_are_reclaimed() {
    // One arena's worth of retention, so the first stays and the rest go.
    let cfg = reclaiming_pool_config(64 * GRANULE as u64);
    let (backend, pool) = mock_pool(cfg);

    // Fill the pool to its ceiling, so growth really produces several arenas
    // rather than however many a fixed allocation count happens to force.
    let mut held = Vec::new();
    while let Ok(buf) = pool.alloc(32 * GRANULE).await {
        held.push(buf);
        if held.len() > 256 {
            panic!("the budget never refused an allocation");
        }
    }
    let grown = pool.arena_count();
    assert!(grown >= 3, "expected several arenas, got {grown}");
    let peak = pool.registered_bytes();
    assert!(peak > 0);

    // Still in use: nothing is reclaimable, whatever the timer says.
    assert_eq!(
        pool.reclaim_idle().await,
        0,
        "an arena with live suballocations was reclaimed"
    );
    assert_eq!(pool.registered_bytes(), peak);

    held.clear();
    let reclaimed = pool.reclaim_idle().await;
    assert!(
        reclaimed > 0,
        "nothing was reclaimed after the pool emptied"
    );
    assert_eq!(backend.unmap_calls(), reclaimed);
    assert_eq!(
        backend.live(),
        pool.arena_count(),
        "the backend and the pool disagree about what is mapped"
    );
    assert!(
        pool.registered_bytes() < peak,
        "the budget was not released: {} vs {peak}",
        pool.registered_bytes()
    );
}

/// The retention floor is the whole point of the knob: an idle pool keeps a warm
/// arena so the next allocation does not pay a fresh registration.
#[tokio::test]
async fn retention_keeps_a_warm_arena_mapped() {
    const RETAIN: u64 = 64 * GRANULE as u64;
    let (backend, pool) = mock_pool(reclaiming_pool_config(RETAIN));

    let mut held = Vec::new();
    while let Ok(buf) = pool.alloc(32 * GRANULE).await {
        held.push(buf);
        if held.len() > 256 {
            panic!("the budget never refused an allocation");
        }
    }
    assert!(
        pool.arena_count() >= 3,
        "the pool did not grow enough to test retention"
    );
    held.clear();
    pool.reclaim_idle().await;

    assert!(pool.arena_count() >= 1, "retention kept nothing at all");
    assert!(
        pool.registered_bytes() >= RETAIN,
        "the pool fell below its retention floor: {} < {RETAIN}",
        pool.registered_bytes()
    );
    let after = pool.registered_bytes();
    // Idempotent: a second sweep over a pool already at its floor does nothing.
    assert_eq!(pool.reclaim_idle().await, 0);
    assert_eq!(pool.registered_bytes(), after);

    // And the floor is warm — the next allocation comes out of what was kept.
    let before = backend.live();
    let _buf = pool.alloc(4 * GRANULE).await.expect("alloc from the floor");
    assert_eq!(
        backend.live(),
        before,
        "the retained arena did not serve the next allocation"
    );
}

/// Zero retention means every empty pooled arena goes, which is what a
/// memory-tight deployment asks for.
#[tokio::test]
async fn zero_retention_reclaims_every_empty_arena() {
    let (backend, pool) = mock_pool(reclaiming_pool_config(0));

    let mut held = Vec::new();
    for _ in 0..4 {
        held.push(pool.alloc(32 * GRANULE).await.expect("alloc"));
    }
    held.clear();
    pool.reclaim_idle().await;

    assert_eq!(
        pool.arena_count(),
        0,
        "an empty arena survived zero retention"
    );
    assert_eq!(backend.live(), 0);
    assert_eq!(
        pool.registered_bytes(),
        0,
        "the budget was not fully released"
    );
}

/// The Phase-2 gap, closed: a dedicated arena is reclaimed as soon as its single
/// suballocation drops, with no timer and no retention — because it can never
/// serve another request, so keeping it mapped is pure charge against the
/// budget.
///
/// The old behaviour is what this asserts against: sixteen oversize stagings at
/// the production defaults exhausted the registered-bytes budget for the life of
/// the process, and the documented mitigation was to raise
/// `dedicated_arena_min` past the sizes a hot path used.
#[tokio::test]
async fn a_dedicated_arena_is_reclaimed_when_its_buffer_drops() {
    // Reclamation *disabled*, to show dedicated arenas do not depend on it.
    let cfg = small_pool_config();
    assert_eq!(cfg.arena_reclaim_after, None);
    let dedicated_min = cfg.dedicated_arena_min as usize;
    let (backend, pool) = mock_pool(cfg);

    // Enough cycles that the old never-reclaim behaviour would exhaust the
    // budget: the config allows eight arenas of this size at most.
    let mut peak_arenas = 0;
    for cycle in 0..24 {
        let buf = pool
            .alloc(dedicated_min)
            .await
            .unwrap_or_else(|e| panic!("cycle {cycle}: oversize alloc refused: {e}"));
        peak_arenas = peak_arenas.max(pool.arena_count());
        drop(buf);
        pool.reclaim_idle().await;
        assert_eq!(
            pool.registered_bytes(),
            0,
            "cycle {cycle}: the dedicated arena's budget was not released"
        );
        assert_eq!(backend.live(), 0, "cycle {cycle}: it stayed mapped");
    }
    assert_eq!(
        peak_arenas, 1,
        "dedicated arenas accumulated instead of being reclaimed each cycle"
    );
}

/// A dedicated arena is not reclaimed while a transfer is still writing into it,
/// even though its `PinnedBuf` is gone — a `TransferHold` outliving its buffer
/// is exactly the cancelled-`get_pinned` case, and unmapping under it would hand
/// the NIC a deregistered range.
#[tokio::test]
async fn a_hold_outliving_its_buffer_blocks_reclaim() {
    let cfg = reclaiming_pool_config(0);
    let (backend, pool) = mock_pool(cfg);

    let buf = pool.alloc(8 * GRANULE).await.expect("alloc");
    let hold = buf.hold();
    drop(buf);

    assert_eq!(
        pool.reclaim_idle().await,
        0,
        "an arena with a transfer still holding it was reclaimed"
    );
    assert_eq!(backend.live(), 1);

    drop(hold);
    assert_eq!(pool.reclaim_idle().await, 1);
    assert_eq!(backend.live(), 0);
}

/// Reclaim racing allocation: whatever the interleaving, an arena is either
/// handed out or unmapped, never both.
///
/// The assertion that matters is `backend.live() == arena_count()` — the pool
/// and the backend agreeing about what is mapped — plus every buffer handed out
/// being readable and writable, which an arena unmapped underneath one would not
/// be. Counting reclaims would only prove the sweep ran.
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn reclaim_racing_alloc_never_hands_out_a_reclaimed_arena() {
    const TASKS: usize = 8;
    const ROUNDS: usize = 40;
    let (backend, pool) = mock_pool(reclaiming_pool_config(0));
    let pool = Arc::new(pool);

    let mut tasks = Vec::new();
    for task in 0..TASKS {
        let pool = Arc::clone(&pool);
        tasks.push(tokio::spawn(async move {
            for round in 0..ROUNDS {
                if let Ok(mut buf) = pool.alloc(4 * GRANULE).await {
                    let byte = (task * ROUNDS + round) as u8;
                    buf.fill(byte);
                    tokio::task::yield_now().await;
                    assert!(
                        buf.iter().all(|b| *b == byte),
                        "task {task} round {round}: the buffer was overwritten"
                    );
                }
            }
        }));
    }
    let sweeper = {
        let pool = Arc::clone(&pool);
        tokio::spawn(async move {
            for _ in 0..ROUNDS * 4 {
                pool.reclaim_idle().await;
                tokio::task::yield_now().await;
            }
        })
    };

    for task in tasks {
        task.await.expect("allocator task must not panic");
    }
    sweeper.await.expect("sweeper must not panic");

    assert_eq!(
        backend.live(),
        pool.arena_count(),
        "the backend and the pool disagree about what is mapped"
    );
    // Everything that was reclaimed gave its bytes back, and everything still
    // mapped is still charged.
    let mapped: u64 = pool.registered_bytes();
    assert_eq!(
        mapped == 0,
        pool.arena_count() == 0,
        "registered bytes and arena count disagree: {mapped} B over {} arenas",
        pool.arena_count()
    );
}

/// The registry gate, not the pool policy: once shutdown has closed admission a
/// reclaim is refused outright, so it cannot be unmapping arenas while the sweep
/// walks them or leave one in flight when `latch_all_deregistered` asks the
/// backend whether anything is still registered.
#[tokio::test]
async fn reclaim_is_refused_once_shutdown_has_gated() {
    let cfg = RdmaConfig {
        pool: reclaiming_pool_config(0),
        ..RdmaConfig::default()
    };
    let (backend, registry) = mock_registry(cfg);

    let buf = registry.alloc_pinned(8 * GRANULE).await.expect("alloc");
    drop(buf);
    assert_eq!(backend.live(), 1);

    registry.shutdown(Duration::from_secs(5)).await;
    assert_eq!(backend.live(), 0, "the sweep left an arena mapped");

    // The gate is closed: a late tick reclaims nothing and, in particular, does
    // not touch the backend.
    let calls = backend.unmap_calls();
    assert_eq!(registry.reclaim_idle_arenas().await, 0);
    assert_eq!(
        backend.unmap_calls(),
        calls,
        "a reclaim ran after the shutdown gate closed"
    );
}

/// A reclaim the backend will not confirm must not credit the budget: `Err` from
/// an unmap means *unknown*, not *unmapped*. The pages and their charge are held
/// until the end of velo shutdown, which is the same rule the shutdown sweep
/// follows and the same machinery.
#[tokio::test]
async fn an_unconfirmed_reclaim_holds_its_budget() {
    let (backend, pool) = mock_pool(reclaiming_pool_config(0));

    let buf = pool.alloc(8 * GRANULE).await.expect("alloc");
    let charged = pool.registered_bytes();
    assert!(charged > 0);
    drop(buf);

    backend.refuse_unmap.store(true, Ordering::SeqCst);
    assert_eq!(
        pool.reclaim_idle().await,
        0,
        "an unconfirmed unmap was counted as a reclaim"
    );
    assert_eq!(
        pool.registered_bytes(),
        charged,
        "the budget was credited for pages that may still be pinned"
    );
    assert_eq!(pool.arena_count(), 0, "the arena stayed in the live set");

    // The end-of-shutdown release is what eventually gives them back, exactly
    // once — no double credit with the reclaim above.
    pool.release_unconfirmed();
    assert_eq!(pool.registered_bytes(), 0);
}

/// The wiring, end to end: nothing in this test calls `reclaim_idle` — the
/// rendezvous tick has to do it, over a real UCX backend, from configuration
/// alone.
///
/// Every other reclamation test drives the sweep directly, so all of them would
/// pass with `set_rdma_context` never having been taught to call it. This is the
/// one that would not.
#[tokio::test(flavor = "multi_thread")]
async fn the_rendezvous_tick_reclaims_arenas_without_being_asked() {
    let transport = Arc::new(
        crate::transports::ucx::UcxTransportBuilder::new()
            .tls("tcp")
            .build()
            .expect("build ucx transport"),
    );
    let velo = crate::Velo::builder()
        .add_ucx_transport(Arc::clone(&transport))
        .rdma_config(RdmaConfig {
            pool: RdmaPoolConfig {
                // Small enough that one staging maps one arena, and every
                // staging above `dedicated_arena_min` gets its own.
                initial_arena_bytes: 64 * GRANULE as u64,
                max_arena_bytes: 64 * GRANULE as u64,
                dedicated_arena_min: 16 * GRANULE as u64,
                registered_bytes_budget: 1024 * GRANULE as u64,
                arena_reclaim_after: Some(Duration::from_millis(10)),
                retain_arena_bytes: 0,
            },
            ..RdmaConfig::default()
        })
        .build()
        .await
        .expect("build velo");

    // A pooled staging and an oversize one, so both policies are on the hook.
    let registry = velo.rdma().expect("the ucx transport gives us a registry");
    let small = registry
        .alloc_pinned(4 * GRANULE)
        .await
        .expect("pooled alloc");
    let big = registry
        .alloc_pinned(16 * GRANULE)
        .await
        .expect("dedicated alloc");
    assert!(velo.rdma_registered_bytes() > 0);
    assert!(transport.live_regions() >= 2, "expected two arenas");

    drop(small);
    drop(big);

    // No explicit sweep: the tick is the only thing that can do this.
    let deadline = std::time::Instant::now() + T;
    while std::time::Instant::now() < deadline {
        if velo.rdma_registered_bytes() == 0 && transport.live_regions() == 0 {
            break;
        }
        tokio::time::sleep(Duration::from_millis(5)).await;
    }
    assert_eq!(
        velo.rdma_registered_bytes(),
        0,
        "the rendezvous tick never reclaimed the empty arenas"
    );
    assert_eq!(transport.live_regions(), 0);

    velo.graceful_shutdown(velo_ext::ShutdownPolicy::Timeout(T))
        .await;
    assert_eq!(transport.live_rkeys(), 0);
}

/// The soak the plan names, on the mock backend so the cycle count can be high:
/// allocate, transfer, release, sweep — and assert the registered-byte total and
/// the arena count come back to the same place every time.
///
/// A leak that grows by one arena per cycle and one that appears once are
/// different bugs; only a per-cycle assertion tells them apart.
#[tokio::test]
async fn pool_lifecycle_soak() {
    const CYCLES: usize = 200;
    let cfg = reclaiming_pool_config(0);
    let dedicated_min = cfg.dedicated_arena_min as usize;
    let (backend, pool) = mock_pool(cfg);

    for cycle in 0..CYCLES {
        // A pooled allocation and an oversize one, so both reclamation policies
        // run every cycle.
        let mut small = pool
            .alloc(4 * GRANULE)
            .await
            .unwrap_or_else(|e| panic!("cycle {cycle}: pooled alloc: {e}"));
        let big = pool
            .alloc(dedicated_min)
            .await
            .unwrap_or_else(|e| panic!("cycle {cycle}: dedicated alloc: {e}"));

        // Stand in for a transfer into the pooled buffer: a hold taken and
        // released around the write, as `get_pinned` does.
        let hold = small.hold();
        small.fill(cycle as u8);
        assert!(small.iter().all(|b| *b == cycle as u8));
        drop(hold);

        drop(small);
        drop(big);
        pool.reclaim_idle().await;

        assert_eq!(
            pool.registered_bytes(),
            0,
            "cycle {cycle}: registered bytes did not return to zero"
        );
        assert_eq!(
            pool.arena_count(),
            0,
            "cycle {cycle}: an arena survived the sweep"
        );
        assert_eq!(
            backend.live(),
            0,
            "cycle {cycle}: the backend still holds a registration"
        );
    }
}

// ---------------------------------------------------------------------------
// RegionGuard lifecycle
// ---------------------------------------------------------------------------

/// The happy path: register, observe, unregister, and only then does
/// `deregistered()` resolve.
#[tokio::test]
async fn region_unregister_latches_deregistered() {
    let (backend, registry) = mock_registry(RdmaConfig::default());
    let guard = registry
        .register_owned(vec![0u8; 8192].into_boxed_slice())
        .await
        .expect("register");

    assert_eq!(backend.live(), 1);
    assert_eq!(registry.region_count(), 1);
    // The budget charges the page-enclosing range, not the requested length: a
    // heap `Box` is byte-aligned, so 8192 bytes generally straddles three pages
    // and that is what the kernel pins.
    let charged = registry.registered_bytes();
    assert!(
        charged >= 8192 && charged % GRANULE as u64 == 0,
        "expected a page-enclosing charge, got {charged}"
    );
    assert!(
        !guard.is_deregistered(),
        "a live registration must not claim to be deregistered"
    );

    // The latch is not resolved yet, so awaiting it must not complete.
    let watch = guard.watch();
    assert!(
        tokio::time::timeout(Duration::from_millis(50), watch.deregistered())
            .await
            .is_err(),
        "deregistered() resolved while the memory was still registered"
    );

    assert_eq!(
        guard.unregister(T).await.expect("unregister"),
        Deregistered::Drained
    );

    assert_eq!(backend.live(), 0);
    assert_eq!(registry.region_count(), 0);
    assert_eq!(
        registry.registered_bytes(),
        0,
        "the budget was not credited back"
    );
    assert!(watch.is_deregistered());
    tokio::time::timeout(T, watch.deregistered())
        .await
        .expect("the latch must be resolved for every observer, not just the unregisterer");
}

/// An unmap the backend could not confirm must **not** latch.
///
/// `ShuttingDown` from an unmap means *unknown*, not *unmapped*; latching on it
/// would tell the caller it may free memory that is still pinned. The latch is
/// closed later instead, at the end of velo shutdown, once teardown has made
/// the release a fact — see
/// `shutdown_latches_regions_whose_unmap_was_never_confirmed`.
#[tokio::test]
async fn unconfirmed_unmap_does_not_latch() {
    let (backend, registry) = mock_registry(RdmaConfig::default());
    let guard = registry
        .register_owned(vec![0u8; 4096].into_boxed_slice())
        .await
        .expect("register");
    let watch = guard.watch();
    let charged = registry.registered_bytes();

    backend.refuse_unmap.store(true, Ordering::SeqCst);
    let err = guard
        .unregister(T)
        .await
        .expect_err("the unmap was refused");
    assert_eq!(err, RdmaError::ShuttingDown);

    assert!(
        !watch.is_deregistered(),
        "an unconfirmed unmap latched deregistered(): a caller would now free pinned memory"
    );
    assert_eq!(
        registry.region_count(),
        1,
        "an unconfirmed region must stay tracked so the shutdown sweep asks again"
    );
    assert_eq!(
        registry.registered_bytes(),
        charged,
        "the budget was credited back for memory that may still be pinned"
    );

    // And once the backend cooperates, the sweep does resolve it.
    backend.refuse_unmap.store(false, Ordering::SeqCst);
    registry.shutdown(T).await;
    assert!(watch.is_deregistered());
    assert_eq!(registry.registered_bytes(), 0);
}

/// Dropping the guard without awaiting deregisters in the background — it does
/// not leak, and it does not block the dropping thread either.
///
/// The observable is the backend actually being asked to unmap, reached through
/// a `RegionWatch` taken before the drop. `Drop` returning is explicitly *not*
/// the point at which the memory is free.
#[tokio::test]
async fn dropped_guard_deregisters_in_the_background() {
    let (backend, registry) = mock_registry(RdmaConfig::default());
    let guard = registry
        .register_owned(vec![0u8; 4096].into_boxed_slice())
        .await
        .expect("register");
    let watch = guard.watch();

    drop(guard);

    tokio::time::timeout(T, watch.deregistered())
        .await
        .expect("a dropped guard must still deregister");
    assert_eq!(
        backend.live(),
        0,
        "the backend still holds the registration"
    );
    assert_eq!(registry.region_count(), 0);
    assert_eq!(registry.registered_bytes(), 0);
}

/// Dropping a guard on a plain thread with no ambient runtime still works,
/// because the registry captured a runtime handle at construction rather than
/// reading one from the environment at drop time.
#[tokio::test(flavor = "multi_thread")]
async fn guard_dropped_off_runtime_still_deregisters() {
    let (backend, registry) = mock_registry(RdmaConfig::default());
    let guard = registry
        .register_owned(vec![0u8; 4096].into_boxed_slice())
        .await
        .expect("register");
    let watch = guard.watch();

    std::thread::spawn(move || drop(guard))
        .join()
        .expect("dropper thread panicked");

    tokio::time::timeout(T, watch.deregistered())
        .await
        .expect("a guard dropped off the runtime must still deregister");
    assert_eq!(backend.live(), 0);
}

/// `unregister` waits for the region in-flight count to drain before unmapping.
///
/// Phase 3 acquires one of these guards per RDMA lease; this is the mechanism
/// that stops a registration being pulled out from under a transfer. Asserted
/// as "pending while held, resolves after release" rather than by sleeping past
/// a guessed interval.
#[tokio::test(flavor = "multi_thread")]
async fn unregister_waits_for_in_flight() {
    let (backend, registry) = mock_registry(RdmaConfig::default());
    let guard = registry
        .register_owned(vec![0u8; 4096].into_boxed_slice())
        .await
        .expect("register");

    let lease = guard.in_flight().acquire();
    let watch = guard.watch();

    let mut unregistering = tokio::spawn(async move { guard.unregister(T).await });

    assert!(
        tokio::time::timeout(Duration::from_millis(100), &mut unregistering)
            .await
            .is_err(),
        "unregister completed while an operation was still in flight"
    );
    assert_eq!(
        backend.unmap_calls(),
        0,
        "the backend was asked to unmap before the region had drained"
    );

    drop(lease);

    let outcome = tokio::time::timeout(T, unregistering)
        .await
        .expect("unregister must resolve once the last in-flight guard is released")
        .expect("task panicked")
        .expect("unregister");
    assert_eq!(
        outcome,
        Deregistered::Drained,
        "the drain completed, so this is not a timed-out deregistration"
    );
    assert!(watch.is_deregistered());
    assert_eq!(backend.live(), 0);
}

/// A drain that outlasts the budget still unmaps, and says so with `Timeout`.
///
/// Waiting forever on a peer that may have crashed is the worse failure, so the
/// bounded wait force-unmaps — and the latch does resolve, because the unmap
/// itself was confirmed. Only the *drain* was cut short.
#[tokio::test(flavor = "multi_thread")]
async fn unregister_timeout_still_unmaps() {
    let (backend, registry) = mock_registry(RdmaConfig::default());
    let guard = registry
        .register_owned(vec![0u8; 4096].into_boxed_slice())
        .await
        .expect("register");
    let watch = guard.watch();

    // Never released: this stands in for a lease whose holder has gone away.
    let _stuck = guard.in_flight().acquire();

    let outcome = guard
        .unregister(Duration::from_millis(100))
        .await
        .expect("a confirmed unmap is Ok even when the drain was cut short");
    assert_eq!(
        outcome,
        Deregistered::DrainTimedOut,
        "the caller must be able to tell that in-flight work was not waited for"
    );

    assert!(
        watch.is_deregistered(),
        "the unmap was confirmed, so the latch must resolve even though the drain timed out"
    );
    assert_eq!(
        backend.live(),
        0,
        "a timed-out drain must still force the unmap"
    );
}

/// `watch()` observes the same states as the guard, and keeps working after the
/// guard is gone. Holding one neither keeps the registration alive nor ends it.
#[tokio::test]
async fn watch_observes_without_owning() {
    let (_backend, registry) = mock_registry(RdmaConfig::default());
    let guard = registry
        .register_owned(vec![0u8; 4096].into_boxed_slice())
        .await
        .expect("register");

    let watch = guard.watch();
    let second = watch.clone();
    assert!(!watch.is_shutting_down());
    assert!(!watch.is_deregistered());

    registry.shutdown(T).await;

    assert!(watch.is_shutting_down(), "a watch must see shutdown begin");
    assert!(second.is_deregistered());
    tokio::time::timeout(T, second.shutdown_initiated())
        .await
        .expect("shutdown_initiated must resolve once shutdown has begun");
    drop(guard);
}

/// `register_owned` hands the buffer back on a confirmed deregistration, and
/// keeps it otherwise — a `Box` returned while the pages may still be pinned is
/// exactly the free-while-mapped hazard.
#[tokio::test]
async fn register_owned_returns_the_buffer() {
    let (backend, registry) = mock_registry(RdmaConfig::default());

    let mut buf = vec![0u8; 4096].into_boxed_slice();
    buf[0] = 0x5A;
    let guard = registry.register_owned(buf).await.expect("register");
    assert_eq!(guard.len(), 4096);

    let (returned, outcome) = guard.unregister_owned(T).await.expect("unregister_owned");
    assert_eq!(outcome, Deregistered::Drained);
    assert_eq!(returned.len(), 4096);
    assert_eq!(
        returned[0], 0x5A,
        "the buffer that came back is not the one that went in"
    );
    assert_eq!(backend.live(), 0);
}

/// The unsafe path, exercised against an allocation deliberately leaked for the
/// duration — which is the honest way to satisfy the safety contract in a test.
#[tokio::test]
async fn register_external_memory_smoke() {
    let (backend, registry) = mock_registry(RdmaConfig::default());

    // Leaked on purpose: the contract requires the allocation to outlive the
    // registration, and this test asserts on the registration, not on reclaim.
    let leaked: &'static mut [u8] = Box::leak(vec![7u8; 8192].into_boxed_slice());
    let ptr = std::ptr::NonNull::new(leaked.as_mut_ptr()).expect("non-null");

    // SAFETY: `leaked` is a live 8192-byte allocation that is never freed, so it
    // outlives the registration unconditionally.
    let guard = unsafe { registry.register_external(ptr, leaked.len()) }
        .await
        .expect("register external");

    assert_eq!(guard.addr(), ptr.as_ptr() as u64);
    assert_eq!(guard.len(), 8192);
    let (eff_addr, eff_len) = guard.effective_range();
    assert!(
        eff_addr <= guard.addr() && eff_len >= guard.len(),
        "the effective range must cover the requested one"
    );
    assert_ne!(
        guard.generation(),
        0,
        "every registration gets a generation"
    );

    assert_eq!(
        guard.unregister(T).await.expect("unregister"),
        Deregistered::Drained
    );
    assert_eq!(backend.live(), 0);
}

/// Degenerate arguments are refused at the boundary rather than reaching the
/// backend as a map of nothing.
#[tokio::test]
async fn register_external_refuses_degenerate_ranges() {
    let (backend, registry) = mock_registry(RdmaConfig::default());
    let mut byte = 0u8;
    let ptr = std::ptr::NonNull::new(&mut byte as *mut u8).expect("non-null");

    // SAFETY: `ptr` is valid; the call is refused before it is ever used.
    let err = unsafe { registry.register_external(ptr, 0) }
        .await
        .unwrap_err();
    assert_eq!(err, RdmaError::OutOfRange);
    assert_eq!(
        backend.live(),
        0,
        "a refused registration must not reach the backend"
    );
}

// ---------------------------------------------------------------------------
// Registry shutdown (D8 steps 1 to 3)
// ---------------------------------------------------------------------------

/// The sweep unmaps everything — external regions and pool arenas alike — and
/// resolves every latch.
#[tokio::test]
async fn shutdown_deregisters_regions_and_arenas() {
    let cfg = RdmaConfig {
        pool: small_pool_config(),
        ..RdmaConfig::default()
    };
    let (backend, registry) = mock_registry(cfg);

    let guard = registry
        .register_owned(vec![0u8; 8192].into_boxed_slice())
        .await
        .expect("register");
    let watch = guard.watch();
    let buf = registry.alloc_pinned(4096).await.expect("alloc pinned");
    assert!(
        backend.live() >= 2,
        "expected an external region and an arena"
    );

    // The buffer outlives the sweep on purpose: shutdown must not depend on
    // every caller having tidied up first.
    registry.shutdown(T).await;

    assert_eq!(backend.live(), 0, "shutdown left something registered");
    assert!(
        watch.is_deregistered(),
        "shutdown did not resolve the latch"
    );
    assert_eq!(registry.registered_bytes(), 0);
    assert_eq!(registry.pool().arena_count(), 0);
    drop(buf);
    drop(guard);
}

/// After the gate closes, both registration paths refuse. This is what stops a
/// registration landing behind the sweep with no tracking entry and no latch.
#[tokio::test]
async fn shutdown_gates_new_registrations() {
    let (_backend, registry) = mock_registry(RdmaConfig::default());
    registry.shutdown(T).await;

    let refused = registry
        .register_owned(vec![0u8; 4096].into_boxed_slice())
        .await
        .expect_err("a gated registry must refuse");
    assert_eq!(refused.cause, RdmaError::ShuttingDown);
    assert_eq!(
        refused.buffer.map(|b| b.len()),
        Some(4096),
        "a refused registration must hand the caller buffer back"
    );
    assert_eq!(
        registry.alloc_pinned(4096).await.err(),
        Some(RdmaError::ShuttingDown),
        "pool allocation must go through the same gate as external registration"
    );
}

/// Shutdown is idempotent: a second sweep over an empty registry is a no-op
/// rather than a double-unmap or a double budget credit.
#[tokio::test]
async fn shutdown_is_idempotent() {
    let (backend, registry) = mock_registry(RdmaConfig::default());
    let guard = registry
        .register_owned(vec![0u8; 4096].into_boxed_slice())
        .await
        .expect("register");
    drop(guard);

    registry.shutdown(T).await;
    let calls = backend.unmap_calls();
    registry.shutdown(T).await;
    assert_eq!(
        backend.unmap_calls(),
        calls,
        "a second shutdown re-issued unmaps"
    );
    assert_eq!(registry.registered_bytes(), 0);
}

/// A registration already past the gate must land before the sweep enumerates.
///
/// This is the race the admission counter exists for: a token-check gate would
/// let this registration pass the check, map after step 3 had walked the
/// region map, and leave pinned memory with no entry and no latch. The
/// backend delay widens the window from instructions to milliseconds so the
/// test is a detector rather than a coin flip.
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn registration_in_flight_is_not_missed_by_shutdown() {
    let (backend, registry) = mock_registry(RdmaConfig::default());
    *backend.map_delay.lock() = Some(Duration::from_millis(200));

    let registering = {
        let registry = Arc::clone(&registry);
        tokio::spawn(async move {
            registry
                .register_owned(vec![0u8; 4096].into_boxed_slice())
                .await
        })
    };

    // Let the registration get past the gate and into the backend map.
    tokio::time::sleep(Duration::from_millis(50)).await;
    registry.shutdown(T).await;

    let outcome = tokio::time::timeout(T, registering)
        .await
        .expect("the registration must resolve")
        .expect("task panicked");

    match outcome {
        Ok(guard) => {
            // It got in. The sweep must have waited for it and unmapped it.
            assert!(
                guard.is_deregistered(),
                "a registration admitted before the gate closed was missed by the sweep"
            );
        }
        Err(e) if e.cause == RdmaError::ShuttingDown => {
            // It was refused at the gate, which is equally correct.
            assert!(e.buffer.is_some(), "a refusal must return the buffer");
        }
        Err(e) => panic!("unexpected registration failure: {}", e.cause),
    }
    assert_eq!(
        backend.live(),
        0,
        "shutdown returned with memory still registered"
    );
    assert_eq!(registry.registered_bytes(), 0);
}

// ---------------------------------------------------------------------------
// Against a real UCX transport, over UCX_TLS=tcp
// ---------------------------------------------------------------------------

/// A started [`UcxTransport`] plus a registry over its RMA endpoint.
///
/// `_streams` is held because dropping the receivers would tear the inbound
/// channels down under a transport that is still running.
struct UcxHarness {
    transport: Arc<crate::transports::ucx::UcxTransport>,
    registry: Arc<RdmaRegistry>,
    _streams: crate::transports::DataStreams,
}

impl UcxHarness {
    async fn start(cfg: RdmaConfig) -> Self {
        use velo_ext::InstanceId;

        let transport = Arc::new(
            crate::transports::ucx::UcxTransportBuilder::new()
                .tls("tcp")
                .build()
                .expect("build ucx transport"),
        );
        let (adapter, streams) = crate::transports::make_channels();
        tokio::time::timeout(
            T,
            transport.start(
                InstanceId::new_v4(),
                adapter,
                tokio::runtime::Handle::current(),
            ),
        )
        .await
        .expect("ucx startup must not hang")
        .expect("start ucx transport");

        let registry = Arc::new(RdmaRegistry::new(
            UcxBackend::new(transport.rdma_endpoint()),
            cfg,
            tokio::runtime::Handle::current(),
            None,
        ));
        Self {
            transport,
            registry,
            _streams: streams,
        }
    }

    /// Regions the progress thread itself believes it holds. The authoritative
    /// count: this module cannot fake it, which is the point of asserting on it
    /// rather than on `registry.region_count()`.
    fn live_regions(&self) -> usize {
        self.transport.live_regions()
    }
}

/// The backend really registers with UCX, and really releases it.
///
/// Asserted through the progress thread's own `live_regions`, so a registration
/// the registry has forgotten but UCX still holds cannot pass.
#[tokio::test(flavor = "multi_thread")]
async fn ucx_backend_maps_and_unmaps() {
    let harness = UcxHarness::start(RdmaConfig::default()).await;
    assert_eq!(harness.live_regions(), 0);
    assert_eq!(harness.registry.backend_key(), "ucx");

    let guard = harness
        .registry
        .register_owned(vec![0u8; 256 * 1024].into_boxed_slice())
        .await
        .expect("register with ucx");

    assert_eq!(harness.live_regions(), 1, "ucx did not register the range");
    let remote = guard.remote();
    assert!(
        !remote.packed_key.is_empty(),
        "a real registration must produce a packed key"
    );
    let (eff_addr, eff_len) = guard.effective_range();
    assert!(
        eff_addr <= guard.addr() && eff_len >= guard.len(),
        "ucx reported an effective range that does not cover the request"
    );

    assert_eq!(
        guard.unregister(T).await.expect("unregister"),
        Deregistered::Drained
    );
    assert_eq!(harness.live_regions(), 0, "ucx still holds the region");
    assert_eq!(harness.registry.registered_bytes(), 0);
    harness.transport.shutdown();
}

/// The pool over a real backend: an arena is one UCX registration however many
/// buffers come out of it, and the shutdown sweep releases it.
#[tokio::test(flavor = "multi_thread")]
async fn ucx_pool_arena_is_one_registration() {
    let cfg = RdmaConfig {
        pool: small_pool_config(),
        ..RdmaConfig::default()
    };
    let harness = UcxHarness::start(cfg).await;

    let a = harness.registry.alloc_pinned(4096).await.expect("alloc a");
    let b = harness.registry.alloc_pinned(4096).await.expect("alloc b");
    assert_eq!(
        harness.live_regions(),
        1,
        "two suballocations from one arena must be one ucx registration"
    );
    assert_eq!(a.backend_region_id(), b.backend_region_id());
    assert_ne!(a.arena_offset(), b.arena_offset());
    drop((a, b));

    harness.registry.shutdown(T).await;
    assert_eq!(harness.live_regions(), 0, "the arena was not unmapped");
    harness.transport.shutdown();
}

/// The ordering D8 exists for, asserted end to end through the real
/// `Velo::graceful_shutdown`.
///
/// The load-bearing claim is not "the memory is eventually released" but
/// "it is released *before* graceful_shutdown returns, and before the transport
/// is torn down". So the assertion is taken at the moment shutdown returns,
/// against the progress thread's own region count — the one number that cannot be
/// satisfied by bookkeeping in the layer under test. If the registry sweep ran
/// after transport teardown, or not at all, `live_regions` would be non-zero
/// here or the unmap would have been a forced one from teardown rather than an
/// orderly one from the sweep.
#[tokio::test(flavor = "multi_thread")]
async fn velo_graceful_shutdown_deregisters_before_transport_teardown() {
    let transport = Arc::new(
        crate::transports::ucx::UcxTransportBuilder::new()
            .tls("tcp")
            .build()
            .expect("build ucx transport"),
    );
    let velo = crate::Velo::builder()
        .add_ucx_transport(Arc::clone(&transport))
        .build()
        .await
        .expect("build velo");

    let guard = velo
        .register_owned(vec![0u8; 128 * 1024].into_boxed_slice())
        .await
        .expect("register through the velo facade");
    let watch = guard.watch();
    assert_eq!(transport.live_regions(), 1);
    assert!(
        velo.rdma_registered_bytes() >= 128 * 1024,
        "the budget must charge at least the requested length"
    );

    velo.graceful_shutdown(velo_ext::ShutdownPolicy::Timeout(T))
        .await;

    assert!(
        watch.is_deregistered(),
        "graceful_shutdown returned without resolving the deregistered() latch"
    );
    assert_eq!(
        transport.live_regions(),
        0,
        "graceful_shutdown returned with memory still registered with ucx"
    );
    assert_eq!(
        transport.live_rkeys(),
        0,
        "an unpacked rkey outlived shutdown"
    );
    assert_eq!(velo.rdma_registered_bytes(), 0);
    drop(guard);
}

/// A guard dropped without awaiting is deregistered in the background against
/// the real backend too — the failure mode being ruled out is a warn that
/// announces a deregistration nobody actually performs.
#[tokio::test(flavor = "multi_thread")]
async fn ucx_dropped_guard_deregisters() {
    let harness = UcxHarness::start(RdmaConfig::default()).await;
    let guard = harness
        .registry
        .register_owned(vec![0u8; 64 * 1024].into_boxed_slice())
        .await
        .expect("register");
    let watch = guard.watch();
    assert_eq!(harness.live_regions(), 1);

    drop(guard);

    tokio::time::timeout(T, watch.deregistered())
        .await
        .expect("a dropped guard must deregister against ucx too");
    assert_eq!(
        harness.live_regions(),
        0,
        "the warn claimed a background deregistration that never reached ucx"
    );
    harness.transport.shutdown();
}

/// Without a UCX transport the facade refuses rather than panicking, and the
/// rest of Velo is unaffected. Adding the transport with `add_transport`
/// instead of `add_ucx_transport` is a legitimate messaging-only setup.
#[tokio::test(flavor = "multi_thread")]
async fn velo_without_ucx_transport_refuses_registration() {
    let listener = std::net::TcpListener::bind("127.0.0.1:0").expect("bind");
    let transport = Arc::new(
        crate::transports::tcp::TcpTransportBuilder::new()
            .from_listener(listener)
            .expect("listener")
            .build()
            .expect("build tcp transport"),
    );
    let velo = crate::Velo::builder()
        .add_transport(transport)
        .build()
        .await
        .expect("build velo");

    let err = velo
        .register_owned(vec![0u8; 4096].into_boxed_slice())
        .await
        .expect_err("registration must be refused without a ucx transport");
    assert_eq!(
        err.cause,
        RdmaError::NotConfigured,
        "a missing backend is a permanent configuration fact, not a retryable backend error"
    );
    assert!(
        err.buffer.is_some(),
        "the caller buffer must come back from a refusal"
    );
    assert_eq!(velo.rdma_registered_bytes(), 0);

    // And shutdown still works, with no registry to sweep.
    velo.graceful_shutdown(velo_ext::ShutdownPolicy::Timeout(T))
        .await;
}

// ---------------------------------------------------------------------------
// Budget accounting under cancellation and awkward sizes
// ---------------------------------------------------------------------------

/// A registration whose future is dropped at the map must give its budget claim
/// back.
///
/// Wrapping a registration in a `timeout` is an ordinary thing for a caller to
/// write, and a claim released only on the error arm survives it: the arm never
/// runs. The failure is silent and permanent — enough cancellations and every
/// later registration answers `BudgetExceeded` for the life of the process,
/// which Phase 3 reads as "stage chunked" and never reports as broken.
#[tokio::test(flavor = "multi_thread")]
async fn cancelled_registration_returns_its_budget() {
    let (backend, registry) = mock_registry(RdmaConfig::default());
    *backend.map_delay.lock() = Some(Duration::from_millis(500));

    for _ in 0..4 {
        let attempt = registry.register_owned(vec![0u8; 64 * 1024].into_boxed_slice());
        assert!(
            tokio::time::timeout(Duration::from_millis(50), attempt)
                .await
                .is_err(),
            "the registration was supposed to be cancelled mid-map"
        );
    }

    assert_eq!(
        registry.registered_bytes(),
        0,
        "cancelled registrations leaked their budget claim"
    );

    // And the budget is genuinely usable again, not merely reported as zero.
    *backend.map_delay.lock() = None;
    let guard = registry
        .register_owned(vec![0u8; 64 * 1024].into_boxed_slice())
        .await
        .expect("a registration after cancellations must still be admitted");
    assert_eq!(
        guard.unregister(T).await.expect("unregister"),
        Deregistered::Drained
    );
    assert_eq!(registry.registered_bytes(), 0);
}

/// The same property for the pool path, which claims its budget in
/// `map_arena`.
#[tokio::test(flavor = "multi_thread")]
async fn cancelled_pool_alloc_returns_its_budget() {
    let cfg = RdmaConfig {
        pool: small_pool_config(),
        ..RdmaConfig::default()
    };
    let (backend, registry) = mock_registry(cfg);
    *backend.map_delay.lock() = Some(Duration::from_millis(500));

    for _ in 0..4 {
        let attempt = registry.alloc_pinned(4096);
        assert!(
            tokio::time::timeout(Duration::from_millis(50), attempt)
                .await
                .is_err(),
            "the allocation was supposed to be cancelled mid-map"
        );
    }

    assert_eq!(
        registry.registered_bytes(),
        0,
        "cancelled pool allocations leaked their budget claim"
    );
    assert_eq!(
        registry.pool().arena_count(),
        0,
        "a cancelled map left an arena in the set"
    );

    *backend.map_delay.lock() = None;
    let buf = registry
        .alloc_pinned(4096)
        .await
        .expect("the pool must still be usable after cancellations");
    drop(buf);
}

/// Arena sizes that are not granule multiples must still balance.
///
/// `initial_arena_bytes` is a public field wired through
/// `VeloBuilder::rdma_config`, so nothing stops a caller passing 100_000. If
/// the claim were taken on the requested size and the release on the
/// page-rounded one, every arena would under-release by the difference; the
/// counter is unsigned, so the drift accumulates until the pool refuses
/// everything, and `publish` saturates so the gauge would keep reading zero.
/// Every other test in this file uses granule multiples, which is exactly why
/// this one does not.
#[tokio::test]
async fn unaligned_arena_sizes_balance_the_budget() {
    let cfg = RdmaConfig {
        pool: RdmaPoolConfig {
            initial_arena_bytes: 100_000,
            max_arena_bytes: 300_000,
            dedicated_arena_min: 1 << 30,
            registered_bytes_budget: 4_000_000,
            ..small_pool_config()
        },
        ..RdmaConfig::default()
    };
    let (_backend, registry) = mock_registry(cfg);

    let mut held = Vec::new();
    for _ in 0..6 {
        held.push(registry.alloc_pinned(30_000).await.expect("alloc"));
    }
    let registered = registry.registered_bytes();
    assert!(registered > 0, "nothing was accounted as registered");
    assert_eq!(
        registered % GRANULE as u64,
        0,
        "the budget claim must be the page-rounded length that is actually mapped"
    );

    drop(held);
    registry.shutdown(T).await;
    assert_eq!(
        registry.registered_bytes(),
        0,
        "reserve and release disagreed on the length; the budget is now permanently skewed"
    );
}

/// The budget ceiling holds under concurrent pressure.
///
/// The CAS loop in `try_reserve` is the only thing standing between N tasks
/// each reading an under-budget total and each concluding it has room. A plain
/// "read, compare, add" passes every sequential test in this file.
///
/// The assertion is on **how many registrations were admitted**, not on the
/// counter. A lost update makes the counter under-report, so a bare
/// "counter stays under the ceiling" check passes with the bug in place — the
/// counter is exactly the thing the bug corrupts. What cannot be faked is the
/// number of registrations simultaneously alive: every admitted registration
/// pins at least its requested length, so admitting more than the ceiling
/// allows means the ceiling did not hold, whatever the counter says.
#[tokio::test(flavor = "multi_thread", worker_threads = 8)]
async fn concurrent_registration_never_overshoots_the_budget() {
    const TASKS: usize = 24;
    const ROUNDS: usize = 8;
    const CHUNK: usize = 16 * GRANULE;
    /// Room for eight concurrent registrations, so tasks genuinely race at the
    /// ceiling rather than all fitting or all being refused.
    const BUDGET: u64 = (8 * 16 * GRANULE) as u64;

    let cfg = RdmaConfig {
        pool: RdmaPoolConfig {
            registered_bytes_budget: BUDGET,
            ..small_pool_config()
        },
        ..RdmaConfig::default()
    };
    let (_backend, registry) = mock_registry(cfg);
    let admitted = Arc::new(parking_lot::Mutex::new(Vec::new()));

    let mut tasks = Vec::new();
    for _ in 0..TASKS {
        let registry = Arc::clone(&registry);
        let admitted = Arc::clone(&admitted);
        tasks.push(tokio::spawn(async move {
            for _ in 0..ROUNDS {
                // Guards are kept, never released, so the peak is the total.
                if let Ok(guard) = registry
                    .register_owned(vec![0u8; CHUNK].into_boxed_slice())
                    .await
                {
                    admitted.lock().push(guard);
                }
                tokio::task::yield_now().await;
            }
        }));
    }
    for task in tasks {
        // Bounded: a regressed CAS loop can livelock, and a hung test binary
        // reports nothing. A timeout reports which invariant stopped holding.
        tokio::time::timeout(T, task)
            .await
            .expect("a registration task did not finish; the budget CAS loop may be livelocked")
            .expect("task panicked");
    }

    let live = admitted.lock().len();
    assert!(
        live > 0,
        "nothing was admitted at all; the test proves nothing"
    );
    assert!(
        (live * CHUNK) as u64 <= BUDGET,
        "{live} concurrent registrations of {CHUNK} B were admitted against a {BUDGET} B \
         budget: the ceiling did not hold under concurrency"
    );

    admitted.lock().clear();
    registry.shutdown(T).await;
    assert_eq!(registry.registered_bytes(), 0, "the budget did not balance");
}

/// An unaligned, non-granule external registration charges the enclosing pages.
///
/// The budget exists to be the `RLIMIT_MEMLOCK` valve, so it has to count what
/// the kernel pins. A 4097-byte buffer at an arbitrary heap address spans three
/// pages; charging its requested length would undercount by most of a factor of
/// two, and the ceiling would let through roughly twice the memory the operator
/// asked it to allow.
#[tokio::test]
async fn external_registration_charges_page_enclosing_bytes() {
    let (_backend, registry) = mock_registry(RdmaConfig::default());

    // Deliberately odd, and deliberately from the heap so it is not page-aligned.
    let leaked: &'static mut [u8] = Box::leak(vec![0u8; 4097].into_boxed_slice());
    let ptr = std::ptr::NonNull::new(leaked.as_mut_ptr()).expect("non-null");

    // SAFETY: a leaked allocation is never freed, so it outlives the
    // registration unconditionally.
    let guard = unsafe { registry.register_external(ptr, leaked.len()) }
        .await
        .expect("register");

    let charged = registry.registered_bytes();
    assert!(
        charged >= 4097,
        "the charge must cover at least the requested range: {charged}"
    );
    assert_eq!(
        charged % GRANULE as u64,
        0,
        "the charge must be a whole number of pages: {charged}"
    );

    assert_eq!(
        guard.unregister(T).await.expect("unregister"),
        Deregistered::Drained
    );
    assert_eq!(
        registry.registered_bytes(),
        0,
        "reserve and release disagreed; the budget is now permanently skewed"
    );
}

/// `deregistered()` resolves at the end of velo shutdown even when the unmap
/// itself was never confirmed.
///
/// This is what makes the future a signal a caller can actually wait on. A
/// backend that answers `ShuttingDown` — a transport already going down, a
/// wedged progress thread — leaves the sweep unable to latch honestly; but by
/// the time `graceful_shutdown` returns, transport teardown has force-unmapped
/// everything, so the region really is released. Without the final latch the
/// future would stay pending forever and a caller waiting on it before freeing
/// would wait for the life of the process.
#[tokio::test(flavor = "multi_thread")]
async fn shutdown_latches_regions_whose_unmap_was_never_confirmed() {
    let (backend, registry) = mock_registry(RdmaConfig::default());
    let guard = registry
        .register_owned(vec![0u8; 8192].into_boxed_slice())
        .await
        .expect("register");
    let watch = guard.watch();

    // The sweep will not be able to confirm anything.
    backend.refuse_unmap.store(true, Ordering::SeqCst);
    registry.shutdown(Duration::from_millis(200)).await;
    assert!(
        !watch.is_deregistered(),
        "an unconfirmed unmap must not latch during the sweep; that is the point of the sweep \
         being honest about what it knows"
    );

    // Standing in for the transport teardown that force-unmaps everything, and
    // then for the end of `Velo::graceful_shutdown`.
    backend.force_teardown();
    registry.latch_all_deregistered();

    assert!(
        watch.is_deregistered(),
        "the end of velo shutdown must resolve every surviving latch"
    );
    tokio::time::timeout(T, watch.deregistered())
        .await
        .expect("deregistered() must resolve once shutdown has completed");
    assert_eq!(
        registry.registered_bytes(),
        0,
        "regions released at the end of shutdown must give their budget back"
    );
    drop(guard);
}

/// The owned buffer survives a runtime abandoned without `shutdown`.
///
/// Plain drop glue on `RegionInner` would free the `Box` here while the backend
/// still had the pages mapped, and a peer holding the key then reads or writes
/// freed heap. The leak is the correct outcome.
///
/// Reproducing it faithfully needs the *last* `Arc<RegionInner>` to go while
/// the registration is unconfirmed. Dropping the guard spawns a background
/// deregistration that holds one, so the scenario is a runtime that dies before
/// that task ever runs — a panicking process, a `Runtime` dropped out from
/// under its tasks. Dropping the runtime cancels the task, releases the last
/// reference, and runs the destructor under test.
///
/// The assertion is on the recorded decision rather than on the memory: proving
/// a `Box` was not freed by reading it is the very use-after-free being
/// prevented, so only Miri or ASan could see it directly.
#[test]
fn abandoned_runtime_leaks_owned_buffers_rather_than_freeing_them() {
    // Taken before any runtime exists, so a blocking acquire is safe here.
    let _serialised = LEAK_COUNTER.blocking_lock();
    let before = super::region::LEAKED_BUFFERS.load(Ordering::SeqCst);
    let backend = MockBackend::new();

    let runtime = tokio::runtime::Builder::new_current_thread()
        .enable_all()
        .build()
        .expect("runtime");
    runtime.block_on({
        let backend = Arc::clone(&backend);
        async move {
            let registry = Arc::new(RdmaRegistry::new(
                backend as Arc<dyn RdmaBackend>,
                RdmaConfig::default(),
                tokio::runtime::Handle::current(),
                None,
            ));
            let guard = registry
                .register_owned(vec![0xC5u8; 8192].into_boxed_slice())
                .await
                .expect("register");
            // The background deregistration is spawned and never polled.
            drop(guard);
            drop(registry);
        }
    });
    // Cancels the pending deregistration, releasing the last reference.
    drop(runtime);

    assert_eq!(
        backend.unmap_calls(),
        0,
        "the deregistration was supposed to never run; the scenario is not what it claims"
    );
    assert_eq!(
        backend.live(),
        1,
        "the backend still holds the registration, so the pages are still pinned"
    );
    assert!(
        super::region::LEAKED_BUFFERS.load(Ordering::SeqCst) > before,
        "the owned buffer was freed while its pages were still pinned"
    );
}

/// `wait_deregistered` must not lose the wakeup when the latch closes while it
/// is between reading the flag and parking.
///
/// `notify_waiters()` stores no permit, so a `Notified` created *after* the
/// latch never hears it — the future would hang forever on a region that is
/// already released, and a caller waiting before freeing would wait for the
/// life of the process. The fix is to create the future before reading the
/// flag; this scans the window rather than hoping to hit it, mirroring
/// `velo_ext`'s `wait_for_drain_survives_guard_dropped_at_the_check`, whose
/// discipline the implementation cites.
///
/// A lost wakeup is permanent, so the per-iteration bound is short and the
/// first hit fails. It is a detector, not a deadline: a runner that fails to
/// schedule the latcher inside it looks identical from here, so the latcher is
/// joined — making the latch a fact — and the wait re-awaited under a generous
/// grace window, which costs no detection power.
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn wait_deregistered_survives_a_latch_at_the_check() {
    const ITERATIONS: usize = 4096;
    const SPIN_SWEEP: usize = 512;
    const WAITER_LEAD: usize = 600;
    const GRACE: Duration = Duration::from_secs(2);

    // `black_box` keeps LLVM from folding the busy-wait away and deleting the
    // delay the scan depends on.
    fn burn(rounds: usize) {
        let mut sink = 0usize;
        for k in 0..rounds {
            sink = std::hint::black_box(sink.wrapping_add(k));
        }
    }

    let (_backend, registry) = mock_registry(RdmaConfig::default());

    for iteration in 0..ITERATIONS {
        let guard = registry
            .register_owned(vec![0u8; GRANULE].into_boxed_slice())
            .await
            .expect("register");
        let watch = guard.watch();
        let inner = guard.watch();

        let armed = Arc::new(AtomicBool::new(false));
        let spins = iteration % SPIN_SWEEP;

        let latcher_armed = Arc::clone(&armed);
        let latcher = std::thread::spawn(move || {
            while !latcher_armed.load(Ordering::Acquire) {
                std::hint::spin_loop();
            }
            burn(spins);
            inner.latch_for_test();
        });

        let mut waiter = tokio::spawn(async move {
            armed.store(true, Ordering::Release);
            burn(WAITER_LEAD);
            watch.deregistered().await;
        });

        let finished = tokio::time::timeout(Duration::from_millis(200), &mut waiter).await;
        latcher.join().expect("latcher thread panicked");
        let joined = match finished {
            Ok(joined) => joined,
            Err(_) => tokio::time::timeout(GRACE, &mut waiter)
                .await
                .unwrap_or_else(|_| {
                    panic!(
                        "wait_deregistered lost the latch wakeup (iteration {iteration}, \
                         spins {spins})"
                    )
                }),
        };
        joined.expect("waiter task panicked");
        std::mem::forget(guard);
    }
    registry.shutdown(T).await;
}

/// Reclamation is **sweep-driven, never drop-driven**: dropping the last buffer
/// makes an arena a candidate, and nothing more.
///
/// Worth pinning on its own because the difference is invisible from the outside
/// and load-bearing on the inside. `Drop` cannot await, so an unmap issued from
/// one would have to be spawned — a detached task racing the shutdown sweep for
/// the same arena, with no admission ticket and no ordering against
/// `latch_all_deregistered`. Doing the work on the periodic sweep instead is
/// what lets a single gate cover it.
///
/// The complementary half — that a sweep *does* reclaim it — is
/// `a_dedicated_arena_is_reclaimed_when_its_buffer_drops`. This one runs no
/// sweep at all, which is exactly its point.
#[tokio::test]
async fn dropping_a_buffer_does_not_itself_unmap_its_arena() {
    let cfg = small_pool_config();
    let (backend, pool) = mock_pool(cfg.clone());
    let size = cfg.dedicated_arena_min as usize;

    let first = pool.alloc(size).await.expect("first oversize");
    assert_eq!(pool.arena_count(), 1);
    let charged = pool.registered_bytes();
    drop(first);

    assert_eq!(
        pool.live_allocations(),
        0,
        "the suballocation was returned to its arena"
    );
    assert_eq!(
        backend.live(),
        1,
        "dropping a PinnedBuf unmapped its arena; reclamation must happen on the sweep, \
         where the shutdown gate can order it"
    );
    assert_eq!(
        pool.registered_bytes(),
        charged,
        "the budget was released without an unmap having been confirmed"
    );

    // And still not reused meanwhile: a dedicated arena stays out of the general
    // search for as long as it exists, so the next oversize request maps its own
    // rather than being handed one whose reclaim has not run yet.
    let second = pool.alloc(size).await.expect("second oversize");
    assert_eq!(
        pool.arena_count(),
        2,
        "an unreclaimed dedicated arena was handed to another request"
    );
    assert_eq!(
        backend.live(),
        2,
        "each oversize request maps its own arena"
    );
    drop(second);

    // The sweep is what closes it out, which is the contrast this test exists to
    // draw.
    assert_eq!(pool.reclaim_idle().await, 2);
    assert_eq!(pool.registered_bytes(), 0);
}

/// The latch refuses to open while the backend still holds registrations.
///
/// `latch_all_deregistered` closing the latch says two things at once: callers
/// may free their memory, and `RegionInner::drop` may free the buffers velo
/// owns. Both are false if anything is still pinned, so an out-of-order
/// shutdown — the latch reached before teardown finished — must not open it.
/// This is where the two high-severity fixes meet: the leak gate reads exactly
/// the boolean this function sets.
#[tokio::test(flavor = "multi_thread")]
async fn latch_refuses_while_the_backend_still_holds_registrations() {
    let (backend, registry) = mock_registry(RdmaConfig::default());
    let guard = registry
        .register_owned(vec![0u8; 8192].into_boxed_slice())
        .await
        .expect("register");
    let watch = guard.watch();

    backend.refuse_unmap.store(true, Ordering::SeqCst);
    registry.shutdown(Duration::from_millis(100)).await;

    // Teardown has *not* run: the backend still holds the region.
    assert_eq!(backend.live_registrations(), Some(1));
    registry.latch_all_deregistered();
    assert!(
        !watch.is_deregistered(),
        "the latch opened while the backend still had the region pinned; a caller would now \
         free live memory, and the owned buffer would be freed by drop glue"
    );

    // Now teardown really has happened.
    backend.force_teardown();
    registry.latch_all_deregistered();
    assert!(
        watch.is_deregistered(),
        "the latch must open once nothing is pinned"
    );
    drop(guard);
}

/// After the latch opens, the owned buffer is freed normally rather than
/// leaked.
///
/// The other half of the leak gate. `RegionInner::drop` leaks whenever the
/// registration is unconfirmed, so an orderly shutdown has to be able to *close*
/// that gate — otherwise every owned registration in a well-behaved process
/// would leak, and the safety property would have been bought by making the
/// normal path pathological.
#[tokio::test(flavor = "multi_thread")]
async fn latched_regions_free_their_buffers_normally() {
    let _serialised = LEAK_COUNTER.lock().await;
    let before = super::region::LEAKED_BUFFERS.load(Ordering::SeqCst);
    let (backend, registry) = mock_registry(RdmaConfig::default());
    let guard = registry
        .register_owned(vec![0u8; 8192].into_boxed_slice())
        .await
        .expect("register");

    // The awkward path: the sweep cannot confirm, so only the end-of-shutdown
    // latch releases this region.
    backend.refuse_unmap.store(true, Ordering::SeqCst);
    registry.shutdown(Duration::from_millis(100)).await;
    backend.force_teardown();
    registry.latch_all_deregistered();
    assert!(guard.is_deregistered());

    // Everything goes. The buffer must be freed, not leaked.
    drop(guard);
    drop(registry);

    assert_eq!(
        super::region::LEAKED_BUFFERS.load(Ordering::SeqCst),
        before,
        "an orderly shutdown leaked an owned buffer; the leak gate never closes"
    );
}

/// `unregister_owned` on a guard that owns nothing must refuse *before*
/// deregistering.
///
/// The old order deregistered successfully and then returned `NotOwned`,
/// telling the caller the opposite of what happened: the region was gone, and
/// the error said the call had found nothing to do.
///
/// Single-threaded on purpose. `unregister_owned` takes `self`, so the refused
/// guard is dropped as the call returns, and `RegionGuard::drop` legitimately
/// *spawns* a deregistration for the region nobody released. That spawn is not
/// the bug this test is about, but on a multi-threaded runtime it can land on
/// another worker before the assertions below read the counter — which is how
/// this test failed under `llvm-cov`, reporting a deregistration that came from
/// the drop rather than from `unregister_owned`. With one worker the spawned
/// task cannot run until the test next awaits, and there is no await between
/// the call and the assertions, so the counter is read at the only instant that
/// means anything: after the refusal, before the drop's cleanup.
#[tokio::test(flavor = "current_thread")]
async fn unregister_owned_refuses_before_deregistering() {
    let (backend, registry) = mock_registry(RdmaConfig::default());

    // Leaked on purpose: the caller-owned path, where velo holds no buffer.
    let leaked: &'static mut [u8] = Box::leak(vec![0u8; 4096].into_boxed_slice());
    let ptr = std::ptr::NonNull::new(leaked.as_mut_ptr()).expect("non-null");
    // SAFETY: a leaked allocation is never freed, so it outlives the
    // registration unconditionally.
    let guard = unsafe { registry.register_external(ptr, leaked.len()) }
        .await
        .expect("register");
    let watch = guard.watch();

    let err = guard
        .unregister_owned(T)
        .await
        .expect_err("a caller-owned region has no buffer to hand back");
    assert_eq!(err, RdmaError::NotOwned);
    assert_eq!(
        backend.unmap_calls(),
        0,
        "the region was deregistered before the ownership check; the error described the \
         opposite of what happened"
    );
    assert!(!watch.is_deregistered());
}

/// Concurrent `graceful_shutdown` calls are serialised, and the sequence runs
/// once.
///
/// `Velo` is `Clone`, so several clones can arrive together. Only one can take
/// the transport join handle; another would skip the join and run straight on
/// to declaring registrations released while the progress thread was still
/// alive and its pages still pinned.
///
/// The assertion is on the *sweep count*, not on the race outcome. The
/// dangerous interleaving needs a second caller to reach the latch inside the
/// window where the first has force-unmapped but not yet joined — narrow enough
/// that a test which waits for it to happen proves nothing when it does not.
/// What is deterministic is that the sequence ran exactly once, which is the
/// property the lock exists to provide. (The latch precondition is the
/// independent backstop if it ever does interleave: it refuses while the
/// backend still reports registrations.)
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn concurrent_graceful_shutdowns_are_serialised() {
    const CALLERS: usize = 8;

    let transport = Arc::new(
        crate::transports::ucx::UcxTransportBuilder::new()
            .tls("tcp")
            .build()
            .expect("build ucx transport"),
    );
    let velo = crate::Velo::builder()
        .add_ucx_transport(Arc::clone(&transport))
        .build()
        .await
        .expect("build velo");

    let guard = velo
        .register_owned(vec![0u8; 128 * 1024].into_boxed_slice())
        .await
        .expect("register");
    let watch = guard.watch();
    assert_eq!(transport.live_regions(), 1);

    let mut callers = Vec::new();
    for _ in 0..CALLERS {
        let velo = crate::Velo::clone(&velo);
        callers.push(tokio::spawn(async move {
            velo.graceful_shutdown(velo_ext::ShutdownPolicy::Timeout(T))
                .await
        }));
    }
    for caller in callers {
        tokio::time::timeout(T, caller)
            .await
            .expect("every shutdown caller must return; a serialised one must not deadlock")
            .expect("task panicked");
    }

    assert_eq!(
        velo.rdma().expect("registry").sweep_count(),
        1,
        "the shutdown sequence ran more than once; concurrent callers were not serialised"
    );
    assert!(watch.is_deregistered());
    assert_eq!(
        transport.live_regions(),
        0,
        "shutdown returned with memory still registered"
    );
    assert_eq!(transport.live_rkeys(), 0);
    assert_eq!(velo.rdma_registered_bytes(), 0);
    drop(guard);
}

/// `Budget::charge` adds unconditionally and reports whether it stayed within
/// the ceiling.
///
/// Unit-tested directly because it is the one place the budget is *allowed* to
/// exceed its limit: the pages are already pinned by the time it is called, so
/// refusing would mean unmapping a live registration over a rounding delta.
/// The arithmetic and the verdict are what callers act on.
#[test]
fn budget_charge_reports_whether_it_stayed_within_the_ceiling() {
    let budget = Arc::new(Budget::new(1000, None));

    assert!(budget.charge(400), "400 of 1000 is within the ceiling");
    assert_eq!(budget.registered(), 400);

    assert!(
        budget.charge(600),
        "exactly at the ceiling is still within it"
    );
    assert_eq!(budget.registered(), 1000);

    assert!(
        !budget.charge(1),
        "past the ceiling must report the overshoot"
    );
    assert_eq!(
        budget.registered(),
        1001,
        "charge adds unconditionally: the pages are already pinned"
    );

    // Saturating, so an over-release cannot wrap the counter into a permanent
    // refusal of every future registration.
    budget.release(u64::MAX);
    assert_eq!(budget.registered(), 0);
}

/// A backend that pinned more than this side estimated tops the claim up, and
/// the claim is what gets released.
///
/// `raise_to` only ever raises, so this is the path that matters: an
/// under-charged region releases less than it claimed and skews the budget
/// upward for good.
#[tokio::test]
async fn a_backend_that_pins_more_than_estimated_tops_up_the_claim() {
    let (backend, registry) = mock_registry(RdmaConfig::default());
    // Two pages beyond whatever this side computes.
    backend
        .extra_effective
        .store(2 * GRANULE as u64, Ordering::SeqCst);

    let guard = registry
        .register_owned(vec![0u8; 4096].into_boxed_slice())
        .await
        .expect("register");

    let charged = registry.registered_bytes();
    let (_, effective_len) = guard.effective_range();
    assert_eq!(
        charged, effective_len,
        "the claim must match what the backend says it pinned"
    );

    assert_eq!(
        guard.unregister(T).await.expect("unregister"),
        Deregistered::Drained
    );
    assert_eq!(
        registry.registered_bytes(),
        0,
        "the top-up was not released; the budget is now permanently skewed"
    );
}

/// A refused latch retains the pool arenas too.
///
/// The pool side of the same rule: arena pages are freed only on evidence that
/// nothing is pinned, and a backend still reporting registrations is that
/// evidence being absent. Tested through `alloc_pinned` rather than
/// `register_owned`, because the two travel different code paths to the same
/// gate.
#[tokio::test(flavor = "multi_thread")]
async fn a_refused_latch_retains_pool_arenas() {
    let cfg = RdmaConfig {
        pool: small_pool_config(),
        ..RdmaConfig::default()
    };
    let (backend, registry) = mock_registry(cfg);

    let buf = registry.alloc_pinned(4096).await.expect("alloc");
    let charged = registry.registered_bytes();
    assert!(charged > 0);
    drop(buf);

    // The sweep cannot confirm, so the arena joins the unconfirmed list.
    backend.refuse_unmap.store(true, Ordering::SeqCst);
    registry.shutdown(Duration::from_millis(100)).await;
    assert_eq!(
        registry.registered_bytes(),
        charged,
        "an unconfirmed arena must keep its budget charge; the pages may still be pinned"
    );

    // Teardown has not happened: the latch must refuse, and the arena must stay.
    assert_eq!(backend.live_registrations(), Some(1));
    registry.latch_all_deregistered();
    assert_eq!(
        registry.registered_bytes(),
        charged,
        "the pool was released while the backend still held the arena"
    );

    // And once teardown really has happened, the pages are given back.
    backend.force_teardown();
    registry.latch_all_deregistered();
    assert_eq!(
        registry.registered_bytes(),
        0,
        "a successful latch must release the arenas the sweep could not confirm"
    );
}

// ---------------------------------------------------------------------------
// The copy gate
// ---------------------------------------------------------------------------

/// A `RegionInner` over a plain heap buffer, with no backend behind it.
///
/// The gate is pure velo ordering — it has nothing to say to the backend — so
/// exercising it against a real registration would only add scheduling noise
/// to the race this test is trying to lose.
fn bare_region(buffer: Box<[u8]>) -> Arc<RegionInner> {
    let ptr = buffer.as_ptr() as usize;
    let len = buffer.len();
    Arc::new(RegionInner::new(RegionParts {
        id: 1,
        generation: 1,
        backend_region_id: 1,
        ptr,
        len,
        packed_key: Bytes::from_static(b"key"),
        effective_addr: ptr as u64,
        effective_len: len as u64,
        owned: Some(buffer),
        charged: len as u64,
        shutdown: tokio_util::sync::CancellationToken::new(),
    }))
}

/// A copy can never be in progress while the deregistration latch closes.
///
/// The latch is the moment the region's owner is told it may free the memory,
/// and an anchor staged inside the region copies out of it through a raw
/// pointer. A bare `is_deregistered()` check followed by a copy is a
/// check-then-act against exactly that event: nothing stops the latch landing
/// between them, and the in-flight guard does not help, because a drain that
/// times out latches with guards still outstanding.
///
/// So the assertion lives *inside* the copy — every reader let through
/// re-checks the flag while holding the gate. Readers and the latch race
/// deliberately, so a gate that does not exclude them fails within a few
/// iterations rather than needing one unlucky interleaving.
#[test]
fn a_copy_never_overlaps_the_deregistration_latch() {
    const READERS: usize = 8;
    const ROUNDS: usize = 400;

    for _ in 0..16 {
        let region = bare_region(vec![0u8; 4096].into_boxed_slice());
        let start = Arc::new(std::sync::Barrier::new(READERS + 1));

        let readers: Vec<_> = (0..READERS)
            .map(|_| {
                let region = Arc::clone(&region);
                let start = Arc::clone(&start);
                std::thread::spawn(move || {
                    start.wait();
                    for _ in 0..ROUNDS {
                        region.with_live(|| {
                            assert!(
                                !region.is_deregistered(),
                                "the latch closed while a copy was in progress; the region's \
                                 owner may already have freed these bytes"
                            );
                        });
                    }
                })
            })
            .collect();

        start.wait();
        // Latched from this thread, racing the readers rather than after them.
        region.latch_deregistered();
        for reader in readers {
            reader.join().expect("a reader observed the latch mid-copy");
        }

        assert!(region.is_deregistered());
        assert!(
            region.with_live(|| ()).is_none(),
            "every read after the latch must refuse"
        );
        // Latched, so `RegionInner::drop` frees the owned buffer normally
        // rather than leaking it.
        drop(region);
    }
}

/// An anchor refuses to read a region that has been released.
///
/// The other half of the copy gate: the ordering test proves a copy and the
/// latch cannot overlap, and this proves the anchor actually consults it. A
/// `read_at` that skipped the check would return bytes out of memory the
/// region's owner has already been told it may free.
///
/// Driven directly rather than through a shutdown, because after
/// `RendezvousManager::shutdown` drops pinned slots ahead of the registration
/// sweep there is no longer a *normal* path that leaves an anchor alive over a
/// released region — which is the point of that change, and would make this an
/// assertion about nothing.
#[test]
fn an_anchor_refuses_to_read_a_region_that_has_been_released() {
    use super::region::RegionWatch;
    use crate::rendezvous::descriptor::DescriptorBackend;
    use crate::rendezvous::pinned::PinnedSlot;

    // Several copy-gate chunks, so the loop that walks them is exercised
    // rather than short-circuited by a single-chunk read.
    const LEN: u64 = 3 * 512 * 1024 + 977;
    let region = bare_region(vec![0x5Au8; LEN as usize].into_boxed_slice());
    let slot = PinnedSlot::from_region(
        region.in_flight.acquire(),
        RegionWatch::for_test(Arc::clone(&region)),
        DescriptorBackend::Ucx,
        region.ptr as u64,
        LEN,
        1,
        Bytes::from_static(b"packed-key"),
    );

    // While the region is live the anchor reads it, so the refusal below is a
    // change of state rather than a slot that never worked.
    let read = slot
        .read_at(0, LEN as usize)
        .expect("read while registered");
    assert_eq!(read.len(), LEN as usize);
    assert!(read.iter().all(|b| *b == 0x5A));
    assert!(slot.descriptor().is_some());
    assert!(slot.is_live());

    // The latch is the moment the region's owner may free the memory.
    region.latch_deregistered();

    assert!(
        slot.read_at(0, LEN as usize).is_none(),
        "the anchor read a region whose owner has been told it may free it"
    );
    assert!(slot.to_bytes().is_none());
    assert!(
        slot.descriptor().is_none(),
        "a descriptor for a released region would send a peer's NIC at freed memory"
    );
    assert!(!slot.is_live());

    drop(slot);
    drop(region);
}

/// A byte whose value depends on its offset with **no power-of-two period**.
///
/// This matters more than it looks. An arithmetic pattern like `i * 31 + i / 256`
/// repeats every 512 KiB, which is exactly the copy-gate chunk size — so a chunk
/// copied from the wrong offset lands on identical bytes and the test passes on
/// a coincidence. (It did, until this was fixed.) Two coprime periods, neither a
/// factor of the chunk size, make a misplaced chunk observable wherever it came
/// from.
fn chunk_pattern(i: usize) -> u8 {
    ((i % 251) ^ ((i / 251) % 257)) as u8
}

/// A copy spanning several gate chunks reassembles in the right order.
///
/// The gate is taken per chunk so a multi-gigabyte anchor cannot park a
/// deregistration for the length of its copy, which means the copy is a loop —
/// and a loop over offsets is where an off-by-one lives. The pattern depends on
/// position, so a chunk landing at the wrong offset, repeated, or dropped fails
/// rather than passing on a uniform fill.
#[test]
fn a_copy_spanning_several_gate_chunks_reassembles_in_order() {
    // Deliberately not a multiple of the chunk size: the tail is its own case.
    const LEN: usize = 2 * 512 * 1024 + 1_237;
    let backing: Vec<u8> = (0..LEN).map(chunk_pattern).collect();

    let region = bare_region(backing.clone().into_boxed_slice());
    let slot = crate::rendezvous::pinned::PinnedSlot::from_region(
        region.in_flight.acquire(),
        super::region::RegionWatch::for_test(Arc::clone(&region)),
        crate::rendezvous::descriptor::DescriptorBackend::Ucx,
        region.ptr as u64,
        LEN as u64,
        1,
        Bytes::from_static(b"packed-key"),
    );

    let whole = slot.to_bytes().expect("copy the whole anchor");
    assert_eq!(whole.len(), LEN);
    assert_eq!(
        &whole[..],
        &backing[..],
        "the chunk loop reassembled the anchor out of order"
    );

    // And a sub-range that straddles a chunk boundary.
    let straddle = slot
        .read_at(512 * 1024 - 7, 64)
        .expect("read across a chunk boundary");
    assert_eq!(&straddle[..], &backing[512 * 1024 - 7..512 * 1024 - 7 + 64]);

    region.latch_deregistered();
    drop(slot);
    drop(region);
}

/// A GET is refused once the registration layer has closed its gate.
///
/// The arena sweep samples each arena's in-flight transfer count and waits for
/// it to reach zero. Without a gate on `get` that sample is a check-then-act: a
/// `PinnedBuf` allocated before shutdown could still raise a hold and submit a
/// transfer while the sweep was awaiting an unmap, and the count would rise
/// again behind it. Gated, the count only falls, and the wait converges.
#[tokio::test(flavor = "multi_thread")]
async fn a_transfer_is_refused_once_the_registration_gate_is_closed() {
    let (backend, registry) = mock_registry(RdmaConfig::default());
    let buf = registry.alloc_pinned(4096).await.expect("alloc");
    let req = BackendGet {
        peer: velo_ext::InstanceId::new_v4(),
        remote_addr: 0x1000,
        packed_key: Bytes::from_static(b"key"),
        local_region_id: buf.backend_region_id(),
        local_offset: buf.arena_offset(),
        len: 4096,
    };

    // Open: the mock accepts it.
    registry
        .get(req.clone())
        .await
        .expect("get before shutdown");

    registry.shutdown(T).await;

    assert_eq!(
        registry.get(req).await,
        Err(RdmaError::ShuttingDown),
        "a transfer submitted after the gate closed would make the arena sweep's in-flight \
         sample a check-then-act"
    );
    drop(buf);
    let _ = backend;
}

/// The arena sweep waits for a transfer hold, and proceeds when it releases.
///
/// The integration test asserts this end to end through `graceful_shutdown`;
/// this one asserts it against the sweep itself, so a change to the wait fails
/// here without needing a UCX transport to notice.
#[tokio::test(flavor = "multi_thread")]
async fn the_arena_sweep_waits_for_a_transfer_hold() {
    let (backend, registry) = mock_registry(RdmaConfig::default());
    let buf = registry.alloc_pinned(4096).await.expect("alloc");
    let hold = buf.hold();

    // Released after the sweep has certainly started waiting.
    let releaser = tokio::spawn(async move {
        tokio::time::sleep(Duration::from_millis(300)).await;
        drop(hold);
    });

    let started = std::time::Instant::now();
    registry.shutdown(T).await;
    let waited = started.elapsed();

    releaser.await.expect("releaser");
    assert!(
        waited >= Duration::from_millis(200),
        "the sweep unmapped after {waited:?}, without waiting for the transfer holding one of \
         its arenas"
    );
    assert_eq!(
        backend.live(),
        0,
        "the arena should still have been unmapped"
    );
    drop(buf);
}

/// Shutdown demotes a pinned slot instead of dropping it, so a pull already in
/// flight still finishes.
///
/// The sweep needs the pool suballocation and the region in-flight guard back
/// before it can drain anything — that is why the slot is touched at all — but
/// the messenger has only just gated *new* requests. A chunked pull admitted
/// before that gate is still entitled to complete, and dropping its slot turns
/// a completion into "chunk not found" halfway through a transfer.
///
/// So the test asserts both halves at once: the region is released, and the
/// transfer that was already open keeps returning its chunks.
#[test]
fn shutdown_demotes_a_pinned_slot_rather_than_dropping_it() {
    use super::region::RegionWatch;
    use crate::rendezvous::pinned::PinnedSlot;
    use crate::rendezvous::store::{DEFAULT_CHUNK_SIZE, DataStore, SlotBody, StageMode};

    const LEN: usize = 3 * 512 * 1024 + 41;
    let backing: Vec<u8> = (0..LEN).map(|i| (i.wrapping_mul(17)) as u8).collect();
    let region = bare_region(backing.clone().into_boxed_slice());
    let store = DataStore::new();

    let local_id = store.register_body(
        SlotBody::Pinned(PinnedSlot::from_region(
            region.in_flight.acquire(),
            RegionWatch::for_test(Arc::clone(&region)),
            crate::rendezvous::descriptor::DescriptorBackend::Ucx,
            region.ptr as u64,
            LEN as u64,
            1,
            Bytes::from_static(b"packed-key"),
        )),
        None,
    );
    assert_eq!(store.stage_mode(local_id), Some(StageMode::Pinned));
    assert_eq!(
        region.in_flight.in_flight_count(),
        1,
        "the anchor should be holding the region open"
    );

    // A pull that is already under way when shutdown begins.
    let lease = store.acquire_read_lock(local_id).expect("lease");
    let (transfer, chunk_size, chunks) = store
        .create_transfer(local_id, lease, DEFAULT_CHUNK_SIZE)
        .expect("transfer");
    let first = store.get_chunk(transfer, 0).expect("the first chunk");
    assert_eq!(&first[..], &backing[..chunk_size as usize]);

    let (demoted, dropped) = store.demote_pinned_slots();
    assert_eq!((demoted, dropped), (1, 0));

    assert_eq!(
        store.stage_mode(local_id),
        Some(StageMode::InMemory),
        "the slot should have moved to the heap, not vanished"
    );
    assert_eq!(
        region.in_flight.in_flight_count(),
        0,
        "demotion did not release the region guard the sweep is waiting on"
    );

    // The rest of the pull still completes, byte for byte.
    for index in 1..chunks {
        let at = index as usize * chunk_size as usize;
        let chunk = store
            .get_chunk(transfer, index)
            .unwrap_or_else(|| panic!("chunk {index} of an admitted pull was dropped"));
        assert_eq!(&chunk[..], &backing[at..at + chunk.len()]);
    }
    assert_eq!(
        &store.get_data(local_id).expect("the whole slot")[..],
        &backing[..]
    );

    region.latch_deregistered();
    drop(store);
    drop(region);
}

/// A slot whose region is already gone is dropped rather than demoted.
///
/// There is nothing left to copy, and a pull against it was going to fail
/// whatever happened — so the sweep takes the slot away rather than leaving a
/// body it could not fill.
#[test]
fn shutdown_drops_a_slot_whose_region_has_already_gone() {
    use super::region::RegionWatch;
    use crate::rendezvous::pinned::PinnedSlot;
    use crate::rendezvous::store::{DataStore, SlotBody};

    let region = bare_region(vec![3u8; 4096].into_boxed_slice());
    let store = DataStore::new();
    let local_id = store.register_body(
        SlotBody::Pinned(PinnedSlot::from_region(
            region.in_flight.acquire(),
            RegionWatch::for_test(Arc::clone(&region)),
            crate::rendezvous::descriptor::DescriptorBackend::Ucx,
            region.ptr as u64,
            4096,
            1,
            Bytes::from_static(b"packed-key"),
        )),
        None,
    );

    region.latch_deregistered();

    assert_eq!(store.demote_pinned_slots(), (0, 1));
    assert!(
        store.metadata(local_id).is_none(),
        "a slot over a released region has nothing to demote to"
    );
    drop(region);
}