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>>>,
order: Vec<usize>,
}
impl CacheInner {
fn new() -> Self {
Self {
cells: HashMap::new(),
order: Vec::new(),
}
}
}
pub(crate) struct LazyCache {
inner: RwLock<CacheInner>,
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())
}
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);
}
pub fn get_or_init_with<F>(&self, key: usize, factory: F) -> AnyService
where
F: FnOnce() -> AnyService,
{
{
let inner = self.inner.read().unwrap();
if let Some(cell) = inner.cells.get(&key) {
if let Some(instance) = cell.get() {
return instance.clone();
}
let cell = cell.clone();
drop(inner);
return cell.get_or_init(factory).clone();
}
}
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
};
cell.get_or_init(factory).clone()
}
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,
{
if let Some(cached) = self.get(key) {
return Ok(cached);
}
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;
if let Some(ref instance) = *guard {
return Ok(instance.clone());
}
if let Some(cached) = self.get(key) {
*guard = Some(cached.clone());
return Ok(cached);
}
let instance = factory().await?;
*guard = Some(instance.clone());
self.prefill(key, instance.clone());
Ok(instance)
}
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());
}
}
}
drained
}
}
impl Default for LazyCache {
fn default() -> Self {
Self::new()
}
}
struct KeyedInner {
cells: HashMap<(usize, ServiceKey), Arc<OnceLock<AnyService>>>,
order: Vec<(usize, ServiceKey)>,
}
pub(crate) struct KeyedInstanceCache {
inner: RwLock<KeyedInner>,
}
impl KeyedInstanceCache {
pub fn new() -> Self {
Self {
inner: RwLock::new(KeyedInner {
cells: HashMap::new(),
order: Vec::new(),
}),
}
}
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())
}
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);
}
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());
{
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();
}
}
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()
}
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);
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(), "二次释放为空");
}
}