use std::any::TypeId;
use std::sync::mpsc::{self, Receiver, Sender};
use std::sync::Mutex;
#[cfg(not(feature = "tokio"))]
use std::thread::JoinHandle;
use std::time::Duration;
#[cfg(feature = "critical")]
use std::sync::OnceLock;
use crate::cache::global;
use crate::traits::ImcCacheable;
#[cfg(feature = "critical")]
static REDIS_URL: OnceLock<String> = OnceLock::new();
#[cfg(feature = "critical")]
pub(crate) fn get_redis_url() -> Option<&'static str> {
REDIS_URL.get().map(|s| s.as_str())
}
pub enum CacheCmd {
Remove(TypeId, u64),
Clear(TypeId),
ClearAll,
Shutdown,
}
#[cfg(not(feature = "tokio"))]
type WorkerJoinHandle = JoinHandle<()>;
#[cfg(feature = "tokio")]
type WorkerJoinHandle = tokio::task::JoinHandle<()>;
#[cfg(not(feature = "tokio"))]
fn spawn_worker_thread(name: &'static str, f: impl FnOnce() + Send + 'static) -> WorkerJoinHandle {
std::thread::Builder::new()
.name(name.into())
.spawn(f)
.expect("failed to spawn imc worker thread")
}
#[cfg(feature = "tokio")]
fn spawn_worker_thread(_name: &'static str, f: impl FnOnce() + Send + 'static) -> WorkerJoinHandle {
tokio::task::spawn_blocking(f)
}
pub(crate) static WORKER_TX: Mutex<Option<Sender<CacheCmd>>> = Mutex::new(None);
#[derive(Clone, Debug)]
pub struct WorkerConfig {
pub sweep_interval: Duration,
#[cfg(feature = "metrics-prometheus")]
#[cfg_attr(docsrs, doc(cfg(feature = "metrics-prometheus")))]
pub metrics_listen_addr: Option<String>,
#[cfg(feature = "invalidation-redis")]
#[cfg_attr(docsrs, doc(cfg(feature = "invalidation-redis")))]
pub redis_connection_string: Option<String>,
}
impl Default for WorkerConfig {
fn default() -> Self {
Self {
sweep_interval: Duration::from_secs(10),
#[cfg(feature = "metrics-prometheus")]
metrics_listen_addr: None,
#[cfg(feature = "invalidation-redis")]
redis_connection_string: None,
}
}
}
pub struct CacheWorker {
tx: Sender<CacheCmd>,
_handle: WorkerJoinHandle,
#[cfg(feature = "invalidation-redis")]
_invalidation_handle: Option<std::thread::JoinHandle<()>>,
#[cfg(feature = "critical")]
_critical_handle: Option<std::thread::JoinHandle<()>>,
}
impl CacheWorker {
pub fn spawn() -> Self {
Self::spawn_with_config(WorkerConfig::default())
}
pub fn spawn_with_config(config: WorkerConfig) -> Self {
let (tx, rx): (Sender<CacheCmd>, Receiver<CacheCmd>) = mpsc::channel();
{
let mut guard = WORKER_TX.lock().unwrap_or_else(|e| e.into_inner());
assert!(guard.is_none(), "only one imc worker may run at a time");
*guard = Some(tx.clone());
}
#[cfg(feature = "invalidation-redis")]
let invalidation_handle: Option<std::thread::JoinHandle<()>> =
if let Some(ref redis_url) = config.redis_connection_string {
let channels = crate::invalidation::snapshot_channels();
let url = redis_url.clone();
Some(
std::thread::Builder::new()
.name("imc-invalidation".into())
.spawn(move || crate::invalidation::redis_subscriber_loop(&url, channels))
.expect("failed to spawn invalidation thread"),
)
} else {
None
};
#[cfg(feature = "critical")]
let critical_handle: Option<std::thread::JoinHandle<()>> =
if let Some(ref redis_url) = config.redis_connection_string {
REDIS_URL.set(redis_url.clone()).ok();
let url = redis_url.clone();
Some(
std::thread::Builder::new()
.name("imc-critical".into())
.spawn(move || crate::critical::subscriber_loop(&url))
.expect("failed to spawn critical subscriber thread"),
)
} else {
None
};
let _handle = spawn_worker_thread("imc-worker", move || {
crate::log_event!(INFO, crate::log::WORKER, crate::log::START,
"imc worker started");
worker_loop(rx, config);
crate::log_event!(INFO, crate::log::WORKER, crate::log::STOP,
"imc worker stopped");
});
Self {
tx,
_handle,
#[cfg(feature = "invalidation-redis")]
_invalidation_handle: invalidation_handle,
#[cfg(feature = "critical")]
_critical_handle: critical_handle,
}
}
pub fn remove<T: ImcCacheable>(&self, id: &T::Id) -> Result<(), mpsc::SendError<CacheCmd>> {
let id_hash = crate::hasher::hash_value(id);
self.tx.send(CacheCmd::Remove(TypeId::of::<T>(), id_hash))
}
pub fn send(&self, cmd: CacheCmd) -> Result<(), mpsc::SendError<CacheCmd>> {
self.tx.send(cmd)
}
}
impl Drop for CacheWorker {
fn drop(&mut self) {
let _ = self.tx.send(CacheCmd::Shutdown);
*WORKER_TX.lock().unwrap_or_else(|e| e.into_inner()) = None;
crate::metrics::set_entries(0);
}
}
fn worker_loop(rx: Receiver<CacheCmd>, config: WorkerConfig) {
loop {
let result = rx.recv_timeout(config.sweep_interval);
match result {
Ok(CacheCmd::Remove(type_id, id_hash)) => {
let mut stores = global().stores.write().unwrap();
if let Some(cache) = stores.get_mut(&type_id) {
cache.remove_data(id_hash);
}
}
Ok(CacheCmd::Clear(type_id)) => {
let mut stores = global().stores.write().unwrap();
if let Some(cache) = stores.get_mut(&type_id) {
cache.clear();
}
}
Ok(CacheCmd::ClearAll) => {
let mut stores = global().stores.write().unwrap();
stores.clear();
}
Ok(CacheCmd::Shutdown) | Err(mpsc::RecvTimeoutError::Disconnected) => break,
Err(mpsc::RecvTimeoutError::Timeout) => {}
}
sweep_all(&global().stores);
let total: usize = {
let stores = global().stores.read().unwrap();
stores.values().map(|c| c.len()).sum()
};
crate::metrics::set_entries(total);
}
}
fn sweep_all(
stores: &std::sync::RwLock<std::collections::HashMap<TypeId, crate::cache::PerTypeCache>>,
) {
let mut stores = stores.write().unwrap();
for cache in stores.values_mut() {
cache.remove_expired();
cache.evict_to_max_size();
}
}