rust-dix 0.9.0

//! 懒初始化服务缓存。
//!
//! [`LazyCache`] 用 per-key [`OnceLock`] 保证多线程下工厂只执行一次，
//! 且工厂在无全局锁状态下执行，避免递归解析（Scoped 依赖 Scoped）时死锁。
//! 统一用于 [`crate::provider::ServiceProvider`] 的 singleton 缓存与 scoped 缓存。
//!
//! 除缓存值外，本结构还记录 **首次创建顺序**：容器释放时按创建逆序释放实例
//! （LIFO，先建依赖、后建依赖者，等价于 MEDI 的逆序 Dispose 语义），
//! 使 `Drop` 顺序可预测、可测试。

use futures::lock::Mutex as FuturesMutex;
use std::any::Any;
use std::collections::HashMap;
use std::future::Future;
use std::sync::{Arc, OnceLock, RwLock};

use crate::key::ServiceKey;

/// 类型擦除的服务实例。
pub(crate) type AnyService = Arc<dyn Any + Send + Sync>;

struct CacheInner {
    cells: HashMap<usize, Arc<OnceLock<AnyService>>>,
    /// 首次创建顺序（cache_key 去重，仅记录一次）。
    order: Vec<usize>,
}

impl CacheInner {
    fn new() -> Self {
        Self {
            cells: HashMap::new(),
            order: Vec::new(),
        }
    }
}

/// 懒初始化缓存。
///
/// 结构：`RwLock<CacheInner { cells: HashMap<cache_key, Arc<OnceLock<AnyService>>>, order }>`
///
/// - **读快速路径**：读锁查 cell，`OnceLock::get` 命中即 clone Arc 返回。
/// - **写慢路径**：cell 不存在时，写锁插入空 `OnceLock` 并记录创建顺序，
///   立即释放写锁，再在无锁状态下 `OnceLock::get_or_init` 执行工厂。
/// - **并发安全**：`OnceLock::get_or_init` 保证多线程下工厂只执行一次，
///   彻底消除 check-then-act 竞态（TOCTOU）。
/// - **无死锁**：工厂执行时不持有任何全局锁，递归解析依赖可安全获取
///   相同 `LazyCache` 的其他 key（不同 cell，无锁竞争）。
/// - **确定性释放**：`drain_reverse` 按创建逆序返回实例；调用方必须在
///   释放锁之后再丢弃返回值（本方法自身不持有锁返回）。
pub(crate) struct LazyCache {
    inner: RwLock<CacheInner>,
    /// Per-key async factory 执行锁：保证 async factory 只执行一次。
    /// 仅 `get_or_init_with_async` 使用，sync 路径不受影响。
    async_locks: RwLock<HashMap<usize, Arc<FuturesMutex<Option<AnyService>>>>>,
}

impl LazyCache {
    pub fn new() -> Self {
        Self {
            inner: RwLock::new(CacheInner::new()),
            async_locks: RwLock::new(HashMap::new()),
        }
    }

    /// 查询已缓存的实例，不执行工厂。
    pub fn get(&self, key: usize) -> Option<AnyService> {
        self.inner
            .read()
            .unwrap()
            .cells
            .get(&key)
            .and_then(|cell| cell.get().cloned())
    }

    /// 预填充：build 阶段 validation 已构建实例，直接填入。
    ///
    /// 若 cell 已初始化则保留首个值（validation 单线程执行，不会冲突）。
    /// 新 cell 会记录创建顺序（用于确定性释放）。
    pub fn prefill(&self, key: usize, instance: AnyService) {
        let cell = {
            let mut inner = self.inner.write().unwrap();
            let is_new = !inner.cells.contains_key(&key);
            let cell = inner
                .cells
                .entry(key)
                .or_insert_with(|| Arc::new(OnceLock::new()))
                .clone();
            if is_new {
                inner.order.push(key);
            }
            cell
        };
        let _ = cell.set(instance);
    }

    /// 懒初始化：若已缓存则返回，否则在无锁状态下执行工厂并缓存。
    ///
    /// 多线程同时 miss 时，`OnceLock::get_or_init` 保证只有一个线程
    /// 执行工厂，其余线程等待并复用结果。
    pub fn get_or_init_with<F>(&self, key: usize, factory: F) -> AnyService
    where
        F: FnOnce() -> AnyService,
    {
        // Fast path: 读锁查已初始化的 cell。
        {
            let inner = self.inner.read().unwrap();
            if let Some(cell) = inner.cells.get(&key) {
                if let Some(instance) = cell.get() {
                    return instance.clone();
                }
                // cell 存在但未初始化：另一线程正在初始化。
                // clone Arc 后释放读锁，通过 get_or_init 等待。
                let cell = cell.clone();
                drop(inner);
                return cell.get_or_init(factory).clone();
            }
        }
        // Slow path: cell 不存在，写锁创建空 cell、记录顺序，然后立即释放。
        let cell = {
            let mut inner = self.inner.write().unwrap();
            let is_new = !inner.cells.contains_key(&key);
            let cell = inner
                .cells
                .entry(key)
                .or_insert_with(|| Arc::new(OnceLock::new()))
                .clone();
            if is_new {
                inner.order.push(key);
            }
            cell
        };
        // 无锁状态下 get_or_init：工厂执行期间不持有任何全局锁，
        // 递归解析依赖可安全获取相同 LazyCache 的其他 key。
        cell.get_or_init(factory).clone()
    }

    /// 懒初始化（async 版本）：若已缓存则返回，否则在 `FuturesMutex` 保护下
    /// 执行 async factory 并缓存。
    ///
    /// 与 sync `get_or_init_with` 的区别：`OnceLock::get_or_init` 不支持 async
    /// 闭包，故用 per-key `FuturesMutex<Option<AnyService>>` 替代。多 future
    /// 同时 miss 时，只有一个 future 执行 factory，其余 await 等待结果。
    ///
    /// factory 返回 `Result`：成功时缓存并返回 `Ok`，失败时不缓存（允许重试）。
    pub async fn get_or_init_with_async<F, Fut, E>(
        &self,
        key: usize,
        factory: F,
    ) -> Result<AnyService, E>
    where
        F: FnOnce() -> Fut + Send,
        Fut: Future<Output = Result<AnyService, E>> + Send,
        E: Send,
    {
        // Fast path: 查 OnceLock（sync 缓存）
        if let Some(cached) = self.get(key) {
            return Ok(cached);
        }

        // Slow path: 获取 per-key async 锁
        let lock = {
            let mut locks = self.async_locks.write().unwrap();
            locks
                .entry(key)
                .or_insert_with(|| Arc::new(FuturesMutex::new(None)))
                .clone()
        };

        let mut guard = lock.lock().await;
        // 双重检查：其他 future 可能已填充
        if let Some(ref instance) = *guard {
            return Ok(instance.clone());
        }
        // 再次检查 OnceLock（可能被 sync 路径填充）
        if let Some(cached) = self.get(key) {
            *guard = Some(cached.clone());
            return Ok(cached);
        }

        // 执行 async factory
        let instance = factory().await?;
        *guard = Some(instance.clone());
        // 填入 OnceLock，供后续 sync/async fast path 使用
        self.prefill(key, instance.clone());
        Ok(instance)
    }

    /// 按创建逆序取出全部缓存实例并清空缓存。
    ///
    /// 返回的 `Vec` 由调用方持有并通过 `drop` 触发实例释放；本方法在返回前
    /// 已释放内部锁，调用方在锁外丢弃即可，避免用户 `Drop` 回调重入缓存。
    ///
    /// 注意：与并发解析不构成同步点。容器释放期间不应并发解析服务
    /// （与 MEDI 相同：已释放的 provider 不可再用）。
    pub(crate) fn drain_reverse(&self) -> Vec<AnyService> {
        let (order, cells) = {
            let mut inner = self.inner.write().unwrap();
            let order = std::mem::take(&mut inner.order);
            let cells = std::mem::take(&mut inner.cells);
            (order, cells)
        };
        self.async_locks.write().unwrap().clear();

        let mut drained = Vec::with_capacity(order.len());
        for key in order.into_iter().rev() {
            if let Some(cell) = cells.get(&key) {
                if let Some(instance) = cell.get() {
                    drained.push(instance.clone());
                }
            }
        }
        // `cells`（仍持有 Arc 引用）在函数返回时释放；`drained` 交给调用方。
        drained
    }
}

impl Default for LazyCache {
    fn default() -> Self {
        Self::new()
    }
}

struct KeyedInner {
    /// `(registration cache_key, requested key)` → instance cell.
    ///
    /// The registration dimension is required: two different `AnyKey` services
    /// resolved with the same requested key must not share a cache slot.
    cells: HashMap<(usize, ServiceKey), Arc<OnceLock<AnyService>>>,
    /// 首次创建顺序（复合键去重，仅记录一次）。
    order: Vec<(usize, ServiceKey)>,
}

/// Per-`(registration, requested key)` instance cache for `AnyKey` fallback
/// registrations.
///
/// An `AnyKey` singleton/scoped registration produces one instance **per
/// requested key**, and different `AnyKey` registrations are independent even
/// for the same requested key. Instances are tracked in creation order so
/// disposal can release them in reverse (LIFO) order.
pub(crate) struct KeyedInstanceCache {
    inner: RwLock<KeyedInner>,
}

impl KeyedInstanceCache {
    pub fn new() -> Self {
        Self {
            inner: RwLock::new(KeyedInner {
                cells: HashMap::new(),
                order: Vec::new(),
            }),
        }
    }

    /// Cached instance for `(registration, key)`, without executing a factory.
    pub(crate) fn get(&self, cache_key: usize, key: &ServiceKey) -> Option<AnyService> {
        self.inner
            .read()
            .unwrap()
            .cells
            .get(&(cache_key, key.clone()))
            .and_then(|cell| cell.get().cloned())
    }

    /// Pre-populate an instance (used by build-time validation so the instance
    /// created while validating a dependent is reused at runtime).
    pub(crate) fn prefill(&self, cache_key: usize, key: &ServiceKey, instance: AnyService) {
        let composite = (cache_key, key.clone());
        let cell = {
            let mut inner = self.inner.write().unwrap();
            let is_new = !inner.cells.contains_key(&composite);
            let cell = inner
                .cells
                .entry(composite.clone())
                .or_insert_with(|| Arc::new(OnceLock::new()))
                .clone();
            if is_new {
                inner.order.push(composite);
            }
            cell
        };
        let _ = cell.set(instance);
    }

    /// Single-flight initialization for `(registration, key)`.
    pub fn get_or_init_with<F>(&self, cache_key: usize, key: &ServiceKey, factory: F) -> AnyService
    where
        F: FnOnce() -> AnyService,
    {
        let composite = (cache_key, key.clone());
        // Fast path: initialized cell.
        {
            let inner = self.inner.read().unwrap();
            if let Some(cell) = inner.cells.get(&composite) {
                if let Some(instance) = cell.get() {
                    return instance.clone();
                }
                let cell = cell.clone();
                drop(inner);
                return cell.get_or_init(factory).clone();
            }
        }
        // Slow path: create the cell (recording creation order) then init lock-free.
        let cell = {
            let mut inner = self.inner.write().unwrap();
            let is_new = !inner.cells.contains_key(&composite);
            let cell = inner
                .cells
                .entry(composite.clone())
                .or_insert_with(|| Arc::new(OnceLock::new()))
                .clone();
            if is_new {
                inner.order.push(composite);
            }
            cell
        };
        cell.get_or_init(factory).clone()
    }

    /// Drain instances in reverse creation order (used by container disposal).
    pub(crate) fn drain_reverse(&self) -> Vec<AnyService> {
        let (order, cells) = {
            let mut inner = self.inner.write().unwrap();
            let order = std::mem::take(&mut inner.order);
            let cells = std::mem::take(&mut inner.cells);
            (order, cells)
        };
        let mut drained = Vec::with_capacity(order.len());
        for composite in order.into_iter().rev() {
            if let Some(cell) = cells.get(&composite) {
                if let Some(instance) = cell.get() {
                    drained.push(instance.clone());
                }
            }
        }
        drained
    }
}

impl Default for KeyedInstanceCache {
    fn default() -> Self {
        Self::new()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::sync::atomic::{AtomicU64, Ordering};
    use std::thread;

    #[test]
    fn get_or_init_executes_factory_once_under_concurrency() {
        static CALLS: AtomicU64 = AtomicU64::new(0);
        let cache = LazyCache::new();
        let cache = Arc::new(cache);

        let mut handles = Vec::new();
        for _ in 0..8 {
            let c = cache.clone();
            handles.push(thread::spawn(move || {
                c.get_or_init_with(42, || {
                    CALLS.fetch_add(1, Ordering::SeqCst);
                    Arc::new(99_u64) as AnyService
                })
            }));
        }
        let results: Vec<_> = handles.into_iter().map(|h| h.join().unwrap()).collect();

        assert_eq!(CALLS.load(Ordering::SeqCst), 1, "工厂应只执行一次");
        for r in &results {
            let v = r.downcast_ref::<u64>().unwrap();
            assert_eq!(*v, 99);
        }
    }

    #[test]
    fn prefill_skips_reinit() {
        let cache = LazyCache::new();
        cache.prefill(1, Arc::new(7_u64) as AnyService);
        // 再次 prefill 不覆盖
        cache.prefill(1, Arc::new(999_u64) as AnyService);
        let instance = cache.get(1).unwrap();
        let v = instance.downcast_ref::<u64>().unwrap();
        assert_eq!(*v, 7);
    }

    #[test]
    fn get_returns_none_for_missing() {
        let cache = LazyCache::new();
        assert!(cache.get(0).is_none());
    }

    #[test]
    fn drain_reverse_follows_creation_order() {
        let cache = LazyCache::new();
        cache.get_or_init_with(10, || Arc::new(1_u64) as AnyService);
        cache.get_or_init_with(20, || Arc::new(2_u64) as AnyService);
        cache.prefill(30, Arc::new(3_u64) as AnyService);
        // 重复解析不重复记录顺序
        let _ = cache.get_or_init_with(10, || Arc::new(99_u64) as AnyService);

        let drained: Vec<u64> = cache
            .drain_reverse()
            .iter()
            .map(|a| *a.downcast_ref::<u64>().unwrap())
            .collect();
        assert_eq!(drained, vec![3, 2, 1], "按创建逆序释放");
        assert!(cache.get(10).is_none(), "释放后缓存清空");
        assert!(cache.drain_reverse().is_empty(), "二次释放为空");
    }
}