aisimulate-core 0.12.0

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

use std::cmp::Ordering;
use std::collections::{BTreeSet, VecDeque};
use std::sync::Arc;

use super::{HostOffloadObservation, HostOffloadObservationData, HostOffloadObserver};
use crate::engine::common::hashing::SequenceHash;
use anyhow::{Result, bail};
use rustc_hash::{FxHashMap, FxHashSet};
use uuid::Uuid;

/// Logical identity shared by native G1 and the framework-selected host tier.
#[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
#[repr(transparent)]
pub(crate) struct HostBlockKey(SequenceHash);

impl HostBlockKey {
    pub(crate) const fn new(sequence_hash: SequenceHash) -> Self {
        Self(sequence_hash)
    }

    pub(crate) const fn sequence_hash(self) -> SequenceHash {
        self.0
    }
}

/// Rank-local transfer identity used to correlate scheduler state and events.
#[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
#[repr(transparent)]
pub(crate) struct TransferId(u64);

impl TransferId {
    pub(crate) const fn new(value: u64) -> Self {
        Self(value)
    }

    pub(crate) const fn get(self) -> u64 {
        self.0
    }
}

/// Physical parameters for one rank-local host tier.
#[derive(Clone, Copy, Debug, PartialEq)]
pub(crate) struct HostTierConfig {
    pub(crate) capacity_blocks: usize,
    pub(crate) block_bytes: usize,
    /// `0.0` models an instantaneous transfer.
    pub(crate) d2h_bandwidth_gbps: f64,
    /// `0.0` models an instantaneous transfer.
    pub(crate) h2d_bandwidth_gbps: f64,
}

impl HostTierConfig {
    fn validate(self) -> Result<Self> {
        if self.capacity_blocks == 0 {
            bail!("host capacity must be positive");
        }
        if self.block_bytes == 0 {
            bail!("host block size must be positive");
        }
        let max_bytes = self
            .capacity_blocks
            .checked_mul(self.block_bytes)
            .ok_or_else(|| anyhow::anyhow!("host capacity in bytes overflowed"))?;
        validate_bandwidth("D2H", self.d2h_bandwidth_gbps, max_bytes)?;
        validate_bandwidth("H2D", self.h2d_bandwidth_gbps, max_bytes)?;
        Ok(self)
    }
}

/// Visibility of one logical host block.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub(crate) enum Lookup {
    Miss,
    Pending { transfer_id: TransferId },
    Hit,
}

/// Result of atomically admitting one framework-selected store cohort.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub(crate) enum StoreOutcome {
    Prepared {
        transfer_id: TransferId,
        stored_blocks: usize,
    },
    AlreadyPresent,
    RetryCapacity {
        structurally_unfittable: bool,
    },
}

/// Result of scheduling an H2D selected by a framework adapter.
#[derive(Clone, Copy, Debug, PartialEq)]
pub(crate) enum LoadOutcome {
    Queued(TransferId),
    Miss,
}

/// Successful transfer completion returned to the framework adapter.
#[derive(Clone, Debug, Eq, PartialEq)]
pub(crate) enum CompletedTransfer {
    Store {
        request_id: Uuid,
        transfer_id: TransferId,
        blocks: Vec<HostBlockKey>,
    },
    Load {
        request_id: Uuid,
        transfer_id: TransferId,
        blocks: Vec<HostBlockKey>,
    },
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
enum EntryState {
    PendingStore { transfer_id: TransferId },
    Resident,
}

#[derive(Clone, Copy, Debug)]
struct Entry {
    state: EntryState,
    load_pins: usize,
}

#[derive(Clone, Debug)]
enum Transfer {
    Store {
        request_id: Uuid,
        blocks: Vec<HostBlockKey>,
        completes_at_ms: Option<f64>,
    },
    Load {
        request_id: Uuid,
        blocks: Vec<HostBlockKey>,
        completes_at_ms: f64,
    },
}

impl Transfer {
    fn deadline(&self) -> Option<f64> {
        match self {
            Self::Store {
                completes_at_ms, ..
            } => *completes_at_ms,
            Self::Load {
                completes_at_ms, ..
            } => Some(*completes_at_ms),
        }
    }
}

#[derive(Clone, Copy, Debug)]
struct Deadline {
    at_ms: f64,
    completion_order: u8,
    transfer_id: TransferId,
}

impl PartialEq for Deadline {
    fn eq(&self, other: &Self) -> bool {
        self.at_ms.to_bits() == other.at_ms.to_bits()
            && self.completion_order == other.completion_order
            && self.transfer_id == other.transfer_id
    }
}

impl Eq for Deadline {}

impl PartialOrd for Deadline {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

impl Ord for Deadline {
    fn cmp(&self, other: &Self) -> Ordering {
        self.at_ms
            .total_cmp(&other.at_ms)
            .then_with(|| self.completion_order.cmp(&other.completion_order))
            .then_with(|| self.transfer_id.cmp(&other.transfer_id))
    }
}

/// Recency for resident, currently unpinned blocks.
///
/// A framework operation may make several blocks evictable at once. They share
/// one epoch because vLLM exposes no stable policy order within its Python-set
/// completion cohort; `HostBlockKey` is only a deterministic tie-break.
#[derive(Default)]
struct Lru {
    clock: u64,
    touched_at: FxHashMap<HostBlockKey, u64>,
    oldest_first: BTreeSet<(u64, HostBlockKey)>,
}

impl Lru {
    fn next_epoch(&mut self) -> u64 {
        self.clock = self.clock.checked_add(1).expect("host LRU clock overflow");
        self.clock
    }

    fn touch_at(&mut self, key: HostBlockKey, epoch: u64) {
        self.remove(key);
        self.touched_at.insert(key, epoch);
        assert!(self.oldest_first.insert((epoch, key)));
    }

    fn touch(&mut self, key: HostBlockKey) {
        let epoch = self.next_epoch();
        self.touch_at(key, epoch);
    }

    fn touch_cohort(&mut self, keys: &[HostBlockKey]) {
        if keys.is_empty() {
            return;
        }
        let epoch = self.next_epoch();
        for key in keys {
            self.touch_at(*key, epoch);
        }
    }

    fn remove(&mut self, key: HostBlockKey) {
        if let Some(timestamp) = self.touched_at.remove(&key) {
            assert!(self.oldest_first.remove(&(timestamp, key)));
        }
    }

    fn victims(
        &self,
        count: usize,
        mut is_eligible: impl FnMut(HostBlockKey) -> bool,
    ) -> Vec<HostBlockKey> {
        self.oldest_first
            .iter()
            .map(|(_, key)| *key)
            .filter(|key| is_eligible(*key))
            .take(count)
            .collect()
    }
}

#[derive(Clone, Copy, Debug)]
struct FifoLane {
    bytes_per_ms: f64,
    tail_ms: f64,
}

impl FifoLane {
    fn new(gbps: f64) -> Self {
        Self {
            bytes_per_ms: if gbps == 0.0 {
                f64::INFINITY
            } else {
                gbps * 1_000_000.0
            },
            tail_ms: 0.0,
        }
    }

    fn submit(&mut self, not_before_ms: f64, bytes: usize) -> f64 {
        let starts_at_ms = self.tail_ms.max(not_before_ms);
        let duration_ms = if self.bytes_per_ms.is_infinite() {
            0.0
        } else {
            bytes as f64 / self.bytes_per_ms
        };
        self.tail_ms = starts_at_ms + duration_ms;
        self.tail_ms
    }
}

/// Success-only G2 residency and virtual-transfer state.
pub(crate) struct HostTier {
    config: HostTierConfig,
    entries: FxHashMap<HostBlockKey, Entry>,
    lru: Lru,
    transfers: FxHashMap<TransferId, Transfer>,
    deadlines: BTreeSet<Deadline>,
    prepared_stores: VecDeque<TransferId>,
    d2h: FifoLane,
    h2d: FifoLane,
    next_transfer_id: u64,
    current_time_ms: f64,
    warned_structurally_unfittable: bool,
    /// Installed only by detailed replay artifacts; ordinary runs retain no
    /// host-event buffer and allocate no observation payloads.
    observer: Option<Arc<dyn HostOffloadObserver>>,
}

impl HostTier {
    pub(crate) fn new(config: HostTierConfig) -> Result<Self> {
        let config = config.validate()?;
        Ok(Self {
            config,
            entries: FxHashMap::default(),
            lru: Lru::default(),
            transfers: FxHashMap::default(),
            deadlines: BTreeSet::new(),
            prepared_stores: VecDeque::new(),
            d2h: FifoLane::new(config.d2h_bandwidth_gbps),
            h2d: FifoLane::new(config.h2d_bandwidth_gbps),
            next_transfer_id: 0,
            current_time_ms: 0.0,
            warned_structurally_unfittable: false,
            observer: None,
        })
    }

    pub(crate) fn set_observer(&mut self, observer: Arc<dyn HostOffloadObserver>) {
        self.observer = Some(observer);
    }

    /// Atomically admit the missing subset of one per-request store cohort.
    pub(crate) fn prepare_store(
        &mut self,
        request_id: Uuid,
        blocks: &[HostBlockKey],
        now_ms: f64,
    ) -> StoreOutcome {
        let now_ms = self.prepare_mutation(now_ms);
        let mut protected = FxHashSet::default();
        let mut missing = Vec::new();
        for key in blocks.iter().copied() {
            if protected.insert(key) && !self.entries.contains_key(&key) {
                missing.push(key);
            }
        }
        if missing.is_empty() {
            return StoreOutcome::AlreadyPresent;
        }

        let free = self.config.capacity_blocks - self.entries.len();
        let needed_victims = missing.len().saturating_sub(free);
        let entries = &self.entries;
        let victims = self.lru.victims(needed_victims, |key| {
            !protected.contains(&key)
                && entries.get(&key).is_some_and(|entry| {
                    entry.state == EntryState::Resident && entry.load_pins == 0
                })
        });
        if victims.len() != needed_victims {
            let structurally_unfittable = protected.len() > self.config.capacity_blocks;
            if structurally_unfittable && !self.warned_structurally_unfittable {
                tracing::warn!(
                    cohort_blocks = protected.len(),
                    missing_blocks = missing.len(),
                    host_capacity_blocks = self.config.capacity_blocks,
                    "native host-offload store cohort cannot fit; retrying without advancing the store cursor"
                );
                self.warned_structurally_unfittable = true;
            }
            self.observe(HostOffloadObservation {
                request_id,
                event: HostOffloadObservationData::CapacityRetry {
                    at_ms: now_ms,
                    blocks: &missing,
                    structurally_unfittable,
                },
            });
            return StoreOutcome::RetryCapacity {
                structurally_unfittable,
            };
        }

        for victim in victims {
            let entry = self
                .entries
                .remove(&victim)
                .expect("selected host victim disappeared");
            assert_eq!(entry.state, EntryState::Resident);
            assert_eq!(entry.load_pins, 0);
            self.lru.remove(victim);
            self.observe(HostOffloadObservation {
                request_id,
                event: HostOffloadObservationData::Evicted {
                    at_ms: now_ms,
                    block: victim,
                },
            });
        }

        let transfer_id = self.allocate_transfer_id();
        for key in &missing {
            assert!(
                self.entries
                    .insert(
                        *key,
                        Entry {
                            state: EntryState::PendingStore { transfer_id },
                            load_pins: 0,
                        },
                    )
                    .is_none()
            );
        }
        let stored_blocks = missing.len();
        self.observe(HostOffloadObservation {
            request_id,
            event: HostOffloadObservationData::StorePrepared {
                at_ms: now_ms,
                transfer_id,
                blocks: &missing,
            },
        });
        self.transfers.insert(
            transfer_id,
            Transfer::Store {
                request_id,
                blocks: missing,
                completes_at_ms: None,
            },
        );
        self.prepared_stores.push_back(transfer_id);
        StoreOutcome::Prepared {
            transfer_id,
            stored_blocks,
        }
    }

    /// Submit stores prepared during the preceding engine step.
    pub(crate) fn submit_prepared_stores(&mut self, now_ms: f64) -> usize {
        self.prepare_mutation(now_ms);
        let mut submitted = 0usize;
        while let Some(transfer_id) = self.prepared_stores.pop_front() {
            let block_count = match self.transfers.get(&transfer_id) {
                Some(Transfer::Store {
                    blocks,
                    completes_at_ms: None,
                    ..
                }) => blocks.len(),
                _ => panic!("prepared store queue lost its transfer"),
            };
            let completes_at_ms = self.d2h.submit(now_ms, self.transfer_bytes(block_count));
            let Some(Transfer::Store {
                completes_at_ms: deadline,
                ..
            }) = self.transfers.get_mut(&transfer_id)
            else {
                unreachable!()
            };
            *deadline = Some(completes_at_ms);
            assert!(self.deadlines.insert(Deadline {
                at_ms: completes_at_ms,
                completion_order: 0,
                transfer_id,
            }));
            if let Some(Transfer::Store {
                request_id, blocks, ..
            }) = self.transfers.get(&transfer_id)
            {
                self.observe(HostOffloadObservation {
                    request_id: *request_id,
                    event: HostOffloadObservationData::StoreSubmitted {
                        at_ms: now_ms,
                        completes_at_ms,
                        transfer_id,
                        blocks,
                    },
                });
            }
            submitted += 1;
        }
        submitted
    }

    pub(crate) fn lookup(&self, key: HostBlockKey) -> Lookup {
        match self.entries.get(&key).map(|entry| entry.state) {
            None => Lookup::Miss,
            Some(EntryState::PendingStore { transfer_id }) => Lookup::Pending { transfer_id },
            Some(EntryState::Resident) => Lookup::Hit,
        }
    }

    /// Apply framework-defined recency order. Not-ready or pinned blocks are
    /// outside vLLM's evictable LRU and ignore ordinary touches.
    pub(crate) fn touch(&mut self, key: HostBlockKey) {
        if self
            .entries
            .get(&key)
            .is_some_and(|entry| entry.state == EntryState::Resident && entry.load_pins == 0)
        {
            self.lru.touch(key);
        }
    }

    /// Queue one H2D after any G1 source-reuse fence selected by the adapter.
    pub(crate) fn schedule_load(
        &mut self,
        request_id: Uuid,
        blocks: &[HostBlockKey],
        now_ms: f64,
        not_before_ms: f64,
    ) -> LoadOutcome {
        let now_ms = self.prepare_mutation(now_ms);
        assert_valid_time("host load dependency time", not_before_ms);
        let mut seen = FxHashSet::default();
        if blocks.is_empty()
            || blocks.len() > self.config.capacity_blocks
            || blocks.iter().any(|key| !seen.insert(*key))
            || blocks.iter().any(|key| {
                !self
                    .entries
                    .get(key)
                    .is_some_and(|entry| entry.state == EntryState::Resident)
            })
        {
            return LoadOutcome::Miss;
        }

        for key in blocks {
            let entry = self
                .entries
                .get_mut(key)
                .expect("validated host load source disappeared");
            if entry.load_pins == 0 {
                self.lru.remove(*key);
            }
            entry.load_pins = entry
                .load_pins
                .checked_add(1)
                .expect("host load pin count overflow");
        }

        let transfer_id = self.allocate_transfer_id();
        let completes_at_ms = self
            .h2d
            .submit(now_ms.max(not_before_ms), self.transfer_bytes(blocks.len()));
        self.transfers.insert(
            transfer_id,
            Transfer::Load {
                request_id,
                blocks: blocks.to_vec(),
                completes_at_ms,
            },
        );
        assert!(self.deadlines.insert(Deadline {
            at_ms: completes_at_ms,
            completion_order: 1,
            transfer_id,
        }));
        self.observe(HostOffloadObservation {
            request_id,
            event: HostOffloadObservationData::LoadQueued {
                at_ms: now_ms,
                completes_at_ms,
                transfer_id,
                blocks,
            },
        });
        LoadOutcome::Queued(transfer_id)
    }

    /// Cancel an H2D and release its G2 pins without shortening the FIFO lane.
    /// The observation may carry the same or a later command timestamp without
    /// advancing the HostTier clock past a transfer hidden until the next boundary.
    pub(crate) fn cancel_load(
        &mut self,
        transfer_id: TransferId,
        mutation_time_ms: f64,
        observed_at_ms: f64,
    ) -> bool {
        let deadline = match self.transfers.get(&transfer_id) {
            Some(Transfer::Load {
                completes_at_ms, ..
            }) => Deadline {
                at_ms: *completes_at_ms,
                completion_order: 1,
                transfer_id,
            },
            _ => return false,
        };
        assert_valid_time("host load cancellation mutation time", mutation_time_ms);
        assert_valid_time("host load cancellation observation time", observed_at_ms);
        assert!(
            observed_at_ms >= mutation_time_ms,
            "host load cancellation observation cannot precede its mutation"
        );
        self.prepare_mutation(mutation_time_ms);
        let Some(Transfer::Load {
            request_id, blocks, ..
        }) = self.transfers.remove(&transfer_id)
        else {
            unreachable!()
        };
        assert!(
            self.deadlines.remove(&deadline),
            "cancelled host load lost its deadline"
        );
        self.release_load_pins(&blocks);
        self.observe(HostOffloadObservation {
            request_id,
            event: HostOffloadObservationData::LoadCancelled {
                at_ms: observed_at_ms,
                transfer_id,
                blocks: &blocks,
            },
        });
        true
    }

    /// Complete every submitted transfer due by `now_ms` in stable order.
    pub(crate) fn tick(&mut self, now_ms: f64) -> Vec<CompletedTransfer> {
        assert_valid_time("host transfer time", now_ms);
        assert!(
            now_ms >= self.current_time_ms,
            "host transfer clock cannot move backwards"
        );
        self.current_time_ms = now_ms;
        let mut completed = Vec::new();
        while self
            .deadlines
            .first()
            .is_some_and(|deadline| deadline.at_ms <= now_ms)
        {
            let deadline = self
                .deadlines
                .pop_first()
                .expect("selected deadline disappeared");
            let transfer_id = deadline.transfer_id;
            let transfer = self
                .transfers
                .remove(&transfer_id)
                .expect("due host transfer disappeared");
            assert_eq!(transfer.deadline(), Some(deadline.at_ms));
            match transfer {
                Transfer::Store {
                    request_id, blocks, ..
                } => {
                    for key in &blocks {
                        let entry = self
                            .entries
                            .get_mut(key)
                            .expect("completed store lost its pending entry");
                        assert_eq!(entry.state, EntryState::PendingStore { transfer_id });
                        entry.state = EntryState::Resident;
                    }
                    self.lru.touch_cohort(&blocks);
                    self.observe(HostOffloadObservation {
                        request_id,
                        event: HostOffloadObservationData::StoreCompleted {
                            at_ms: deadline.at_ms,
                            transfer_id,
                            blocks: &blocks,
                        },
                    });
                    completed.push(CompletedTransfer::Store {
                        request_id,
                        transfer_id,
                        blocks,
                    });
                }
                Transfer::Load {
                    request_id, blocks, ..
                } => {
                    self.release_load_pins(&blocks);
                    self.observe(HostOffloadObservation {
                        request_id,
                        event: HostOffloadObservationData::LoadCompleted {
                            at_ms: deadline.at_ms,
                            transfer_id,
                            blocks: &blocks,
                        },
                    });
                    completed.push(CompletedTransfer::Load {
                        request_id,
                        transfer_id,
                        blocks,
                    });
                }
            }
        }
        completed
    }

    pub(crate) fn next_deadline(&self) -> Option<f64> {
        self.deadlines.first().map(|deadline| deadline.at_ms)
    }

    pub(crate) fn current_time_ms(&self) -> f64 {
        self.current_time_ms
    }

    pub(crate) fn transfer_deadline(&self, transfer_id: TransferId) -> Option<f64> {
        self.transfers.get(&transfer_id)?.deadline()
    }

    pub(crate) fn has_transfer(&self, transfer_id: TransferId) -> bool {
        self.transfers.contains_key(&transfer_id)
    }

    #[cfg(test)]
    pub(crate) fn needs_engine_boundary(&self) -> bool {
        !self.prepared_stores.is_empty()
    }

    pub(crate) fn has_pending_work(&self) -> bool {
        !self.transfers.is_empty()
    }

    fn observe(&self, observation: HostOffloadObservation<'_>) {
        if let Some(observer) = &self.observer {
            observer.record(observation);
        }
    }

    #[cfg(test)]
    pub(crate) fn used_blocks(&self) -> usize {
        self.entries.len()
    }

    #[cfg(test)]
    pub(crate) fn resident_blocks(&self) -> usize {
        self.entries
            .values()
            .filter(|entry| entry.state == EntryState::Resident)
            .count()
    }

    #[cfg(test)]
    pub(crate) fn is_resident(&self, key: HostBlockKey) -> bool {
        self.lookup(key) == Lookup::Hit
    }

    #[cfg(test)]
    pub(crate) fn resident_snapshot(&self) -> Vec<HostBlockKey> {
        let mut blocks: Vec<_> = self
            .entries
            .iter()
            .filter_map(|(key, entry)| (entry.state == EntryState::Resident).then_some(*key))
            .collect();
        blocks.sort_unstable();
        blocks
    }

    fn release_load_pins(&mut self, blocks: &[HostBlockKey]) {
        let mut newly_unpinned = Vec::new();
        for key in blocks {
            let entry = self
                .entries
                .get_mut(key)
                .expect("terminal host load lost its source");
            entry.load_pins = entry
                .load_pins
                .checked_sub(1)
                .expect("terminal host load did not hold a pin");
            if entry.load_pins == 0 {
                newly_unpinned.push(*key);
            }
        }
        self.lru.touch_cohort(&newly_unpinned);
    }

    fn prepare_mutation(&mut self, now_ms: f64) -> f64 {
        assert_valid_time("host mutation time", now_ms);
        let now_ms = now_ms.max(self.current_time_ms);
        assert!(
            self.next_deadline()
                .is_none_or(|deadline| deadline >= now_ms),
            "host completions must be ticked before a later mutation"
        );
        self.current_time_ms = now_ms;
        now_ms
    }

    fn allocate_transfer_id(&mut self) -> TransferId {
        let id = TransferId::new(self.next_transfer_id);
        self.next_transfer_id = self
            .next_transfer_id
            .checked_add(1)
            .expect("host transfer id overflow");
        id
    }

    fn transfer_bytes(&self, blocks: usize) -> usize {
        self.config
            .block_bytes
            .checked_mul(blocks)
            .expect("host transfer byte count overflow")
    }
}

fn validate_bandwidth(direction: &str, gbps: f64, max_bytes: usize) -> Result<()> {
    if !gbps.is_finite() || gbps < 0.0 {
        bail!("host {direction} bandwidth must be finite and non-negative");
    }
    let bytes_per_ms = gbps * 1_000_000.0;
    if !bytes_per_ms.is_finite() || (gbps > 0.0 && !(max_bytes as f64 / bytes_per_ms).is_finite()) {
        bail!("host {direction} transfer timing overflowed");
    }
    Ok(())
}

fn assert_valid_time(label: &str, time_ms: f64) {
    assert!(
        time_ms.is_finite() && time_ms >= 0.0,
        "{label} must be finite and non-negative"
    );
}

#[cfg(test)]
mod tests {
    use std::sync::Mutex;

    use super::*;

    #[derive(Debug, PartialEq)]
    struct Cancellation {
        request_id: Uuid,
        at_ms: f64,
        transfer_id: u64,
        block_hashes: Vec<u64>,
    }

    #[derive(Default)]
    struct CancellationObserver {
        cancelled: Mutex<Option<Cancellation>>,
    }

    impl HostOffloadObserver for CancellationObserver {
        fn record(&self, observation: HostOffloadObservation<'_>) {
            let HostOffloadObservation {
                request_id,
                event:
                    HostOffloadObservationData::LoadCancelled {
                        at_ms,
                        transfer_id,
                        blocks,
                    },
            } = observation
            else {
                return;
            };
            *self
                .cancelled
                .lock()
                .unwrap_or_else(|poisoned| poisoned.into_inner()) = Some(Cancellation {
                request_id,
                at_ms,
                transfer_id: transfer_id.get(),
                block_hashes: blocks.iter().map(|block| block.sequence_hash()).collect(),
            });
        }
    }

    fn key(value: u64) -> HostBlockKey {
        HostBlockKey::new(value)
    }

    fn request_id() -> Uuid {
        Uuid::from_u128(1)
    }

    fn tier(capacity_blocks: usize) -> HostTier {
        HostTier::new(HostTierConfig {
            capacity_blocks,
            block_bytes: 1_000_000,
            d2h_bandwidth_gbps: 1.0,
            h2d_bandwidth_gbps: 1.0,
        })
        .unwrap()
    }

    fn make_resident(tier: &mut HostTier, blocks: &[HostBlockKey], at_ms: f64) -> f64 {
        let StoreOutcome::Prepared { transfer_id, .. } =
            tier.prepare_store(request_id(), blocks, at_ms)
        else {
            panic!("test host cohort must be absent")
        };
        assert_eq!(tier.submit_prepared_stores(at_ms), 1);
        let deadline = tier.transfer_deadline(transfer_id).unwrap();
        tier.tick(deadline);
        deadline
    }

    #[test]
    fn store_is_pending_until_the_next_boundary_and_deadline() {
        let mut tier = tier(2);
        assert_eq!(
            tier.prepare_store(request_id(), &[key(1), key(2)], 5.0),
            StoreOutcome::Prepared {
                transfer_id: TransferId::new(0),
                stored_blocks: 2,
            }
        );
        assert!(tier.needs_engine_boundary());
        assert_eq!(tier.next_deadline(), None);
        assert_eq!(
            tier.lookup(key(1)),
            Lookup::Pending {
                transfer_id: TransferId::new(0)
            }
        );
        assert_eq!(
            tier.prepare_store(request_id(), &[key(1), key(2)], 5.0),
            StoreOutcome::AlreadyPresent,
            "pending stores participate in the missing-only filter"
        );
        assert_eq!(tier.next_transfer_id, 1);

        assert_eq!(tier.submit_prepared_stores(8.0), 1);
        assert_eq!(tier.transfer_deadline(TransferId::new(0)), Some(10.0));
        assert!(!tier.needs_engine_boundary());
        assert!(tier.tick(9.0).is_empty());
        assert_eq!(tier.tick(10.0).len(), 1);
        assert_eq!(tier.resident_snapshot(), vec![key(1), key(2)]);
    }

    #[test]
    fn pending_touches_are_ignored_and_completion_shares_one_epoch() {
        let mut tier = tier(2);
        let StoreOutcome::Prepared { transfer_id, .. } =
            tier.prepare_store(request_id(), &[key(2), key(1)], 0.0)
        else {
            panic!("initial store must prepare")
        };

        // vLLM excludes ref_cnt=-1 entries from its evictable LRU.
        tier.touch(key(2));
        tier.touch(key(1));
        assert!(tier.lru.touched_at.is_empty());
        tier.submit_prepared_stores(0.0);
        tier.tick(tier.transfer_deadline(transfer_id).unwrap());
        assert_eq!(tier.lru.touched_at[&key(1)], tier.lru.touched_at[&key(2)]);

        assert!(matches!(
            tier.prepare_store(request_id(), &[key(3)], 2.0),
            StoreOutcome::Prepared { .. }
        ));
        // Key order is only the simulator's deterministic tie-break within
        // the completion cohort, not a framework policy claim.
        assert_eq!(tier.lookup(key(1)), Lookup::Miss);
        assert!(tier.is_resident(key(2)));
    }

    #[test]
    fn older_completion_cohort_is_exhausted_before_newer_cohort() {
        let mut tier = tier(3);
        let mut now = make_resident(&mut tier, &[key(1), key(2)], 0.0);
        let older_epoch = tier.lru.touched_at[&key(1)];
        assert_eq!(older_epoch, tier.lru.touched_at[&key(2)]);
        now = make_resident(&mut tier, &[key(3)], now);
        assert!(older_epoch < tier.lru.touched_at[&key(3)]);

        let StoreOutcome::Prepared { transfer_id, .. } =
            tier.prepare_store(request_id(), &[key(4)], now)
        else {
            panic!("older cohort should provide capacity")
        };
        assert_eq!(tier.lookup(key(1)), Lookup::Miss);
        assert!(tier.is_resident(key(2)) && tier.is_resident(key(3)));
        tier.submit_prepared_stores(now);
        now = tier.transfer_deadline(transfer_id).unwrap();
        tier.tick(now);

        assert!(matches!(
            tier.prepare_store(request_id(), &[key(5)], now),
            StoreOutcome::Prepared { .. }
        ));
        assert_eq!(tier.lookup(key(2)), Lookup::Miss);
        assert!(tier.is_resident(key(3)) && tier.is_resident(key(4)));
    }

    #[test]
    fn missing_only_store_is_atomic_and_protects_its_present_inputs() {
        let mut tier = tier(2);
        let now = make_resident(&mut tier, &[key(1), key(2)], 0.0);
        tier.touch(key(1));
        assert_eq!(
            tier.prepare_store(request_id(), &[key(1), key(3)], now),
            StoreOutcome::Prepared {
                transfer_id: TransferId::new(1),
                stored_blocks: 1,
            }
        );
        assert!(tier.is_resident(key(1)));
        assert_eq!(tier.lookup(key(2)), Lookup::Miss);
        assert!(matches!(tier.lookup(key(3)), Lookup::Pending { .. }));
        assert_eq!(tier.used_blocks(), 2);
    }

    #[test]
    fn structurally_unfittable_cohort_retries_without_mutation_or_id_use() {
        let mut tier = tier(1);
        assert_eq!(
            tier.prepare_store(request_id(), &[key(1), key(2)], 0.0),
            StoreOutcome::RetryCapacity {
                structurally_unfittable: true
            }
        );
        assert!(tier.warned_structurally_unfittable);
        assert_eq!(tier.used_blocks(), 0);
        assert_eq!(tier.next_transfer_id, 0);
        assert_eq!(
            tier.prepare_store(request_id(), &[key(3), key(4)], 0.0),
            StoreOutcome::RetryCapacity {
                structurally_unfittable: true
            }
        );
        assert_eq!(tier.next_transfer_id, 0);
    }

    #[test]
    fn pinned_load_source_is_not_an_eviction_victim() {
        let mut tier = tier(2);
        let now = make_resident(&mut tier, &[key(1), key(2)], 0.0);
        let LoadOutcome::Queued(_) = tier.schedule_load(request_id(), &[key(1)], now, now) else {
            panic!("resident block should load")
        };
        assert!(matches!(
            tier.prepare_store(request_id(), &[key(3)], now),
            StoreOutcome::Prepared { .. }
        ));
        assert!(tier.is_resident(key(1)));
        assert_eq!(tier.lookup(key(2)), Lookup::Miss);
    }

    #[test]
    fn cancelled_pinned_batch_reenters_as_one_newer_cohort() {
        let mut tier = tier(3);
        let mut now = make_resident(&mut tier, &[key(1), key(2)], 0.0);
        now = make_resident(&mut tier, &[key(3)], now);
        let third_epoch = tier.lru.touched_at[&key(3)];
        let LoadOutcome::Queued(load) =
            tier.schedule_load(request_id(), &[key(1), key(2)], now, now)
        else {
            panic!("resident blocks should load")
        };
        assert!(!tier.lru.touched_at.contains_key(&key(1)));
        assert!(!tier.lru.touched_at.contains_key(&key(2)));
        tier.touch(key(1));
        assert!(tier.cancel_load(load, now, now));
        let reentry_epoch = tier.lru.touched_at[&key(1)];
        assert_eq!(reentry_epoch, tier.lru.touched_at[&key(2)]);
        assert!(third_epoch < reentry_epoch);

        assert!(matches!(
            tier.prepare_store(request_id(), &[key(4)], now),
            StoreOutcome::Prepared { .. }
        ));
        assert_eq!(tier.lookup(key(3)), Lookup::Miss);
        assert!(tier.is_resident(key(1)) && tier.is_resident(key(2)));
    }

    #[test]
    fn directional_lanes_overlap_and_each_remains_fifo() {
        let mut tier = tier(4);
        let now = make_resident(&mut tier, &[key(1), key(2)], 0.0);
        let StoreOutcome::Prepared {
            transfer_id: first_store,
            ..
        } = tier.prepare_store(request_id(), &[key(3)], now)
        else {
            panic!("first store should prepare")
        };
        let StoreOutcome::Prepared {
            transfer_id: second_store,
            ..
        } = tier.prepare_store(request_id(), &[key(4)], now)
        else {
            panic!("second store should prepare")
        };
        assert_eq!(tier.submit_prepared_stores(now), 2);
        let LoadOutcome::Queued(first_load) = tier.schedule_load(request_id(), &[key(1)], now, now)
        else {
            panic!("first load should queue")
        };
        let LoadOutcome::Queued(second_load) =
            tier.schedule_load(request_id(), &[key(2)], now, now)
        else {
            panic!("second load should queue")
        };

        assert_eq!(tier.transfer_deadline(first_store), Some(now + 1.0));
        assert_eq!(tier.transfer_deadline(second_store), Some(now + 2.0));
        assert_eq!(tier.transfer_deadline(first_load), Some(now + 1.0));
        assert_eq!(tier.transfer_deadline(second_load), Some(now + 2.0));
    }

    #[test]
    fn cancelling_a_load_removes_its_exact_deadline_without_reflowing_the_lane() {
        let mut tier = tier(2);
        let observer = Arc::new(CancellationObserver::default());
        tier.set_observer(observer.clone());
        let now = make_resident(&mut tier, &[key(1), key(2)], 0.0);
        let LoadOutcome::Queued(first) = tier.schedule_load(request_id(), &[key(1)], now, now)
        else {
            panic!("first load should queue")
        };
        let first_deadline = tier.transfer_deadline(first).unwrap();
        let LoadOutcome::Queued(second) = tier.schedule_load(request_id(), &[key(2)], now, now)
        else {
            panic!("second load should queue")
        };
        let second_deadline = tier.transfer_deadline(second).unwrap();
        assert_eq!(second_deadline, first_deadline + 1.0);
        assert_eq!(tier.next_deadline(), Some(first_deadline));
        assert_eq!(tier.deadlines.len(), 2);

        let observed_at_ms = now + 0.5;
        assert!(tier.cancel_load(first, now, observed_at_ms));
        assert_eq!(tier.transfer_deadline(second), Some(second_deadline));
        assert_eq!(tier.next_deadline(), Some(second_deadline));
        assert_eq!(tier.deadlines.len(), 1);
        assert!(tier.tick(first_deadline).is_empty());
        assert_eq!(tier.next_deadline(), Some(second_deadline));
        assert!(matches!(
            tier.prepare_store(request_id(), &[key(3)], now),
            StoreOutcome::Prepared { .. }
        ));
        assert_eq!(
            *observer
                .cancelled
                .lock()
                .unwrap_or_else(|poisoned| poisoned.into_inner()),
            Some(Cancellation {
                request_id: request_id(),
                at_ms: observed_at_ms,
                transfer_id: first.get(),
                block_hashes: vec![1],
            })
        );
    }
}