use std::future::Future;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Mutex;
use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};
use super::{AsyncLock, InMemoryAsyncLock, LockError};
#[derive(Debug, Clone)]
pub(crate) struct LeaseConfig {
pub lease_ttl: Duration,
pub retry_interval: Duration,
pub max_wait: Option<Duration>,
}
impl Default for LeaseConfig {
fn default() -> Self {
Self {
lease_ttl: Duration::from_secs(30),
retry_interval: Duration::from_millis(50),
max_wait: None,
}
}
}
pub(crate) struct LockShared {
pub gate: InMemoryAsyncLock,
pub token: Mutex<Option<String>>,
}
impl LockShared {
pub fn new() -> Self {
Self {
gate: InMemoryAsyncLock::new(),
token: Mutex::new(None),
}
}
}
pub(crate) trait LeaseBackend: Send + Sync {
fn shared(&self) -> &LockShared;
fn config(&self) -> &LeaseConfig;
fn mint_token(&self) -> String;
fn db_acquire(&self, token: &str) -> impl Future<Output = Result<bool, LockError>> + Send;
fn db_release(&self, token: &str) -> impl Future<Output = Result<(), LockError>> + Send;
}
struct GateGuard<'a> {
gate: &'a InMemoryAsyncLock,
armed: bool,
}
impl<'a> GateGuard<'a> {
fn new(gate: &'a InMemoryAsyncLock) -> Self {
Self { gate, armed: true }
}
fn disarm(&mut self) {
self.armed = false;
}
}
impl Drop for GateGuard<'_> {
fn drop(&mut self) {
if self.armed {
let _ = self.gate.unlock_core();
}
}
}
pub(crate) async fn lease_lock(backend: &impl LeaseBackend) -> Result<(), LockError> {
let started = Instant::now();
let max_wait = backend.config().max_wait;
match max_wait {
Some(max) => match tokio::time::timeout(max, backend.shared().gate.lock()).await {
Ok(result) => result?,
Err(_elapsed) => {
return Err(LockError::AcquireFailed(format!(
"lease acquire timed out after {max:?} waiting for the in-process gate"
)))
}
},
None => backend.shared().gate.lock().await?,
}
let mut guard = GateGuard::new(&backend.shared().gate);
let token = backend.mint_token();
loop {
match backend.db_acquire(&token).await {
Ok(true) => {
store_token(backend.shared(), Some(token));
guard.disarm(); return Ok(());
}
Ok(false) => {}
Err(err) => return Err(err), }
if let Some(max) = max_wait {
if started.elapsed() >= max {
return Err(LockError::AcquireFailed(format!(
"lease acquire timed out after {max:?}"
)));
}
}
tokio::time::sleep(jittered(backend.config().retry_interval, &token)).await;
}
}
pub(crate) async fn lease_try_lock(backend: &impl LeaseBackend) -> Result<bool, LockError> {
if !backend.shared().gate.try_lock().await? {
return Ok(false);
}
let mut guard = GateGuard::new(&backend.shared().gate);
let token = backend.mint_token();
match backend.db_acquire(&token).await {
Ok(true) => {
store_token(backend.shared(), Some(token));
guard.disarm();
Ok(true)
}
Ok(false) => Ok(false), Err(err) => Err(err), }
}
pub(crate) async fn lease_unlock(backend: &impl LeaseBackend) -> Result<(), LockError> {
let _guard = GateGuard::new(&backend.shared().gate);
let token = store_token(backend.shared(), None);
if let Some(token) = token {
let _ = backend.db_release(&token).await;
}
Ok(())
}
fn store_token(shared: &LockShared, next: Option<String>) -> Option<String> {
let mut slot = shared
.token
.lock()
.unwrap_or_else(|poisoned| poisoned.into_inner());
std::mem::replace(&mut slot, next)
}
pub(crate) fn default_owner_id(prefix: &str) -> String {
static MANAGER_SEQ: AtomicU64 = AtomicU64::new(0);
let pid = std::process::id();
let nanos = SystemTime::now()
.duration_since(UNIX_EPOCH)
.map(|d| d.as_nanos())
.unwrap_or(0);
let seq = MANAGER_SEQ.fetch_add(1, Ordering::Relaxed);
format!("{prefix}-{pid}-{nanos}-{seq}")
}
pub(crate) fn mint_token(owner_id: &str, seq: &AtomicU64) -> String {
format!("{owner_id}:{}", seq.fetch_add(1, Ordering::Relaxed))
}
fn jittered(base: Duration, token: &str) -> Duration {
use std::hash::{Hash, Hasher};
let mut hasher = std::collections::hash_map::DefaultHasher::new();
token.hash(&mut hasher);
let frac = (hasher.finish() % 1000) as f64 / 1000.0;
base + base.mul_f64(0.25 * frac)
}
pub(crate) fn lease_acquire_error(err: sqlx::Error) -> LockError {
LockError::AcquireFailed(format!("sqlx lease acquire failed: {err}"))
}
pub(crate) fn lease_release_error(err: sqlx::Error) -> LockError {
LockError::ReleaseFailed(format!("sqlx lease release failed: {err}"))
}