use std::collections::HashMap;
use std::sync::Arc;
use std::time::{Duration, Instant};
use async_trait::async_trait;
use tokio::sync::RwLock;
use crate::error::{PhotonError, Result};
#[derive(Debug, Clone)]
pub struct ConsumerLease {
pub group_id: String,
pub shard_id: u32,
pub instance_id: String,
pub ttl_secs: u64,
}
struct LeaseRecord {
instance_id: String,
expires_at: Instant,
}
#[async_trait]
pub trait LeaseStore: Send + Sync {
async fn claim(&self, lease: ConsumerLease) -> Result<()>;
async fn renew(&self, group_id: &str, instance_id: &str, ttl_secs: u64) -> Result<()>;
async fn release(&self, group_id: &str, instance_id: &str) -> Result<()>;
async fn list_for_instance(&self, group_id: &str, instance_id: &str) -> Result<Vec<u32>>;
}
pub struct MemoryLeaseStore {
inner: Arc<RwLock<HashMap<(String, u32), LeaseRecord>>>,
}
impl Default for MemoryLeaseStore {
fn default() -> Self {
Self::new()
}
}
impl MemoryLeaseStore {
#[must_use]
pub fn new() -> Self {
Self {
inner: Arc::new(RwLock::new(HashMap::new())),
}
}
fn key(group_id: &str, shard_id: u32) -> (String, u32) {
(group_id.to_string(), shard_id)
}
}
#[async_trait]
impl LeaseStore for MemoryLeaseStore {
#[allow(clippy::significant_drop_tightening)]
async fn claim(&self, lease: ConsumerLease) -> Result<()> {
let expires_at = Instant::now() + Duration::from_secs(lease.ttl_secs.max(1));
let k = Self::key(&lease.group_id, lease.shard_id);
let mut guard = self.inner.write().await;
if let Some(existing) = guard.get(&k) {
if existing.expires_at > Instant::now() && existing.instance_id != lease.instance_id {
return Err(PhotonError::Internal(format!(
"shard {} already leased by {}",
lease.shard_id, existing.instance_id
)));
}
}
guard.insert(
k,
LeaseRecord {
instance_id: lease.instance_id,
expires_at,
},
);
Ok(())
}
#[allow(clippy::significant_drop_tightening)]
async fn renew(&self, group_id: &str, instance_id: &str, ttl_secs: u64) -> Result<()> {
let expires_at = Instant::now() + Duration::from_secs(ttl_secs.max(1));
let mut guard = self.inner.write().await;
for ((g, _shard), rec) in guard.iter_mut() {
if g == group_id && rec.instance_id == instance_id {
rec.expires_at = expires_at;
}
}
Ok(())
}
async fn release(&self, group_id: &str, instance_id: &str) -> Result<()> {
self.inner
.write()
.await
.retain(|(g, _), rec| !(g == group_id && rec.instance_id == instance_id));
Ok(())
}
async fn list_for_instance(&self, group_id: &str, instance_id: &str) -> Result<Vec<u32>> {
let now = Instant::now();
let mut shards: Vec<u32> = {
let guard = self.inner.read().await;
guard
.iter()
.filter(|((g, _), rec)| {
g == group_id && rec.instance_id == instance_id && rec.expires_at > now
})
.map(|((_, shard), _)| *shard)
.collect()
};
shards.sort_unstable();
Ok(shards)
}
}