use std::time::{Duration, Instant};
use async_trait::async_trait;
use dashmap::DashMap;
use crate::error::{Error, Result};
#[cfg(feature = "etcd")]
pub mod etcd;
#[async_trait]
pub trait LockBackend: Send + Sync {
async fn try_acquire(&self, key: &str, ttl: Duration) -> Result<(u64, Instant)>;
async fn try_renew(&self, key: &str, token: u64, ttl: Duration) -> Result<Instant>;
async fn release(&self, key: &str, token: u64) -> Result<()>;
fn name(&self) -> &'static str {
"unknown"
}
}
struct Entry {
token: u64,
held: bool,
deadline: Instant,
}
#[derive(Default)]
pub struct InMemoryBackend {
entries: DashMap<String, Entry>,
}
impl InMemoryBackend {
pub fn new() -> Self {
Self {
entries: DashMap::new(),
}
}
pub fn tracked_keys(&self) -> usize {
self.entries.len()
}
}
#[async_trait]
impl LockBackend for InMemoryBackend {
async fn try_acquire(&self, key: &str, ttl: Duration) -> Result<(u64, Instant)> {
let now = Instant::now();
let deadline = now + ttl;
let mut e = self.entries.entry(key.to_string()).or_insert(Entry {
token: 0,
held: false,
deadline: now,
});
let expired = now >= e.deadline;
if e.held && !expired {
return Err(Error::Contended {
resource: key.to_string(),
});
}
e.token = e
.token
.checked_add(1)
.expect("fencing token overflow (u64)");
e.held = true;
e.deadline = deadline;
Ok((e.token, deadline))
}
async fn try_renew(&self, key: &str, token: u64, ttl: Duration) -> Result<Instant> {
let now = Instant::now();
let mut e = self.entries.get_mut(key).ok_or_else(|| Error::StaleToken {
resource: key.to_string(),
token,
current: 0,
})?;
if !e.held || now >= e.deadline {
return Err(Error::LeaseExpired {
resource: key.to_string(),
});
}
if e.token != token {
return Err(Error::StaleToken {
resource: key.to_string(),
token,
current: e.token,
});
}
let deadline = now + ttl;
e.deadline = deadline;
Ok(deadline)
}
async fn release(&self, key: &str, token: u64) -> Result<()> {
if let Some(mut e) = self.entries.get_mut(key) {
if e.held && e.token == token {
e.held = false;
}
}
Ok(())
}
fn name(&self) -> &'static str {
"in-memory"
}
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn acquire_is_exclusive_and_monotonic() {
let b = InMemoryBackend::new();
let (t1, _) = b.try_acquire("k", Duration::from_secs(10)).await.unwrap();
assert_eq!(t1, 1);
assert!(matches!(
b.try_acquire("k", Duration::from_secs(10)).await,
Err(Error::Contended { .. })
));
b.release("k", t1).await.unwrap();
let (t2, _) = b.try_acquire("k", Duration::from_secs(10)).await.unwrap();
assert!(t2 > t1, "tokens must strictly increase across acquisitions");
}
#[tokio::test]
async fn stale_renew_is_rejected() {
let b = InMemoryBackend::new();
let (t, _) = b.try_acquire("k", Duration::from_secs(10)).await.unwrap();
assert!(
b.try_renew("k", t + 999, Duration::from_secs(10))
.await
.is_err()
);
assert!(b.try_renew("k", t, Duration::from_secs(10)).await.is_ok());
}
}