radarmp/mp_core/src/datapool/mod.rs
2024-11-11 15:47:07 +08:00

110 lines
2.6 KiB
Rust

use super::data::Data;
use crate::errors::DataError;
use crate::plugin_system::PluginManager;
use once_cell::sync::Lazy;
use quick_cache::sync::Cache;
use std::collections::HashMap;
use std::fmt::Display;
use std::ops::Deref;
use std::path::PathBuf;
use std::sync::Arc;
static RUNTIME: Lazy<tokio::runtime::Runtime> =
Lazy::new(|| tokio::runtime::Runtime::new().unwrap());
pub type DataValue = Value<Data>;
#[derive(Clone, Debug)]
pub struct Value<T: Display>(Arc<HashMap<usize, Arc<T>>>);
impl Value<Data> {
pub fn new(data: Vec<Data>) -> Self {
let mut hashmap = HashMap::new();
for data in data {
hashmap.insert(data.id, Arc::new(data));
}
Self(Arc::new(hashmap))
}
pub fn get(&self, key: usize) -> Option<&Arc<Data>> {
self.0.get(&key)
}
pub fn iter(&self) -> std::collections::hash_map::Iter<'_, usize, Arc<Data>> {
self.0.iter()
}
}
impl<T: Display> Value<T> {
pub fn len(&self) -> usize {
self.0.len()
}
}
impl<T: Display> Deref for Value<T> {
type Target = HashMap<usize, Arc<T>>;
fn deref(&self) -> &Self::Target {
&self.0
}
}
pub struct DataPool {
plugin_manager: &'static PluginManager,
pool: Cache<PathBuf, Value<Data>>,
}
impl DataPool {
pub fn new(plugin_manager: &'static PluginManager, cap: usize) -> Self {
Self {
plugin_manager,
pool: Cache::new(cap),
}
}
pub fn get_or_load(&self, path: impl Into<PathBuf>) -> Result<Value<Data>, DataError> {
let path = path.into();
self.pool.get_or_insert_with(&path, || {
self.plugin_manager.try_load_data(&path).map(Value::new)
})
}
pub fn len(&self) -> usize {
self.pool.len()
}
pub async fn get_or_load_async(
&self,
path: impl Into<PathBuf>,
) -> Result<Value<Data>, DataError> {
let path = path.into();
let plugin_manager = self.plugin_manager;
self.pool
.get_or_insert_async(&path.clone(), async move {
RUNTIME
.spawn_blocking(move || plugin_manager.try_load_data(&path))
.await
.unwrap()
.map(Value::new)
})
.await
}
pub async fn get_or_insert_batch_async<A: Into<PathBuf>, P: IntoIterator<Item = A>>(
&self,
paths: P,
) -> Vec<Result<Value<Data>, DataError>> {
use futures::future::join_all;
let tasks = paths
.into_iter()
.map(|path| self.get_or_load_async(path.into()));
let results = join_all(tasks).await;
results
}
}