photon_backend/consumer_group/
lease_store.rs

1//! Durable shard lease metadata.
2
3use std::collections::HashMap;
4use std::sync::Arc;
5use std::time::{Duration, Instant};
6
7use async_trait::async_trait;
8use tokio::sync::RwLock;
9
10use crate::error::{PhotonError, Result};
11
12/// One shard lease held by a group member.
13#[derive(Debug, Clone)]
14pub struct ConsumerLease {
15    /// Consumer group id.
16    pub group_id: String,
17    /// Virtual shard being leased.
18    pub shard_id: u32,
19    /// Member instance holding the lease.
20    pub instance_id: String,
21    /// Lease time-to-live in seconds.
22    pub ttl_secs: u64,
23}
24
25struct LeaseRecord {
26    instance_id: String,
27    expires_at: Instant,
28}
29
30/// Persistence for `(group, shard) → instance` assignments.
31#[async_trait]
32pub trait LeaseStore: Send + Sync {
33    /// Claim a shard lease for an instance (fails if held by another live member).
34    async fn claim(&self, lease: ConsumerLease) -> Result<()>;
35    /// Renew all leases held by an instance in the group.
36    async fn renew(&self, group_id: &str, instance_id: &str, ttl_secs: u64) -> Result<()>;
37    /// Release all leases held by an instance in the group.
38    async fn release(&self, group_id: &str, instance_id: &str) -> Result<()>;
39    /// List live shard ids leased by an instance.
40    async fn list_for_instance(&self, group_id: &str, instance_id: &str) -> Result<Vec<u32>>;
41}
42
43/// In-process lease table for lab and unit tests.
44pub struct MemoryLeaseStore {
45    inner: Arc<RwLock<HashMap<(String, u32), LeaseRecord>>>,
46}
47
48impl Default for MemoryLeaseStore {
49    fn default() -> Self {
50        Self::new()
51    }
52}
53
54impl MemoryLeaseStore {
55    /// Create an empty in-memory lease table.
56    #[must_use]
57    pub fn new() -> Self {
58        Self {
59            inner: Arc::new(RwLock::new(HashMap::new())),
60        }
61    }
62
63    fn key(group_id: &str, shard_id: u32) -> (String, u32) {
64        (group_id.to_string(), shard_id)
65    }
66}
67
68#[async_trait]
69impl LeaseStore for MemoryLeaseStore {
70    #[allow(clippy::significant_drop_tightening)]
71    async fn claim(&self, lease: ConsumerLease) -> Result<()> {
72        let expires_at = Instant::now() + Duration::from_secs(lease.ttl_secs.max(1));
73        let k = Self::key(&lease.group_id, lease.shard_id);
74        let mut guard = self.inner.write().await;
75        if let Some(existing) = guard.get(&k) {
76            if existing.expires_at > Instant::now() && existing.instance_id != lease.instance_id {
77                return Err(PhotonError::Internal(format!(
78                    "shard {} already leased by {}",
79                    lease.shard_id, existing.instance_id
80                )));
81            }
82        }
83        guard.insert(
84            k,
85            LeaseRecord {
86                instance_id: lease.instance_id,
87                expires_at,
88            },
89        );
90        Ok(())
91    }
92
93    #[allow(clippy::significant_drop_tightening)]
94    async fn renew(&self, group_id: &str, instance_id: &str, ttl_secs: u64) -> Result<()> {
95        let expires_at = Instant::now() + Duration::from_secs(ttl_secs.max(1));
96        let mut guard = self.inner.write().await;
97        for ((g, _shard), rec) in guard.iter_mut() {
98            if g == group_id && rec.instance_id == instance_id {
99                rec.expires_at = expires_at;
100            }
101        }
102        Ok(())
103    }
104
105    async fn release(&self, group_id: &str, instance_id: &str) -> Result<()> {
106        self.inner
107            .write()
108            .await
109            .retain(|(g, _), rec| !(g == group_id && rec.instance_id == instance_id));
110        Ok(())
111    }
112
113    async fn list_for_instance(&self, group_id: &str, instance_id: &str) -> Result<Vec<u32>> {
114        let now = Instant::now();
115        let mut shards: Vec<u32> = {
116            let guard = self.inner.read().await;
117            guard
118                .iter()
119                .filter(|((g, _), rec)| {
120                    g == group_id && rec.instance_id == instance_id && rec.expires_at > now
121                })
122                .map(|((_, shard), _)| *shard)
123                .collect()
124        };
125        shards.sort_unstable();
126        Ok(shards)
127    }
128}
photon_backend/consumer_group/lease_store.rs

photon_backend/consumer_group/
lease_store.rs
