1mod db;
10mod execution;
11mod schema;
12
13pub use execution::{
14 HostLeaseCargoExecutionContext, HostLeaseExecutionContext, HostLeaseOperationKind,
15 HostLeasePathIdentity, HostLeaseProcessExit, HostLeaseRunLaunchFailure, HostLeaseRunReceipt,
16 HostLeaseRunReleaseOutcome, HostLeaseRunStartFailure, HostLeaseRunState,
17};
18
19use std::collections::BTreeMap;
20use std::path::{Path, PathBuf};
21use std::sync::{mpsc, Arc};
22use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};
23
24use notify::{RecursiveMode, Watcher};
25use rusqlite::{params, ErrorCode, OptionalExtension, Transaction, TransactionBehavior};
26use serde::{Deserialize, Serialize};
27use sysinfo::System;
28use uuid::Uuid;
29
30use self::schema::{
31 add_domain_key, add_execution_context_column, create_current_lease_table, lease_table_layout,
32 migrate_legacy_lease_table, LeaseTableLayout,
33};
34
35pub const HOST_LEASE_ROOT_ENV: &str = "HARN_HOST_LEASE_ROOT";
37const HARN_HOME_ENV: &str = "HARN_HOME";
38const LEASE_DB_FILE: &str = "host-leases.sqlite";
39const LEASE_WAKE_FILE: &str = "host-leases.wake";
40const RUN_RECEIPTS_DIR: &str = "receipts";
41const SQLITE_MUTATION_BUSY_TIMEOUT: Duration = Duration::from_secs(5);
47const REGISTRY_BUSY_RETRY_INTERVAL: Duration = Duration::from_millis(250);
48const PROCESS_LIVENESS_RECHECK_INTERVAL: Duration = Duration::from_secs(5);
49const SCHEMA_VERSION: u32 = 3;
50const RUN_RECEIPT_SCHEMA_VERSION: u32 = 3;
51const WHOLE_MACHINE_RESOURCE_CLASS: &str = "whole-machine";
52pub const DEFAULT_HOST_LEASE_DOMAIN: &str = "default";
54
55#[derive(Debug, thiserror::Error)]
57pub enum HostLeaseError {
58 #[error("invalid host lease request: {0}")]
60 InvalidRequest(String),
61 #[error("host lease state I/O failed: {0}")]
63 Io(#[from] std::io::Error),
64 #[error("host lease database failed: {0}")]
66 Database(#[from] rusqlite::Error),
67 #[error("host lease metadata serialization failed: {0}")]
69 Serialization(#[from] serde_json::Error),
70 #[error("host lease watcher failed: {0}")]
72 Watch(String),
73 #[error("system clock is before the Unix epoch")]
75 Clock,
76}
77
78#[derive(Clone, Copy, Debug, PartialEq, Eq)]
79enum ProcessObservation {
80 Alive { identity: u64 },
81 Dead,
82 Unknown,
83}
84
85trait ProcessInspector: std::fmt::Debug + Send + Sync {
86 fn observe(&self, pid: u32) -> ProcessObservation;
87}
88
89#[derive(Debug)]
90struct SystemProcessInspector;
91
92impl ProcessInspector for SystemProcessInspector {
93 fn observe(&self, pid: u32) -> ProcessObservation {
94 match crate::process_liveness::process_liveness(pid) {
95 crate::process_liveness::ProcessLiveness::Dead => ProcessObservation::Dead,
96 crate::process_liveness::ProcessLiveness::Unknown => ProcessObservation::Unknown,
97 crate::process_liveness::ProcessLiveness::Alive => {
98 crate::process_liveness::process_identity(pid)
99 .map_or(ProcessObservation::Unknown, |identity| {
100 ProcessObservation::Alive { identity }
101 })
102 }
103 }
104 }
105}
106
107#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
109#[serde(rename_all = "kebab-case")]
110pub enum HostLeasePriorityClass {
111 Interactive,
113 Measurement,
115 CiVerify,
117 #[default]
118 Deferrable,
120}
121
122impl HostLeasePriorityClass {
123 pub const fn as_str(self) -> &'static str {
125 match self {
126 Self::Interactive => "interactive",
127 Self::Measurement => "measurement",
128 Self::CiVerify => "ci-verify",
129 Self::Deferrable => "deferrable",
130 }
131 }
132
133 fn parse(raw: &str) -> Result<Self, HostLeaseError> {
134 match raw {
135 "interactive" => Ok(Self::Interactive),
136 "measurement" => Ok(Self::Measurement),
137 "ci-verify" => Ok(Self::CiVerify),
138 "deferrable" => Ok(Self::Deferrable),
139 other => Err(HostLeaseError::InvalidRequest(format!(
140 "unknown priority class `{other}`"
141 ))),
142 }
143 }
144}
145
146#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
153#[serde(rename_all = "kebab-case")]
154pub enum HostLeaseResourceClass {
155 #[default]
157 WholeMachine,
158 RustHeavy,
160}
161
162#[derive(Clone, Copy, Debug, PartialEq, Eq)]
164pub struct HostLeaseResourceDefinition {
165 pub name: &'static str,
167 pub capacity: u16,
169}
170
171const HOST_LEASE_RESOURCE_DEFINITIONS: [HostLeaseResourceDefinition; 2] = [
172 HostLeaseResourceDefinition {
173 name: WHOLE_MACHINE_RESOURCE_CLASS,
174 capacity: 1,
175 },
176 HostLeaseResourceDefinition {
177 name: "rust-heavy",
178 capacity: 1,
179 },
180];
181
182impl HostLeaseResourceClass {
183 pub const fn definition(self) -> &'static HostLeaseResourceDefinition {
185 match self {
186 Self::WholeMachine => &HOST_LEASE_RESOURCE_DEFINITIONS[0],
187 Self::RustHeavy => &HOST_LEASE_RESOURCE_DEFINITIONS[1],
188 }
189 }
190
191 pub const fn as_str(self) -> &'static str {
193 self.definition().name
194 }
195
196 pub const fn capacity(self) -> u16 {
201 self.definition().capacity
202 }
203
204 fn parse(raw: &str) -> Result<Self, HostLeaseError> {
205 match raw {
206 WHOLE_MACHINE_RESOURCE_CLASS => Ok(Self::WholeMachine),
207 "rust-heavy" => Ok(Self::RustHeavy),
208 other => Err(HostLeaseError::InvalidRequest(format!(
209 "unknown resource class `{other}`"
210 ))),
211 }
212 }
213}
214
215#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
217pub struct HostLeaseResourceKey {
218 pub machine: String,
221 pub resource_class: HostLeaseResourceClass,
223 #[serde(default = "default_host_lease_domain")]
224 pub domain: String,
226}
227
228impl HostLeaseResourceKey {
229 fn normalize(
230 machine: &str,
231 resource_class: HostLeaseResourceClass,
232 domain: &str,
233 ) -> Result<Self, HostLeaseError> {
234 Ok(Self {
235 machine: normalize_component("host", machine)?,
236 resource_class,
237 domain: normalize_domain(domain)?,
238 })
239 }
240}
241
242fn default_host_lease_domain() -> String {
243 DEFAULT_HOST_LEASE_DOMAIN.to_string()
244}
245
246#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
248pub struct HostLeaseRequest {
249 pub host: String,
251 #[serde(default)]
252 pub resource_class: HostLeaseResourceClass,
255 #[serde(default = "default_host_lease_domain")]
256 pub domain: String,
258 #[serde(default)]
259 pub execution_context: Option<HostLeaseExecutionContext>,
261 pub owner: String,
263 #[serde(default)]
264 pub priority_class: HostLeasePriorityClass,
266 #[serde(default)]
270 pub ttl_ms: Option<u64>,
271 #[serde(default)]
272 pub owner_pid: Option<u32>,
274 #[serde(default)]
275 pub reason: Option<String>,
277 #[serde(default)]
278 pub metadata: BTreeMap<String, String>,
280}
281
282#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
284pub struct HostLeaseHandle {
285 pub schema_version: u32,
287 pub host: String,
289 #[serde(default)]
290 pub resource_class: HostLeaseResourceClass,
292 #[serde(default = "default_host_lease_domain")]
293 pub domain: String,
295 #[serde(default)]
296 pub execution_context: Option<HostLeaseExecutionContext>,
298 pub lease_id: String,
300 pub owner: String,
302 pub priority_class: HostLeasePriorityClass,
304 pub acquired_at_ms: i64,
306 pub updated_at_ms: i64,
308 #[serde(default)]
309 pub expires_at_ms: Option<i64>,
311 #[serde(default)]
312 pub owner_pid: Option<u32>,
314 #[serde(default)]
315 pub owner_process_identity: Option<u64>,
317 #[serde(default)]
318 pub reason: Option<String>,
320 #[serde(default)]
321 pub metadata: BTreeMap<String, String>,
323}
324
325#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
327#[serde(rename_all = "snake_case")]
328pub enum HostLeaseAcquireStatus {
329 Acquired,
331 Deferred,
333}
334
335#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
337pub enum HostLeaseDeferReason {
338 #[serde(rename = "host_lease_contended")]
340 Contended,
341 #[serde(rename = "host_lease_registry_busy")]
343 RegistryBusy,
344}
345
346impl HostLeaseDeferReason {
347 pub const fn as_str(self) -> &'static str {
349 match self {
350 Self::Contended => "host_lease_contended",
351 Self::RegistryBusy => "host_lease_registry_busy",
352 }
353 }
354}
355
356#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
358pub struct HostLeaseDeferReceipt {
359 pub host: String,
361 #[serde(default)]
362 pub resource_class: HostLeaseResourceClass,
364 #[serde(default = "default_host_lease_domain")]
365 pub domain: String,
367 pub deferred_reason: HostLeaseDeferReason,
369 pub observed_at_ms: i64,
371 #[serde(default)]
372 pub next_wake_at_ms: Option<i64>,
374 #[serde(default)]
375 pub deadline_at_ms: Option<i64>,
377 #[serde(default)]
380 pub active: Option<HostLeaseHandle>,
381}
382
383#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
385pub struct HostLeaseAcquireReceipt {
386 pub schema_version: u32,
388 pub status: HostLeaseAcquireStatus,
390 pub observed_at_ms: i64,
392 pub waited_ms: u64,
394 #[serde(default)]
395 pub handle: Option<HostLeaseHandle>,
397 #[serde(default)]
398 pub defer: Option<HostLeaseDeferReceipt>,
400 pub recovered_stale_lease: bool,
402 #[serde(default)]
403 pub recovered: Option<HostLeaseHandle>,
405}
406
407#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
409pub struct HostLeaseState {
410 pub schema_version: u32,
412 pub host: String,
414 #[serde(default)]
415 pub resource_class: HostLeaseResourceClass,
417 #[serde(default = "default_host_lease_domain")]
418 pub domain: String,
420 pub observed_at_ms: i64,
422 #[serde(default)]
423 pub active: Option<HostLeaseHandle>,
425 pub recovered_stale_lease: bool,
427 #[serde(default)]
428 pub recovered: Option<HostLeaseHandle>,
430}
431
432#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
434pub struct HostLeaseRenewReceipt {
435 pub schema_version: u32,
437 pub renewed: bool,
439 pub observed_at_ms: i64,
441 #[serde(default)]
442 pub handle: Option<HostLeaseHandle>,
444}
445
446#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
448pub struct HostLeaseMetadataUpdateReceipt {
449 pub schema_version: u32,
451 pub updated: bool,
453 pub observed_at_ms: i64,
455 #[serde(default)]
456 pub handle: Option<HostLeaseHandle>,
458}
459
460#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
462pub struct HostLeaseReleaseReceipt {
463 pub schema_version: u32,
465 pub released: bool,
467 pub host: String,
469 #[serde(default)]
470 pub resource_class: HostLeaseResourceClass,
472 #[serde(default = "default_host_lease_domain")]
473 pub domain: String,
475 pub lease_id: String,
477 pub observed_at_ms: i64,
479}
480
481#[derive(Clone, Debug)]
483pub struct HostLeaseStore {
484 root: PathBuf,
485 db_path: PathBuf,
486 wake_path: PathBuf,
487 process_inspector: Arc<dyn ProcessInspector>,
488 #[cfg(test)]
489 busy_handler: Option<fn(i32) -> bool>,
490}
491
492impl HostLeaseStore {
493 pub fn from_env() -> Result<Self, HostLeaseError> {
495 let root = if let Some(path) = std::env::var_os(HOST_LEASE_ROOT_ENV) {
496 PathBuf::from(path)
497 } else if let Some(path) = std::env::var_os(HARN_HOME_ENV) {
498 PathBuf::from(path).join("host-leases")
499 } else {
500 harn_vm::user_dirs::home_dir()
501 .ok_or_else(|| {
502 HostLeaseError::InvalidRequest(
503 "cannot resolve a home directory; set HARN_HOST_LEASE_ROOT".to_string(),
504 )
505 })?
506 .join(".harn/host-leases")
507 };
508 Self::for_root(root)
509 }
510
511 pub fn for_root(root: impl Into<PathBuf>) -> Result<Self, HostLeaseError> {
513 Self::for_root_with_inspector(root, Arc::new(SystemProcessInspector))
514 }
515
516 fn for_root_with_inspector(
517 root: impl Into<PathBuf>,
518 process_inspector: Arc<dyn ProcessInspector>,
519 ) -> Result<Self, HostLeaseError> {
520 let root = root.into();
521 std::fs::create_dir_all(&root)?;
522 let store = Self {
523 db_path: root.join(LEASE_DB_FILE),
524 wake_path: root.join(LEASE_WAKE_FILE),
525 root,
526 process_inspector,
527 #[cfg(test)]
528 busy_handler: None,
529 };
530 store.initialize()?;
531 std::fs::OpenOptions::new()
532 .create(true)
533 .append(true)
534 .open(&store.wake_path)?;
535 Ok(store)
536 }
537
538 pub fn root(&self) -> &Path {
540 &self.root
541 }
542
543 pub fn begin_run(
545 &self,
546 owner: &str,
547 priority_class: HostLeasePriorityClass,
548 resource: HostLeaseResourceKey,
549 execution_context: HostLeaseExecutionContext,
550 wait_limit_ms: u64,
551 ) -> Result<HostLeaseRunReceipt, HostLeaseError> {
552 let resource = HostLeaseResourceKey::normalize(
553 &resource.machine,
554 resource.resource_class,
555 &resource.domain,
556 )?;
557 let receipt = HostLeaseRunReceipt {
558 schema_version: RUN_RECEIPT_SCHEMA_VERSION,
559 run_id: Uuid::now_v7().to_string(),
560 owner: normalize_component("owner", owner)?,
561 priority_class,
562 wait_limit_ms,
563 resource,
564 execution_context,
565 status: HostLeaseRunState::Pending {
566 requested_at_ms: unix_now_ms()?,
567 },
568 };
569 let path = self.run_receipt_path(&receipt.run_id)?;
570 let bytes = serde_json::to_vec_pretty(&receipt)?;
571 harn_vm::atomic_io::atomic_write(&path, &bytes)?;
572 Ok(receipt)
573 }
574
575 pub fn load_run(&self, run_id: &str) -> Result<HostLeaseRunReceipt, HostLeaseError> {
577 let path = self.run_receipt_path(run_id)?;
578 Ok(serde_json::from_slice(&std::fs::read(path)?)?)
579 }
580
581 pub fn transition_run(
583 &self,
584 run_id: &str,
585 status: HostLeaseRunState,
586 ) -> Result<HostLeaseRunReceipt, HostLeaseError> {
587 let mut receipt = self.load_run(run_id)?;
588 if !receipt.status.may_transition_to(&status) {
589 return Err(HostLeaseError::InvalidRequest(format!(
590 "invalid run receipt transition from {:?} to {:?}",
591 receipt.status, status
592 )));
593 }
594 receipt.status = status;
595 let path = self.run_receipt_path(run_id)?;
596 let bytes = serde_json::to_vec_pretty(&receipt)?;
597 harn_vm::atomic_io::atomic_write(&path, &bytes)?;
598 Ok(receipt)
599 }
600
601 pub fn run_receipt_path(&self, run_id: &str) -> Result<PathBuf, HostLeaseError> {
603 let run_id = normalize_component("run_id", run_id)?;
604 Ok(self
605 .root
606 .join(RUN_RECEIPTS_DIR)
607 .join(format!("{run_id}.json")))
608 }
609
610 pub fn default_host() -> String {
612 System::host_name()
613 .filter(|name| !name.trim().is_empty())
614 .unwrap_or_else(|| "local".to_string())
615 }
616
617 pub fn try_acquire(
619 &self,
620 request: HostLeaseRequest,
621 ) -> Result<HostLeaseAcquireReceipt, HostLeaseError> {
622 self.try_acquire_once(request, None, None)
623 }
624
625 pub fn acquire_wait(
627 &self,
628 request: HostLeaseRequest,
629 wait_timeout: Duration,
630 ) -> Result<HostLeaseAcquireReceipt, HostLeaseError> {
631 if wait_timeout.is_zero() {
632 return self.try_acquire(request);
633 }
634 let started_at_ms = unix_now_ms()?;
635 let started_at = Instant::now();
636 let deadline = started_at.checked_add(wait_timeout).ok_or_else(|| {
637 HostLeaseError::InvalidRequest("wait timeout exceeds the monotonic clock".to_string())
638 })?;
639 let deadline_at_ms = started_at_ms.saturating_add(duration_ms_i64(wait_timeout));
640 let (tx, rx) = mpsc::channel();
641 let mut watcher = notify::recommended_watcher(move |event| {
642 let _ = tx.send(event);
643 })
644 .map_err(|error| HostLeaseError::Watch(error.to_string()))?;
645 watcher
646 .watch(&self.wake_path, RecursiveMode::NonRecursive)
647 .map_err(|error| HostLeaseError::Watch(error.to_string()))?;
648
649 loop {
650 let receipt =
651 self.try_acquire_once(request.clone(), Some(started_at), Some(deadline_at_ms))?;
652 if receipt.status == HostLeaseAcquireStatus::Acquired || Instant::now() >= deadline {
653 return Ok(receipt);
654 }
655 let wake_at = receipt
656 .defer
657 .as_ref()
658 .and_then(|defer| defer.next_wake_at_ms)
659 .map(|wake| wake.min(deadline_at_ms))
660 .unwrap_or(deadline_at_ms);
661 let wake_duration =
662 Duration::from_millis(wake_at.saturating_sub(receipt.observed_at_ms).max(1) as u64);
663 let remaining = deadline.saturating_duration_since(Instant::now());
664 if remaining.is_zero() {
665 return Ok(receipt);
666 }
667 match rx.recv_timeout(wake_duration.min(remaining)) {
668 Ok(Ok(_)) | Err(mpsc::RecvTimeoutError::Timeout) => {}
669 Ok(Err(error)) => return Err(HostLeaseError::Watch(error.to_string())),
670 Err(mpsc::RecvTimeoutError::Disconnected) => {
671 return Err(HostLeaseError::Watch(
672 "host lease watcher disconnected".to_string(),
673 ));
674 }
675 }
676 }
677 }
678
679 pub fn status(&self, host: &str) -> Result<HostLeaseState, HostLeaseError> {
681 self.status_for_resource(host, HostLeaseResourceClass::WholeMachine)
682 }
683
684 pub fn status_for_resource(
686 &self,
687 host: &str,
688 resource_class: HostLeaseResourceClass,
689 ) -> Result<HostLeaseState, HostLeaseError> {
690 self.status_for_domain(host, resource_class, DEFAULT_HOST_LEASE_DOMAIN)
691 }
692
693 pub fn status_for_domain(
695 &self,
696 host: &str,
697 resource_class: HostLeaseResourceClass,
698 domain: &str,
699 ) -> Result<HostLeaseState, HostLeaseError> {
700 let resource = HostLeaseResourceKey::normalize(host, resource_class, domain)?;
701 let mut conn = self.connection(SQLITE_MUTATION_BUSY_TIMEOUT)?;
702 let tx = conn.transaction_with_behavior(TransactionBehavior::Immediate)?;
703 let now = unix_now_ms()?;
704 self.status_in_transaction(
705 tx,
706 &resource.machine,
707 resource.resource_class,
708 &resource.domain,
709 now,
710 )
711 }
712
713 pub fn renew(
715 &self,
716 host: &str,
717 lease_id: &str,
718 ttl_ms: u64,
719 ) -> Result<HostLeaseRenewReceipt, HostLeaseError> {
720 self.renew_for_resource(host, HostLeaseResourceClass::WholeMachine, lease_id, ttl_ms)
721 }
722
723 pub fn renew_for_resource(
725 &self,
726 host: &str,
727 resource_class: HostLeaseResourceClass,
728 lease_id: &str,
729 ttl_ms: u64,
730 ) -> Result<HostLeaseRenewReceipt, HostLeaseError> {
731 self.renew_for_domain(
732 host,
733 resource_class,
734 DEFAULT_HOST_LEASE_DOMAIN,
735 lease_id,
736 ttl_ms,
737 )
738 }
739
740 pub fn renew_for_domain(
742 &self,
743 host: &str,
744 resource_class: HostLeaseResourceClass,
745 domain: &str,
746 lease_id: &str,
747 ttl_ms: u64,
748 ) -> Result<HostLeaseRenewReceipt, HostLeaseError> {
749 let resource = HostLeaseResourceKey::normalize(host, resource_class, domain)?;
750 let lease_id = normalize_component("lease_id", lease_id)?;
751 validate_ttl(Some(ttl_ms))?;
752 let mut conn = self.connection(SQLITE_MUTATION_BUSY_TIMEOUT)?;
753 let tx = conn.transaction_with_behavior(TransactionBehavior::Immediate)?;
754 let now = unix_now_ms()?;
755 let (active, _) = active_handle(
756 &tx,
757 &resource.machine,
758 resource.resource_class,
759 &resource.domain,
760 now,
761 self.process_inspector.as_ref(),
762 )?;
763 let Some(mut handle) = active.filter(|handle| handle.lease_id == lease_id) else {
764 tx.commit()?;
765 return Ok(HostLeaseRenewReceipt {
766 schema_version: SCHEMA_VERSION,
767 renewed: false,
768 observed_at_ms: now,
769 handle: None,
770 });
771 };
772 handle.updated_at_ms = now;
773 handle.expires_at_ms = Some(now.saturating_add(u64_ms_i64(ttl_ms)));
774 write_handle(&tx, &handle)?;
775 tx.commit()?;
776 self.signal_waiters();
777 Ok(HostLeaseRenewReceipt {
778 schema_version: SCHEMA_VERSION,
779 renewed: true,
780 observed_at_ms: now,
781 handle: Some(handle),
782 })
783 }
784
785 pub fn update_metadata_for_domain(
791 &self,
792 host: &str,
793 resource_class: HostLeaseResourceClass,
794 domain: &str,
795 lease_id: &str,
796 metadata: BTreeMap<String, String>,
797 ) -> Result<HostLeaseMetadataUpdateReceipt, HostLeaseError> {
798 self.update_metadata_at_domain(
799 host,
800 resource_class,
801 domain,
802 lease_id,
803 metadata,
804 unix_now_ms()?,
805 )
806 }
807
808 fn update_metadata_at_domain(
809 &self,
810 host: &str,
811 resource_class: HostLeaseResourceClass,
812 domain: &str,
813 lease_id: &str,
814 metadata: BTreeMap<String, String>,
815 now: i64,
816 ) -> Result<HostLeaseMetadataUpdateReceipt, HostLeaseError> {
817 let resource = HostLeaseResourceKey::normalize(host, resource_class, domain)?;
818 let lease_id = normalize_component("lease_id", lease_id)?;
819 let mut conn = self.connection(SQLITE_MUTATION_BUSY_TIMEOUT)?;
820 let tx = conn.transaction_with_behavior(TransactionBehavior::Immediate)?;
821 let (active, _) = active_handle(
822 &tx,
823 &resource.machine,
824 resource.resource_class,
825 &resource.domain,
826 now,
827 self.process_inspector.as_ref(),
828 )?;
829 let Some(mut handle) = active.filter(|handle| handle.lease_id == lease_id) else {
830 tx.commit()?;
831 return Ok(HostLeaseMetadataUpdateReceipt {
832 schema_version: SCHEMA_VERSION,
833 updated: false,
834 observed_at_ms: now,
835 handle: None,
836 });
837 };
838 handle.updated_at_ms = now;
839 handle.metadata = metadata;
840 write_handle(&tx, &handle)?;
841 tx.commit()?;
842 Ok(HostLeaseMetadataUpdateReceipt {
843 schema_version: SCHEMA_VERSION,
844 updated: true,
845 observed_at_ms: now,
846 handle: Some(handle),
847 })
848 }
849
850 pub fn release(
852 &self,
853 host: &str,
854 lease_id: &str,
855 ) -> Result<HostLeaseReleaseReceipt, HostLeaseError> {
856 self.release_for_resource(host, HostLeaseResourceClass::WholeMachine, lease_id)
857 }
858
859 pub fn release_for_resource(
861 &self,
862 host: &str,
863 resource_class: HostLeaseResourceClass,
864 lease_id: &str,
865 ) -> Result<HostLeaseReleaseReceipt, HostLeaseError> {
866 self.release_for_domain(host, resource_class, DEFAULT_HOST_LEASE_DOMAIN, lease_id)
867 }
868
869 pub fn release_for_domain(
871 &self,
872 host: &str,
873 resource_class: HostLeaseResourceClass,
874 domain: &str,
875 lease_id: &str,
876 ) -> Result<HostLeaseReleaseReceipt, HostLeaseError> {
877 let resource = HostLeaseResourceKey::normalize(host, resource_class, domain)?;
878 let lease_id = normalize_component("lease_id", lease_id)?;
879 let conn = self.connection(SQLITE_MUTATION_BUSY_TIMEOUT)?;
880 let released = conn.execute(
881 "DELETE FROM host_leases
882 WHERE host = ?1 AND resource_class = ?2 AND domain = ?3 AND lease_id = ?4",
883 params![
884 &resource.machine,
885 resource.resource_class.as_str(),
886 &resource.domain,
887 &lease_id,
888 ],
889 )? == 1;
890 let now = unix_now_ms()?;
891 if released {
892 self.signal_waiters();
893 }
894 Ok(HostLeaseReleaseReceipt {
895 schema_version: SCHEMA_VERSION,
896 released,
897 host: resource.machine,
898 resource_class: resource.resource_class,
899 domain: resource.domain,
900 lease_id,
901 observed_at_ms: now,
902 })
903 }
904
905 fn initialize(&self) -> Result<(), HostLeaseError> {
906 let mut conn = self.connection(SQLITE_MUTATION_BUSY_TIMEOUT)?;
907 let tx = conn.transaction_with_behavior(TransactionBehavior::Immediate)?;
908 match lease_table_layout(&tx)? {
909 LeaseTableLayout::Missing => create_current_lease_table(&tx)?,
910 LeaseTableLayout::LegacyWholeMachine => migrate_legacy_lease_table(&tx)?,
911 LeaseTableLayout::ResourceClassWithoutExecutionContext => {
912 add_execution_context_column(&tx)?;
913 add_domain_key(&tx)?;
914 }
915 LeaseTableLayout::CurrentWithoutDomain => add_domain_key(&tx)?,
916 LeaseTableLayout::Current => {}
917 }
918 tx.commit()?;
919 Ok(())
920 }
921
922 fn signal_waiters(&self) {
931 let _ = std::fs::write(&self.wake_path, Uuid::now_v7().to_string());
932 }
933
934 #[cfg(test)]
935 fn with_busy_handler(mut self, handler: fn(i32) -> bool) -> Self {
936 self.busy_handler = Some(handler);
937 self
938 }
939
940 fn try_acquire_once(
941 &self,
942 request: HostLeaseRequest,
943 started_at: Option<Instant>,
944 deadline_at_ms: Option<i64>,
945 ) -> Result<HostLeaseAcquireReceipt, HostLeaseError> {
946 self.try_acquire_once_with_registry_timeout(
947 request,
948 started_at,
949 deadline_at_ms,
950 SQLITE_MUTATION_BUSY_TIMEOUT,
951 )
952 }
953
954 fn try_acquire_once_with_registry_timeout(
955 &self,
956 request: HostLeaseRequest,
957 started_at: Option<Instant>,
958 deadline_at_ms: Option<i64>,
959 registry_timeout: Duration,
960 ) -> Result<HostLeaseAcquireReceipt, HostLeaseError> {
961 let request = normalize_request(request)?;
962 let mut conn = self.connection(registry_timeout)?;
968 let tx = match conn.transaction_with_behavior(TransactionBehavior::Immediate) {
969 Ok(tx) => tx,
970 Err(error) if sqlite_is_busy(&error) => {
971 let now = unix_now_ms()?;
972 return Ok(registry_busy_receipt(
973 request.host,
974 request.resource_class,
975 request.domain,
976 now,
977 started_at.map(|started| duration_ms_u64(started.elapsed())),
978 deadline_at_ms,
979 ));
980 }
981 Err(error) => return Err(error.into()),
982 };
983 let now = unix_now_ms()?;
984 let waited_ms = started_at
985 .map(|started| duration_ms_u64(started.elapsed()))
986 .unwrap_or(0);
987 let host = request.host.clone();
988 let resource_class = request.resource_class;
989 let domain = request.domain.clone();
990 match self.acquire_in_transaction(tx, request, now, deadline_at_ms, waited_ms) {
991 Err(HostLeaseError::Database(error)) if sqlite_is_busy(&error) => {
992 Ok(registry_busy_receipt(
993 host,
994 resource_class,
995 domain,
996 now,
997 Some(waited_ms),
998 deadline_at_ms,
999 ))
1000 }
1001 result => result,
1002 }
1003 }
1004
1005 #[cfg(test)]
1006 fn try_acquire_at(
1007 &self,
1008 request: HostLeaseRequest,
1009 now: i64,
1010 deadline_at_ms: Option<i64>,
1011 waited_ms: u64,
1012 ) -> Result<HostLeaseAcquireReceipt, HostLeaseError> {
1013 let request = normalize_request(request)?;
1014 let mut conn = self.connection(SQLITE_MUTATION_BUSY_TIMEOUT)?;
1015 let tx = conn.transaction_with_behavior(TransactionBehavior::Immediate)?;
1016 self.acquire_in_transaction(tx, request, now, deadline_at_ms, waited_ms)
1017 }
1018
1019 fn acquire_in_transaction(
1020 &self,
1021 tx: Transaction<'_>,
1022 request: HostLeaseRequest,
1023 now: i64,
1024 deadline_at_ms: Option<i64>,
1025 waited_ms: u64,
1026 ) -> Result<HostLeaseAcquireReceipt, HostLeaseError> {
1027 let owner_process_identity = request
1028 .owner_pid
1029 .map(|pid| match self.process_inspector.observe(pid) {
1030 ProcessObservation::Alive { identity } => Ok(identity),
1031 ProcessObservation::Dead => Err(HostLeaseError::InvalidRequest(
1032 "owner_pid is not a live local process".to_string(),
1033 )),
1034 ProcessObservation::Unknown => Err(HostLeaseError::InvalidRequest(
1035 "owner_pid liveness could not be verified".to_string(),
1036 )),
1037 })
1038 .transpose()?;
1039 let (active, recovered) = active_handle(
1040 &tx,
1041 &request.host,
1042 request.resource_class,
1043 &request.domain,
1044 now,
1045 self.process_inspector.as_ref(),
1046 )?;
1047 if let Some(active) = active {
1048 let defer = HostLeaseDeferReceipt {
1049 host: request.host,
1050 resource_class: request.resource_class,
1051 domain: request.domain,
1052 deferred_reason: HostLeaseDeferReason::Contended,
1053 observed_at_ms: now,
1054 next_wake_at_ms: Some(next_lease_wake_at(&active, now, deadline_at_ms)),
1055 deadline_at_ms,
1056 active: Some(active),
1057 };
1058 tx.commit()?;
1059 return Ok(HostLeaseAcquireReceipt {
1060 schema_version: SCHEMA_VERSION,
1061 status: HostLeaseAcquireStatus::Deferred,
1062 observed_at_ms: now,
1063 waited_ms,
1064 handle: None,
1065 defer: Some(defer),
1066 recovered_stale_lease: recovered.is_some(),
1067 recovered,
1068 });
1069 }
1070
1071 let handle = HostLeaseHandle {
1072 schema_version: SCHEMA_VERSION,
1073 host: request.host,
1074 resource_class: request.resource_class,
1075 domain: request.domain,
1076 execution_context: request.execution_context,
1077 lease_id: Uuid::now_v7().to_string(),
1078 owner: request.owner,
1079 priority_class: request.priority_class,
1080 acquired_at_ms: now,
1081 updated_at_ms: now,
1082 expires_at_ms: request
1083 .ttl_ms
1084 .map(|ttl| now.saturating_add(u64_ms_i64(ttl))),
1085 owner_pid: request.owner_pid,
1086 owner_process_identity,
1087 reason: request.reason,
1088 metadata: request.metadata,
1089 };
1090 write_handle(&tx, &handle)?;
1091 tx.commit()?;
1092 Ok(HostLeaseAcquireReceipt {
1093 schema_version: SCHEMA_VERSION,
1094 status: HostLeaseAcquireStatus::Acquired,
1095 observed_at_ms: now,
1096 waited_ms,
1097 handle: Some(handle),
1098 defer: None,
1099 recovered_stale_lease: recovered.is_some(),
1100 recovered,
1101 })
1102 }
1103
1104 #[cfg(test)]
1105 fn status_at(
1106 &self,
1107 host: &str,
1108 resource_class: HostLeaseResourceClass,
1109 now: i64,
1110 ) -> Result<HostLeaseState, HostLeaseError> {
1111 self.status_at_domain(host, resource_class, DEFAULT_HOST_LEASE_DOMAIN, now)
1112 }
1113
1114 #[cfg(test)]
1115 fn status_at_domain(
1116 &self,
1117 host: &str,
1118 resource_class: HostLeaseResourceClass,
1119 domain: &str,
1120 now: i64,
1121 ) -> Result<HostLeaseState, HostLeaseError> {
1122 let mut conn = self.connection(SQLITE_MUTATION_BUSY_TIMEOUT)?;
1123 let tx = conn.transaction_with_behavior(TransactionBehavior::Immediate)?;
1124 self.status_in_transaction(tx, host, resource_class, domain, now)
1125 }
1126
1127 fn status_in_transaction(
1128 &self,
1129 tx: Transaction<'_>,
1130 host: &str,
1131 resource_class: HostLeaseResourceClass,
1132 domain: &str,
1133 now: i64,
1134 ) -> Result<HostLeaseState, HostLeaseError> {
1135 let (active, recovered) = active_handle(
1136 &tx,
1137 host,
1138 resource_class,
1139 domain,
1140 now,
1141 self.process_inspector.as_ref(),
1142 )?;
1143 tx.commit()?;
1144 if recovered.is_some() {
1145 self.signal_waiters();
1146 }
1147 Ok(HostLeaseState {
1148 schema_version: SCHEMA_VERSION,
1149 host: host.to_string(),
1150 resource_class,
1151 domain: domain.to_string(),
1152 observed_at_ms: now,
1153 active,
1154 recovered_stale_lease: recovered.is_some(),
1155 recovered,
1156 })
1157 }
1158}
1159
1160fn normalize_request(mut request: HostLeaseRequest) -> Result<HostLeaseRequest, HostLeaseError> {
1161 let resource =
1162 HostLeaseResourceKey::normalize(&request.host, request.resource_class, &request.domain)?;
1163 request.host = resource.machine;
1164 request.resource_class = resource.resource_class;
1165 request.domain = resource.domain;
1166 request.owner = normalize_component("owner", &request.owner)?;
1167 validate_ttl(request.ttl_ms)?;
1168 if request.ttl_ms.is_none() && request.owner_pid.is_none() {
1169 return Err(HostLeaseError::InvalidRequest(
1170 "a non-expiring lease requires owner_pid for crash recovery".to_string(),
1171 ));
1172 }
1173 request.reason = request.reason.and_then(|reason| {
1174 let trimmed = reason.trim();
1175 (!trimmed.is_empty()).then(|| trimmed.to_string())
1176 });
1177 Ok(request)
1178}
1179
1180fn normalize_component(name: &str, value: &str) -> Result<String, HostLeaseError> {
1181 let value = value.trim();
1182 if value.is_empty() {
1183 return Err(HostLeaseError::InvalidRequest(format!(
1184 "{name} cannot be empty"
1185 )));
1186 }
1187 if !value
1188 .chars()
1189 .all(|ch| ch.is_ascii_alphanumeric() || matches!(ch, '.' | '_' | '-'))
1190 {
1191 return Err(HostLeaseError::InvalidRequest(format!(
1192 "{name} may contain only ASCII letters, digits, '.', '_' and '-'"
1193 )));
1194 }
1195 Ok(value.to_string())
1196}
1197
1198fn normalize_domain(value: &str) -> Result<String, HostLeaseError> {
1199 let value = value.trim();
1200 if value.is_empty() {
1201 return Err(HostLeaseError::InvalidRequest(
1202 "domain cannot be empty".to_string(),
1203 ));
1204 }
1205 if value.len() > 128 {
1206 return Err(HostLeaseError::InvalidRequest(
1207 "domain cannot exceed 128 bytes".to_string(),
1208 ));
1209 }
1210 if matches!(value, "." | "..")
1211 || !value
1212 .chars()
1213 .all(|ch| ch.is_ascii_alphanumeric() || matches!(ch, '.' | '_' | '-'))
1214 {
1215 return Err(HostLeaseError::InvalidRequest(
1216 "domain may contain only ASCII letters, digits, '.', '_' and '-' and must not be path-like"
1217 .to_string(),
1218 ));
1219 }
1220 Ok(value.to_string())
1221}
1222
1223fn validate_ttl(ttl_ms: Option<u64>) -> Result<(), HostLeaseError> {
1224 if ttl_ms == Some(0) {
1225 return Err(HostLeaseError::InvalidRequest(
1226 "ttl_ms must be greater than zero when supplied".to_string(),
1227 ));
1228 }
1229 Ok(())
1230}
1231
1232fn active_handle(
1233 tx: &Transaction<'_>,
1234 host: &str,
1235 resource_class: HostLeaseResourceClass,
1236 domain: &str,
1237 now: i64,
1238 process_inspector: &dyn ProcessInspector,
1239) -> Result<(Option<HostLeaseHandle>, Option<HostLeaseHandle>), HostLeaseError> {
1240 let handle = read_handle(tx, host, resource_class, domain)?;
1241 let Some(handle) = handle else {
1242 return Ok((None, None));
1243 };
1244 let expired = handle.expires_at_ms.is_some_and(|expiry| expiry <= now);
1245 let owner_dead = match (handle.owner_pid, handle.owner_process_identity) {
1246 (Some(pid), Some(expected_identity)) => match process_inspector.observe(pid) {
1247 ProcessObservation::Alive { identity } => identity != expected_identity,
1248 ProcessObservation::Dead => true,
1249 ProcessObservation::Unknown => false,
1250 },
1251 _ => false,
1252 };
1253 if expired || owner_dead {
1254 tx.execute(
1255 "DELETE FROM host_leases
1256 WHERE host = ?1 AND resource_class = ?2 AND domain = ?3 AND lease_id = ?4",
1257 params![host, resource_class.as_str(), domain, handle.lease_id],
1258 )?;
1259 return Ok((None, Some(handle)));
1260 }
1261 Ok((Some(handle), None))
1262}
1263
1264fn next_lease_wake_at(active: &HostLeaseHandle, now: i64, deadline_at_ms: Option<i64>) -> i64 {
1265 let mut wake_at = deadline_at_ms.unwrap_or(i64::MAX);
1266 if let Some(expiry) = active.expires_at_ms {
1267 wake_at = wake_at.min(expiry);
1268 }
1269 if active.owner_pid.is_some() {
1270 wake_at =
1271 wake_at.min(now.saturating_add(duration_ms_i64(PROCESS_LIVENESS_RECHECK_INTERVAL)));
1272 }
1273 wake_at
1274}
1275
1276fn registry_busy_receipt(
1277 host: String,
1278 resource_class: HostLeaseResourceClass,
1279 domain: String,
1280 now: i64,
1281 waited_ms: Option<u64>,
1282 deadline_at_ms: Option<i64>,
1283) -> HostLeaseAcquireReceipt {
1284 let next_wake_at_ms = now
1285 .saturating_add(duration_ms_i64(REGISTRY_BUSY_RETRY_INTERVAL))
1286 .min(deadline_at_ms.unwrap_or(i64::MAX));
1287 HostLeaseAcquireReceipt {
1288 schema_version: SCHEMA_VERSION,
1289 status: HostLeaseAcquireStatus::Deferred,
1290 observed_at_ms: now,
1291 waited_ms: waited_ms.unwrap_or(0),
1292 handle: None,
1293 defer: Some(HostLeaseDeferReceipt {
1294 host,
1295 resource_class,
1296 domain,
1297 deferred_reason: HostLeaseDeferReason::RegistryBusy,
1298 observed_at_ms: now,
1299 next_wake_at_ms: Some(next_wake_at_ms),
1300 deadline_at_ms,
1301 active: None,
1302 }),
1303 recovered_stale_lease: false,
1304 recovered: None,
1305 }
1306}
1307
1308fn sqlite_is_busy(error: &rusqlite::Error) -> bool {
1309 matches!(
1310 error,
1311 rusqlite::Error::SqliteFailure(inner, _)
1312 if matches!(inner.code, ErrorCode::DatabaseBusy | ErrorCode::DatabaseLocked)
1313 )
1314}
1315
1316fn read_handle(
1317 tx: &Transaction<'_>,
1318 host: &str,
1319 resource_class: HostLeaseResourceClass,
1320 domain: &str,
1321) -> Result<Option<HostLeaseHandle>, HostLeaseError> {
1322 tx.query_row(
1323 "SELECT resource_class, domain, lease_id, owner, priority_class, acquired_at_ms, updated_at_ms,
1324 expires_at_ms, owner_pid, owner_process_identity, reason, metadata_json,
1325 execution_context_json
1326 FROM host_leases WHERE host = ?1 AND resource_class = ?2 AND domain = ?3",
1327 params![host, resource_class.as_str(), domain],
1328 |row| {
1329 let priority: String = row.get(4)?;
1330 let metadata_json: String = row.get(11)?;
1331 let execution_context_json: Option<String> = row.get(12)?;
1332 let owner_pid_i64: Option<i64> = row.get(8)?;
1333 let owner_identity_i64: Option<i64> = row.get(9)?;
1334 Ok((
1335 row.get::<_, String>(2)?,
1336 row.get::<_, String>(3)?,
1337 priority,
1338 row.get::<_, i64>(5)?,
1339 row.get::<_, i64>(6)?,
1340 row.get::<_, Option<i64>>(7)?,
1341 owner_pid_i64,
1342 owner_identity_i64,
1343 row.get::<_, Option<String>>(10)?,
1344 metadata_json,
1345 row.get::<_, String>(0)?,
1346 row.get::<_, String>(1)?,
1347 execution_context_json,
1348 ))
1349 },
1350 )
1351 .optional()?
1352 .map(
1353 |(
1354 lease_id,
1355 owner,
1356 priority,
1357 acquired_at_ms,
1358 updated_at_ms,
1359 expires_at_ms,
1360 owner_pid,
1361 owner_process_identity,
1362 reason,
1363 metadata_json,
1364 stored_resource_class,
1365 stored_domain,
1366 execution_context_json,
1367 )| {
1368 let owner_pid = owner_pid
1369 .map(|pid| {
1370 u32::try_from(pid).map_err(|_| {
1371 HostLeaseError::InvalidRequest(
1372 "persisted owner_pid is outside the u32 range".to_string(),
1373 )
1374 })
1375 })
1376 .transpose()?;
1377 let owner_process_identity = owner_process_identity
1378 .map(|identity| {
1379 u64::try_from(identity).map_err(|_| {
1380 HostLeaseError::InvalidRequest(
1381 "persisted process identity is negative".to_string(),
1382 )
1383 })
1384 })
1385 .transpose()?;
1386 Ok(HostLeaseHandle {
1387 schema_version: SCHEMA_VERSION,
1388 host: host.to_string(),
1389 resource_class: HostLeaseResourceClass::parse(&stored_resource_class)?,
1390 domain: stored_domain,
1391 execution_context: execution_context_json
1392 .map(|encoded| serde_json::from_str(&encoded))
1393 .transpose()?,
1394 lease_id,
1395 owner,
1396 priority_class: HostLeasePriorityClass::parse(&priority)?,
1397 acquired_at_ms,
1398 updated_at_ms,
1399 expires_at_ms,
1400 owner_pid,
1401 owner_process_identity,
1402 reason,
1403 metadata: serde_json::from_str(&metadata_json)?,
1404 })
1405 },
1406 )
1407 .transpose()
1408}
1409
1410fn write_handle(tx: &Transaction<'_>, handle: &HostLeaseHandle) -> Result<(), HostLeaseError> {
1411 let metadata_json = serde_json::to_string(&handle.metadata)?;
1412 let execution_context_json = handle
1413 .execution_context
1414 .as_ref()
1415 .map(serde_json::to_string)
1416 .transpose()?;
1417 let owner_process_identity = handle
1418 .owner_process_identity
1419 .map(|value| {
1420 i64::try_from(value).map_err(|_| {
1421 HostLeaseError::InvalidRequest(
1422 "owner process identity is outside the SQLite integer range".to_string(),
1423 )
1424 })
1425 })
1426 .transpose()?;
1427 tx.execute(
1428 "INSERT INTO host_leases (
1429 host, resource_class, domain, lease_id, owner, priority_class, acquired_at_ms, updated_at_ms,
1430 expires_at_ms, owner_pid, owner_process_identity, reason, metadata_json,
1431 execution_context_json
1432 ) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14)
1433 ON CONFLICT(host, resource_class, domain) DO UPDATE SET
1434 lease_id = excluded.lease_id,
1435 owner = excluded.owner,
1436 priority_class = excluded.priority_class,
1437 acquired_at_ms = excluded.acquired_at_ms,
1438 updated_at_ms = excluded.updated_at_ms,
1439 expires_at_ms = excluded.expires_at_ms,
1440 owner_pid = excluded.owner_pid,
1441 owner_process_identity = excluded.owner_process_identity,
1442 reason = excluded.reason,
1443 metadata_json = excluded.metadata_json,
1444 execution_context_json = excluded.execution_context_json",
1445 params![
1446 handle.host,
1447 handle.resource_class.as_str(),
1448 handle.domain,
1449 handle.lease_id,
1450 handle.owner,
1451 handle.priority_class.as_str(),
1452 handle.acquired_at_ms,
1453 handle.updated_at_ms,
1454 handle.expires_at_ms,
1455 handle.owner_pid.map(i64::from),
1456 owner_process_identity,
1457 handle.reason,
1458 metadata_json,
1459 execution_context_json,
1460 ],
1461 )?;
1462 Ok(())
1463}
1464
1465fn unix_now_ms() -> Result<i64, HostLeaseError> {
1466 let millis = SystemTime::now()
1467 .duration_since(UNIX_EPOCH)
1468 .map_err(|_| HostLeaseError::Clock)?
1469 .as_millis();
1470 Ok(millis.min(i64::MAX as u128) as i64)
1471}
1472
1473fn duration_ms_i64(duration: Duration) -> i64 {
1474 duration.as_millis().min(i64::MAX as u128) as i64
1475}
1476
1477fn duration_ms_u64(duration: Duration) -> u64 {
1478 duration.as_millis().min(u64::MAX as u128) as u64
1479}
1480
1481fn u64_ms_i64(value: u64) -> i64 {
1482 value.min(i64::MAX as u64) as i64
1483}
1484
1485#[cfg(test)]
1486#[path = "host_lease/tests.rs"]
1487mod tests;