26af08e
use std::time::Duration;
use uuid::Uuid;
use crate::error::{CacheError, CacheResult};
use crate::kv::Cache;
pub const RETRY_INTERVAL: Duration = Duration::from_millis(50);
/// TTL 既是互斥租约也是故障自愈上限：持锁方必须在 TTL 内完成临界区，
/// 超时后锁自动失效，临界区可能与其他持有者并发执行；长任务请调用
/// `refresh()` 续期或改用更长的 TTL。
pub struct CacheLock {
cache: Cache,
key: String,
token: String,
ttl: Duration,
released: bool,
}
pub struct CacheLocks {
cache: Cache,
keys: Vec<String>,
token: String,
ttl: Duration,
released: bool,
}
impl Cache {
pub async fn try_lock(&self, key: &str, ttl: Duration) -> CacheResult<Option<CacheLock>> {
let token = { Uuid::now_v7().to_string() };
if self.set_nx_ex(key, token.as_str(), ttl).await? {
Ok(Some(CacheLock::new(self.clone(), key, token, ttl)))
} else {
Ok(None)
}
}
pub async fn lock(
&self,
key: &str,
ttl: Duration,
wait: Option<Duration>,
) -> CacheResult<CacheLock> {
let Some(wait) = wait else {
return self
.try_lock(key, ttl)
.await?
.ok_or_else(|| CacheError::lock(key, "already locked"));
};
let deadline = tokio::time::Instant::now() + wait;
loop {
if let Some(lock) = self.try_lock(key, ttl).await? {
return Ok(lock);
}
let now = tokio::time::Instant::now();
if now >= deadline {
return Err(CacheError::lock(
key,
"timed out while waiting for the lock",
));
}
tokio::time::sleep((deadline - now).min(RETRY_INTERVAL)).await;
}
}
}
impl CacheLock {
pub(crate) fn new(cache: Cache, key: &str, token: String, ttl: Duration) -> Self {
Self {
cache,
key: key.to_string(),
token,
ttl,
released: false,
}
}
pub fn key(&self) -> &str {
&self.key
}
pub fn token(&self) -> &str {
&self.token
}
pub fn ttl(&self) -> Duration {
self.ttl
}
pub async fn refresh(&self) -> CacheResult<bool> {
self.refresh_with(self.ttl).await
}
pub async fn refresh_with(&self, ttl: Duration) -> CacheResult<bool> {
self.cache
.scripts()
.refresh_lock(&self.key, &self.token, ttl)
.await
}
pub async fn release(&mut self) -> CacheResult<bool> {
if self.released {
return Ok(false);
}
let released = self
.cache
.scripts()
.release_lock(&self.key, &self.token)
.await?;
if released {
self.released = true;
}
Ok(released)
}
}
impl CacheLocks {
pub fn keys(&self) -> &[String] {
&self.keys
}
pub fn token(&self) -> &str {
&self.token
}
pub fn ttl(&self) -> Duration {
self.ttl
}
pub async fn release(&mut self) -> CacheResult<usize> {
if self.released {
return Ok(0);
}
let script = self.cache.scripts();
let releases = futures_util::future::join_all(
self.keys
.iter()
.map(|key| script.release_lock(key, &self.token)),
)
.await;
let mut released = 0;
for result in releases {
match result {
Ok(true) => released += 1,
Ok(false) => {}
Err(error) => {
tracing::warn!(%error, "distributed lock release failed; key expires via TTL");
}
}
}
self.released = true;
Ok(released)
}
}
impl Drop for CacheLock {
fn drop(&mut self) {
if self.released {
return;
}
{
let cache = self.cache.clone();
let key = self.key.clone();
let token = self.token.clone();
if let Ok(handle) = tokio::runtime::Handle::try_current() {
handle.spawn(async move {
let _ = cache.scripts().release_lock(&key, &token).await;
});
}
}
}
}
impl Drop for CacheLocks {
fn drop(&mut self) {
if self.released {
return;
}
let cache = self.cache.clone();
let keys = self.keys.clone();
let token = self.token.clone();
if let Ok(handle) = tokio::runtime::Handle::try_current() {
handle.spawn(async move {
for key in keys {
let _ = cache.scripts().release_lock(&key, &token).await;
}
});
}
}
}