use crate::manifest::object_key;
use object_store::path::Path as ObjectPath;
use object_store::{ObjectStore, PutMode, PutPayload, PutResult, UpdateVersion};
use serde::{Deserialize, Serialize};
use std::sync::{Arc, Mutex};
use std::time::{Duration, SystemTime};
use time::format_description::well_known::Rfc3339;
use time::OffsetDateTime;
use tokio::time::Instant;
const ACQUIRE_RACES: usize = 3;
pub const MAX_SANE_TTL_S: u64 = 3600;
const PROBE_ETAG: &str = "\"mushroomdb-lease-probe-not-an-etag\"";
pub fn new_instance_id() -> String {
rand::random::<[u8; 16]>()
.iter()
.map(|b| format!("{b:02x}"))
.collect()
}
#[derive(Debug, Default)]
pub struct LeaseWatch {
seen: Option<(String, Instant)>,
#[cfg(test)]
pub(crate) probe_control: Option<Arc<ProbeControl>>,
}
#[cfg(test)]
#[derive(Debug, Default)]
pub(crate) struct ProbeControl {
pub(crate) block_probe: usize,
pub(crate) entered: tokio::sync::Notify,
pub(crate) resume: tokio::sync::Notify,
pub(crate) shutdown: tokio::sync::Notify,
pub(crate) failed: tokio::sync::Notify,
pub(crate) acquisitions: std::sync::atomic::AtomicUsize,
pub(crate) probes: std::sync::atomic::AtomicUsize,
pub(crate) renewals: std::sync::atomic::AtomicUsize,
}
pub fn lease_key(prefix: &str) -> String {
object_key(prefix, "LEASE")
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct LeaseRecord {
pub instance_id: String,
pub acquired_at: String,
pub renewed_at: String,
pub ttl_s: u64,
}
impl LeaseRecord {
pub fn is_fresh(&self, now: OffsetDateTime) -> bool {
let Ok(renewed) = OffsetDateTime::parse(&self.renewed_at, &Rfc3339) else {
return false;
};
let ttl = time::Duration::seconds(i64::try_from(self.ttl_s).unwrap_or(i64::MAX));
now - renewed < ttl
}
pub fn is_suspect(&self, now: OffsetDateTime) -> bool {
if self.ttl_s > MAX_SANE_TTL_S {
return true;
}
let Ok(renewed) = OffsetDateTime::parse(&self.renewed_at, &Rfc3339) else {
return false;
};
renewed - now > time::Duration::seconds(self.ttl_s as i64)
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Exclusion {
Conditional,
BestEffort,
}
#[derive(Debug)]
pub enum LeaseError {
ForeignActive { instance_id: String },
Lost { instance_id: Option<String> },
PreconditionsIgnored { header: &'static str },
Store(object_store::Error),
}
impl std::fmt::Display for LeaseError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
LeaseError::ForeignActive { instance_id } => write!(
f,
"the backup bucket's lease is held by another instance ({instance_id})"
),
LeaseError::Lost {
instance_id: Some(id),
} => write!(f, "the bucket lease was taken over by {id}"),
LeaseError::Lost { instance_id: None } => write!(f, "the bucket lease was lost"),
LeaseError::PreconditionsIgnored { header } => write!(
f,
"the backup service ignores {header} on PUT, so the bucket lease cannot keep a \
second writer out; backup is not started. Set \
MUSHROOMDB_BACKUP_LEASE_CONDITIONAL=off to run with a best-effort lease"
),
LeaseError::Store(e) => {
write!(f, "object store: {}", crate::redact(&e.to_string()))
}
}
}
}
impl std::error::Error for LeaseError {}
impl From<object_store::Error> for LeaseError {
fn from(e: object_store::Error) -> Self {
LeaseError::Store(e)
}
}
#[derive(Debug)]
struct Held {
version: UpdateVersion,
confirmed: Instant,
confirmed_wall: SystemTime,
exclusion: Exclusion,
lost: bool,
}
#[derive(Debug)]
pub struct BucketLease {
store: Arc<dyn ObjectStore>,
key: ObjectPath,
instance_id: String,
ttl_s: u64,
acquired_at: String,
held: Mutex<Held>,
#[cfg(test)]
probe_control: Option<Arc<ProbeControl>>,
}
impl BucketLease {
pub async fn acquire(
store: Arc<dyn ObjectStore>,
prefix: &str,
ttl_s: u64,
) -> Result<BucketLease, LeaseError> {
let id = new_instance_id();
Self::acquire_as(store, prefix, ttl_s, &id, &mut LeaseWatch::default()).await
}
pub async fn acquire_as(
store: Arc<dyn ObjectStore>,
prefix: &str,
ttl_s: u64,
instance_id: &str,
watch: &mut LeaseWatch,
) -> Result<BucketLease, LeaseError> {
let ttl_s = ttl_s.max(1);
let key = ObjectPath::from(lease_key(prefix));
for _ in 0..ACQUIRE_RACES {
let current = read(&store, &key).await?;
if let Some((Some(rec), version)) = ¤t {
if rec.instance_id != instance_id {
foreign_blocks(rec, version, ttl_s, watch)?;
}
}
let (confirmed, confirmed_wall) = (Instant::now(), SystemTime::now());
let stamp = rfc3339(confirmed_wall);
let record = LeaseRecord {
instance_id: instance_id.to_string(),
acquired_at: stamp.clone(),
renewed_at: stamp.clone(),
ttl_s,
};
let mode = match ¤t {
None => PutMode::Create,
Some((_, v)) => PutMode::Update(v.clone()),
};
let (result, exclusion) = match put(&store, &key, &record, mode).await {
Ok(r) => (r, Exclusion::Conditional),
Err(object_store::Error::AlreadyExists { .. })
| Err(object_store::Error::Precondition { .. }) => continue,
Err(object_store::Error::NotImplemented) => (
put(&store, &key, &record, PutMode::Overwrite).await?,
Exclusion::BestEffort,
),
Err(e) => return Err(e.into()),
};
let lease = BucketLease {
store,
key,
instance_id: instance_id.to_string(),
ttl_s,
acquired_at: stamp,
held: Mutex::new(Held {
version: version_of(result),
confirmed,
confirmed_wall,
exclusion,
lost: false,
}),
#[cfg(test)]
probe_control: watch.probe_control.clone(),
};
#[cfg(test)]
if let Some(control) = &lease.probe_control {
control
.acquisitions
.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
}
if exclusion == Exclusion::Conditional {
if let Err(e) = lease.probe(&record).await {
if matches!(e, LeaseError::PreconditionsIgnored { .. }) {
let _ =
release_own_record(&lease.store, &lease.key, &lease.instance_id).await;
}
return Err(e);
}
}
*watch = LeaseWatch::default();
return Ok(lease);
}
let holder = read(&store, &key)
.await?
.and_then(|(r, _)| r)
.map_or_else(|| "unknown".to_string(), |r| r.instance_id);
Err(LeaseError::ForeignActive {
instance_id: holder,
})
}
async fn probe(&self, record: &LeaseRecord) -> Result<(), LeaseError> {
let fake = UpdateVersion {
e_tag: Some(PROBE_ETAG.to_string()),
version: None,
};
for (mode, header) in [
(PutMode::Create, "If-None-Match"),
(PutMode::Update(fake), "If-Match"),
] {
#[cfg(test)]
if let Some(control) = &self.probe_control {
let index = control.probes.load(std::sync::atomic::Ordering::SeqCst);
if index == control.block_probe {
control.entered.notify_one();
control.resume.notified().await;
}
control
.probes
.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
}
let answer = put(&self.store, &self.key, record, mode).await;
match probe_verdict(answer, header)? {
Exclusion::Conditional => {}
Exclusion::BestEffort => self.lock().exclusion = Exclusion::BestEffort,
}
}
Ok(())
}
pub(crate) fn mark_lost(&self) {
self.lock().lost = true;
}
fn lock(&self) -> std::sync::MutexGuard<'_, Held> {
self.held.lock().unwrap_or_else(|e| e.into_inner())
}
pub fn instance_id(&self) -> &str {
&self.instance_id
}
pub fn ttl_s(&self) -> u64 {
self.ttl_s
}
pub fn renew_interval(&self) -> Duration {
Duration::from_millis(self.ttl_s * 1000 / 3)
}
pub fn exclusion(&self) -> Exclusion {
self.lock().exclusion
}
pub fn is_lost(&self) -> bool {
self.lock().lost
}
pub fn write_deadline(&self) -> Option<Instant> {
let h = self.lock();
let fence = Duration::from_millis(self.ttl_s * 2000 / 3);
let wall_ok = match SystemTime::now().duration_since(h.confirmed_wall) {
Ok(elapsed) => elapsed < fence,
Err(_) => true,
};
let deadline = h.confirmed + fence;
(!h.lost && wall_ok && Instant::now() < deadline).then_some(deadline)
}
pub fn may_write(&self) -> bool {
self.write_deadline().is_some()
}
pub async fn renew(&self) -> Result<(), LeaseError> {
let (exclusion, version) = {
let h = self.lock();
if h.lost {
return Err(LeaseError::Lost { instance_id: None });
}
(h.exclusion, h.version.clone())
};
let (confirmed, confirmed_wall) = (Instant::now(), SystemTime::now());
let record = LeaseRecord {
instance_id: self.instance_id.clone(),
acquired_at: self.acquired_at.clone(),
renewed_at: rfc3339(confirmed_wall),
ttl_s: self.ttl_s,
};
let written = match exclusion {
Exclusion::Conditional => {
match put(&self.store, &self.key, &record, PutMode::Update(version)).await {
Err(object_store::Error::NotImplemented) => {
self.lock().exclusion = Exclusion::BestEffort;
self.renew_best_effort(&record).await
}
Err(object_store::Error::Precondition { .. })
| Err(object_store::Error::NotFound { .. }) => self.reconcile(&record).await,
other => other.map_err(LeaseError::from),
}
}
Exclusion::BestEffort => self.renew_best_effort(&record).await,
}?;
let mut h = self.lock();
h.version = version_of(written);
h.confirmed = confirmed;
h.confirmed_wall = confirmed_wall;
#[cfg(test)]
if let Some(control) = &self.probe_control {
control
.renewals
.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
}
Ok(())
}
async fn reconcile(&self, record: &LeaseRecord) -> Result<PutResult, LeaseError> {
match read(&self.store, &self.key).await? {
Some((Some(r), v)) if r.instance_id == self.instance_id => {
match put(&self.store, &self.key, record, PutMode::Update(v)).await {
Ok(w) => Ok(w),
Err(object_store::Error::Precondition { .. })
| Err(object_store::Error::NotFound { .. }) => Err(self.lose().await),
Err(e) => Err(e.into()),
}
}
_ => Err(self.lose().await),
}
}
async fn renew_best_effort(&self, record: &LeaseRecord) -> Result<PutResult, LeaseError> {
match read(&self.store, &self.key).await? {
Some((Some(r), _)) if r.instance_id == self.instance_id => {
Ok(put(&self.store, &self.key, record, PutMode::Overwrite).await?)
}
_ => Err(self.lose().await),
}
}
async fn lose(&self) -> LeaseError {
self.lock().lost = true;
let holder = read(&self.store, &self.key)
.await
.ok()
.flatten()
.and_then(|(r, _)| r)
.map(|r| r.instance_id)
.filter(|id| *id != self.instance_id);
LeaseError::Lost {
instance_id: holder,
}
}
pub async fn release(self) -> Result<(), LeaseError> {
if self.is_lost() {
return Ok(());
}
release_own_record(&self.store, &self.key, &self.instance_id).await
}
}
pub(crate) async fn release_own_record(
store: &Arc<dyn ObjectStore>,
key: &ObjectPath,
instance_id: &str,
) -> Result<(), LeaseError> {
match read(store, key).await? {
Some((Some(record), _)) if record.instance_id == instance_id => {
match store.delete(key).await {
Ok(()) | Err(object_store::Error::NotFound { .. }) => Ok(()),
Err(error) => Err(error.into()),
}
}
_ => Ok(()),
}
}
async fn read(
store: &Arc<dyn ObjectStore>,
key: &ObjectPath,
) -> Result<Option<(Option<LeaseRecord>, UpdateVersion)>, LeaseError> {
let got = match store.get(key).await {
Ok(g) => g,
Err(object_store::Error::NotFound { .. }) => return Ok(None),
Err(e) => return Err(e.into()),
};
let version = UpdateVersion {
e_tag: got.meta.e_tag.clone(),
version: got.meta.version.clone(),
};
let bytes = got.bytes().await?;
Ok(Some((serde_json::from_slice(&bytes).ok(), version)))
}
async fn put(
store: &Arc<dyn ObjectStore>,
key: &ObjectPath,
record: &LeaseRecord,
mode: PutMode,
) -> Result<PutResult, object_store::Error> {
let body = serde_json::to_vec(record).expect("a lease record always serialises");
store
.put_opts(key, PutPayload::from(body), mode.into())
.await
}
fn probe_verdict(
answer: Result<PutResult, object_store::Error>,
header: &'static str,
) -> Result<Exclusion, LeaseError> {
match answer {
Err(object_store::Error::AlreadyExists { .. })
| Err(object_store::Error::Precondition { .. }) => Ok(Exclusion::Conditional),
Err(object_store::Error::NotImplemented) => Ok(Exclusion::BestEffort),
Ok(_) => Err(LeaseError::PreconditionsIgnored { header }),
Err(e) => Err(e.into()),
}
}
fn foreign_blocks(
rec: &LeaseRecord,
version: &UpdateVersion,
ttl_s: u64,
watch: &mut LeaseWatch,
) -> Result<(), LeaseError> {
let now = OffsetDateTime::now_utc();
let blocked = || LeaseError::ForeignActive {
instance_id: rec.instance_id.clone(),
};
if !rec.is_fresh(now) {
return Ok(());
}
if !rec.is_suspect(now) {
watch.seen = None;
return Err(blocked());
}
let tag = format!(
"{}|{}|{}",
version.e_tag.as_deref().unwrap_or(""),
version.version.as_deref().unwrap_or(""),
rec.renewed_at
);
match &watch.seen {
Some((seen, since)) if *seen == tag => {
if since.elapsed() >= Duration::from_secs(ttl_s) {
tracing::warn!(
holder = %rec.instance_id,
renewed_at = %rec.renewed_at,
ttl_s = rec.ttl_s,
"taking over a suspect bucket lease that has not changed for a TTL"
);
return Ok(());
}
}
_ => {
tracing::warn!(
holder = %rec.instance_id,
renewed_at = %rec.renewed_at,
ttl_s = rec.ttl_s,
"the bucket lease is stamped far in the future or has an absurd TTL; \
it is taken over if it stays unchanged for a TTL"
);
watch.seen = Some((tag, Instant::now()));
}
}
Err(blocked())
}
fn version_of(r: PutResult) -> UpdateVersion {
UpdateVersion {
e_tag: r.e_tag,
version: r.version,
}
}
fn rfc3339(t: SystemTime) -> String {
OffsetDateTime::from(t).format(&Rfc3339).unwrap_or_default()
}
#[cfg(test)]
mod tests {
use super::*;
fn record(renewed: OffsetDateTime, ttl_s: u64) -> LeaseRecord {
let s = renewed.format(&Rfc3339).unwrap();
LeaseRecord {
instance_id: "x".into(),
acquired_at: s.clone(),
renewed_at: s,
ttl_s,
}
}
#[test]
fn freshness_is_strictly_less_than_the_ttl() {
let now = OffsetDateTime::now_utc();
assert!(record(now - time::Duration::seconds(29), 30).is_fresh(now));
assert!(!record(now - time::Duration::seconds(30), 30).is_fresh(now));
assert!(record(now + time::Duration::seconds(5), 30).is_fresh(now));
let mut garbage = record(now, 30);
garbage.renewed_at = "yesterday".into();
assert!(!garbage.is_fresh(now));
assert!(record(now, u64::MAX).is_fresh(now));
}
#[test]
fn the_probe_fails_closed_when_a_precondition_is_ignored() {
let accepted = || {
Ok(PutResult {
e_tag: None,
version: None,
})
};
for header in ["If-None-Match", "If-Match"] {
match probe_verdict(accepted(), header) {
Err(e @ LeaseError::PreconditionsIgnored { .. }) => {
let shown = e.to_string();
assert!(shown.contains(header), "{shown}");
assert!(
shown.contains("MUSHROOMDB_BACKUP_LEASE_CONDITIONAL=off"),
"{shown}"
);
}
other => panic!("{header}: expected PreconditionsIgnored, got {other:?}"),
}
}
let refused = |e| probe_verdict(Err(e), "If-Match").unwrap();
assert_eq!(
refused(object_store::Error::Precondition {
path: "p".into(),
source: "no".into()
}),
Exclusion::Conditional
);
assert_eq!(
refused(object_store::Error::AlreadyExists {
path: "p".into(),
source: "no".into()
}),
Exclusion::Conditional
);
assert_eq!(
refused(object_store::Error::NotImplemented),
Exclusion::BestEffort
);
assert!(matches!(
probe_verdict(
Err(object_store::Error::Generic {
store: "S3",
source: "503".into()
}),
"If-Match"
),
Err(LeaseError::Store(_))
));
}
#[test]
fn far_future_stamps_and_absurd_ttls_are_suspect() {
let now = OffsetDateTime::now_utc();
assert!(!record(now, 30).is_suspect(now));
assert!(!record(now + time::Duration::seconds(20), 30).is_suspect(now));
assert!(record(now + time::Duration::seconds(31), 30).is_suspect(now));
assert!(record(now, MAX_SANE_TTL_S + 1).is_suspect(now));
assert!(!record(now, MAX_SANE_TTL_S).is_suspect(now));
}
#[tokio::test(start_paused = true)]
async fn a_suspect_lease_unchanged_for_a_ttl_is_taken_over() {
let store: Arc<dyn ObjectStore> = Arc::new(object_store::memory::InMemory::new());
let ahead = OffsetDateTime::now_utc() + time::Duration::days(1);
let mut rec = record(ahead, 30);
rec.instance_id = "clock-a-day-ahead".into();
store
.put(
&ObjectPath::from(lease_key("p")),
PutPayload::from(serde_json::to_vec(&rec).unwrap()),
)
.await
.unwrap();
let mut watch = LeaseWatch::default();
let me = new_instance_id();
for _ in 0..2 {
assert!(matches!(
BucketLease::acquire_as(Arc::clone(&store), "p", 30, &me, &mut watch).await,
Err(LeaseError::ForeignActive { .. })
));
tokio::time::advance(Duration::from_secs(10)).await;
}
tokio::time::advance(Duration::from_secs(11)).await;
let l = BucketLease::acquire_as(Arc::clone(&store), "p", 30, &me, &mut watch)
.await
.expect("unchanged for 31 s of a 30 s TTL");
assert_eq!(l.instance_id(), me);
}
#[tokio::test]
async fn a_lease_naming_our_id_is_taken_back_at_once() {
let store: Arc<dyn ObjectStore> = Arc::new(object_store::memory::InMemory::new());
let me = new_instance_id();
let mut watch = LeaseWatch::default();
let first = BucketLease::acquire_as(Arc::clone(&store), "p", 30, &me, &mut watch)
.await
.unwrap();
drop(first); let again = BucketLease::acquire_as(Arc::clone(&store), "p", 30, &me, &mut watch)
.await
.expect("our own fresh lease is ours");
again.renew().await.unwrap();
}
#[test]
fn the_key_sits_at_the_prefix_root() {
assert_eq!(lease_key("a/b"), "a/b/LEASE");
assert_eq!(lease_key(""), "LEASE");
}
}