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