use crate::{
ClientError, Error, Result,
client::{Client, PreparedCommand},
commands::{
BitFieldSubCommand, BitRange, BitmapCommands, ClientTrackingOptions, ClientTrackingStatus,
ConnectionCommands, HashCommands, ListCommands, SetCommands, SortedSetCommands,
StringCommands, ZRangeOptions,
},
network::{JoinHandle, spawn},
resp::{
BulkString, Command, CommandArgsMut, FastPathCommandBuilder, RespDeserializer,
RespResponse, Response,
},
};
use bytes::Bytes;
use dashmap::DashMap;
use futures_util::StreamExt;
use serde::{Serialize, de::DeserializeOwned};
use std::{
sync::{
Arc,
atomic::{AtomicU64, Ordering},
},
time::Duration,
};
pub use moka::future::CacheBuilder;
type SubCache = DashMap<Bytes, RespResponse>;
type MokaCache = moka::future::Cache<BulkString, Arc<SubCache>>;
type MokaCacheBuilder = moka::future::CacheBuilder<BulkString, Arc<SubCache>, MokaCache>;
pub struct Cache {
cache: Arc<MokaCache>,
client: Client,
generation_counter: Arc<AtomicU64>,
key_generations: Arc<DashMap<BulkString, u64>>,
flush_generation: Arc<AtomicU64>,
#[allow(dead_code)]
invalidation_task: JoinHandle<()>,
#[allow(dead_code)]
reconnection_task: JoinHandle<()>,
}
impl Cache {
#[allow(clippy::type_complexity)]
#[expect(
clippy::arithmetic_side_effects,
reason = "`dropped` is a monotonic counter and `dropped_seen` its last \
observed value, so the difference cannot go below zero. The \
generation counter counts cache flushes over the life of a \
client."
)]
pub(crate) async fn from_builder(
client: Client,
builder: MokaCacheBuilder,
tracking_opts: ClientTrackingOptions,
) -> Result<Arc<Self>> {
client
.client_tracking(ClientTrackingStatus::On, tracking_opts.clone())
.await?;
let stream = client.create_client_tracking_invalidation_stream()?;
let cache = Arc::new(builder.build());
let cache_clone = cache.clone();
let generation_counter = Arc::new(AtomicU64::new(0));
let key_generations: Arc<DashMap<BulkString, u64>> = Arc::new(DashMap::new());
let flush_generation = Arc::new(AtomicU64::new(0));
let connection_tag = client.connection_tag().to_owned();
let counter_clone = generation_counter.clone();
let key_generations_clone = key_generations.clone();
let flush_generation_clone = flush_generation.clone();
let invalidation_task = spawn(async move {
let mut stream = stream;
let mut dropped_seen = 0usize;
while let Some(keys) = stream.next().await {
let dropped = stream.dropped_messages();
if dropped != dropped_seen {
tracing::warn!(
tag = %connection_tag,
"Dropped {} invalidation message(s) under backpressure; \
invalidating the whole client cache",
dropped - dropped_seen
);
dropped_seen = dropped;
let generation = counter_clone.fetch_add(1, Ordering::SeqCst) + 1;
flush_generation_clone.store(generation, Ordering::SeqCst);
cache_clone.invalidate_all();
}
for key in keys {
tracing::debug!(
tag = %connection_tag,
"Invalidating key `{key}` from client cache"
);
let generation = counter_clone.fetch_add(1, Ordering::SeqCst) + 1;
key_generations_clone.insert(key.clone(), generation);
cache_clone.invalidate(&key).await;
}
}
});
let cache_clone = cache.clone();
let client_clone = client.clone();
let connection_tag = client.connection_tag().to_owned();
let mut on_reconnect = client.on_reconnect();
let counter_clone = generation_counter.clone();
let flush_generation_clone = flush_generation.clone();
let reconnection_task = spawn(async move {
while on_reconnect.recv().await.is_ok() {
tracing::debug!(tag = %connection_tag, "Re-enabling client tracking after reconnection");
let generation = counter_clone.fetch_add(1, Ordering::SeqCst) + 1;
flush_generation_clone.store(generation, Ordering::SeqCst);
cache_clone.invalidate_all();
if let Err(e) = client_clone
.client_tracking(ClientTrackingStatus::On, tracking_opts.clone())
.await
{
tracing::error!(
tag = %connection_tag,
"Cannot re-enable client tracking after reconnection: {e}"
);
}
}
});
Ok(Arc::new(Self {
cache,
client,
generation_counter,
key_generations,
flush_generation,
invalidation_task,
reconnection_task,
}))
}
pub async fn new(
client: Client,
ttl_secs: u64,
tracking_opts: ClientTrackingOptions,
) -> Result<Arc<Self>> {
let builder = MokaCache::builder()
.time_to_live(Duration::from_secs(ttl_secs))
.max_capacity(10_000);
Self::from_builder(client, builder, tracking_opts).await
}
#[cfg(test)]
pub(crate) fn flush_generation(&self) -> u64 {
self.flush_generation.load(Ordering::SeqCst)
}
pub async fn get<R: Response + DeserializeOwned>(&self, key: impl Serialize) -> Result<R> {
self.process_prepared_command(key_to_bulk_string(&key)?, self.client.get(key))
.await
}
pub async fn mget<R: Response + DeserializeOwned>(&self, keys: impl Serialize) -> Result<R> {
let prepared_command = self.client.mget::<R>(keys);
let mut responses = Vec::with_capacity(prepared_command.command.num_args());
let mut missing_indices = Vec::new();
let mut missing_keys = Vec::new();
let mut missing_subcache_keys = Vec::new();
for (i, arg) in prepared_command.command.args().enumerate() {
let key = BulkString::from(arg.clone());
let subcache_key = get_subcache_key(&key);
if let Some(values) = self.cache.get(&key).await
&& let Some(response) = values.get(&subcache_key)
{
tracing::debug!(
tag = %self.client.connection_tag(),
"Cache hit on key `{key}`"
);
responses.push(response.clone());
} else {
tracing::debug!(
tag = %self.client.connection_tag(),
"Cache miss on key `{key}`"
);
responses.push(RespResponse::null());
missing_indices.push(i);
missing_keys.push(key);
missing_subcache_keys.push(subcache_key);
}
}
if !missing_keys.is_empty() {
let missing_prepared_command = self.client.mget::<R>(missing_keys);
let response = self
.client
.internal_send(missing_prepared_command.command, None)
.await?;
let Ok(collection_iter) = response.clone().into_collection_iter() else {
return Err(Error::Client(ClientError::ExpectedArrayForMGet));
};
for (idx_in_missing, response) in collection_iter.enumerate() {
let response = response?;
let original_idx = missing_indices[idx_in_missing];
let Some(key) = prepared_command
.command
.get_arg(original_idx)
.map(BulkString::from)
else {
break;
};
self.cache
.entry(key)
.or_insert_with(async { Arc::new(DashMap::new()) })
.await
.value()
.insert(
missing_subcache_keys[idx_in_missing].clone(),
response.compact(),
);
responses[original_idx] = response;
}
} else {
tracing::debug!(tag = %self.client.connection_tag(), "Cache hit on mget");
}
let response = RespResponse::owned_array(responses);
let deserializer = RespDeserializer::new(response.view()?);
R::deserialize(deserializer)
}
pub async fn getrange<R: Response + DeserializeOwned>(
&self,
key: impl Serialize,
start: isize,
end: isize,
) -> Result<R> {
self.process_prepared_command(
key_to_bulk_string(&key)?,
self.client.getrange(key, start, end),
)
.await
}
pub async fn strlen(&self, key: impl Serialize) -> Result<usize> {
self.process_prepared_command(key_to_bulk_string(&key)?, self.client.strlen(key))
.await
}
pub async fn hexists(&self, key: impl Serialize, field: impl Serialize) -> Result<bool> {
self.process_prepared_command(
key_to_bulk_string(&key)?,
self.client.hexists(key_to_bulk_string(&key)?, field),
)
.await
}
pub async fn hget<R: Response + DeserializeOwned>(
&self,
key: impl Serialize,
field: impl Serialize,
) -> Result<R> {
self.process_prepared_command(key_to_bulk_string(&key)?, self.client.hget(key, field))
.await
}
pub async fn hgetall<R: Response + DeserializeOwned>(&self, key: impl Serialize) -> Result<R> {
self.process_prepared_command(key_to_bulk_string(&key)?, self.client.hgetall(key))
.await
}
pub async fn hlen(&self, key: impl Serialize) -> Result<usize> {
self.process_prepared_command(key_to_bulk_string(&key)?, self.client.hlen(key))
.await
}
pub async fn hkeys<R: Response + DeserializeOwned>(&self, key: impl Serialize) -> Result<R> {
self.process_prepared_command(key_to_bulk_string(&key)?, self.client.hkeys(key))
.await
}
pub async fn hvals<R: Response + DeserializeOwned>(&self, key: impl Serialize) -> Result<R> {
self.process_prepared_command(key_to_bulk_string(&key)?, self.client.hvals(key))
.await
}
pub async fn hstrlen(&self, key: impl Serialize, field: impl Serialize) -> Result<usize> {
self.process_prepared_command(key_to_bulk_string(&key)?, self.client.hstrlen(key, field))
.await
}
pub async fn hmget<R: Response + DeserializeOwned>(
&self,
key: impl Serialize,
fields: impl Serialize,
) -> Result<R> {
self.process_prepared_command(key_to_bulk_string(&key)?, self.client.hmget(key, fields))
.await
}
pub async fn lrange<R: Response + DeserializeOwned>(
&self,
key: impl Serialize,
start: isize,
stop: isize,
) -> Result<R> {
self.process_prepared_command(
key_to_bulk_string(&key)?,
self.client.lrange(key, start, stop),
)
.await
}
pub async fn llen(&self, key: impl Serialize) -> Result<usize> {
self.process_prepared_command(key_to_bulk_string(&key)?, self.client.llen(key))
.await
}
pub async fn lindex<R: Response + DeserializeOwned>(
&self,
key: impl Serialize,
index: isize,
) -> Result<R> {
self.process_prepared_command(key_to_bulk_string(&key)?, self.client.lindex(key, index))
.await
}
pub async fn smembers<R: Response + DeserializeOwned>(&self, key: impl Serialize) -> Result<R> {
self.process_prepared_command(key_to_bulk_string(&key)?, self.client.smembers(key))
.await
}
pub async fn scard(&self, key: impl Serialize) -> Result<usize> {
self.process_prepared_command(key_to_bulk_string(&key)?, self.client.scard(key))
.await
}
pub async fn sismember(&self, key: impl Serialize, member: impl Serialize) -> Result<bool> {
self.process_prepared_command(
key_to_bulk_string(&key)?,
self.client.sismember(key, member),
)
.await
}
pub async fn zcard(&self, key: impl Serialize) -> Result<usize> {
self.process_prepared_command(key_to_bulk_string(&key)?, self.client.zcard(key))
.await
}
pub async fn zcount(
&self,
key: impl Serialize,
min: impl Serialize,
max: impl Serialize,
) -> Result<usize> {
self.process_prepared_command(key_to_bulk_string(&key)?, self.client.zcount(key, min, max))
.await
}
pub async fn zlexcount(
&self,
key: impl Serialize,
min: impl Serialize,
max: impl Serialize,
) -> Result<usize> {
self.process_prepared_command(
key_to_bulk_string(&key)?,
self.client.zlexcount(key, min, max),
)
.await
}
pub async fn zrange<R: Response + DeserializeOwned>(
&self,
key: impl Serialize,
start: impl Serialize,
stop: impl Serialize,
options: ZRangeOptions,
) -> Result<R> {
self.process_prepared_command(
key_to_bulk_string(&key)?,
self.client.zrange(key, start, stop, options),
)
.await
}
pub async fn zrank(
&self,
key: impl Serialize,
member: impl Serialize,
) -> Result<Option<usize>> {
self.process_prepared_command(key_to_bulk_string(&key)?, self.client.zrank(key, member))
.await
}
pub async fn zrevrank(
&self,
key: impl Serialize,
member: impl Serialize,
) -> Result<Option<usize>> {
self.process_prepared_command(key_to_bulk_string(&key)?, self.client.zrevrank(key, member))
.await
}
pub async fn zscore(&self, key: impl Serialize, member: impl Serialize) -> Result<Option<f64>> {
self.process_prepared_command(key_to_bulk_string(&key)?, self.client.zscore(key, member))
.await
}
pub async fn bitcount(&self, key: impl Serialize, range: BitRange) -> Result<usize> {
self.process_prepared_command(key_to_bulk_string(&key)?, self.client.bitcount(key, range))
.await
}
pub async fn bitpos(&self, key: impl Serialize, bit: u64, range: BitRange) -> Result<usize> {
self.process_prepared_command(
key_to_bulk_string(&key)?,
self.client.bitpos(key, bit, range),
)
.await
}
pub async fn getbit(&self, key: impl Serialize, offset: u64) -> Result<u64> {
self.process_prepared_command(key_to_bulk_string(&key)?, self.client.getbit(key, offset))
.await
}
pub async fn bitfield_readonly<'a>(
&self,
key: impl Serialize,
sub_commands: impl IntoIterator<Item = BitFieldSubCommand<'a>> + Serialize,
) -> Result<Vec<u64>> {
self.process_prepared_command(
key_to_bulk_string(&key)?,
self.client.bitfield_readonly(key, sub_commands),
)
.await
}
async fn process_prepared_command<'a, R>(
&self,
key: BulkString,
prepared_command: PreparedCommand<'a, &'a Client, R>,
) -> Result<R>
where
R: Response + DeserializeOwned,
{
self.process_command(key, prepared_command.command).await
}
async fn process_command<R>(&self, key: BulkString, command: Command) -> Result<R>
where
R: Response + DeserializeOwned,
{
if let Some(values) = self.cache.get(&key).await
&& let Some(response) = values.get(command.bytes())
{
tracing::debug!(
tag = %self.client.connection_tag(),
"Cache hit on key `{key}`"
);
let deserializer = RespDeserializer::new(response.view()?);
return R::deserialize(deserializer);
}
tracing::debug!(
tag = %self.client.connection_tag(),
"Cache miss on key `{key}`"
);
let generation_before = self.generation_counter.load(Ordering::SeqCst);
let command_bytes = command.bytes().clone();
let response = self.client.internal_send(command, None).await?;
let deserializer = RespDeserializer::new(response.view()?);
let deserialized = R::deserialize(deserializer)?;
let key_for_check = key.clone();
self.cache
.entry(key)
.or_insert_with(async { Arc::new(DashMap::new()) })
.await
.value()
.insert(command_bytes, response.compact());
let recorded = self.key_generations.get(&key_for_check).map(|g| *g);
let flushed_at = self.flush_generation.load(Ordering::SeqCst);
match post_insert_action(recorded, generation_before, flushed_at) {
PostInsertAction::DropStale => {
self.cache.invalidate(&key_for_check).await;
}
PostInsertAction::PruneGeneration => {
self.key_generations.remove(&key_for_check);
}
PostInsertAction::Keep => {}
}
Ok(deserialized)
}
}
fn get_subcache_key(key: &BulkString) -> Bytes {
FastPathCommandBuilder::get(key.clone()).bytes().clone()
}
fn key_to_bulk_string(key: &impl Serialize) -> Result<BulkString> {
let args = CommandArgsMut::default().arg(key).freeze();
args.into_iter()
.next()
.map(Into::into)
.ok_or_else(|| Error::Client(ClientError::InvalidCacheKey))
}
#[derive(Debug, PartialEq, Eq)]
enum PostInsertAction {
DropStale,
PruneGeneration,
Keep,
}
fn post_insert_action(
recorded_generation: Option<u64>,
sampled_before: u64,
flushed_at: u64,
) -> PostInsertAction {
if sampled_before < flushed_at {
return PostInsertAction::DropStale;
}
match recorded_generation {
Some(generation) if generation > sampled_before => PostInsertAction::DropStale,
Some(_) => PostInsertAction::PruneGeneration,
None => PostInsertAction::Keep,
}
}
#[cfg(test)]
mod tests {
#![allow(
clippy::unwrap_used,
clippy::expect_used,
clippy::panic,
clippy::unreachable,
clippy::indexing_slicing,
reason = "test code: a panic is how a test reports failure"
)]
use super::{PostInsertAction, post_insert_action};
#[test]
fn no_invalidation_recorded_keeps_entry() {
assert_eq!(PostInsertAction::Keep, post_insert_action(None, 5, 0));
}
#[test]
fn invalidation_after_sample_drops_stale_entry() {
assert_eq!(
PostInsertAction::DropStale,
post_insert_action(Some(6), 5, 0)
);
}
#[test]
fn a_flush_drops_an_entry_fetched_before_it_whatever_its_key_record() {
assert_eq!(PostInsertAction::DropStale, post_insert_action(None, 5, 6));
assert_eq!(
PostInsertAction::DropStale,
post_insert_action(Some(3), 5, 6)
);
}
#[test]
fn a_flush_leaves_a_later_fetch_alone() {
assert_eq!(PostInsertAction::Keep, post_insert_action(None, 6, 6));
assert_eq!(PostInsertAction::Keep, post_insert_action(None, 7, 6));
}
#[test]
fn invalidation_at_or_before_sample_is_stale_record_pruned() {
assert_eq!(
PostInsertAction::PruneGeneration,
post_insert_action(Some(5), 5, 0)
);
assert_eq!(
PostInsertAction::PruneGeneration,
post_insert_action(Some(4), 5, 0)
);
}
}