f2e745a
use std::future::Future;
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 = new_token();
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;
}
}
pub async fn with_lock<T, F, Fut>(
&self,
key: &str,
ttl: Duration,
wait: Option<Duration>,
operation: F,
) -> CacheResult<T>
where
F: FnOnce() -> Fut,
Fut: Future<Output = CacheResult<T>>,
{
let mut lock = self.lock(key, ttl, wait).await?;
let result = operation().await;
let _ = lock.release().await;
result
}
pub async fn try_lock_many(
&self,
keys: &[String],
ttl: Duration,
) -> CacheResult<Option<CacheLocks>> {
if keys.is_empty() {
return Err(CacheError::lock("", "the key list must not be empty"));
}
let token = new_token();
if self.scripts().lock_many(keys, &token, ttl).await? {
Ok(Some(CacheLocks::new(self.clone(), keys, token, ttl)))
} else {
Ok(None)
}
}
}
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 fn is_released(&self) -> bool {
self.released
}
pub async fn is_held(&self) -> CacheResult<bool> {
Ok(
matches!(self.cache.get_bytes(&self.key).await?, Some(value) if value == self.token.as_bytes()),
)
}
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(crate) fn new(cache: Cache, keys: &[String], token: String, ttl: Duration) -> Self {
Self {
cache,
keys: keys.to_vec(),
token,
ttl,
released: false,
}
}
pub fn keys(&self) -> &[String] {
&self.keys
}
pub fn token(&self) -> &str {
&self.token
}
pub fn ttl(&self) -> Duration {
self.ttl
}
pub fn is_released(&self) -> bool {
self.released
}
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;
}
release_in_background(self.cache.clone(), self.key.clone(), self.token.clone());
}
}
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;
}
});
}
}
}
fn release_in_background(cache: Cache, key: String, token: String) {
if let Ok(handle) = tokio::runtime::Handle::try_current() {
handle.spawn(async move {
let _ = cache.scripts().release_lock(&key, &token).await;
});
}
}
fn new_token() -> String {
Uuid::now_v7().to_string()
}