use corium_store::{DbRoot, FORMAT_VERSION, RootStore, StoreError, db_root_name};
use thiserror::Error;
#[must_use]
pub fn lease_root(db: &str) -> String {
format!("lease:{db}")
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct Lease {
pub owner: String,
pub version: u64,
pub expires_unix_ms: i64,
pub endpoint: String,
}
#[derive(Debug, Error)]
pub enum LeaseError {
#[error("lease held by {owner} until {expires_unix_ms}")]
Held {
owner: String,
expires_unix_ms: i64,
},
#[error("lease lost to another owner")]
Lost,
#[error(transparent)]
Store(#[from] StoreError),
}
fn decode_legacy(bytes: &[u8]) -> Option<(String, u64, i64)> {
let text = std::str::from_utf8(bytes).ok()?;
let mut lines = text.lines();
Some((
lines.next()?.to_owned(),
lines.next()?.parse().ok()?,
lines.next()?.parse().ok()?,
))
}
fn holder(root: Option<&DbRoot>, legacy: Option<&(String, u64, i64)>) -> (String, u64, i64) {
match root {
Some(root) if !root.owner.is_empty() => (
root.owner.clone(),
root.lease_version,
root.lease_expires_unix_ms,
),
Some(root) => match legacy {
Some((owner, version, expiry)) => {
(owner.clone(), (*version).max(root.lease_version), *expiry)
}
None => (String::new(), root.lease_version, 0),
},
None => match legacy {
Some((owner, version, expiry)) => (owner.clone(), *version, *expiry),
None => (String::new(), 0, 0),
},
}
}
pub async fn acquire(
store: &dyn RootStore,
db: &str,
owner: &str,
endpoint: &str,
ttl_ms: i64,
now_unix_ms: i64,
) -> Result<Lease, LeaseError> {
let name = db_root_name(db);
let legacy_name = lease_root(db);
loop {
let current_bytes = store.get_root(&name).await?;
let current = current_bytes.as_deref().and_then(DbRoot::decode);
let legacy = store
.get_root(&legacy_name)
.await?
.as_deref()
.and_then(decode_legacy);
let (cur_owner, cur_version, cur_expiry) = holder(current.as_ref(), legacy.as_ref());
let version = if !cur_owner.is_empty() && cur_owner == owner {
cur_version
} else if cur_expiry > now_unix_ms {
return Err(LeaseError::Held {
owner: cur_owner,
expires_unix_ms: cur_expiry,
});
} else {
cur_version + 1
};
let next = DbRoot {
format_version: FORMAT_VERSION,
lease_version: version,
owner: owner.to_owned(),
lease_expires_unix_ms: now_unix_ms + ttl_ms,
owner_endpoint: endpoint.to_owned(),
index_basis_t: current.as_ref().map_or(0, |root| root.index_basis_t),
roots: current.as_ref().and_then(|root| root.roots.clone()),
next_entity_id: current.as_ref().map_or(0, |root| root.next_entity_id),
last_tx_instant: current
.as_ref()
.map_or(i64::MIN, |root| root.last_tx_instant),
};
match store
.cas_root(&name, current_bytes.as_deref(), &next.encode())
.await
{
Ok(()) => {
if legacy.is_some() {
let _ = store.delete_root(&legacy_name).await;
}
return Ok(Lease {
owner: owner.to_owned(),
version,
expires_unix_ms: next.lease_expires_unix_ms,
endpoint: endpoint.to_owned(),
});
}
Err(StoreError::CasFailed { .. }) => {}
Err(error) => return Err(error.into()),
}
}
}
pub async fn renew(
store: &dyn RootStore,
db: &str,
held: &Lease,
ttl_ms: i64,
now_unix_ms: i64,
) -> Result<Lease, LeaseError> {
let name = db_root_name(db);
loop {
let current_bytes = store.get_root(&name).await?;
let Some(root) = current_bytes.as_deref().and_then(DbRoot::decode) else {
return Err(LeaseError::Lost);
};
if root.owner != held.owner || root.lease_version != held.version {
return Err(LeaseError::Lost);
}
let next = DbRoot {
lease_expires_unix_ms: now_unix_ms + ttl_ms,
owner_endpoint: held.endpoint.clone(),
..root
};
match store
.cas_root(&name, current_bytes.as_deref(), &next.encode())
.await
{
Ok(()) => {
return Ok(Lease {
expires_unix_ms: next.lease_expires_unix_ms,
..held.clone()
});
}
Err(StoreError::CasFailed { .. }) => {}
Err(error) => return Err(error.into()),
}
}
}
pub async fn verify(store: &dyn RootStore, db: &str, held: &Lease) -> Result<(), LeaseError> {
let root = store
.get_root(&db_root_name(db))
.await?
.as_deref()
.and_then(DbRoot::decode);
match root {
Some(root) if root.owner == held.owner && root.lease_version == held.version => Ok(()),
_ => Err(LeaseError::Lost),
}
}
pub async fn release(store: &dyn RootStore, db: &str, held: &Lease) -> Result<(), LeaseError> {
let name = db_root_name(db);
loop {
let current_bytes = store.get_root(&name).await?;
let Some(root) = current_bytes.as_deref().and_then(DbRoot::decode) else {
return Err(LeaseError::Lost);
};
if root.owner != held.owner || root.lease_version != held.version {
return Err(LeaseError::Lost);
}
let next = DbRoot {
lease_expires_unix_ms: 0,
..root
};
match store
.cas_root(&name, current_bytes.as_deref(), &next.encode())
.await
{
Ok(()) => return Ok(()),
Err(StoreError::CasFailed { .. }) => {}
Err(error) => return Err(error.into()),
}
}
}