1use corium_store::{DbRoot, FORMAT_VERSION, RootStore, StoreError, db_root_name};
18use thiserror::Error;
19
20#[must_use]
22pub fn lease_root(db: &str) -> String {
23 format!("lease:{db}")
24}
25
26#[derive(Clone, Debug, Eq, PartialEq)]
28pub struct Lease {
29 pub owner: String,
31 pub version: u64,
33 pub expires_unix_ms: i64,
35 pub endpoint: String,
37}
38
39#[derive(Debug, Error)]
41pub enum LeaseError {
42 #[error("lease held by {owner} until {expires_unix_ms}")]
44 Held {
45 owner: String,
47 expires_unix_ms: i64,
49 },
50 #[error("lease lost to another owner")]
52 Lost,
53 #[error(transparent)]
55 Store(#[from] StoreError),
56}
57
58fn decode_legacy(bytes: &[u8]) -> Option<(String, u64, i64)> {
60 let text = std::str::from_utf8(bytes).ok()?;
61 let mut lines = text.lines();
62 Some((
63 lines.next()?.to_owned(),
64 lines.next()?.parse().ok()?,
65 lines.next()?.parse().ok()?,
66 ))
67}
68
69fn holder(root: Option<&DbRoot>, legacy: Option<&(String, u64, i64)>) -> (String, u64, i64) {
71 match root {
72 Some(root) if !root.owner.is_empty() => (
73 root.owner.clone(),
74 root.lease_version,
75 root.lease_expires_unix_ms,
76 ),
77 Some(root) => match legacy {
78 Some((owner, version, expiry)) => {
79 (owner.clone(), (*version).max(root.lease_version), *expiry)
80 }
81 None => (String::new(), root.lease_version, 0),
82 },
83 None => match legacy {
84 Some((owner, version, expiry)) => (owner.clone(), *version, *expiry),
85 None => (String::new(), 0, 0),
86 },
87 }
88}
89
90pub async fn acquire(
101 store: &dyn RootStore,
102 db: &str,
103 owner: &str,
104 endpoint: &str,
105 ttl_ms: i64,
106 now_unix_ms: i64,
107) -> Result<Lease, LeaseError> {
108 let name = db_root_name(db);
109 let legacy_name = lease_root(db);
110 loop {
111 let current_bytes = store.get_root(&name).await?;
112 let current = current_bytes.as_deref().and_then(DbRoot::decode);
113 let legacy = store
114 .get_root(&legacy_name)
115 .await?
116 .as_deref()
117 .and_then(decode_legacy);
118 let (cur_owner, cur_version, cur_expiry) = holder(current.as_ref(), legacy.as_ref());
119 let version = if !cur_owner.is_empty() && cur_owner == owner {
120 cur_version
121 } else if cur_expiry > now_unix_ms {
122 return Err(LeaseError::Held {
123 owner: cur_owner,
124 expires_unix_ms: cur_expiry,
125 });
126 } else {
127 cur_version + 1
128 };
129 let next = DbRoot {
130 format_version: FORMAT_VERSION,
131 lease_version: version,
132 owner: owner.to_owned(),
133 lease_expires_unix_ms: now_unix_ms + ttl_ms,
134 owner_endpoint: endpoint.to_owned(),
135 index_basis_t: current.as_ref().map_or(0, |root| root.index_basis_t),
136 roots: current.as_ref().and_then(|root| root.roots.clone()),
137 };
138 match store
139 .cas_root(&name, current_bytes.as_deref(), &next.encode())
140 .await
141 {
142 Ok(()) => {
143 if legacy.is_some() {
144 let _ = store.delete_root(&legacy_name).await;
146 }
147 return Ok(Lease {
148 owner: owner.to_owned(),
149 version,
150 expires_unix_ms: next.lease_expires_unix_ms,
151 endpoint: endpoint.to_owned(),
152 });
153 }
154 Err(StoreError::CasFailed { .. }) => {}
155 Err(error) => return Err(error.into()),
156 }
157 }
158}
159
160pub async fn renew(
167 store: &dyn RootStore,
168 db: &str,
169 held: &Lease,
170 ttl_ms: i64,
171 now_unix_ms: i64,
172) -> Result<Lease, LeaseError> {
173 let name = db_root_name(db);
174 loop {
175 let current_bytes = store.get_root(&name).await?;
176 let Some(root) = current_bytes.as_deref().and_then(DbRoot::decode) else {
177 return Err(LeaseError::Lost);
178 };
179 if root.owner != held.owner || root.lease_version != held.version {
180 return Err(LeaseError::Lost);
181 }
182 let next = DbRoot {
183 lease_expires_unix_ms: now_unix_ms + ttl_ms,
184 owner_endpoint: held.endpoint.clone(),
185 ..root
186 };
187 match store
188 .cas_root(&name, current_bytes.as_deref(), &next.encode())
189 .await
190 {
191 Ok(()) => {
192 return Ok(Lease {
193 expires_unix_ms: next.lease_expires_unix_ms,
194 ..held.clone()
195 });
196 }
197 Err(StoreError::CasFailed { .. }) => {}
201 Err(error) => return Err(error.into()),
202 }
203 }
204}
205
206pub async fn verify(store: &dyn RootStore, db: &str, held: &Lease) -> Result<(), LeaseError> {
215 let root = store
216 .get_root(&db_root_name(db))
217 .await?
218 .as_deref()
219 .and_then(DbRoot::decode);
220 match root {
221 Some(root) if root.owner == held.owner && root.lease_version == held.version => Ok(()),
222 _ => Err(LeaseError::Lost),
223 }
224}
225
226pub async fn release(store: &dyn RootStore, db: &str, held: &Lease) -> Result<(), LeaseError> {
232 let name = db_root_name(db);
233 loop {
234 let current_bytes = store.get_root(&name).await?;
235 let Some(root) = current_bytes.as_deref().and_then(DbRoot::decode) else {
236 return Err(LeaseError::Lost);
237 };
238 if root.owner != held.owner || root.lease_version != held.version {
239 return Err(LeaseError::Lost);
240 }
241 let next = DbRoot {
242 lease_expires_unix_ms: 0,
243 ..root
244 };
245 match store
246 .cas_root(&name, current_bytes.as_deref(), &next.encode())
247 .await
248 {
249 Ok(()) => return Ok(()),
250 Err(StoreError::CasFailed { .. }) => {}
251 Err(error) => return Err(error.into()),
252 }
253 }
254}