use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
use std::time::{Duration, Instant};
use redis::aio::ConnectionManager;
use tokio::task::JoinHandle;
use tokio::time::{MissedTickBehavior, interval};
use crate::DistkitError;
use crate::lock::{LockError, LockGuardState, LockOptions, mutex_backend};
#[derive(Debug)]
pub struct Mutex {
connection_manager: ConnectionManager,
full_key: String,
owner: String,
ttl_ms: i64,
ttl_duration: Duration,
max_wait: Option<Duration>,
retry_interval: Duration,
}
impl Mutex {
pub fn new(options: LockOptions) -> Arc<Self> {
let LockOptions {
key,
connection_manager,
namespace,
ttl,
owner_id,
max_wait,
retry_interval,
..
} = options;
let full_key = format!("{}:{}", *namespace, *key);
let owner = owner_id.unwrap_or_else(|| uuid::Uuid::new_v4().to_string());
Arc::new(Self {
connection_manager,
full_key,
owner,
ttl_ms: ttl.as_millis() as i64,
ttl_duration: ttl,
max_wait,
retry_interval,
})
}
pub async fn lock(&self) -> Result<MutexGuard, DistkitError> {
self.acquire_core(self.max_wait, None, self.retry_interval)
.await
}
pub async fn try_lock(&self) -> Result<MutexGuard, DistkitError> {
self.acquire_core(Some(Duration::ZERO), None, Duration::ZERO)
.await
}
pub async fn try_lock_with_timeout(&self, timeout: Duration) -> Result<MutexGuard, DistkitError> {
self.acquire_core(Some(timeout), None, self.retry_interval)
.await
}
pub async fn try_lock_with_retries(
&self,
max_retries: usize,
) -> Result<MutexGuard, DistkitError> {
self.acquire_core(None, Some(max_retries), self.retry_interval)
.await
}
#[deprecated(
since = "0.5.3",
note = "use `try_lock_with_timeout`, which sources the retry interval from the lock's configuration"
)]
pub async fn try_lock_for(
&self,
timeout: Duration,
retry_interval: Duration,
) -> Result<MutexGuard, DistkitError> {
self.acquire_core(Some(timeout), None, retry_interval).await
}
async fn acquire_core(
&self,
timeout: Option<Duration>,
max_retries: Option<usize>,
retry_interval: Duration,
) -> Result<MutexGuard, DistkitError> {
let start = Instant::now();
let mut connection = self.connection_manager.clone();
let mut retry_interval = if retry_interval.is_zero() {
None
} else {
let mut retry_interval = interval(retry_interval);
retry_interval.set_missed_tick_behavior(MissedTickBehavior::Delay);
Some(retry_interval)
};
let mut attempt = 0usize;
loop {
if let Some(retry_interval) = &mut retry_interval {
retry_interval.tick().await;
}
let acquired =
mutex_backend::acquire(&mut connection, &self.full_key, &self.owner, self.ttl_ms)
.await?;
if acquired {
let lost = Arc::new(AtomicBool::new(false));
let refresh_handle = self.spawn_refresh(lost.clone());
return Ok(MutexGuard {
connection_manager: self.connection_manager.clone(),
full_key: self.full_key.clone(),
owner: self.owner.clone(),
refresh_handle: Some(refresh_handle),
lost,
on_attempt: attempt,
});
}
if retry_interval.is_none() {
return Err(LockError::AcquireFail.into());
}
if let Some(ttl) = timeout {
if ttl.is_zero() {
return Err(LockError::AcquireFail.into());
}
let waited = start.elapsed();
if waited >= ttl {
return Err(LockError::Timeout { waited }.into());
}
}
if let Some(max_retries) = max_retries
&& attempt >= max_retries
{
return Err(LockError::RetriesExhausted { retries: attempt }.into());
}
attempt += 1;
}
}
fn spawn_refresh(&self, lost: Arc<AtomicBool>) -> JoinHandle<()> {
let mut connection_manager = self.connection_manager.clone();
let full_key = self.full_key.clone();
let owner = self.owner.clone();
let ttl_ms = self.ttl_ms;
let ttl_duration = self.ttl_duration;
tokio::spawn(async move {
let mut ticker = interval(ttl_duration / 3);
ticker.set_missed_tick_behavior(MissedTickBehavior::Delay);
ticker.tick().await;
loop {
ticker.tick().await;
match mutex_backend::refresh(&mut connection_manager, &full_key, &owner, ttl_ms)
.await
{
Ok(true) => {
if lost.swap(false, Ordering::AcqRel) {
tracing::debug!(
full_key,
owner,
"Lost distributed lock reaquired during refresh"
);
}
}
Ok(false) => {
tracing::debug!(
full_key,
owner,
"Lost distributed lock lease during refresh"
);
lost.store(true, Ordering::Release);
}
Err(error) => {
tracing::debug!(
?error,
full_key,
owner,
"Error refreshing distributed lock lease"
);
lost.store(true, Ordering::Release);
}
}
}
})
} }
#[derive(Debug)]
pub struct MutexGuard {
connection_manager: ConnectionManager,
full_key: String,
owner: String,
refresh_handle: Option<JoinHandle<()>>,
lost: Arc<AtomicBool>,
on_attempt: usize,
}
impl MutexGuard {
pub async fn get_state(&self) -> LockGuardState {
if self.refresh_handle.is_none() {
return LockGuardState::Released;
}
if self.lost.load(Ordering::Acquire) {
return LockGuardState::Lost;
}
LockGuardState::Acquired
}
pub async fn get_on_attempt(&self) -> usize {
self.on_attempt
}
pub async fn release(mut self) -> Result<LockGuardState, DistkitError> {
if let Some(handle) = self.refresh_handle.take() {
handle.abort();
}
if self.lost.load(Ordering::Acquire) {
return Ok(LockGuardState::Lost);
}
let mut connection = self.connection_manager.clone();
mutex_backend::release(&mut connection, &self.full_key, &self.owner).await?;
Ok(LockGuardState::Released)
}
}
impl Drop for MutexGuard {
fn drop(&mut self) {
let Some(handle) = self.refresh_handle.take() else {
return;
};
handle.abort();
let mut connection = self.connection_manager.clone();
let full_key = self.full_key.clone();
let owner = self.owner.clone();
tokio::spawn(async move {
if let Err(error) = mutex_backend::release(&mut connection, &full_key, &owner).await {
tracing::error!(?error, full_key, "Error releasing distributed lock on drop");
}
});
}
}