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