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}
483
484impl HostLeaseStore {
485 pub fn from_env() -> Result<Self, HostLeaseError> {
487 let root = if let Some(path) = std::env::var_os(HOST_LEASE_ROOT_ENV) {
488 PathBuf::from(path)
489 } else if let Some(path) = std::env::var_os(HARN_HOME_ENV) {
490 PathBuf::from(path).join("host-leases")
491 } else {
492 harn_vm::user_dirs::home_dir()
493 .ok_or_else(|| {
494 HostLeaseError::InvalidRequest(
495 "cannot resolve a home directory; set HARN_HOST_LEASE_ROOT".to_string(),
496 )
497 })?
498 .join(".harn/host-leases")
499 };
500 Self::for_root(root)
501 }
502
503 pub fn for_root(root: impl Into<PathBuf>) -> Result<Self, HostLeaseError> {
505 Self::for_root_with_inspector(root, Arc::new(SystemProcessInspector))
506 }
507
508 fn for_root_with_inspector(
509 root: impl Into<PathBuf>,
510 process_inspector: Arc<dyn ProcessInspector>,
511 ) -> Result<Self, HostLeaseError> {
512 let root = root.into();
513 std::fs::create_dir_all(&root)?;
514 let store = Self {
515 db_path: root.join(LEASE_DB_FILE),
516 root,
517 process_inspector,
518 };
519 store.initialize()?;
520 Ok(store)
521 }
522
523 pub fn root(&self) -> &Path {
525 &self.root
526 }
527
528 pub fn begin_run(
530 &self,
531 owner: &str,
532 priority_class: HostLeasePriorityClass,
533 resource: HostLeaseResourceKey,
534 execution_context: HostLeaseExecutionContext,
535 wait_limit_ms: u64,
536 ) -> Result<HostLeaseRunReceipt, HostLeaseError> {
537 let resource = HostLeaseResourceKey::normalize(
538 &resource.machine,
539 resource.resource_class,
540 &resource.domain,
541 )?;
542 let receipt = HostLeaseRunReceipt {
543 schema_version: RUN_RECEIPT_SCHEMA_VERSION,
544 run_id: Uuid::now_v7().to_string(),
545 owner: normalize_component("owner", owner)?,
546 priority_class,
547 wait_limit_ms,
548 resource,
549 execution_context,
550 status: HostLeaseRunState::Pending {
551 requested_at_ms: unix_now_ms()?,
552 },
553 };
554 let path = self.run_receipt_path(&receipt.run_id)?;
555 let bytes = serde_json::to_vec_pretty(&receipt)?;
556 harn_vm::atomic_io::atomic_write(&path, &bytes)?;
557 Ok(receipt)
558 }
559
560 pub fn load_run(&self, run_id: &str) -> Result<HostLeaseRunReceipt, HostLeaseError> {
562 let path = self.run_receipt_path(run_id)?;
563 Ok(serde_json::from_slice(&std::fs::read(path)?)?)
564 }
565
566 pub fn transition_run(
568 &self,
569 run_id: &str,
570 status: HostLeaseRunState,
571 ) -> Result<HostLeaseRunReceipt, HostLeaseError> {
572 let mut receipt = self.load_run(run_id)?;
573 if !receipt.status.may_transition_to(&status) {
574 return Err(HostLeaseError::InvalidRequest(format!(
575 "invalid run receipt transition from {:?} to {:?}",
576 receipt.status, status
577 )));
578 }
579 receipt.status = status;
580 let path = self.run_receipt_path(run_id)?;
581 let bytes = serde_json::to_vec_pretty(&receipt)?;
582 harn_vm::atomic_io::atomic_write(&path, &bytes)?;
583 Ok(receipt)
584 }
585
586 pub fn run_receipt_path(&self, run_id: &str) -> Result<PathBuf, HostLeaseError> {
588 let run_id = normalize_component("run_id", run_id)?;
589 Ok(self
590 .root
591 .join(RUN_RECEIPTS_DIR)
592 .join(format!("{run_id}.json")))
593 }
594
595 pub fn default_host() -> String {
597 System::host_name()
598 .filter(|name| !name.trim().is_empty())
599 .unwrap_or_else(|| "local".to_string())
600 }
601
602 pub fn try_acquire(
604 &self,
605 request: HostLeaseRequest,
606 ) -> Result<HostLeaseAcquireReceipt, HostLeaseError> {
607 self.try_acquire_once(request, None, None)
608 }
609
610 pub fn acquire_wait(
612 &self,
613 request: HostLeaseRequest,
614 wait_timeout: Duration,
615 ) -> Result<HostLeaseAcquireReceipt, HostLeaseError> {
616 if wait_timeout.is_zero() {
617 return self.try_acquire(request);
618 }
619 let started_at_ms = unix_now_ms()?;
620 let started_at = Instant::now();
621 let deadline = started_at.checked_add(wait_timeout).ok_or_else(|| {
622 HostLeaseError::InvalidRequest("wait timeout exceeds the monotonic clock".to_string())
623 })?;
624 let deadline_at_ms = started_at_ms.saturating_add(duration_ms_i64(wait_timeout));
625 let (tx, rx) = mpsc::channel();
626 let mut watcher = notify::recommended_watcher(move |event| {
627 let _ = tx.send(event);
628 })
629 .map_err(|error| HostLeaseError::Watch(error.to_string()))?;
630 watcher
631 .watch(&self.root, RecursiveMode::NonRecursive)
632 .map_err(|error| HostLeaseError::Watch(error.to_string()))?;
633
634 loop {
635 let receipt =
636 self.try_acquire_once(request.clone(), Some(started_at), Some(deadline_at_ms))?;
637 if receipt.status == HostLeaseAcquireStatus::Acquired || Instant::now() >= deadline {
638 return Ok(receipt);
639 }
640 let wake_at = receipt
641 .defer
642 .as_ref()
643 .and_then(|defer| defer.next_wake_at_ms)
644 .map(|wake| wake.min(deadline_at_ms))
645 .unwrap_or(deadline_at_ms);
646 let wake_duration =
647 Duration::from_millis(wake_at.saturating_sub(receipt.observed_at_ms).max(1) as u64);
648 let remaining = deadline.saturating_duration_since(Instant::now());
649 if remaining.is_zero() {
650 return Ok(receipt);
651 }
652 match rx.recv_timeout(wake_duration.min(remaining)) {
653 Ok(Ok(_)) | Err(mpsc::RecvTimeoutError::Timeout) => {}
654 Ok(Err(error)) => return Err(HostLeaseError::Watch(error.to_string())),
655 Err(mpsc::RecvTimeoutError::Disconnected) => {
656 return Err(HostLeaseError::Watch(
657 "host lease watcher disconnected".to_string(),
658 ));
659 }
660 }
661 }
662 }
663
664 pub fn status(&self, host: &str) -> Result<HostLeaseState, HostLeaseError> {
666 self.status_for_resource(host, HostLeaseResourceClass::WholeMachine)
667 }
668
669 pub fn status_for_resource(
671 &self,
672 host: &str,
673 resource_class: HostLeaseResourceClass,
674 ) -> Result<HostLeaseState, HostLeaseError> {
675 self.status_for_domain(host, resource_class, DEFAULT_HOST_LEASE_DOMAIN)
676 }
677
678 pub fn status_for_domain(
680 &self,
681 host: &str,
682 resource_class: HostLeaseResourceClass,
683 domain: &str,
684 ) -> Result<HostLeaseState, HostLeaseError> {
685 let resource = HostLeaseResourceKey::normalize(host, resource_class, domain)?;
686 let mut conn = self.connection(SQLITE_MUTATION_BUSY_TIMEOUT)?;
687 let tx = conn.transaction_with_behavior(TransactionBehavior::Immediate)?;
688 let now = unix_now_ms()?;
689 self.status_in_transaction(
690 tx,
691 &resource.machine,
692 resource.resource_class,
693 &resource.domain,
694 now,
695 )
696 }
697
698 pub fn renew(
700 &self,
701 host: &str,
702 lease_id: &str,
703 ttl_ms: u64,
704 ) -> Result<HostLeaseRenewReceipt, HostLeaseError> {
705 self.renew_for_resource(host, HostLeaseResourceClass::WholeMachine, lease_id, ttl_ms)
706 }
707
708 pub fn renew_for_resource(
710 &self,
711 host: &str,
712 resource_class: HostLeaseResourceClass,
713 lease_id: &str,
714 ttl_ms: u64,
715 ) -> Result<HostLeaseRenewReceipt, HostLeaseError> {
716 self.renew_for_domain(
717 host,
718 resource_class,
719 DEFAULT_HOST_LEASE_DOMAIN,
720 lease_id,
721 ttl_ms,
722 )
723 }
724
725 pub fn renew_for_domain(
727 &self,
728 host: &str,
729 resource_class: HostLeaseResourceClass,
730 domain: &str,
731 lease_id: &str,
732 ttl_ms: u64,
733 ) -> Result<HostLeaseRenewReceipt, HostLeaseError> {
734 let resource = HostLeaseResourceKey::normalize(host, resource_class, domain)?;
735 let lease_id = normalize_component("lease_id", lease_id)?;
736 validate_ttl(Some(ttl_ms))?;
737 let mut conn = self.connection(SQLITE_MUTATION_BUSY_TIMEOUT)?;
738 let tx = conn.transaction_with_behavior(TransactionBehavior::Immediate)?;
739 let now = unix_now_ms()?;
740 let (active, _) = active_handle(
741 &tx,
742 &resource.machine,
743 resource.resource_class,
744 &resource.domain,
745 now,
746 self.process_inspector.as_ref(),
747 )?;
748 let Some(mut handle) = active.filter(|handle| handle.lease_id == lease_id) else {
749 tx.commit()?;
750 return Ok(HostLeaseRenewReceipt {
751 schema_version: SCHEMA_VERSION,
752 renewed: false,
753 observed_at_ms: now,
754 handle: None,
755 });
756 };
757 handle.updated_at_ms = now;
758 handle.expires_at_ms = Some(now.saturating_add(u64_ms_i64(ttl_ms)));
759 write_handle(&tx, &handle)?;
760 tx.commit()?;
761 Ok(HostLeaseRenewReceipt {
762 schema_version: SCHEMA_VERSION,
763 renewed: true,
764 observed_at_ms: now,
765 handle: Some(handle),
766 })
767 }
768
769 pub fn update_metadata_for_domain(
775 &self,
776 host: &str,
777 resource_class: HostLeaseResourceClass,
778 domain: &str,
779 lease_id: &str,
780 metadata: BTreeMap<String, String>,
781 ) -> Result<HostLeaseMetadataUpdateReceipt, HostLeaseError> {
782 self.update_metadata_at_domain(
783 host,
784 resource_class,
785 domain,
786 lease_id,
787 metadata,
788 unix_now_ms()?,
789 )
790 }
791
792 fn update_metadata_at_domain(
793 &self,
794 host: &str,
795 resource_class: HostLeaseResourceClass,
796 domain: &str,
797 lease_id: &str,
798 metadata: BTreeMap<String, String>,
799 now: i64,
800 ) -> Result<HostLeaseMetadataUpdateReceipt, HostLeaseError> {
801 let resource = HostLeaseResourceKey::normalize(host, resource_class, domain)?;
802 let lease_id = normalize_component("lease_id", lease_id)?;
803 let mut conn = self.connection(SQLITE_MUTATION_BUSY_TIMEOUT)?;
804 let tx = conn.transaction_with_behavior(TransactionBehavior::Immediate)?;
805 let (active, _) = active_handle(
806 &tx,
807 &resource.machine,
808 resource.resource_class,
809 &resource.domain,
810 now,
811 self.process_inspector.as_ref(),
812 )?;
813 let Some(mut handle) = active.filter(|handle| handle.lease_id == lease_id) else {
814 tx.commit()?;
815 return Ok(HostLeaseMetadataUpdateReceipt {
816 schema_version: SCHEMA_VERSION,
817 updated: false,
818 observed_at_ms: now,
819 handle: None,
820 });
821 };
822 handle.updated_at_ms = now;
823 handle.metadata = metadata;
824 write_handle(&tx, &handle)?;
825 tx.commit()?;
826 Ok(HostLeaseMetadataUpdateReceipt {
827 schema_version: SCHEMA_VERSION,
828 updated: true,
829 observed_at_ms: now,
830 handle: Some(handle),
831 })
832 }
833
834 pub fn release(
836 &self,
837 host: &str,
838 lease_id: &str,
839 ) -> Result<HostLeaseReleaseReceipt, HostLeaseError> {
840 self.release_for_resource(host, HostLeaseResourceClass::WholeMachine, lease_id)
841 }
842
843 pub fn release_for_resource(
845 &self,
846 host: &str,
847 resource_class: HostLeaseResourceClass,
848 lease_id: &str,
849 ) -> Result<HostLeaseReleaseReceipt, HostLeaseError> {
850 self.release_for_domain(host, resource_class, DEFAULT_HOST_LEASE_DOMAIN, lease_id)
851 }
852
853 pub fn release_for_domain(
855 &self,
856 host: &str,
857 resource_class: HostLeaseResourceClass,
858 domain: &str,
859 lease_id: &str,
860 ) -> Result<HostLeaseReleaseReceipt, HostLeaseError> {
861 let resource = HostLeaseResourceKey::normalize(host, resource_class, domain)?;
862 let lease_id = normalize_component("lease_id", lease_id)?;
863 let conn = self.connection(SQLITE_MUTATION_BUSY_TIMEOUT)?;
864 let released = conn.execute(
865 "DELETE FROM host_leases
866 WHERE host = ?1 AND resource_class = ?2 AND domain = ?3 AND lease_id = ?4",
867 params![
868 &resource.machine,
869 resource.resource_class.as_str(),
870 &resource.domain,
871 &lease_id,
872 ],
873 )? == 1;
874 let now = unix_now_ms()?;
875 Ok(HostLeaseReleaseReceipt {
876 schema_version: SCHEMA_VERSION,
877 released,
878 host: resource.machine,
879 resource_class: resource.resource_class,
880 domain: resource.domain,
881 lease_id,
882 observed_at_ms: now,
883 })
884 }
885
886 fn initialize(&self) -> Result<(), HostLeaseError> {
887 let mut conn = self.connection(SQLITE_MUTATION_BUSY_TIMEOUT)?;
888 let tx = conn.transaction_with_behavior(TransactionBehavior::Immediate)?;
889 match lease_table_layout(&tx)? {
890 LeaseTableLayout::Missing => create_current_lease_table(&tx)?,
891 LeaseTableLayout::LegacyWholeMachine => migrate_legacy_lease_table(&tx)?,
892 LeaseTableLayout::ResourceClassWithoutExecutionContext => {
893 add_execution_context_column(&tx)?;
894 add_domain_key(&tx)?;
895 }
896 LeaseTableLayout::CurrentWithoutDomain => add_domain_key(&tx)?,
897 LeaseTableLayout::Current => {}
898 }
899 tx.commit()?;
900 Ok(())
901 }
902
903 fn connection(&self, busy_timeout: Duration) -> Result<Connection, HostLeaseError> {
904 let conn = Connection::open(&self.db_path)?;
905 conn.busy_timeout(busy_timeout)?;
906 Ok(conn)
907 }
908
909 fn try_acquire_once(
910 &self,
911 request: HostLeaseRequest,
912 started_at: Option<Instant>,
913 deadline_at_ms: Option<i64>,
914 ) -> Result<HostLeaseAcquireReceipt, HostLeaseError> {
915 let request = normalize_request(request)?;
916 let mut conn = self.connection(Duration::ZERO)?;
917 let tx = match conn.transaction_with_behavior(TransactionBehavior::Immediate) {
918 Ok(tx) => tx,
919 Err(error) if sqlite_is_busy(&error) => {
920 let now = unix_now_ms()?;
921 return Ok(registry_busy_receipt(
922 request.host,
923 request.resource_class,
924 request.domain,
925 now,
926 started_at.map(|started| duration_ms_u64(started.elapsed())),
927 deadline_at_ms,
928 ));
929 }
930 Err(error) => return Err(error.into()),
931 };
932 let now = unix_now_ms()?;
933 let waited_ms = started_at
934 .map(|started| duration_ms_u64(started.elapsed()))
935 .unwrap_or(0);
936 let host = request.host.clone();
937 let resource_class = request.resource_class;
938 let domain = request.domain.clone();
939 match self.acquire_in_transaction(tx, request, now, deadline_at_ms, waited_ms) {
940 Err(HostLeaseError::Database(error)) if sqlite_is_busy(&error) => {
941 Ok(registry_busy_receipt(
942 host,
943 resource_class,
944 domain,
945 now,
946 Some(waited_ms),
947 deadline_at_ms,
948 ))
949 }
950 result => result,
951 }
952 }
953
954 #[cfg(test)]
955 fn try_acquire_at(
956 &self,
957 request: HostLeaseRequest,
958 now: i64,
959 deadline_at_ms: Option<i64>,
960 waited_ms: u64,
961 ) -> Result<HostLeaseAcquireReceipt, HostLeaseError> {
962 let request = normalize_request(request)?;
963 let mut conn = self.connection(SQLITE_MUTATION_BUSY_TIMEOUT)?;
964 let tx = conn.transaction_with_behavior(TransactionBehavior::Immediate)?;
965 self.acquire_in_transaction(tx, request, now, deadline_at_ms, waited_ms)
966 }
967
968 fn acquire_in_transaction(
969 &self,
970 tx: Transaction<'_>,
971 request: HostLeaseRequest,
972 now: i64,
973 deadline_at_ms: Option<i64>,
974 waited_ms: u64,
975 ) -> Result<HostLeaseAcquireReceipt, HostLeaseError> {
976 let owner_process_identity = request
977 .owner_pid
978 .map(|pid| match self.process_inspector.observe(pid) {
979 ProcessObservation::Alive { identity } => Ok(identity),
980 ProcessObservation::Dead => Err(HostLeaseError::InvalidRequest(
981 "owner_pid is not a live local process".to_string(),
982 )),
983 ProcessObservation::Unknown => Err(HostLeaseError::InvalidRequest(
984 "owner_pid liveness could not be verified".to_string(),
985 )),
986 })
987 .transpose()?;
988 let (active, recovered) = active_handle(
989 &tx,
990 &request.host,
991 request.resource_class,
992 &request.domain,
993 now,
994 self.process_inspector.as_ref(),
995 )?;
996 if let Some(active) = active {
997 let defer = HostLeaseDeferReceipt {
998 host: request.host,
999 resource_class: request.resource_class,
1000 domain: request.domain,
1001 deferred_reason: HostLeaseDeferReason::Contended,
1002 observed_at_ms: now,
1003 next_wake_at_ms: Some(next_lease_wake_at(&active, now, deadline_at_ms)),
1004 deadline_at_ms,
1005 active: Some(active),
1006 };
1007 tx.commit()?;
1008 return Ok(HostLeaseAcquireReceipt {
1009 schema_version: SCHEMA_VERSION,
1010 status: HostLeaseAcquireStatus::Deferred,
1011 observed_at_ms: now,
1012 waited_ms,
1013 handle: None,
1014 defer: Some(defer),
1015 recovered_stale_lease: recovered.is_some(),
1016 recovered,
1017 });
1018 }
1019
1020 let handle = HostLeaseHandle {
1021 schema_version: SCHEMA_VERSION,
1022 host: request.host,
1023 resource_class: request.resource_class,
1024 domain: request.domain,
1025 execution_context: request.execution_context,
1026 lease_id: Uuid::now_v7().to_string(),
1027 owner: request.owner,
1028 priority_class: request.priority_class,
1029 acquired_at_ms: now,
1030 updated_at_ms: now,
1031 expires_at_ms: request
1032 .ttl_ms
1033 .map(|ttl| now.saturating_add(u64_ms_i64(ttl))),
1034 owner_pid: request.owner_pid,
1035 owner_process_identity,
1036 reason: request.reason,
1037 metadata: request.metadata,
1038 };
1039 write_handle(&tx, &handle)?;
1040 tx.commit()?;
1041 Ok(HostLeaseAcquireReceipt {
1042 schema_version: SCHEMA_VERSION,
1043 status: HostLeaseAcquireStatus::Acquired,
1044 observed_at_ms: now,
1045 waited_ms,
1046 handle: Some(handle),
1047 defer: None,
1048 recovered_stale_lease: recovered.is_some(),
1049 recovered,
1050 })
1051 }
1052
1053 #[cfg(test)]
1054 fn status_at(
1055 &self,
1056 host: &str,
1057 resource_class: HostLeaseResourceClass,
1058 now: i64,
1059 ) -> Result<HostLeaseState, HostLeaseError> {
1060 self.status_at_domain(host, resource_class, DEFAULT_HOST_LEASE_DOMAIN, now)
1061 }
1062
1063 #[cfg(test)]
1064 fn status_at_domain(
1065 &self,
1066 host: &str,
1067 resource_class: HostLeaseResourceClass,
1068 domain: &str,
1069 now: i64,
1070 ) -> Result<HostLeaseState, HostLeaseError> {
1071 let mut conn = self.connection(SQLITE_MUTATION_BUSY_TIMEOUT)?;
1072 let tx = conn.transaction_with_behavior(TransactionBehavior::Immediate)?;
1073 self.status_in_transaction(tx, host, resource_class, domain, now)
1074 }
1075
1076 fn status_in_transaction(
1077 &self,
1078 tx: Transaction<'_>,
1079 host: &str,
1080 resource_class: HostLeaseResourceClass,
1081 domain: &str,
1082 now: i64,
1083 ) -> Result<HostLeaseState, HostLeaseError> {
1084 let (active, recovered) = active_handle(
1085 &tx,
1086 host,
1087 resource_class,
1088 domain,
1089 now,
1090 self.process_inspector.as_ref(),
1091 )?;
1092 tx.commit()?;
1093 Ok(HostLeaseState {
1094 schema_version: SCHEMA_VERSION,
1095 host: host.to_string(),
1096 resource_class,
1097 domain: domain.to_string(),
1098 observed_at_ms: now,
1099 active,
1100 recovered_stale_lease: recovered.is_some(),
1101 recovered,
1102 })
1103 }
1104}
1105
1106fn normalize_request(mut request: HostLeaseRequest) -> Result<HostLeaseRequest, HostLeaseError> {
1107 let resource =
1108 HostLeaseResourceKey::normalize(&request.host, request.resource_class, &request.domain)?;
1109 request.host = resource.machine;
1110 request.resource_class = resource.resource_class;
1111 request.domain = resource.domain;
1112 request.owner = normalize_component("owner", &request.owner)?;
1113 validate_ttl(request.ttl_ms)?;
1114 if request.ttl_ms.is_none() && request.owner_pid.is_none() {
1115 return Err(HostLeaseError::InvalidRequest(
1116 "a non-expiring lease requires owner_pid for crash recovery".to_string(),
1117 ));
1118 }
1119 request.reason = request.reason.and_then(|reason| {
1120 let trimmed = reason.trim();
1121 (!trimmed.is_empty()).then(|| trimmed.to_string())
1122 });
1123 Ok(request)
1124}
1125
1126fn normalize_component(name: &str, value: &str) -> Result<String, HostLeaseError> {
1127 let value = value.trim();
1128 if value.is_empty() {
1129 return Err(HostLeaseError::InvalidRequest(format!(
1130 "{name} cannot be empty"
1131 )));
1132 }
1133 if !value
1134 .chars()
1135 .all(|ch| ch.is_ascii_alphanumeric() || matches!(ch, '.' | '_' | '-'))
1136 {
1137 return Err(HostLeaseError::InvalidRequest(format!(
1138 "{name} may contain only ASCII letters, digits, '.', '_' and '-'"
1139 )));
1140 }
1141 Ok(value.to_string())
1142}
1143
1144fn normalize_domain(value: &str) -> Result<String, HostLeaseError> {
1145 let value = value.trim();
1146 if value.is_empty() {
1147 return Err(HostLeaseError::InvalidRequest(
1148 "domain cannot be empty".to_string(),
1149 ));
1150 }
1151 if value.len() > 128 {
1152 return Err(HostLeaseError::InvalidRequest(
1153 "domain cannot exceed 128 bytes".to_string(),
1154 ));
1155 }
1156 if matches!(value, "." | "..")
1157 || !value
1158 .chars()
1159 .all(|ch| ch.is_ascii_alphanumeric() || matches!(ch, '.' | '_' | '-'))
1160 {
1161 return Err(HostLeaseError::InvalidRequest(
1162 "domain may contain only ASCII letters, digits, '.', '_' and '-' and must not be path-like"
1163 .to_string(),
1164 ));
1165 }
1166 Ok(value.to_string())
1167}
1168
1169fn validate_ttl(ttl_ms: Option<u64>) -> Result<(), HostLeaseError> {
1170 if ttl_ms == Some(0) {
1171 return Err(HostLeaseError::InvalidRequest(
1172 "ttl_ms must be greater than zero when supplied".to_string(),
1173 ));
1174 }
1175 Ok(())
1176}
1177
1178fn active_handle(
1179 tx: &Transaction<'_>,
1180 host: &str,
1181 resource_class: HostLeaseResourceClass,
1182 domain: &str,
1183 now: i64,
1184 process_inspector: &dyn ProcessInspector,
1185) -> Result<(Option<HostLeaseHandle>, Option<HostLeaseHandle>), HostLeaseError> {
1186 let handle = read_handle(tx, host, resource_class, domain)?;
1187 let Some(handle) = handle else {
1188 return Ok((None, None));
1189 };
1190 let expired = handle.expires_at_ms.is_some_and(|expiry| expiry <= now);
1191 let owner_dead = match (handle.owner_pid, handle.owner_process_identity) {
1192 (Some(pid), Some(expected_identity)) => match process_inspector.observe(pid) {
1193 ProcessObservation::Alive { identity } => identity != expected_identity,
1194 ProcessObservation::Dead => true,
1195 ProcessObservation::Unknown => false,
1196 },
1197 _ => false,
1198 };
1199 if expired || owner_dead {
1200 tx.execute(
1201 "DELETE FROM host_leases
1202 WHERE host = ?1 AND resource_class = ?2 AND domain = ?3 AND lease_id = ?4",
1203 params![host, resource_class.as_str(), domain, handle.lease_id],
1204 )?;
1205 return Ok((None, Some(handle)));
1206 }
1207 Ok((Some(handle), None))
1208}
1209
1210fn next_lease_wake_at(active: &HostLeaseHandle, now: i64, deadline_at_ms: Option<i64>) -> i64 {
1211 let mut wake_at = deadline_at_ms.unwrap_or(i64::MAX);
1212 if let Some(expiry) = active.expires_at_ms {
1213 wake_at = wake_at.min(expiry);
1214 }
1215 if active.owner_pid.is_some() {
1216 wake_at =
1217 wake_at.min(now.saturating_add(duration_ms_i64(PROCESS_LIVENESS_RECHECK_INTERVAL)));
1218 }
1219 wake_at
1220}
1221
1222fn registry_busy_receipt(
1223 host: String,
1224 resource_class: HostLeaseResourceClass,
1225 domain: String,
1226 now: i64,
1227 waited_ms: Option<u64>,
1228 deadline_at_ms: Option<i64>,
1229) -> HostLeaseAcquireReceipt {
1230 let next_wake_at_ms = now
1231 .saturating_add(duration_ms_i64(REGISTRY_BUSY_RETRY_INTERVAL))
1232 .min(deadline_at_ms.unwrap_or(i64::MAX));
1233 HostLeaseAcquireReceipt {
1234 schema_version: SCHEMA_VERSION,
1235 status: HostLeaseAcquireStatus::Deferred,
1236 observed_at_ms: now,
1237 waited_ms: waited_ms.unwrap_or(0),
1238 handle: None,
1239 defer: Some(HostLeaseDeferReceipt {
1240 host,
1241 resource_class,
1242 domain,
1243 deferred_reason: HostLeaseDeferReason::RegistryBusy,
1244 observed_at_ms: now,
1245 next_wake_at_ms: Some(next_wake_at_ms),
1246 deadline_at_ms,
1247 active: None,
1248 }),
1249 recovered_stale_lease: false,
1250 recovered: None,
1251 }
1252}
1253
1254fn sqlite_is_busy(error: &rusqlite::Error) -> bool {
1255 matches!(
1256 error,
1257 rusqlite::Error::SqliteFailure(inner, _)
1258 if matches!(inner.code, ErrorCode::DatabaseBusy | ErrorCode::DatabaseLocked)
1259 )
1260}
1261
1262fn read_handle(
1263 tx: &Transaction<'_>,
1264 host: &str,
1265 resource_class: HostLeaseResourceClass,
1266 domain: &str,
1267) -> Result<Option<HostLeaseHandle>, HostLeaseError> {
1268 tx.query_row(
1269 "SELECT resource_class, domain, lease_id, owner, priority_class, acquired_at_ms, updated_at_ms,
1270 expires_at_ms, owner_pid, owner_process_identity, reason, metadata_json,
1271 execution_context_json
1272 FROM host_leases WHERE host = ?1 AND resource_class = ?2 AND domain = ?3",
1273 params![host, resource_class.as_str(), domain],
1274 |row| {
1275 let priority: String = row.get(4)?;
1276 let metadata_json: String = row.get(11)?;
1277 let execution_context_json: Option<String> = row.get(12)?;
1278 let owner_pid_i64: Option<i64> = row.get(8)?;
1279 let owner_identity_i64: Option<i64> = row.get(9)?;
1280 Ok((
1281 row.get::<_, String>(2)?,
1282 row.get::<_, String>(3)?,
1283 priority,
1284 row.get::<_, i64>(5)?,
1285 row.get::<_, i64>(6)?,
1286 row.get::<_, Option<i64>>(7)?,
1287 owner_pid_i64,
1288 owner_identity_i64,
1289 row.get::<_, Option<String>>(10)?,
1290 metadata_json,
1291 row.get::<_, String>(0)?,
1292 row.get::<_, String>(1)?,
1293 execution_context_json,
1294 ))
1295 },
1296 )
1297 .optional()?
1298 .map(
1299 |(
1300 lease_id,
1301 owner,
1302 priority,
1303 acquired_at_ms,
1304 updated_at_ms,
1305 expires_at_ms,
1306 owner_pid,
1307 owner_process_identity,
1308 reason,
1309 metadata_json,
1310 stored_resource_class,
1311 stored_domain,
1312 execution_context_json,
1313 )| {
1314 let owner_pid = owner_pid
1315 .map(|pid| {
1316 u32::try_from(pid).map_err(|_| {
1317 HostLeaseError::InvalidRequest(
1318 "persisted owner_pid is outside the u32 range".to_string(),
1319 )
1320 })
1321 })
1322 .transpose()?;
1323 let owner_process_identity = owner_process_identity
1324 .map(|identity| {
1325 u64::try_from(identity).map_err(|_| {
1326 HostLeaseError::InvalidRequest(
1327 "persisted process identity is negative".to_string(),
1328 )
1329 })
1330 })
1331 .transpose()?;
1332 Ok(HostLeaseHandle {
1333 schema_version: SCHEMA_VERSION,
1334 host: host.to_string(),
1335 resource_class: HostLeaseResourceClass::parse(&stored_resource_class)?,
1336 domain: stored_domain,
1337 execution_context: execution_context_json
1338 .map(|encoded| serde_json::from_str(&encoded))
1339 .transpose()?,
1340 lease_id,
1341 owner,
1342 priority_class: HostLeasePriorityClass::parse(&priority)?,
1343 acquired_at_ms,
1344 updated_at_ms,
1345 expires_at_ms,
1346 owner_pid,
1347 owner_process_identity,
1348 reason,
1349 metadata: serde_json::from_str(&metadata_json)?,
1350 })
1351 },
1352 )
1353 .transpose()
1354}
1355
1356fn write_handle(tx: &Transaction<'_>, handle: &HostLeaseHandle) -> Result<(), HostLeaseError> {
1357 let metadata_json = serde_json::to_string(&handle.metadata)?;
1358 let execution_context_json = handle
1359 .execution_context
1360 .as_ref()
1361 .map(serde_json::to_string)
1362 .transpose()?;
1363 let owner_process_identity = handle
1364 .owner_process_identity
1365 .map(|value| {
1366 i64::try_from(value).map_err(|_| {
1367 HostLeaseError::InvalidRequest(
1368 "owner process identity is outside the SQLite integer range".to_string(),
1369 )
1370 })
1371 })
1372 .transpose()?;
1373 tx.execute(
1374 "INSERT INTO host_leases (
1375 host, resource_class, domain, lease_id, owner, priority_class, acquired_at_ms, updated_at_ms,
1376 expires_at_ms, owner_pid, owner_process_identity, reason, metadata_json,
1377 execution_context_json
1378 ) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14)
1379 ON CONFLICT(host, resource_class, domain) DO UPDATE SET
1380 lease_id = excluded.lease_id,
1381 owner = excluded.owner,
1382 priority_class = excluded.priority_class,
1383 acquired_at_ms = excluded.acquired_at_ms,
1384 updated_at_ms = excluded.updated_at_ms,
1385 expires_at_ms = excluded.expires_at_ms,
1386 owner_pid = excluded.owner_pid,
1387 owner_process_identity = excluded.owner_process_identity,
1388 reason = excluded.reason,
1389 metadata_json = excluded.metadata_json,
1390 execution_context_json = excluded.execution_context_json",
1391 params![
1392 handle.host,
1393 handle.resource_class.as_str(),
1394 handle.domain,
1395 handle.lease_id,
1396 handle.owner,
1397 handle.priority_class.as_str(),
1398 handle.acquired_at_ms,
1399 handle.updated_at_ms,
1400 handle.expires_at_ms,
1401 handle.owner_pid.map(i64::from),
1402 owner_process_identity,
1403 handle.reason,
1404 metadata_json,
1405 execution_context_json,
1406 ],
1407 )?;
1408 Ok(())
1409}
1410
1411fn unix_now_ms() -> Result<i64, HostLeaseError> {
1412 let millis = SystemTime::now()
1413 .duration_since(UNIX_EPOCH)
1414 .map_err(|_| HostLeaseError::Clock)?
1415 .as_millis();
1416 Ok(millis.min(i64::MAX as u128) as i64)
1417}
1418
1419fn duration_ms_i64(duration: Duration) -> i64 {
1420 duration.as_millis().min(i64::MAX as u128) as i64
1421}
1422
1423fn duration_ms_u64(duration: Duration) -> u64 {
1424 duration.as_millis().min(u64::MAX as u128) as u64
1425}
1426
1427fn u64_ms_i64(value: u64) -> i64 {
1428 value.min(i64::MAX as u64) as i64
1429}
1430
1431#[cfg(test)]
1432#[path = "host_lease/tests.rs"]
1433mod tests;