use async_trait::async_trait;
use chrono::{DateTime, Utc};
use std::sync::Arc;
use std::time::Duration;
use uuid::Uuid;
use crate::error::StateResult;
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct LockToken(String);
impl LockToken {
pub fn new() -> Self {
Self(Uuid::new_v4().to_string())
}
pub fn value(&self) -> &str {
&self.0
}
pub fn from_string(s: String) -> Self {
Self(s)
}
}
impl Default for LockToken {
fn default() -> Self {
Self::new()
}
}
impl std::fmt::Display for LockToken {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "{}", self.0)
}
}
#[derive(Debug)]
pub struct LockGuard {
resource: String,
token: LockToken,
acquired_at: DateTime<Utc>,
ttl: Duration,
backend: Arc<dyn DistributedLock>,
}
impl LockGuard {
pub fn new(
resource: String,
token: LockToken,
acquired_at: DateTime<Utc>,
ttl: Duration,
backend: Arc<dyn DistributedLock>,
) -> Self {
Self {
resource,
token,
acquired_at,
ttl,
backend,
}
}
pub fn resource(&self) -> &str {
&self.resource
}
pub fn token(&self) -> &LockToken {
&self.token
}
pub fn acquired_at(&self) -> DateTime<Utc> {
self.acquired_at
}
pub fn ttl(&self) -> Duration {
self.ttl
}
pub fn is_expired(&self) -> bool {
let elapsed = Utc::now().signed_duration_since(self.acquired_at);
elapsed.to_std().unwrap_or_default() >= self.ttl
}
pub fn time_remaining(&self) -> Duration {
let elapsed = Utc::now().signed_duration_since(self.acquired_at);
let elapsed_std = elapsed.to_std().unwrap_or_default();
if elapsed_std >= self.ttl {
Duration::from_secs(0)
} else {
self.ttl - elapsed_std
}
}
pub async fn release(self) -> StateResult<()> {
drop(self);
Ok(())
}
}
impl Drop for LockGuard {
fn drop(&mut self) {
let backend = self.backend.clone();
let resource = self.resource.clone();
let token = self.token.clone();
tokio::spawn(async move {
let _ = backend.release_internal(&resource, &token).await;
});
}
}
#[async_trait]
pub trait DistributedLock: Send + Sync + std::fmt::Debug {
async fn acquire(&self, resource: &str, ttl: Duration) -> StateResult<LockGuard>;
async fn try_acquire(&self, resource: &str, ttl: Duration) -> StateResult<Option<LockGuard>>;
async fn release(&self, guard: LockGuard) -> StateResult<()> {
self.release_internal(guard.resource(), guard.token()).await
}
async fn release_internal(&self, resource: &str, token: &LockToken) -> StateResult<()>;
async fn is_locked(&self, resource: &str) -> StateResult<bool>;
async fn extend(&self, guard: &mut LockGuard, ttl: Duration) -> StateResult<()>;
async fn get_lock_info(&self, resource: &str) -> StateResult<Option<LockToken>>;
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_lock_token_creation() {
let token1 = LockToken::new();
let token2 = LockToken::new();
assert_ne!(token1, token2);
assert!(!token1.value().is_empty());
}
#[test]
fn test_lock_token_from_string() {
let original = "test-token-123";
let token = LockToken::from_string(original.to_string());
assert_eq!(token.value(), original);
}
#[test]
fn test_lock_token_display() {
let token = LockToken::from_string("test-123".to_string());
assert_eq!(format!("{}", token), "test-123");
}
}