use std::path::PathBuf;
use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};
use crate::api::{RedDBError, RedDBResult};
use crate::replication::lease::{LeaseError, LeaseStore, WriterLease};
use crate::runtime::audit_log::{
AuditAuthSource, AuditEvent, AuditFieldEscaper, AuditLogger, Outcome,
};
use crate::runtime::write_gate::{LeaseGateState, WriteGate};
pub type MarkDraining = Arc<dyn Fn() + Send + Sync>;
pub type CurrentTerm = Arc<dyn Fn() -> u64 + Send + Sync>;
pub struct LeaseLifecycle {
store: Arc<LeaseStore>,
write_gate: Arc<WriteGate>,
audit_log: Arc<AuditLogger>,
mark_draining: MarkDraining,
current_term: CurrentTerm,
holder_id: String,
database_key: String,
ttl_ms: u64,
current: Mutex<Option<WriterLease>>,
transition: Mutex<()>,
snapshot_publication:
Option<Arc<crate::runtime::serverless_snapshot_publication::SnapshotPublicationStore>>,
initial_snapshot_path: Option<PathBuf>,
}
impl LeaseLifecycle {
pub fn new(
store: Arc<LeaseStore>,
write_gate: Arc<WriteGate>,
audit_log: Arc<AuditLogger>,
mark_draining: MarkDraining,
current_term: CurrentTerm,
holder_id: String,
database_key: String,
ttl_ms: u64,
) -> Self {
write_gate.set_lease_state(LeaseGateState::NotHeld);
Self {
store,
write_gate,
audit_log,
mark_draining,
current_term,
holder_id,
database_key,
ttl_ms,
current: Mutex::new(None),
transition: Mutex::new(()),
snapshot_publication: None,
initial_snapshot_path: None,
}
}
pub fn with_snapshot_publication(
mut self,
publication: Arc<crate::runtime::serverless_snapshot_publication::SnapshotPublicationStore>,
initial_snapshot_path: impl Into<PathBuf>,
) -> Self {
self.snapshot_publication = Some(publication);
self.initial_snapshot_path = Some(initial_snapshot_path.into());
self
}
pub fn holder_id(&self) -> &str {
&self.holder_id
}
pub fn database_key(&self) -> &str {
&self.database_key
}
pub fn ttl_ms(&self) -> u64 {
self.ttl_ms
}
pub(crate) fn refresh_delay(&self) -> Duration {
self.write_gate.lease_remaining() / 3
}
pub fn current_lease(&self) -> Option<WriterLease> {
self.current.lock().expect("poisoned lease mutex").clone()
}
pub fn try_acquire(&self) -> RedDBResult<()> {
let _transition = self.transition.lock().expect("poisoned lease transition");
self.write_gate.set_lease_state(LeaseGateState::NotHeld);
let deadline = self.request_deadline()?;
let term = (self.current_term)();
match self.store.try_acquire_for_term(
&self.database_key,
&self.holder_id,
self.ttl_ms,
term,
) {
Ok(lease) => {
if let Some(publication) = &self.snapshot_publication {
let local_snapshot =
self.initial_snapshot_path.as_deref().ok_or_else(|| {
RedDBError::Internal(
"snapshot publication configured without a local snapshot path"
.to_string(),
)
})?;
if let Err(err) = publication.claim(&lease, local_snapshot) {
let _ = self.store.release(&lease);
self.audit_log.record_event(
AuditEvent::builder("lease/acquire")
.principal(self.holder_id.clone())
.source(AuditAuthSource::System)
.resource(self.database_key.clone())
.outcome(Outcome::Error)
.field(AuditFieldEscaper::field("error", err.to_string()))
.build(),
);
return Err(RedDBError::ReadOnly(format!(
"writer lease acquired but snapshot head claim failed: {err}"
)));
}
}
if Instant::now() >= deadline {
self.on_refresh_lost("acquire response arrived after the local lease deadline");
return Err(RedDBError::ReadOnly(
"writer lease expired during acquire".into(),
));
}
*self.current.lock().expect("poisoned lease mutex") = Some(lease.clone());
self.write_gate.grant_lease_until(deadline);
self.audit_log.record_event(
AuditEvent::builder("lease/acquire")
.principal(self.holder_id.clone())
.source(AuditAuthSource::System)
.resource(self.database_key.clone())
.outcome(Outcome::Success)
.field(AuditFieldEscaper::field(
"generation",
lease.generation as i64,
))
.field(AuditFieldEscaper::field("term", lease.term as i64))
.field(AuditFieldEscaper::field("ttl_ms", self.ttl_ms))
.build(),
);
Ok(())
}
Err(err) => {
self.audit_log.record_event(
AuditEvent::builder("lease/acquire")
.principal(self.holder_id.clone())
.source(AuditAuthSource::System)
.resource(self.database_key.clone())
.outcome(Outcome::Error)
.field(AuditFieldEscaper::field("error", err.to_string()))
.build(),
);
Err(RedDBError::Internal(format!("acquire writer lease: {err}")))
}
}
}
pub fn publish_snapshot(
&self,
prepare_snapshot: impl FnOnce() -> RedDBResult<reddb_file::ServerlessWriterLeaseTempFile>,
) -> RedDBResult<bool> {
let Some(publication) = &self.snapshot_publication else {
return Ok(false);
};
let _transition = self.transition.lock().expect("poisoned lease transition");
if self.write_gate.lease_state() != LeaseGateState::Held {
self.on_refresh_lost("local writer lease expired before snapshot publication");
return Err(RedDBError::ReadOnly(
"writer lease expired before snapshot publication".to_string(),
));
}
let lease = self
.current
.lock()
.expect("poisoned lease mutex")
.clone()
.ok_or_else(|| {
RedDBError::ReadOnly(
"snapshot publication requires an acquired writer lease".to_string(),
)
})?;
let local_snapshot = prepare_snapshot()?;
if let Err(err) = publication.publish(&lease, local_snapshot.path()) {
self.on_refresh_lost(&err);
return Err(RedDBError::ReadOnly(format!(
"writer snapshot publication lost its fence: {err}"
)));
}
Ok(true)
}
pub fn refresh(&self) -> RedDBResult<()> {
let _transition = self.transition.lock().expect("poisoned lease transition");
let deadline = self.request_deadline()?;
let snapshot = match self.current.lock().expect("poisoned lease mutex").clone() {
Some(lease) => lease,
None => {
return Err(RedDBError::Internal(
"refresh called without an acquired lease".to_string(),
));
}
};
if self.write_gate.lease_state() != LeaseGateState::Held {
self.on_refresh_lost("local writer lease deadline expired");
return Err(RedDBError::ReadOnly(
"writer lease expired before refresh".into(),
));
}
let term = (self.current_term)();
match self.store.refresh_for_term(&snapshot, self.ttl_ms, term) {
Ok(updated) => {
if Instant::now() >= deadline {
self.on_refresh_lost("refresh response arrived after the local lease deadline");
return Err(RedDBError::ReadOnly(
"writer lease expired during refresh".into(),
));
}
*self.current.lock().expect("poisoned lease mutex") = Some(updated);
self.write_gate.grant_lease_until(deadline);
Ok(())
}
Err(err) => {
self.on_refresh_lost(err);
Err(RedDBError::Internal("writer lease lost".to_string()))
}
}
}
pub fn release(&self) -> RedDBResult<()> {
let _transition = self.transition.lock().expect("poisoned lease transition");
self.write_gate.set_lease_state(LeaseGateState::NotHeld);
let snapshot = match self.current.lock().expect("poisoned lease mutex").take() {
Some(lease) => lease,
None => {
self.write_gate.set_lease_state(LeaseGateState::NotHeld);
return Ok(());
}
};
let result = self.store.release(&snapshot);
self.write_gate.set_lease_state(LeaseGateState::NotHeld);
match result {
Ok(()) => {
self.audit_log.record_event(
AuditEvent::builder("lease/release")
.principal(self.holder_id.clone())
.source(AuditAuthSource::System)
.resource(self.database_key.clone())
.outcome(Outcome::Success)
.build(),
);
Ok(())
}
Err(err) => {
self.audit_log.record_event(
AuditEvent::builder("lease/release")
.principal(self.holder_id.clone())
.source(AuditAuthSource::System)
.resource(self.database_key.clone())
.outcome(Outcome::Error)
.field(AuditFieldEscaper::field("error", err.to_string()))
.build(),
);
tracing::warn!(
target: "reddb::serverless::lease",
error = %err,
"lease release on shutdown failed"
);
Ok(())
}
}
}
fn request_deadline(&self) -> RedDBResult<Instant> {
let ttl = self
.ttl_ms
.checked_sub(1)
.filter(|ttl| *ttl > 0)
.ok_or_else(|| {
RedDBError::InvalidOperation("writer lease TTL must exceed 1 ms".into())
})?;
Instant::now()
.checked_add(Duration::from_millis(ttl))
.ok_or_else(|| RedDBError::InvalidOperation("writer lease TTL is too large".into()))
}
fn on_refresh_lost(&self, err: impl std::fmt::Display) {
tracing::error!(
target: "reddb::serverless::lease",
error = %err,
holder = %self.holder_id,
database_key = %self.database_key,
"lease refresh failed; flipping to NotHeld + drain"
);
*self.current.lock().expect("poisoned lease mutex") = None;
self.write_gate.set_lease_state(LeaseGateState::NotHeld);
self.audit_log.record_event(
AuditEvent::builder("lease/lost")
.principal(self.holder_id.clone())
.source(AuditAuthSource::System)
.resource(self.database_key.clone())
.outcome(Outcome::Error)
.field(AuditFieldEscaper::field("error", err.to_string()))
.build(),
);
(self.mark_draining)();
}
}
pub fn admin_promote_lease(
store: &LeaseStore,
audit_log: &AuditLogger,
database_key: &str,
holder_id: &str,
ttl_ms: u64,
term: u64,
) -> Result<WriterLease, LeaseError> {
match store.try_acquire_for_term(database_key, holder_id, ttl_ms, term) {
Ok(lease) => {
audit_log.record_event(
AuditEvent::builder("admin/failover/promote")
.principal(lease.holder_id.clone())
.source(AuditAuthSource::System)
.resource(database_key.to_string())
.outcome(Outcome::Success)
.field(AuditFieldEscaper::field(
"holder_id",
lease.holder_id.clone(),
))
.field(AuditFieldEscaper::field(
"generation",
lease.generation as i64,
))
.field(AuditFieldEscaper::field("term", lease.term as i64))
.field(AuditFieldEscaper::field("ttl_ms", ttl_ms))
.build(),
);
Ok(lease)
}
Err(err) => {
audit_log.record_event(
AuditEvent::builder("admin/failover/promote")
.principal(holder_id.to_string())
.source(AuditAuthSource::System)
.resource(database_key.to_string())
.outcome(Outcome::Error)
.field(AuditFieldEscaper::field("error", err.to_string()))
.build(),
);
Err(err)
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::api::RedDBOptions;
use crate::storage::backend::LocalBackend;
use std::path::PathBuf;
use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
fn temp_prefix(tag: &str) -> PathBuf {
let mut p = PathBuf::from(std::env::temp_dir());
p.push(format!(
"reddb-lease-lifecycle-{tag}-{}-{}",
std::process::id(),
crate::utils::now_unix_nanos(),
));
std::fs::create_dir_all(&p).unwrap();
p
}
fn build_lifecycle(
tag: &str,
) -> (
Arc<LeaseLifecycle>,
Arc<WriteGate>,
Arc<AuditLogger>,
Arc<AtomicUsize>,
PathBuf,
) {
let (lifecycle, write_gate, audit_log, drain_counter, _current_term, prefix) =
build_lifecycle_with_term(tag, 1);
(lifecycle, write_gate, audit_log, drain_counter, prefix)
}
fn build_lifecycle_with_term(
tag: &str,
initial_term: u64,
) -> (
Arc<LeaseLifecycle>,
Arc<WriteGate>,
Arc<AuditLogger>,
Arc<AtomicUsize>,
Arc<AtomicU64>,
PathBuf,
) {
let prefix = temp_prefix(tag);
let store = Arc::new(
LeaseStore::new(Arc::new(LocalBackend))
.with_prefix(prefix.to_string_lossy().to_string()),
);
let mut opts = RedDBOptions::default();
opts.read_only = false;
let write_gate = Arc::new(WriteGate::from_options(&opts));
let audit_log = Arc::new(AuditLogger::for_data_path(&prefix.join("data.rdb")));
let drain_counter = Arc::new(AtomicUsize::new(0));
let drain_counter_clone = Arc::clone(&drain_counter);
let mark_draining: MarkDraining = Arc::new(move || {
drain_counter_clone.fetch_add(1, Ordering::SeqCst);
});
let current_term_value = Arc::new(AtomicU64::new(initial_term));
let current_term_clone = Arc::clone(¤t_term_value);
let current_term: CurrentTerm = Arc::new(move || current_term_clone.load(Ordering::SeqCst));
let lifecycle = Arc::new(LeaseLifecycle::new(
store,
Arc::clone(&write_gate),
Arc::clone(&audit_log),
mark_draining,
current_term,
"writer-1".to_string(),
"main".to_string(),
60_000,
));
(
lifecycle,
write_gate,
audit_log,
drain_counter,
current_term_value,
prefix,
)
}
#[test]
fn acquire_flips_gate_to_held_and_records_audit() {
let (lifecycle, gate, audit, drain, prefix) = build_lifecycle("acquire");
assert!(lifecycle.try_acquire().is_ok());
assert_eq!(gate.lease_state(), LeaseGateState::Held);
assert!(lifecycle.current_lease().is_some());
assert_eq!(drain.load(Ordering::SeqCst), 0);
assert!(audit.wait_idle(std::time::Duration::from_secs(2)));
let body = std::fs::read_to_string(audit.path()).unwrap();
assert!(body.contains("lease/acquire"));
assert!(body.contains("\"outcome\":\"success\""));
let _ = std::fs::remove_dir_all(&prefix);
}
#[test]
fn successor_acquire_cannot_leave_expired_predecessor_writable() {
let (mut lifecycle, gate, _audit, drain, prefix) = build_lifecycle("expiry");
Arc::get_mut(&mut lifecycle)
.expect("invariant: test owns the only lifecycle reference")
.ttl_ms = 1_000;
lifecycle
.try_acquire()
.expect("first writer acquires lease");
let first = lifecycle
.current_lease()
.expect("acquire publishes the current lease");
let remaining = first
.expires_at_ms
.saturating_sub(crate::utils::now_unix_millis());
std::thread::sleep(Duration::from_millis(remaining + 20));
let successor = lifecycle
.store
.try_acquire_for_term("main", "writer-2", 60_000, 1)
.expect("successor acquires expired lease");
assert!(successor.generation > first.generation);
assert!(gate
.check(crate::runtime::write_gate::WriteKind::Dml)
.is_err());
assert!(lifecycle.refresh().is_err());
assert_eq!(drain.load(Ordering::SeqCst), 1);
assert!(lifecycle.current_lease().is_none());
lifecycle
.store
.release(&successor)
.expect("successor lease releases cleanly");
let _ = std::fs::remove_dir_all(prefix);
}
#[test]
fn refresh_cannot_reopen_a_released_gate() {
let (lifecycle, gate, _audit, _drain, prefix) = build_lifecycle("release-race");
lifecycle.try_acquire().expect("writer acquires lease");
let worker = lifecycle.clone();
let refresh = std::thread::spawn(move || worker.refresh());
lifecycle.release().expect("lease release succeeds");
let _ = refresh.join().expect("refresh worker does not panic");
assert_eq!(gate.lease_state(), LeaseGateState::NotHeld);
assert!(lifecycle.current_lease().is_none());
let _ = std::fs::remove_dir_all(prefix);
}
#[test]
fn renewal_installs_a_new_local_deadline() {
let (lifecycle, gate, _audit, _drain, prefix) = build_lifecycle("renewal");
lifecycle.try_acquire().expect("writer acquires lease");
gate.grant_lease_until(Instant::now() + Duration::from_millis(100));
lifecycle.refresh().expect("lease refresh succeeds");
std::thread::sleep(Duration::from_millis(120));
assert_eq!(gate.lease_state(), LeaseGateState::Held);
lifecycle.release().expect("lease release succeeds");
let _ = std::fs::remove_dir_all(prefix);
}
#[test]
fn minimum_configured_ttl_schedules_refresh_before_expiry() {
let (mut lifecycle, gate, _audit, _drain, prefix) = build_lifecycle("short-ttl");
Arc::get_mut(&mut lifecycle)
.expect("invariant: test owns the only lifecycle reference")
.ttl_ms = 1_000;
lifecycle.try_acquire().expect("writer acquires lease");
let delay = lifecycle.refresh_delay();
assert!(delay < Duration::from_millis(334));
assert!(delay < gate.lease_remaining());
gate.grant_lease_until(Instant::now() + Duration::from_millis(90));
assert!(lifecycle.refresh_delay() <= Duration::from_millis(30));
lifecycle.release().expect("lease release succeeds");
let _ = std::fs::remove_dir_all(prefix);
}
#[test]
fn local_deadline_accounts_for_backend_timestamp_quantization() {
let (mut lifecycle, _gate, _audit, _drain, prefix) = build_lifecycle("quantization");
Arc::get_mut(&mut lifecycle)
.expect("invariant: test owns the only lifecycle reference")
.ttl_ms = 1;
assert!(lifecycle.request_deadline().is_err());
Arc::get_mut(&mut lifecycle)
.expect("invariant: test owns the only lifecycle reference")
.ttl_ms = 1_000;
let deadline = lifecycle
.request_deadline()
.expect("one-second TTL has a conservative deadline");
assert!(deadline.saturating_duration_since(Instant::now()) <= Duration::from_millis(999));
let _ = std::fs::remove_dir_all(prefix);
}
#[test]
fn delayed_acquire_response_never_opens_admission() {
let (mut lifecycle, gate, _audit, _drain, prefix) = build_lifecycle("delayed-acquire");
let lifecycle_mut = Arc::get_mut(&mut lifecycle)
.expect("invariant: test owns the only lifecycle reference");
lifecycle_mut.ttl_ms = 2;
lifecycle_mut.current_term = Arc::new(|| {
std::thread::sleep(Duration::from_millis(5));
1
});
assert!(lifecycle.try_acquire().is_err());
assert_eq!(gate.lease_state(), LeaseGateState::NotHeld);
assert!(lifecycle.current_lease().is_none());
let _ = std::fs::remove_dir_all(prefix);
}
#[test]
fn acquire_stamps_current_replication_term() {
let (lifecycle, _gate, _audit, _drain, _term, prefix) =
build_lifecycle_with_term("acquire-term", 7);
lifecycle.try_acquire().unwrap();
assert_eq!(lifecycle.current_lease().unwrap().term, 7);
let _ = std::fs::remove_dir_all(&prefix);
}
#[test]
fn release_flips_gate_to_not_held_and_clears_inner_state() {
let (lifecycle, gate, audit, _drain, prefix) = build_lifecycle("release");
lifecycle.try_acquire().unwrap();
assert!(lifecycle.release().is_ok());
assert_eq!(gate.lease_state(), LeaseGateState::NotHeld);
assert!(lifecycle.current_lease().is_none());
assert!(audit.wait_idle(std::time::Duration::from_secs(2)));
let body = std::fs::read_to_string(audit.path()).unwrap();
assert!(body.contains("lease/release"));
let _ = std::fs::remove_dir_all(&prefix);
}
#[test]
fn release_is_idempotent_when_no_lease_held() {
let (lifecycle, gate, _audit, _drain, prefix) = build_lifecycle("release-idem");
assert!(lifecycle.release().is_ok());
assert_eq!(gate.lease_state(), LeaseGateState::NotHeld);
let _ = std::fs::remove_dir_all(&prefix);
}
#[test]
fn refresh_without_acquire_keeps_gate_closed() {
let (lifecycle, gate, _audit, drain, prefix) = build_lifecycle("refresh-noop");
let err = lifecycle.refresh().unwrap_err();
match err {
RedDBError::Internal(msg) => assert!(msg.contains("without an acquired lease")),
other => panic!("expected Internal, got {other:?}"),
}
assert_eq!(gate.lease_state(), LeaseGateState::NotHeld);
assert_eq!(drain.load(Ordering::SeqCst), 0);
let _ = std::fs::remove_dir_all(&prefix);
}
#[test]
fn refresh_fences_when_current_term_advances() {
let (lifecycle, gate, _audit, drain, current_term, prefix) =
build_lifecycle_with_term("refresh-stale-term", 4);
lifecycle.try_acquire().unwrap();
current_term.store(5, Ordering::SeqCst);
let err = lifecycle.refresh().unwrap_err();
match err {
RedDBError::Internal(msg) => assert_eq!(msg, "writer lease lost"),
other => panic!("expected Internal, got {other:?}"),
}
assert_eq!(gate.lease_state(), LeaseGateState::NotHeld);
assert_eq!(drain.load(Ordering::SeqCst), 1);
assert!(lifecycle.current_lease().is_none());
let _ = std::fs::remove_dir_all(&prefix);
}
#[test]
fn admin_promote_lease_audits_success() {
let prefix = temp_prefix("admin-ok");
let store = LeaseStore::new(Arc::new(LocalBackend))
.with_prefix(prefix.to_string_lossy().to_string());
let audit = AuditLogger::for_data_path(&prefix.join("data.rdb"));
let lease = admin_promote_lease(&store, &audit, "main", "promoter-1", 30_000, 9).unwrap();
assert_eq!(lease.holder_id, "promoter-1");
assert_eq!(lease.term, 9);
assert!(audit.wait_idle(std::time::Duration::from_secs(2)));
let body = std::fs::read_to_string(audit.path()).unwrap();
assert!(body.contains("admin/failover/promote"));
assert!(body.contains("\"outcome\":\"success\""));
let _ = std::fs::remove_dir_all(&prefix);
}
#[test]
fn admin_promote_lease_does_not_flip_a_separate_gate() {
let prefix = temp_prefix("admin-no-gate");
let store = LeaseStore::new(Arc::new(LocalBackend))
.with_prefix(prefix.to_string_lossy().to_string());
let audit = AuditLogger::for_data_path(&prefix.join("data.rdb"));
let mut opts = RedDBOptions::default();
opts.read_only = false;
let gate = WriteGate::from_options(&opts);
let _ = admin_promote_lease(&store, &audit, "main", "promoter-2", 30_000, 1).unwrap();
assert_eq!(gate.lease_state(), LeaseGateState::NotRequired);
let _ = std::fs::remove_dir_all(&prefix);
}
}