use crate::oplog::WallClock;
use serde::{Deserialize, Serialize};
use serde_json::{json, Value};
use std::collections::BTreeMap;
use std::fmt;
use std::sync::{Arc, Mutex};
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct Lease {
pub agent_id: String,
pub holder: String,
pub epoch: u64,
pub expires_at_ms: u64,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum LeaseError {
Held {
agent_id: String,
holder: String,
epoch: u64,
expires_at_ms: u64,
},
Lost {
agent_id: String,
claimed_epoch: u64,
current_epoch: u64,
},
Backend(String),
}
impl fmt::Display for LeaseError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
LeaseError::Held {
agent_id,
holder,
epoch,
expires_at_ms,
} => write!(
f,
"lease for {agent_id} is held by {holder} at epoch {epoch} (expires at \
{expires_at_ms}ms) — acquire refused (unexpired)"
),
LeaseError::Lost {
agent_id,
claimed_epoch,
current_epoch,
} => write!(
f,
"lease for {agent_id} claimed at epoch {claimed_epoch} is lost (coordinator is at \
epoch {current_epoch}) — you are no longer the holder"
),
LeaseError::Backend(m) => write!(f, "lease coordinator backend error: {m}"),
}
}
}
impl std::error::Error for LeaseError {}
pub trait LeaseCoordinator {
fn acquire(
&mut self,
agent_id: &str,
device_id: &str,
ttl_ms: u64,
) -> Result<Lease, LeaseError>;
fn renew(
&mut self,
agent_id: &str,
device_id: &str,
epoch: u64,
ttl_ms: u64,
) -> Result<Lease, LeaseError>;
fn release(&mut self, agent_id: &str, device_id: &str, epoch: u64) -> Result<(), LeaseError>;
fn current(&mut self, agent_id: &str) -> Result<Option<Lease>, LeaseError>;
}
#[derive(Debug, Clone, Default)]
struct Slot {
epoch: u64,
holder: Option<String>,
expires_at_ms: u64,
}
#[derive(Clone)]
pub struct InMemoryLeaseCoordinator {
slots: Arc<Mutex<BTreeMap<String, Slot>>>,
wall: WallClock,
}
impl fmt::Debug for InMemoryLeaseCoordinator {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("InMemoryLeaseCoordinator")
.finish_non_exhaustive()
}
}
impl InMemoryLeaseCoordinator {
pub fn new(wall: WallClock) -> Self {
Self {
slots: Arc::new(Mutex::new(BTreeMap::new())),
wall,
}
}
fn lease_of(agent_id: &str, slot: &Slot) -> Option<Lease> {
slot.holder.as_ref().map(|holder| Lease {
agent_id: agent_id.to_string(),
holder: holder.clone(),
epoch: slot.epoch,
expires_at_ms: slot.expires_at_ms,
})
}
}
impl LeaseCoordinator for InMemoryLeaseCoordinator {
fn acquire(
&mut self,
agent_id: &str,
device_id: &str,
ttl_ms: u64,
) -> Result<Lease, LeaseError> {
let now = (self.wall)();
let mut slots = self.slots.lock().expect("lease mutex poisoned");
let slot = slots.entry(agent_id.to_string()).or_default();
let held = slot.holder.is_some() && now < slot.expires_at_ms;
if held {
return Err(LeaseError::Held {
agent_id: agent_id.to_string(),
holder: slot.holder.clone().expect("held ⇒ some holder"),
epoch: slot.epoch,
expires_at_ms: slot.expires_at_ms,
});
}
slot.epoch += 1; slot.holder = Some(device_id.to_string());
slot.expires_at_ms = now.saturating_add(ttl_ms);
Ok(Self::lease_of(agent_id, slot).expect("just set holder"))
}
fn renew(
&mut self,
agent_id: &str,
device_id: &str,
epoch: u64,
ttl_ms: u64,
) -> Result<Lease, LeaseError> {
let now = (self.wall)();
let mut slots = self.slots.lock().expect("lease mutex poisoned");
let slot = slots.entry(agent_id.to_string()).or_default();
let ours = slot.holder.as_deref() == Some(device_id) && slot.epoch == epoch;
if !ours {
return Err(LeaseError::Lost {
agent_id: agent_id.to_string(),
claimed_epoch: epoch,
current_epoch: slot.epoch,
});
}
slot.expires_at_ms = now.saturating_add(ttl_ms);
Ok(Self::lease_of(agent_id, slot).expect("ours ⇒ some holder"))
}
fn release(&mut self, agent_id: &str, device_id: &str, epoch: u64) -> Result<(), LeaseError> {
let mut slots = self.slots.lock().expect("lease mutex poisoned");
let slot = slots.entry(agent_id.to_string()).or_default();
let ours = slot.holder.as_deref() == Some(device_id) && slot.epoch == epoch;
if !ours {
return Err(LeaseError::Lost {
agent_id: agent_id.to_string(),
claimed_epoch: epoch,
current_epoch: slot.epoch,
});
}
slot.holder = None; Ok(())
}
fn current(&mut self, agent_id: &str) -> Result<Option<Lease>, LeaseError> {
let slots = self.slots.lock().expect("lease mutex poisoned");
Ok(slots
.get(agent_id)
.and_then(|slot| Self::lease_of(agent_id, slot)))
}
}
pub const INTENT_FIELD_ID: &str = "id";
pub const INTENT_FIELD_AGENT: &str = "agent_id";
pub const INTENT_FIELD_EPOCH: &str = "epoch";
pub const INTENT_FIELD_STATUS: &str = "status";
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum IntentStatus {
Pending,
Committed,
Failed,
}
impl IntentStatus {
pub fn rank(self) -> u8 {
match self {
IntentStatus::Pending => 0,
IntentStatus::Committed | IntentStatus::Failed => 1,
}
}
fn as_str(self) -> &'static str {
match self {
IntentStatus::Pending => "pending",
IntentStatus::Committed => "committed",
IntentStatus::Failed => "failed",
}
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Intent {
pub agent_id: String,
pub run_id: String,
pub epoch: u64,
pub status: IntentStatus,
}
impl Intent {
pub fn new(
agent_id: impl Into<String>,
run_id: impl Into<String>,
epoch: u64,
status: IntentStatus,
) -> Self {
Self {
agent_id: agent_id.into(),
run_id: run_id.into(),
epoch,
status,
}
}
pub fn payload(&self) -> Value {
json!({
INTENT_FIELD_ID: self.run_id,
INTENT_FIELD_AGENT: self.agent_id,
INTENT_FIELD_EPOCH: self.epoch,
INTENT_FIELD_STATUS: self.status.as_str(),
})
}
pub fn from_payload(payload: &Value) -> Option<Self> {
let status = match payload.get(INTENT_FIELD_STATUS).and_then(Value::as_str)? {
"pending" => IntentStatus::Pending,
"committed" => IntentStatus::Committed,
"failed" => IntentStatus::Failed,
_ => return None,
};
Some(Self {
agent_id: payload
.get(INTENT_FIELD_AGENT)
.and_then(Value::as_str)?
.to_string(),
run_id: payload
.get(INTENT_FIELD_ID)
.and_then(Value::as_str)?
.to_string(),
epoch: payload.get(INTENT_FIELD_EPOCH).and_then(Value::as_u64)?,
status,
})
}
}
pub fn intent_agent(payload: &Value) -> &str {
payload
.get(INTENT_FIELD_AGENT)
.and_then(Value::as_str)
.unwrap_or("")
}
pub fn intent_epoch(payload: &Value) -> u64 {
payload
.get(INTENT_FIELD_EPOCH)
.and_then(Value::as_u64)
.unwrap_or(0)
}
pub fn intent_status_rank(payload: &Value) -> u8 {
match payload.get(INTENT_FIELD_STATUS).and_then(Value::as_str) {
Some("committed") | Some("failed") => 1,
_ => 0,
}
}
pub fn intent_is_committed(payload: &Value) -> bool {
payload.get(INTENT_FIELD_STATUS).and_then(Value::as_str) == Some("committed")
}
#[cfg(test)]
mod tests {
use super::*;
use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
use std::thread;
fn manual_clock() -> (Arc<AtomicU64>, WallClock) {
let t = Arc::new(AtomicU64::new(0));
let reader = t.clone();
(t, Arc::new(move || reader.load(Ordering::SeqCst)))
}
#[test]
fn acquire_is_exclusive_and_bumps_epoch_never_reusing() {
let (_t, wall) = manual_clock(); let mut coord = InMemoryLeaseCoordinator::new(wall);
let a = coord.acquire("agent-x", "dev-a", 100).unwrap();
assert_eq!((a.epoch, a.holder.as_str()), (1, "dev-a"));
let err = coord.acquire("agent-x", "dev-b", 100).unwrap_err();
assert!(
matches!(err, LeaseError::Held { epoch: 1, .. }),
"got {err:?}"
);
coord.release("agent-x", "dev-a", 1).unwrap();
assert_eq!(coord.current("agent-x").unwrap(), None, "released ⇒ unheld");
let b = coord.acquire("agent-x", "dev-b", 100).unwrap();
assert_eq!(
(b.epoch, b.holder.as_str()),
(2, "dev-b"),
"epoch is monotone, never reused"
);
}
#[test]
fn renew_and_release_require_the_current_holder_and_epoch() {
let (t, wall) = manual_clock();
let mut coord = InMemoryLeaseCoordinator::new(wall);
let lease = coord.acquire("a", "dev-a", 100).unwrap();
assert_eq!(lease.expires_at_ms, 100);
t.store(50, Ordering::SeqCst);
let renewed = coord.renew("a", "dev-a", 1, 100).unwrap();
assert_eq!((renewed.epoch, renewed.expires_at_ms), (1, 150));
assert!(matches!(
coord.renew("a", "dev-a", 99, 100),
Err(LeaseError::Lost {
current_epoch: 1,
..
})
));
assert!(matches!(
coord.renew("a", "dev-b", 1, 100),
Err(LeaseError::Lost { .. })
));
assert!(matches!(
coord.release("a", "dev-b", 1),
Err(LeaseError::Lost { .. })
));
coord.release("a", "dev-a", 1).unwrap();
assert!(matches!(
coord.renew("a", "dev-a", 1, 100),
Err(LeaseError::Lost { .. })
));
}
#[test]
fn renew_works_past_expiry_until_actually_stolen() {
let (t, wall) = manual_clock();
let mut coord = InMemoryLeaseCoordinator::new(wall);
coord.acquire("a", "dev-a", 100).unwrap();
t.store(500, Ordering::SeqCst);
assert!(
coord.renew("a", "dev-a", 1, 100).is_ok(),
"renew works until stolen"
);
}
#[test]
fn expired_lease_is_stealable_and_fences_the_old_epoch() {
let (t, wall) = manual_clock();
let mut coord = InMemoryLeaseCoordinator::new(wall);
let old = coord.acquire("a", "dev-a", 100).unwrap();
assert_eq!(old.epoch, 1);
t.store(200, Ordering::SeqCst);
let stolen = coord.acquire("a", "dev-b", 100).unwrap();
assert_eq!((stolen.epoch, stolen.holder.as_str()), (2, "dev-b"));
assert!(matches!(
coord.renew("a", "dev-a", 1, 100),
Err(LeaseError::Lost {
claimed_epoch: 1,
current_epoch: 2,
..
})
));
let cur = coord.current("a").unwrap().unwrap();
assert_eq!((cur.epoch, cur.holder.as_str()), (2, "dev-b"));
}
#[test]
fn reference_impl_is_linearizable_under_concurrent_contention() {
let (_t, wall) = manual_clock(); let coord = InMemoryLeaseCoordinator::new(wall);
let concurrent = Arc::new(AtomicUsize::new(0));
let max_concurrent = Arc::new(AtomicUsize::new(0));
let granted = Arc::new(Mutex::new(Vec::<u64>::new()));
let threads: Vec<_> = (0..8)
.map(|i| {
let mut coord = coord.clone();
let concurrent = concurrent.clone();
let max_concurrent = max_concurrent.clone();
let granted = granted.clone();
let device = format!("dev-{i}");
thread::spawn(move || {
for _ in 0..50 {
let lease = loop {
match coord.acquire("agent", &device, 1_000_000) {
Ok(l) => break l,
Err(LeaseError::Held { .. }) => thread::yield_now(),
Err(e) => panic!("unexpected {e:?}"),
}
};
let now = concurrent.fetch_add(1, Ordering::SeqCst) + 1;
max_concurrent.fetch_max(now, Ordering::SeqCst);
granted.lock().unwrap().push(lease.epoch);
for _ in 0..20 {
std::hint::spin_loop();
}
concurrent.fetch_sub(1, Ordering::SeqCst);
coord.release("agent", &device, lease.epoch).unwrap();
}
})
})
.collect();
for th in threads {
th.join().unwrap();
}
assert_eq!(
max_concurrent.load(Ordering::SeqCst),
1,
"the CAS is linearizable: never two holders at once"
);
let mut epochs = granted.lock().unwrap().clone();
let count = epochs.len();
epochs.sort_unstable();
epochs.dedup();
assert_eq!(epochs.len(), count, "no epoch was ever reused");
assert_eq!(
epochs,
(1..=count as u64).collect::<Vec<_>>(),
"epochs are dense and monotone"
);
}
#[test]
fn intent_payload_round_trips_and_ranks() {
let intent = Intent::new("milo", "run-abc", 3, IntentStatus::Committed);
let payload = intent.payload();
assert_eq!(
payload[INTENT_FIELD_ID],
json!("run-abc"),
"run_id lives under id"
);
assert_eq!(Intent::from_payload(&payload), Some(intent));
assert_eq!(intent_agent(&payload), "milo");
assert_eq!(intent_epoch(&payload), 3);
assert_eq!(intent_status_rank(&payload), 1);
assert_eq!(
intent_status_rank(&Intent::new("m", "r", 1, IntentStatus::Pending).payload()),
0
);
assert!(IntentStatus::Pending.rank() < IntentStatus::Committed.rank());
assert_eq!(IntentStatus::Committed.rank(), IntentStatus::Failed.rank());
assert_eq!(intent_epoch(&json!({})), 0);
assert_eq!(intent_agent(&json!({})), "");
assert_eq!(Intent::from_payload(&json!({"id": "r"})), None);
}
}