use std::error::Error;
use std::future::Future;
use std::sync::atomic::Ordering;
use std::sync::Arc;
use std::time::Instant;
use bytes::Bytes;
use futures_util::FutureExt;
use hydracache_core::{
CacheCodec, CacheDiagnostics, CacheError, CacheEvent, CacheEventKind, CacheEventOptions,
CacheEventOrigin, CacheOptions, CacheStats, PostcardCodec, Result,
};
use moka::future::Cache;
use serde::{de::DeserializeOwned, Serialize};
use crate::builder::HydraCacheBuilder;
use crate::entry::CacheEntry;
use crate::events::{CacheEventListenerHandle, CacheEventSubscriber, EventBus};
use crate::inflight::{InFlightMap, SharedLoadFuture};
use crate::stats::StatsCounters;
use crate::tag_index::{LoadGenerationSnapshot, TagIndex};
use crate::typed::TypedCache;
#[derive(Debug, Clone)]
pub struct HydraCache<C = PostcardCodec>
where
C: CacheCodec,
{
pub(crate) inner: Arc<HydraCacheInner<C>>,
}
#[derive(Debug)]
pub(crate) struct HydraCacheInner<C>
where
C: CacheCodec,
{
pub(crate) store: Cache<String, CacheEntry>,
pub(crate) tag_index: TagIndex,
pub(crate) in_flight: InFlightMap,
pub(crate) codec: C,
pub(crate) default_ttl: std::time::Duration,
pub(crate) stats: Arc<StatsCounters>,
pub(crate) events: EventBus,
}
impl HydraCache<PostcardCodec> {
pub fn local() -> HydraCacheBuilder<PostcardCodec> {
HydraCacheBuilder::default()
}
}
impl<C> HydraCache<C>
where
C: CacheCodec,
{
pub fn typed<T>(&self, namespace: impl Into<String>) -> TypedCache<T, C> {
TypedCache::new(self.clone(), namespace.into())
}
pub fn subscribe(&self, options: CacheEventOptions) -> CacheEventSubscriber {
self.inner
.events
.subscribe(options, self.inner.stats.clone())
}
pub fn subscribe_mutations(&self) -> CacheEventSubscriber {
self.subscribe(CacheEventOptions::mutations())
}
pub fn subscribe_access(&self) -> CacheEventSubscriber {
self.subscribe(CacheEventOptions::access())
}
pub fn subscribe_key(&self, key: impl Into<String>) -> CacheEventSubscriber {
self.subscribe(CacheEventOptions::new().key(key))
}
pub fn subscribe_tag(&self, tag: impl Into<String>) -> CacheEventSubscriber {
self.subscribe(CacheEventOptions::new().tag(tag))
}
pub fn add_listener<F>(
&self,
options: CacheEventOptions,
listener: F,
) -> CacheEventListenerHandle
where
F: Fn(CacheEvent) + Send + 'static,
{
CacheEventListenerHandle::spawn(self.subscribe(options), listener)
}
pub fn on_mutation<F>(&self, listener: F) -> CacheEventListenerHandle
where
F: Fn(CacheEvent) + Send + 'static,
{
self.add_listener(CacheEventOptions::mutations(), listener)
}
pub fn on_access<F>(&self, listener: F) -> CacheEventListenerHandle
where
F: Fn(CacheEvent) + Send + 'static,
{
self.add_listener(CacheEventOptions::access(), listener)
}
pub async fn get<T>(&self, key: &str) -> Result<Option<T>>
where
T: DeserializeOwned,
{
match self.inner.store.get(key).await {
Some(entry) if entry.is_expired() => {
self.remove_expired(key, &entry).await;
self.inner.stats.misses.fetch_add(1, Ordering::Relaxed);
self.publish_key_event(
CacheEventKind::Miss,
key,
CacheEventOrigin::LocalApi,
entry.tags.clone(),
);
Ok(None)
}
Some(entry) => match self.inner.codec.decode::<T>(&entry.value) {
Ok(value) => {
self.inner.stats.hits.fetch_add(1, Ordering::Relaxed);
self.publish_key_event(
CacheEventKind::Hit,
key,
CacheEventOrigin::LocalApi,
entry.tags.clone(),
);
Ok(Some(value))
}
Err(error) => {
self.remove_entry(key, &entry).await;
self.inner.stats.misses.fetch_add(1, Ordering::Relaxed);
self.publish_key_event(
CacheEventKind::Miss,
key,
CacheEventOrigin::LocalApi,
entry.tags.clone(),
);
Err(error)
}
},
None => {
self.inner.stats.misses.fetch_add(1, Ordering::Relaxed);
self.publish_key_event(
CacheEventKind::Miss,
key,
CacheEventOrigin::LocalApi,
Vec::<String>::new(),
);
Ok(None)
}
}
}
pub async fn put<T>(&self, key: &str, value: T, options: CacheOptions) -> Result<()>
where
T: Serialize,
{
let bytes = self.inner.codec.encode(&value)?;
self.put_bytes(key, bytes, options).await
}
pub async fn get_or_load<T, E, F, Fut>(
&self,
key: &str,
options: CacheOptions,
loader: F,
) -> Result<T>
where
T: Serialize + DeserializeOwned,
E: Error + Send + Sync + 'static,
F: FnOnce() -> Fut + Send + 'static,
Fut: Future<Output = std::result::Result<T, E>> + Send + 'static,
{
if let Some(value) = self.get(key).await? {
return Ok(value);
}
let shared = self
.shared_load(key, options, move |cache| async move {
cache.inner.stats.loads.fetch_add(1, Ordering::Relaxed);
let value = loader().await.map_err(CacheError::loader)?;
let bytes = cache.inner.codec.encode(&value)?;
Ok(bytes)
})
.await;
let bytes = shared.await.map_err(|error| (*error).clone())?;
self.inner.codec.decode(&bytes)
}
pub async fn get_or_insert_with<T, F, Fut>(
&self,
key: &str,
options: CacheOptions,
loader: F,
) -> Result<T>
where
T: Serialize + DeserializeOwned,
F: FnOnce() -> Fut + Send + 'static,
Fut: Future<Output = T> + Send + 'static,
{
self.get_or_load(key, options, move || async move {
Ok::<_, std::convert::Infallible>(loader().await)
})
.await
}
pub async fn try_get_or_insert_with<T, E, F, Fut>(
&self,
key: &str,
options: CacheOptions,
loader: F,
) -> Result<T>
where
T: Serialize + DeserializeOwned,
E: Error + Send + Sync + 'static,
F: FnOnce() -> Fut + Send + 'static,
Fut: Future<Output = std::result::Result<T, E>> + Send + 'static,
{
self.get_or_load(key, options, loader).await
}
pub async fn invalidate_key(&self, key: &str) -> Result<bool> {
self.remove_with_event(key, CacheEventKind::KeyInvalidated)
.await
}
pub async fn remove(&self, key: &str) -> Result<bool> {
self.remove_with_event(key, CacheEventKind::Removed).await
}
pub async fn contains_key(&self, key: &str) -> bool {
match self.inner.store.get(key).await {
Some(entry) if entry.is_expired() => {
self.remove_expired(key, &entry).await;
false
}
Some(_) => true,
None => false,
}
}
pub async fn invalidate_tag(&self, tag: &str) -> Result<u64> {
let keys = self.inner.tag_index.take_tag(tag).await;
let mut removed = 0;
for key in keys {
if let Some(entry) = self.inner.store.get(&key).await {
self.remove_entry(&key, &entry).await;
removed += 1;
}
}
if removed > 0 {
self.inner
.stats
.invalidations
.fetch_add(removed, Ordering::Relaxed);
}
self.publish_event(CacheEvent::for_tag(
CacheEventKind::TagInvalidated,
tag,
removed,
CacheEventOrigin::LocalApi,
));
Ok(removed)
}
pub async fn flush(&self) -> Result<()> {
let estimated_entries = self.inner.store.entry_count();
self.inner.store.invalidate_all();
self.inner.tag_index.clear().await;
self.publish_event(CacheEvent::for_cache(
CacheEventKind::Flushed,
Some(estimated_entries),
CacheEventOrigin::LocalApi,
));
Ok(())
}
pub fn stats(&self) -> CacheStats {
self.inner.stats.snapshot()
}
pub async fn diagnostics(&self) -> CacheDiagnostics {
self.inner.store.run_pending_tasks().await;
CacheDiagnostics {
stats: self.stats(),
estimated_entries: self.inner.store.entry_count(),
}
}
pub(crate) async fn put_bytes(
&self,
key: &str,
value: Bytes,
options: CacheOptions,
) -> Result<()> {
self.put_bytes_unchecked(key, value, options, CacheEventOrigin::LocalApi)
.await
}
async fn put_bytes_unchecked(
&self,
key: &str,
value: Bytes,
options: CacheOptions,
origin: CacheEventOrigin,
) -> Result<()> {
let ttl = options.ttl_value().unwrap_or(self.inner.default_ttl);
let tags = options.tags_value().to_vec();
let entry = CacheEntry {
value,
tags: tags.clone(),
expires_at: Instant::now().checked_add(ttl),
};
if let Some(old_entry) = self.inner.store.get(key).await {
self.inner.tag_index.unregister(key, &old_entry.tags).await;
}
self.inner.store.insert(key.to_owned(), entry).await;
self.inner.tag_index.register(key, &tags).await;
self.publish_key_event(CacheEventKind::Stored, key, origin, tags);
Ok(())
}
async fn put_bytes_if_fresh(
&self,
key: &str,
value: Bytes,
options: CacheOptions,
generation: &LoadGenerationSnapshot,
) -> Result<bool> {
if !self.inner.tag_index.is_current(generation).await {
self.inner
.stats
.stale_load_discards
.fetch_add(1, Ordering::Relaxed);
self.publish_key_event(
CacheEventKind::StaleLoadDiscarded,
key,
CacheEventOrigin::Loader,
options.tags_value().to_vec(),
);
return Ok(false);
}
self.put_bytes_unchecked(key, value, options, CacheEventOrigin::Loader)
.await?;
Ok(true)
}
async fn shared_load<F, Fut>(
&self,
key: &str,
options: CacheOptions,
loader: F,
) -> SharedLoadFuture
where
F: FnOnce(Self) -> Fut + Send + 'static,
Fut: Future<Output = Result<Bytes>> + Send + 'static,
{
let generation = self.inner.tag_index.snapshot(options.tags_value()).await;
let event_tags = options.tags_value().to_vec();
let late_join_event_tags = event_tags.clone();
if let Some(shared) = self.inner.in_flight.get_current(key, &generation).await {
self.inner
.stats
.single_flight_joins
.fetch_add(1, Ordering::Relaxed);
self.publish_key_event(
CacheEventKind::SingleFlightJoined,
key,
CacheEventOrigin::SingleFlight,
event_tags,
);
return shared;
}
#[cfg(coverage)]
tokio::task::yield_now().await;
let key_owned = key.to_owned();
let cache = self.clone();
let load_key = key_owned.clone();
let load_generation = generation.clone();
let load_event_tags = event_tags.clone();
let shared = async move {
let result = async {
cache.publish_key_event(
CacheEventKind::LoadStarted,
&load_key,
CacheEventOrigin::Loader,
load_event_tags.clone(),
);
let bytes = loader(cache.clone()).await?;
let accepted = cache
.put_bytes_if_fresh(&load_key, bytes.clone(), options, &load_generation)
.await?;
if accepted {
cache.publish_key_event(
CacheEventKind::LoadCompleted,
&load_key,
CacheEventOrigin::Loader,
load_event_tags.clone(),
);
}
Ok(bytes)
}
.await;
if result.is_err() {
cache.publish_key_event(
CacheEventKind::LoadFailed,
&load_key,
CacheEventOrigin::Loader,
load_event_tags,
);
}
let result = result.map_err(Arc::new);
cache
.inner
.in_flight
.remove_if_generation_matches(&load_key, &load_generation)
.await;
result
}
.boxed()
.shared();
let (shared, inserted) = self
.inner
.in_flight
.insert_or_get_current(key_owned, shared, generation)
.await;
if !inserted {
self.inner
.stats
.single_flight_joins
.fetch_add(1, Ordering::Relaxed);
self.publish_key_event(
CacheEventKind::SingleFlightJoined,
key,
CacheEventOrigin::SingleFlight,
late_join_event_tags,
);
}
shared
}
async fn remove_expired(&self, key: &str, entry: &CacheEntry) {
self.remove_entry(key, entry).await;
self.publish_key_event(
CacheEventKind::Expired,
key,
CacheEventOrigin::Backend,
entry.tags.clone(),
);
}
async fn remove_entry(&self, key: &str, entry: &CacheEntry) {
self.inner.store.invalidate(key).await;
self.inner.tag_index.unregister(key, &entry.tags).await;
}
async fn remove_with_event(&self, key: &str, kind: CacheEventKind) -> Result<bool> {
let Some(entry) = self.inner.store.get(key).await else {
return Ok(false);
};
self.remove_entry(key, &entry).await;
self.inner
.stats
.invalidations
.fetch_add(1, Ordering::Relaxed);
self.publish_key_event(kind, key, CacheEventOrigin::LocalApi, entry.tags.clone());
Ok(true)
}
fn publish_key_event<I, S>(
&self,
kind: CacheEventKind,
key: &str,
origin: CacheEventOrigin,
tags: I,
) where
I: IntoIterator<Item = S>,
S: Into<String>,
{
self.publish_event(CacheEvent::for_key(kind, key, origin, tags));
}
fn publish_event(&self, event: CacheEvent) {
self.inner.events.publish(event, &self.inner.stats);
}
}