1use std::collections::BTreeMap;
10use std::path::{Path, PathBuf};
11use std::sync::{mpsc, Arc};
12use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};
13
14use notify::{RecursiveMode, Watcher};
15use rusqlite::{
16 params, Connection, ErrorCode, OptionalExtension, Transaction, TransactionBehavior,
17};
18use serde::{Deserialize, Serialize};
19use sysinfo::System;
20use uuid::Uuid;
21
22pub const HOST_LEASE_ROOT_ENV: &str = "HARN_HOST_LEASE_ROOT";
24const HARN_HOME_ENV: &str = "HARN_HOME";
25const LEASE_DB_FILE: &str = "host-leases.sqlite";
26const SQLITE_MUTATION_BUSY_TIMEOUT: Duration = Duration::from_secs(1);
27const REGISTRY_BUSY_RETRY_INTERVAL: Duration = Duration::from_millis(250);
28const PROCESS_LIVENESS_RECHECK_INTERVAL: Duration = Duration::from_secs(5);
29const SCHEMA_VERSION: u32 = 1;
30
31#[derive(Debug, thiserror::Error)]
33pub enum HostLeaseError {
34 #[error("invalid host lease request: {0}")]
36 InvalidRequest(String),
37 #[error("host lease state I/O failed: {0}")]
39 Io(#[from] std::io::Error),
40 #[error("host lease database failed: {0}")]
42 Database(#[from] rusqlite::Error),
43 #[error("host lease metadata serialization failed: {0}")]
45 Serialization(#[from] serde_json::Error),
46 #[error("host lease watcher failed: {0}")]
48 Watch(String),
49 #[error("system clock is before the Unix epoch")]
51 Clock,
52}
53
54#[derive(Clone, Copy, Debug, PartialEq, Eq)]
55enum ProcessObservation {
56 Alive { identity: u64 },
57 Dead,
58 Unknown,
59}
60
61trait ProcessInspector: std::fmt::Debug + Send + Sync {
62 fn observe(&self, pid: u32) -> ProcessObservation;
63}
64
65#[derive(Debug)]
66struct SystemProcessInspector;
67
68impl ProcessInspector for SystemProcessInspector {
69 fn observe(&self, pid: u32) -> ProcessObservation {
70 match crate::process_liveness::process_liveness(pid) {
71 crate::process_liveness::ProcessLiveness::Dead => ProcessObservation::Dead,
72 crate::process_liveness::ProcessLiveness::Unknown => ProcessObservation::Unknown,
73 crate::process_liveness::ProcessLiveness::Alive => {
74 crate::process_liveness::process_identity(pid)
75 .map_or(ProcessObservation::Unknown, |identity| {
76 ProcessObservation::Alive { identity }
77 })
78 }
79 }
80 }
81}
82
83#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
85#[serde(rename_all = "kebab-case")]
86pub enum HostLeasePriorityClass {
87 Interactive,
89 Measurement,
91 CiVerify,
93 #[default]
94 Deferrable,
96}
97
98impl HostLeasePriorityClass {
99 pub const fn as_str(self) -> &'static str {
101 match self {
102 Self::Interactive => "interactive",
103 Self::Measurement => "measurement",
104 Self::CiVerify => "ci-verify",
105 Self::Deferrable => "deferrable",
106 }
107 }
108
109 fn parse(raw: &str) -> Result<Self, HostLeaseError> {
110 match raw {
111 "interactive" => Ok(Self::Interactive),
112 "measurement" => Ok(Self::Measurement),
113 "ci-verify" => Ok(Self::CiVerify),
114 "deferrable" => Ok(Self::Deferrable),
115 other => Err(HostLeaseError::InvalidRequest(format!(
116 "unknown priority class `{other}`"
117 ))),
118 }
119 }
120}
121
122#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
124pub struct HostLeaseRequest {
125 pub host: String,
127 pub owner: String,
129 #[serde(default)]
130 pub priority_class: HostLeasePriorityClass,
132 #[serde(default)]
136 pub ttl_ms: Option<u64>,
137 #[serde(default)]
138 pub owner_pid: Option<u32>,
140 #[serde(default)]
141 pub reason: Option<String>,
143 #[serde(default)]
144 pub metadata: BTreeMap<String, String>,
146}
147
148#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
150pub struct HostLeaseHandle {
151 pub schema_version: u32,
153 pub host: String,
155 pub lease_id: String,
157 pub owner: String,
159 pub priority_class: HostLeasePriorityClass,
161 pub acquired_at_ms: i64,
163 pub updated_at_ms: i64,
165 #[serde(default)]
166 pub expires_at_ms: Option<i64>,
168 #[serde(default)]
169 pub owner_pid: Option<u32>,
171 #[serde(default)]
172 pub owner_process_identity: Option<u64>,
174 #[serde(default)]
175 pub reason: Option<String>,
177 #[serde(default)]
178 pub metadata: BTreeMap<String, String>,
180}
181
182#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
184#[serde(rename_all = "snake_case")]
185pub enum HostLeaseAcquireStatus {
186 Acquired,
188 Deferred,
190}
191
192#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
194pub enum HostLeaseDeferReason {
195 #[serde(rename = "host_lease_contended")]
197 Contended,
198 #[serde(rename = "host_lease_registry_busy")]
200 RegistryBusy,
201}
202
203impl HostLeaseDeferReason {
204 pub const fn as_str(self) -> &'static str {
206 match self {
207 Self::Contended => "host_lease_contended",
208 Self::RegistryBusy => "host_lease_registry_busy",
209 }
210 }
211}
212
213#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
215pub struct HostLeaseDeferReceipt {
216 pub host: String,
218 pub deferred_reason: HostLeaseDeferReason,
220 pub observed_at_ms: i64,
222 #[serde(default)]
223 pub next_wake_at_ms: Option<i64>,
225 #[serde(default)]
226 pub deadline_at_ms: Option<i64>,
228 #[serde(default)]
231 pub active: Option<HostLeaseHandle>,
232}
233
234#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
236pub struct HostLeaseAcquireReceipt {
237 pub schema_version: u32,
239 pub status: HostLeaseAcquireStatus,
241 pub observed_at_ms: i64,
243 pub waited_ms: u64,
245 #[serde(default)]
246 pub handle: Option<HostLeaseHandle>,
248 #[serde(default)]
249 pub defer: Option<HostLeaseDeferReceipt>,
251 pub recovered_stale_lease: bool,
253}
254
255#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
257pub struct HostLeaseState {
258 pub schema_version: u32,
260 pub host: String,
262 pub observed_at_ms: i64,
264 #[serde(default)]
265 pub active: Option<HostLeaseHandle>,
267 pub recovered_stale_lease: bool,
269}
270
271#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
273pub struct HostLeaseRenewReceipt {
274 pub schema_version: u32,
276 pub renewed: bool,
278 pub observed_at_ms: i64,
280 #[serde(default)]
281 pub handle: Option<HostLeaseHandle>,
283}
284
285#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
287pub struct HostLeaseReleaseReceipt {
288 pub schema_version: u32,
290 pub released: bool,
292 pub host: String,
294 pub lease_id: String,
296 pub observed_at_ms: i64,
298}
299
300#[derive(Clone, Debug)]
302pub struct HostLeaseStore {
303 root: PathBuf,
304 db_path: PathBuf,
305 process_inspector: Arc<dyn ProcessInspector>,
306}
307
308impl HostLeaseStore {
309 pub fn from_env() -> Result<Self, HostLeaseError> {
311 let root = if let Some(path) = std::env::var_os(HOST_LEASE_ROOT_ENV) {
312 PathBuf::from(path)
313 } else if let Some(path) = std::env::var_os(HARN_HOME_ENV) {
314 PathBuf::from(path).join("host-leases")
315 } else {
316 harn_vm::user_dirs::home_dir()
317 .ok_or_else(|| {
318 HostLeaseError::InvalidRequest(
319 "cannot resolve a home directory; set HARN_HOST_LEASE_ROOT".to_string(),
320 )
321 })?
322 .join(".harn/host-leases")
323 };
324 Self::for_root(root)
325 }
326
327 pub fn for_root(root: impl Into<PathBuf>) -> Result<Self, HostLeaseError> {
329 Self::for_root_with_inspector(root, Arc::new(SystemProcessInspector))
330 }
331
332 fn for_root_with_inspector(
333 root: impl Into<PathBuf>,
334 process_inspector: Arc<dyn ProcessInspector>,
335 ) -> Result<Self, HostLeaseError> {
336 let root = root.into();
337 std::fs::create_dir_all(&root)?;
338 let store = Self {
339 db_path: root.join(LEASE_DB_FILE),
340 root,
341 process_inspector,
342 };
343 store.initialize()?;
344 Ok(store)
345 }
346
347 pub fn root(&self) -> &Path {
349 &self.root
350 }
351
352 pub fn default_host() -> String {
354 System::host_name()
355 .filter(|name| !name.trim().is_empty())
356 .unwrap_or_else(|| "local".to_string())
357 }
358
359 pub fn try_acquire(
361 &self,
362 request: HostLeaseRequest,
363 ) -> Result<HostLeaseAcquireReceipt, HostLeaseError> {
364 self.try_acquire_once(request, None, None)
365 }
366
367 pub fn acquire_wait(
369 &self,
370 request: HostLeaseRequest,
371 wait_timeout: Duration,
372 ) -> Result<HostLeaseAcquireReceipt, HostLeaseError> {
373 if wait_timeout.is_zero() {
374 return self.try_acquire(request);
375 }
376 let started_at_ms = unix_now_ms()?;
377 let started_at = Instant::now();
378 let deadline = started_at.checked_add(wait_timeout).ok_or_else(|| {
379 HostLeaseError::InvalidRequest("wait timeout exceeds the monotonic clock".to_string())
380 })?;
381 let deadline_at_ms = started_at_ms.saturating_add(duration_ms_i64(wait_timeout));
382 let (tx, rx) = mpsc::channel();
383 let mut watcher = notify::recommended_watcher(move |event| {
384 let _ = tx.send(event);
385 })
386 .map_err(|error| HostLeaseError::Watch(error.to_string()))?;
387 watcher
388 .watch(&self.root, RecursiveMode::NonRecursive)
389 .map_err(|error| HostLeaseError::Watch(error.to_string()))?;
390
391 loop {
392 let receipt =
393 self.try_acquire_once(request.clone(), Some(started_at), Some(deadline_at_ms))?;
394 if receipt.status == HostLeaseAcquireStatus::Acquired || Instant::now() >= deadline {
395 return Ok(receipt);
396 }
397 let wake_at = receipt
398 .defer
399 .as_ref()
400 .and_then(|defer| defer.next_wake_at_ms)
401 .map(|wake| wake.min(deadline_at_ms))
402 .unwrap_or(deadline_at_ms);
403 let wake_duration =
404 Duration::from_millis(wake_at.saturating_sub(receipt.observed_at_ms).max(1) as u64);
405 let remaining = deadline.saturating_duration_since(Instant::now());
406 if remaining.is_zero() {
407 return Ok(receipt);
408 }
409 match rx.recv_timeout(wake_duration.min(remaining)) {
410 Ok(Ok(_)) | Err(mpsc::RecvTimeoutError::Timeout) => {}
411 Ok(Err(error)) => return Err(HostLeaseError::Watch(error.to_string())),
412 Err(mpsc::RecvTimeoutError::Disconnected) => {
413 return Err(HostLeaseError::Watch(
414 "host lease watcher disconnected".to_string(),
415 ));
416 }
417 }
418 }
419 }
420
421 pub fn status(&self, host: &str) -> Result<HostLeaseState, HostLeaseError> {
423 let host = normalize_component("host", host)?;
424 let mut conn = self.connection(SQLITE_MUTATION_BUSY_TIMEOUT)?;
425 let tx = conn.transaction_with_behavior(TransactionBehavior::Immediate)?;
426 let now = unix_now_ms()?;
427 self.status_in_transaction(tx, &host, now)
428 }
429
430 pub fn renew(
432 &self,
433 host: &str,
434 lease_id: &str,
435 ttl_ms: u64,
436 ) -> Result<HostLeaseRenewReceipt, HostLeaseError> {
437 let host = normalize_component("host", host)?;
438 let lease_id = normalize_component("lease_id", lease_id)?;
439 validate_ttl(Some(ttl_ms))?;
440 let mut conn = self.connection(SQLITE_MUTATION_BUSY_TIMEOUT)?;
441 let tx = conn.transaction_with_behavior(TransactionBehavior::Immediate)?;
442 let now = unix_now_ms()?;
443 let (active, _) = active_handle(&tx, &host, now, self.process_inspector.as_ref())?;
444 let Some(mut handle) = active.filter(|handle| handle.lease_id == lease_id) else {
445 tx.commit()?;
446 return Ok(HostLeaseRenewReceipt {
447 schema_version: SCHEMA_VERSION,
448 renewed: false,
449 observed_at_ms: now,
450 handle: None,
451 });
452 };
453 handle.updated_at_ms = now;
454 handle.expires_at_ms = Some(now.saturating_add(u64_ms_i64(ttl_ms)));
455 write_handle(&tx, &handle)?;
456 tx.commit()?;
457 Ok(HostLeaseRenewReceipt {
458 schema_version: SCHEMA_VERSION,
459 renewed: true,
460 observed_at_ms: now,
461 handle: Some(handle),
462 })
463 }
464
465 pub fn release(
467 &self,
468 host: &str,
469 lease_id: &str,
470 ) -> Result<HostLeaseReleaseReceipt, HostLeaseError> {
471 let host = normalize_component("host", host)?;
472 let lease_id = normalize_component("lease_id", lease_id)?;
473 let conn = self.connection(SQLITE_MUTATION_BUSY_TIMEOUT)?;
474 let released = conn.execute(
475 "DELETE FROM host_leases WHERE host = ?1 AND lease_id = ?2",
476 params![host, lease_id],
477 )? == 1;
478 let now = unix_now_ms()?;
479 Ok(HostLeaseReleaseReceipt {
480 schema_version: SCHEMA_VERSION,
481 released,
482 host,
483 lease_id,
484 observed_at_ms: now,
485 })
486 }
487
488 fn initialize(&self) -> Result<(), HostLeaseError> {
489 let conn = self.connection(SQLITE_MUTATION_BUSY_TIMEOUT)?;
490 conn.execute_batch(
491 "CREATE TABLE IF NOT EXISTS host_leases (
492 host TEXT PRIMARY KEY NOT NULL,
493 lease_id TEXT NOT NULL,
494 owner TEXT NOT NULL,
495 priority_class TEXT NOT NULL,
496 acquired_at_ms INTEGER NOT NULL,
497 updated_at_ms INTEGER NOT NULL,
498 expires_at_ms INTEGER,
499 owner_pid INTEGER,
500 owner_process_identity INTEGER,
501 reason TEXT,
502 metadata_json TEXT NOT NULL
503 );",
504 )?;
505 Ok(())
506 }
507
508 fn connection(&self, busy_timeout: Duration) -> Result<Connection, HostLeaseError> {
509 let conn = Connection::open(&self.db_path)?;
510 conn.busy_timeout(busy_timeout)?;
511 Ok(conn)
512 }
513
514 fn try_acquire_once(
515 &self,
516 request: HostLeaseRequest,
517 started_at: Option<Instant>,
518 deadline_at_ms: Option<i64>,
519 ) -> Result<HostLeaseAcquireReceipt, HostLeaseError> {
520 let request = normalize_request(request)?;
521 let mut conn = self.connection(Duration::ZERO)?;
522 let tx = match conn.transaction_with_behavior(TransactionBehavior::Immediate) {
523 Ok(tx) => tx,
524 Err(error) if sqlite_is_busy(&error) => {
525 let now = unix_now_ms()?;
526 return Ok(registry_busy_receipt(
527 request.host,
528 now,
529 started_at.map(|started| duration_ms_u64(started.elapsed())),
530 deadline_at_ms,
531 ));
532 }
533 Err(error) => return Err(error.into()),
534 };
535 let now = unix_now_ms()?;
536 let waited_ms = started_at
537 .map(|started| duration_ms_u64(started.elapsed()))
538 .unwrap_or(0);
539 let host = request.host.clone();
540 match self.acquire_in_transaction(tx, request, now, deadline_at_ms, waited_ms) {
541 Err(HostLeaseError::Database(error)) if sqlite_is_busy(&error) => Ok(
542 registry_busy_receipt(host, now, Some(waited_ms), deadline_at_ms),
543 ),
544 result => result,
545 }
546 }
547
548 #[cfg(test)]
549 fn try_acquire_at(
550 &self,
551 request: HostLeaseRequest,
552 now: i64,
553 deadline_at_ms: Option<i64>,
554 waited_ms: u64,
555 ) -> Result<HostLeaseAcquireReceipt, HostLeaseError> {
556 let request = normalize_request(request)?;
557 let mut conn = self.connection(SQLITE_MUTATION_BUSY_TIMEOUT)?;
558 let tx = conn.transaction_with_behavior(TransactionBehavior::Immediate)?;
559 self.acquire_in_transaction(tx, request, now, deadline_at_ms, waited_ms)
560 }
561
562 fn acquire_in_transaction(
563 &self,
564 tx: Transaction<'_>,
565 request: HostLeaseRequest,
566 now: i64,
567 deadline_at_ms: Option<i64>,
568 waited_ms: u64,
569 ) -> Result<HostLeaseAcquireReceipt, HostLeaseError> {
570 let owner_process_identity = request
571 .owner_pid
572 .map(|pid| match self.process_inspector.observe(pid) {
573 ProcessObservation::Alive { identity } => Ok(identity),
574 ProcessObservation::Dead => Err(HostLeaseError::InvalidRequest(
575 "owner_pid is not a live local process".to_string(),
576 )),
577 ProcessObservation::Unknown => Err(HostLeaseError::InvalidRequest(
578 "owner_pid liveness could not be verified".to_string(),
579 )),
580 })
581 .transpose()?;
582 let (active, recovered_stale_lease) =
583 active_handle(&tx, &request.host, now, self.process_inspector.as_ref())?;
584 if let Some(active) = active {
585 let defer = HostLeaseDeferReceipt {
586 host: request.host,
587 deferred_reason: HostLeaseDeferReason::Contended,
588 observed_at_ms: now,
589 next_wake_at_ms: Some(next_lease_wake_at(&active, now, deadline_at_ms)),
590 deadline_at_ms,
591 active: Some(active),
592 };
593 tx.commit()?;
594 return Ok(HostLeaseAcquireReceipt {
595 schema_version: SCHEMA_VERSION,
596 status: HostLeaseAcquireStatus::Deferred,
597 observed_at_ms: now,
598 waited_ms,
599 handle: None,
600 defer: Some(defer),
601 recovered_stale_lease,
602 });
603 }
604
605 let handle = HostLeaseHandle {
606 schema_version: SCHEMA_VERSION,
607 host: request.host,
608 lease_id: Uuid::now_v7().to_string(),
609 owner: request.owner,
610 priority_class: request.priority_class,
611 acquired_at_ms: now,
612 updated_at_ms: now,
613 expires_at_ms: request
614 .ttl_ms
615 .map(|ttl| now.saturating_add(u64_ms_i64(ttl))),
616 owner_pid: request.owner_pid,
617 owner_process_identity,
618 reason: request.reason,
619 metadata: request.metadata,
620 };
621 write_handle(&tx, &handle)?;
622 tx.commit()?;
623 Ok(HostLeaseAcquireReceipt {
624 schema_version: SCHEMA_VERSION,
625 status: HostLeaseAcquireStatus::Acquired,
626 observed_at_ms: now,
627 waited_ms,
628 handle: Some(handle),
629 defer: None,
630 recovered_stale_lease,
631 })
632 }
633
634 #[cfg(test)]
635 fn status_at(&self, host: &str, now: i64) -> Result<HostLeaseState, HostLeaseError> {
636 let mut conn = self.connection(SQLITE_MUTATION_BUSY_TIMEOUT)?;
637 let tx = conn.transaction_with_behavior(TransactionBehavior::Immediate)?;
638 self.status_in_transaction(tx, host, now)
639 }
640
641 fn status_in_transaction(
642 &self,
643 tx: Transaction<'_>,
644 host: &str,
645 now: i64,
646 ) -> Result<HostLeaseState, HostLeaseError> {
647 let (active, recovered_stale_lease) =
648 active_handle(&tx, host, now, self.process_inspector.as_ref())?;
649 tx.commit()?;
650 Ok(HostLeaseState {
651 schema_version: SCHEMA_VERSION,
652 host: host.to_string(),
653 observed_at_ms: now,
654 active,
655 recovered_stale_lease,
656 })
657 }
658}
659
660fn normalize_request(mut request: HostLeaseRequest) -> Result<HostLeaseRequest, HostLeaseError> {
661 request.host = normalize_component("host", &request.host)?;
662 request.owner = normalize_component("owner", &request.owner)?;
663 validate_ttl(request.ttl_ms)?;
664 if request.ttl_ms.is_none() && request.owner_pid.is_none() {
665 return Err(HostLeaseError::InvalidRequest(
666 "a non-expiring lease requires owner_pid for crash recovery".to_string(),
667 ));
668 }
669 request.reason = request.reason.and_then(|reason| {
670 let trimmed = reason.trim();
671 (!trimmed.is_empty()).then(|| trimmed.to_string())
672 });
673 Ok(request)
674}
675
676fn normalize_component(name: &str, value: &str) -> Result<String, HostLeaseError> {
677 let value = value.trim();
678 if value.is_empty() {
679 return Err(HostLeaseError::InvalidRequest(format!(
680 "{name} cannot be empty"
681 )));
682 }
683 if !value
684 .chars()
685 .all(|ch| ch.is_ascii_alphanumeric() || matches!(ch, '.' | '_' | '-'))
686 {
687 return Err(HostLeaseError::InvalidRequest(format!(
688 "{name} may contain only ASCII letters, digits, '.', '_' and '-'"
689 )));
690 }
691 Ok(value.to_string())
692}
693
694fn validate_ttl(ttl_ms: Option<u64>) -> Result<(), HostLeaseError> {
695 if ttl_ms == Some(0) {
696 return Err(HostLeaseError::InvalidRequest(
697 "ttl_ms must be greater than zero when supplied".to_string(),
698 ));
699 }
700 Ok(())
701}
702
703fn active_handle(
704 tx: &Transaction<'_>,
705 host: &str,
706 now: i64,
707 process_inspector: &dyn ProcessInspector,
708) -> Result<(Option<HostLeaseHandle>, bool), HostLeaseError> {
709 let handle = read_handle(tx, host)?;
710 let Some(handle) = handle else {
711 return Ok((None, false));
712 };
713 let expired = handle.expires_at_ms.is_some_and(|expiry| expiry <= now);
714 let owner_dead = match (handle.owner_pid, handle.owner_process_identity) {
715 (Some(pid), Some(expected_identity)) => match process_inspector.observe(pid) {
716 ProcessObservation::Alive { identity } => identity != expected_identity,
717 ProcessObservation::Dead => true,
718 ProcessObservation::Unknown => false,
719 },
720 _ => false,
721 };
722 if expired || owner_dead {
723 tx.execute(
724 "DELETE FROM host_leases WHERE host = ?1 AND lease_id = ?2",
725 params![host, handle.lease_id],
726 )?;
727 return Ok((None, true));
728 }
729 Ok((Some(handle), false))
730}
731
732fn next_lease_wake_at(active: &HostLeaseHandle, now: i64, deadline_at_ms: Option<i64>) -> i64 {
733 let mut wake_at = deadline_at_ms.unwrap_or(i64::MAX);
734 if let Some(expiry) = active.expires_at_ms {
735 wake_at = wake_at.min(expiry);
736 }
737 if active.owner_pid.is_some() {
738 wake_at =
739 wake_at.min(now.saturating_add(duration_ms_i64(PROCESS_LIVENESS_RECHECK_INTERVAL)));
740 }
741 wake_at
742}
743
744fn registry_busy_receipt(
745 host: String,
746 now: i64,
747 waited_ms: Option<u64>,
748 deadline_at_ms: Option<i64>,
749) -> HostLeaseAcquireReceipt {
750 let next_wake_at_ms = now
751 .saturating_add(duration_ms_i64(REGISTRY_BUSY_RETRY_INTERVAL))
752 .min(deadline_at_ms.unwrap_or(i64::MAX));
753 HostLeaseAcquireReceipt {
754 schema_version: SCHEMA_VERSION,
755 status: HostLeaseAcquireStatus::Deferred,
756 observed_at_ms: now,
757 waited_ms: waited_ms.unwrap_or(0),
758 handle: None,
759 defer: Some(HostLeaseDeferReceipt {
760 host,
761 deferred_reason: HostLeaseDeferReason::RegistryBusy,
762 observed_at_ms: now,
763 next_wake_at_ms: Some(next_wake_at_ms),
764 deadline_at_ms,
765 active: None,
766 }),
767 recovered_stale_lease: false,
768 }
769}
770
771fn sqlite_is_busy(error: &rusqlite::Error) -> bool {
772 matches!(
773 error,
774 rusqlite::Error::SqliteFailure(inner, _)
775 if matches!(inner.code, ErrorCode::DatabaseBusy | ErrorCode::DatabaseLocked)
776 )
777}
778
779fn read_handle(
780 tx: &Transaction<'_>,
781 host: &str,
782) -> Result<Option<HostLeaseHandle>, HostLeaseError> {
783 tx.query_row(
784 "SELECT lease_id, owner, priority_class, acquired_at_ms, updated_at_ms,
785 expires_at_ms, owner_pid, owner_process_identity, reason, metadata_json
786 FROM host_leases WHERE host = ?1",
787 [host],
788 |row| {
789 let priority: String = row.get(2)?;
790 let metadata_json: String = row.get(9)?;
791 let owner_pid_i64: Option<i64> = row.get(6)?;
792 let owner_identity_i64: Option<i64> = row.get(7)?;
793 Ok((
794 row.get::<_, String>(0)?,
795 row.get::<_, String>(1)?,
796 priority,
797 row.get::<_, i64>(3)?,
798 row.get::<_, i64>(4)?,
799 row.get::<_, Option<i64>>(5)?,
800 owner_pid_i64,
801 owner_identity_i64,
802 row.get::<_, Option<String>>(8)?,
803 metadata_json,
804 ))
805 },
806 )
807 .optional()?
808 .map(
809 |(
810 lease_id,
811 owner,
812 priority,
813 acquired_at_ms,
814 updated_at_ms,
815 expires_at_ms,
816 owner_pid,
817 owner_process_identity,
818 reason,
819 metadata_json,
820 )| {
821 let owner_pid = owner_pid
822 .map(|pid| {
823 u32::try_from(pid).map_err(|_| {
824 HostLeaseError::InvalidRequest(
825 "persisted owner_pid is outside the u32 range".to_string(),
826 )
827 })
828 })
829 .transpose()?;
830 let owner_process_identity = owner_process_identity
831 .map(|identity| {
832 u64::try_from(identity).map_err(|_| {
833 HostLeaseError::InvalidRequest(
834 "persisted process identity is negative".to_string(),
835 )
836 })
837 })
838 .transpose()?;
839 Ok(HostLeaseHandle {
840 schema_version: SCHEMA_VERSION,
841 host: host.to_string(),
842 lease_id,
843 owner,
844 priority_class: HostLeasePriorityClass::parse(&priority)?,
845 acquired_at_ms,
846 updated_at_ms,
847 expires_at_ms,
848 owner_pid,
849 owner_process_identity,
850 reason,
851 metadata: serde_json::from_str(&metadata_json)?,
852 })
853 },
854 )
855 .transpose()
856}
857
858fn write_handle(tx: &Transaction<'_>, handle: &HostLeaseHandle) -> Result<(), HostLeaseError> {
859 let metadata_json = serde_json::to_string(&handle.metadata)?;
860 let owner_process_identity = handle
861 .owner_process_identity
862 .map(|value| {
863 i64::try_from(value).map_err(|_| {
864 HostLeaseError::InvalidRequest(
865 "owner process identity is outside the SQLite integer range".to_string(),
866 )
867 })
868 })
869 .transpose()?;
870 tx.execute(
871 "INSERT INTO host_leases (
872 host, lease_id, owner, priority_class, acquired_at_ms, updated_at_ms,
873 expires_at_ms, owner_pid, owner_process_identity, reason, metadata_json
874 ) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11)
875 ON CONFLICT(host) DO UPDATE SET
876 lease_id = excluded.lease_id,
877 owner = excluded.owner,
878 priority_class = excluded.priority_class,
879 acquired_at_ms = excluded.acquired_at_ms,
880 updated_at_ms = excluded.updated_at_ms,
881 expires_at_ms = excluded.expires_at_ms,
882 owner_pid = excluded.owner_pid,
883 owner_process_identity = excluded.owner_process_identity,
884 reason = excluded.reason,
885 metadata_json = excluded.metadata_json",
886 params![
887 handle.host,
888 handle.lease_id,
889 handle.owner,
890 handle.priority_class.as_str(),
891 handle.acquired_at_ms,
892 handle.updated_at_ms,
893 handle.expires_at_ms,
894 handle.owner_pid.map(i64::from),
895 owner_process_identity,
896 handle.reason,
897 metadata_json,
898 ],
899 )?;
900 Ok(())
901}
902
903fn unix_now_ms() -> Result<i64, HostLeaseError> {
904 let millis = SystemTime::now()
905 .duration_since(UNIX_EPOCH)
906 .map_err(|_| HostLeaseError::Clock)?
907 .as_millis();
908 Ok(millis.min(i64::MAX as u128) as i64)
909}
910
911fn duration_ms_i64(duration: Duration) -> i64 {
912 duration.as_millis().min(i64::MAX as u128) as i64
913}
914
915fn duration_ms_u64(duration: Duration) -> u64 {
916 duration.as_millis().min(u64::MAX as u128) as u64
917}
918
919fn u64_ms_i64(value: u64) -> i64 {
920 value.min(i64::MAX as u64) as i64
921}
922
923#[cfg(test)]
924mod tests {
925 use std::sync::atomic::{AtomicU64, Ordering};
926 use std::sync::{Arc, Barrier};
927 use std::thread;
928
929 use tempfile::TempDir;
930
931 use super::*;
932
933 fn store(temp: &TempDir) -> HostLeaseStore {
934 HostLeaseStore::for_root(temp.path()).unwrap()
935 }
936
937 #[derive(Debug)]
938 struct ScriptedProcessInspector {
939 observation: AtomicU64,
940 }
941
942 impl ScriptedProcessInspector {
943 fn alive(identity: u64) -> Self {
944 Self {
945 observation: AtomicU64::new(identity.saturating_add(2)),
946 }
947 }
948
949 fn set(&self, observation: ProcessObservation) {
950 let value = match observation {
951 ProcessObservation::Unknown => 0,
952 ProcessObservation::Dead => 1,
953 ProcessObservation::Alive { identity } => identity.saturating_add(2),
954 };
955 self.observation.store(value, Ordering::SeqCst);
956 }
957 }
958
959 impl ProcessInspector for ScriptedProcessInspector {
960 fn observe(&self, _pid: u32) -> ProcessObservation {
961 match self.observation.load(Ordering::SeqCst) {
962 0 => ProcessObservation::Unknown,
963 1 => ProcessObservation::Dead,
964 value => ProcessObservation::Alive {
965 identity: value - 2,
966 },
967 }
968 }
969 }
970
971 fn request(owner: &str) -> HostLeaseRequest {
972 HostLeaseRequest {
973 host: "mac-local".to_string(),
974 owner: owner.to_string(),
975 priority_class: HostLeasePriorityClass::Measurement,
976 ttl_ms: Some(60_000),
977 owner_pid: None,
978 reason: Some("meter run".to_string()),
979 metadata: BTreeMap::new(),
980 }
981 }
982
983 #[test]
984 fn acquire_blocks_second_owner_until_release() {
985 let temp = TempDir::new().unwrap();
986 let store = store(&temp);
987 let first = store
988 .try_acquire_at(request("codex-0"), 1_000, None, 0)
989 .unwrap();
990 assert_eq!(first.status, HostLeaseAcquireStatus::Acquired);
991 let second = store
992 .try_acquire_at(request("codex-1"), 1_001, None, 0)
993 .unwrap();
994 assert_eq!(second.status, HostLeaseAcquireStatus::Deferred);
995 assert_eq!(
996 second.defer.as_ref().unwrap().deferred_reason,
997 HostLeaseDeferReason::Contended
998 );
999
1000 let handle = first.handle.unwrap();
1001 assert!(
1002 store
1003 .release(&handle.host, &handle.lease_id)
1004 .unwrap()
1005 .released
1006 );
1007 let third = store
1008 .try_acquire_at(request("codex-1"), 1_002, None, 0)
1009 .unwrap();
1010 assert_eq!(third.status, HostLeaseAcquireStatus::Acquired);
1011 }
1012
1013 #[test]
1014 fn immediate_transaction_allows_exactly_one_race_winner() {
1015 let temp = TempDir::new().unwrap();
1016 let store = Arc::new(store(&temp));
1017 let worker_count = 12;
1018 let barrier = Arc::new(Barrier::new(worker_count));
1019 let workers = (0..worker_count)
1020 .map(|index| {
1021 let store = Arc::clone(&store);
1022 let barrier = Arc::clone(&barrier);
1023 thread::spawn(move || {
1024 barrier.wait();
1025 store
1026 .try_acquire_at(request(&format!("worker-{index}")), 1_000, None, 0)
1027 .unwrap()
1028 .status
1029 })
1030 })
1031 .collect::<Vec<_>>();
1032 let statuses = workers
1033 .into_iter()
1034 .map(|worker| worker.join().unwrap())
1035 .collect::<Vec<_>>();
1036 assert_eq!(
1037 statuses
1038 .iter()
1039 .filter(|status| **status == HostLeaseAcquireStatus::Acquired)
1040 .count(),
1041 1
1042 );
1043 }
1044
1045 #[test]
1046 fn expiry_is_recovered_inside_the_acquire_transaction() {
1047 let temp = TempDir::new().unwrap();
1048 let store = store(&temp);
1049 let mut short = request("codex-0");
1050 short.ttl_ms = Some(10);
1051 store.try_acquire_at(short, 1_000, None, 0).unwrap();
1052 let next = store
1053 .try_acquire_at(request("codex-1"), 1_011, None, 0)
1054 .unwrap();
1055 assert_eq!(next.status, HostLeaseAcquireStatus::Acquired);
1056 assert!(next.recovered_stale_lease);
1057 }
1058
1059 #[test]
1060 fn wrong_token_cannot_release_an_active_lease() {
1061 let temp = TempDir::new().unwrap();
1062 let store = store(&temp);
1063 let first = store
1064 .try_acquire_at(request("codex-0"), 1_000, None, 0)
1065 .unwrap();
1066 let handle = first.handle.unwrap();
1067 assert!(!store.release(&handle.host, "wrong-token").unwrap().released);
1068 assert_eq!(
1069 store.status_at(&handle.host, 1_001).unwrap().active,
1070 Some(handle)
1071 );
1072 }
1073
1074 #[test]
1075 fn renew_requires_the_active_token() {
1076 let temp = TempDir::new().unwrap();
1077 let store = store(&temp);
1078 let first = store.try_acquire(request("codex-0")).unwrap();
1079 let handle = first.handle.unwrap();
1080
1081 assert!(
1082 !store
1083 .renew(&handle.host, "wrong-token", 120_000)
1084 .unwrap()
1085 .renewed
1086 );
1087 let renewed = store
1088 .renew(&handle.host, &handle.lease_id, 120_000)
1089 .unwrap();
1090 assert!(renewed.renewed);
1091 let renewed_handle = renewed.handle.unwrap();
1092 assert_eq!(renewed_handle.lease_id, handle.lease_id);
1093 assert!(renewed_handle.updated_at_ms >= handle.updated_at_ms);
1094 assert!(renewed_handle.expires_at_ms > Some(renewed_handle.updated_at_ms));
1095 }
1096
1097 #[test]
1098 fn non_expiring_lease_requires_a_live_owner_pid() {
1099 let temp = TempDir::new().unwrap();
1100 let store = store(&temp);
1101 let mut missing = request("codex-0");
1102 missing.ttl_ms = None;
1103 let error = store.try_acquire_at(missing, 1_000, None, 0).unwrap_err();
1104 assert!(error.to_string().contains("requires owner_pid"));
1105
1106 let mut live = request("codex-0");
1107 live.ttl_ms = None;
1108 live.owner_pid = Some(std::process::id());
1109 let acquired = store.try_acquire_at(live, 1_000, None, 0).unwrap();
1110 assert_eq!(acquired.status, HostLeaseAcquireStatus::Acquired);
1111 assert_eq!(acquired.handle.unwrap().expires_at_ms, None);
1112 }
1113
1114 #[test]
1115 fn unknown_process_liveness_preserves_the_active_lease() {
1116 let temp = TempDir::new().unwrap();
1117 let inspector = Arc::new(ScriptedProcessInspector::alive(42));
1118 let store = HostLeaseStore::for_root_with_inspector(
1119 temp.path(),
1120 Arc::clone(&inspector) as Arc<dyn ProcessInspector>,
1121 )
1122 .unwrap();
1123 let mut owner = request("codex-0");
1124 owner.ttl_ms = None;
1125 owner.owner_pid = Some(1234);
1126 store.try_acquire_at(owner, 1_000, None, 0).unwrap();
1127
1128 inspector.set(ProcessObservation::Unknown);
1129 let deferred = store
1130 .try_acquire_at(request("codex-1"), 1_001, None, 0)
1131 .unwrap();
1132 assert_eq!(deferred.status, HostLeaseAcquireStatus::Deferred);
1133 assert!(!deferred.recovered_stale_lease);
1134
1135 inspector.set(ProcessObservation::Dead);
1136 let recovered = store
1137 .try_acquire_at(request("codex-1"), 1_002, None, 0)
1138 .unwrap();
1139 assert_eq!(recovered.status, HostLeaseAcquireStatus::Acquired);
1140 assert!(recovered.recovered_stale_lease);
1141 }
1142
1143 #[test]
1144 fn non_expiring_owner_gets_a_bounded_liveness_wake() {
1145 let mut active = request("codex-0");
1146 active.ttl_ms = None;
1147 active.owner_pid = Some(1234);
1148 let handle = HostLeaseHandle {
1149 schema_version: SCHEMA_VERSION,
1150 host: active.host,
1151 lease_id: "lease-1".to_string(),
1152 owner: active.owner,
1153 priority_class: active.priority_class,
1154 acquired_at_ms: 1_000,
1155 updated_at_ms: 1_000,
1156 expires_at_ms: None,
1157 owner_pid: active.owner_pid,
1158 owner_process_identity: Some(42),
1159 reason: active.reason,
1160 metadata: active.metadata,
1161 };
1162 assert_eq!(next_lease_wake_at(&handle, 1_000, Some(60_000)), 6_000);
1163 }
1164
1165 #[test]
1166 fn registry_write_contention_returns_a_typed_defer_receipt() {
1167 let temp = TempDir::new().unwrap();
1168 let store = store(&temp);
1169 let mut conn = store.connection(SQLITE_MUTATION_BUSY_TIMEOUT).unwrap();
1170 let _tx = conn
1171 .transaction_with_behavior(TransactionBehavior::Immediate)
1172 .unwrap();
1173
1174 let receipt = store.try_acquire(request("codex-0")).unwrap();
1175 assert_eq!(receipt.status, HostLeaseAcquireStatus::Deferred);
1176 let defer = receipt.defer.unwrap();
1177 assert_eq!(defer.deferred_reason, HostLeaseDeferReason::RegistryBusy);
1178 assert!(defer.active.is_none());
1179 }
1180
1181 #[test]
1182 fn wait_rechecks_after_cross_thread_release_without_polling() {
1183 let temp = TempDir::new().unwrap();
1184 let store = Arc::new(store(&temp));
1185 let first = store.try_acquire(request("codex-0")).unwrap();
1186 let handle = first.handle.unwrap();
1187 let waiter = {
1188 let store = Arc::clone(&store);
1189 thread::spawn(move || {
1190 store
1191 .acquire_wait(request("codex-1"), Duration::from_secs(5))
1192 .unwrap()
1193 })
1194 };
1195 assert!(
1196 store
1197 .release(&handle.host, &handle.lease_id)
1198 .unwrap()
1199 .released
1200 );
1201 let receipt = waiter.join().unwrap();
1202 assert_eq!(receipt.status, HostLeaseAcquireStatus::Acquired);
1203 }
1204}