f2e745a
use redis::{AsyncCommands, AsyncIter};
use crate::error::{CacheError, CacheResult};
use crate::kv::Cache;
pub const SEPARATOR: char = ':';
pub const DEFAULT_BATCH: usize = 256;
pub fn join(prefix: &str, segments: &[&str]) -> String {
let mut key = String::with_capacity(prefix.len() + segments.len() * 8);
key.push_str(prefix);
for segment in segments {
if !key.is_empty() {
key.push(SEPARATOR);
}
key.push_str(segment);
}
key
}
pub fn join_tagged(prefix: &str, tag: &str, segments: &[&str]) -> String {
join(
prefix,
&[
&format!("{{{tag}}}"),
&segments.join(&SEPARATOR.to_string()),
],
)
}
pub fn slot_tag(key: &str) -> Option<&str> {
let start = key.find('{')? + 1;
let end = key[start..].find('}')? + start;
if end > start {
Some(&key[start..end])
} else {
None
}
}
pub fn same_slot<S: AsRef<str>>(keys: &[S]) -> bool {
if keys.len() < 2 {
return true;
}
let Some(tag) = slot_tag(keys[0].as_ref()) else {
return false;
};
keys[1..]
.iter()
.all(|key| slot_tag(key.as_ref()) == Some(tag))
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct CacheKey {
prefix: String,
segments: Vec<String>,
}
impl CacheKey {
pub fn new(prefix: impl Into<String>) -> Self {
Self {
prefix: prefix.into(),
segments: Vec::new(),
}
}
pub fn segment(mut self, segment: impl Into<String>) -> Self {
self.segments.push(segment.into());
self
}
pub fn tag(self, tag: impl Into<String>) -> Self {
self.segment(format!("{{{}}}", tag.into()))
}
pub fn build(&self) -> String {
let segments: Vec<&str> = self.segments.iter().map(String::as_str).collect();
join(&self.prefix, &segments)
}
pub fn pattern(&self) -> String {
let mut pattern = self.build();
pattern.push('*');
pattern
}
}
impl std::fmt::Display for CacheKey {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.write_str(&self.build())
}
}
impl Cache {
pub fn key(&self, segments: &[&str]) -> String {
join(self.key_prefix(), segments)
}
pub fn key_tagged(&self, tag: &str, segments: &[&str]) -> String {
join_tagged(self.key_prefix(), tag, segments)
}
pub fn pattern(&self, segments: &[&str]) -> String {
let mut pattern = self.key(segments);
pattern.push('*');
pattern
}
pub async fn scan(&self, pattern: &str) -> CacheResult<Vec<String>> {
self.ensure_scan_supported()?;
let mut conn = self.acquire().await?;
let mut iter: AsyncIter<'_, String> =
conn.scan_match(pattern).await.map_err(CacheError::from)?;
let mut keys = Vec::new();
while let Some(item) = iter.next_item().await {
keys.push(item.map_err(CacheError::from)?);
}
Ok(keys)
}
pub async fn scan_prefix(&self, prefix: &str) -> CacheResult<Vec<String>> {
let pattern = format!("{prefix}*");
self.scan(&pattern).await
}
pub async fn count_prefix(&self, prefix: &str) -> CacheResult<usize> {
Ok(self.scan_prefix(prefix).await?.len())
}
pub async fn exists_prefix(&self, prefix: &str) -> CacheResult<bool> {
self.ensure_scan_supported()?;
let mut conn = self.acquire().await?;
let pattern = format!("{prefix}*");
let mut iter: AsyncIter<'_, String> =
conn.scan_match(pattern).await.map_err(CacheError::from)?;
// SCAN 迭代的错误不能被当成「存在」：否则缓存故障会静默变成 true。
match iter.next_item().await {
Some(Ok(_)) => Ok(true),
Some(Err(error)) => Err(CacheError::from(error)),
None => Ok(false),
}
}
pub async fn delete_prefix(&self, prefix: &str) -> CacheResult<usize> {
let keys = self.scan_prefix(prefix).await?;
if keys.is_empty() {
return Ok(0);
}
let mut removed = 0usize;
for chunk in keys.chunks(DEFAULT_BATCH) {
removed += self.delete_many(chunk.to_vec()).await?;
}
Ok(removed)
}
fn ensure_scan_supported(&self) -> CacheResult<()> {
if self.is_cluster() {
return Err(CacheError::unsupported(
"scan",
"SCAN is not routable over a Redis Cluster connection",
));
}
Ok(())
}
}