1use std::collections::BTreeMap;
10use std::time::Duration;
11
12use harn_vm::{VmResourceGuardHandle, VmValue};
13
14use crate::error::HostlibError;
15use crate::host_lease::{
16 HostLeaseAcquireReceipt, HostLeaseAcquireStatus, HostLeaseDeferReceipt, HostLeaseHandle,
17 HostLeaseMetadataUpdateReceipt, HostLeasePriorityClass, HostLeaseReleaseReceipt,
18 HostLeaseRequest, HostLeaseResourceClass, HostLeaseState, HostLeaseStore,
19 DEFAULT_HOST_LEASE_DOMAIN,
20};
21use crate::registry::{BuiltinRegistry, HostlibCapability};
22use crate::tools::args::{
23 build_dict, dict_arg, optional_int, optional_string, require_string, str_value,
24};
25
26const STATUS_BUILTIN: &str = "hostlib_host_lease_status";
27const ACQUIRE_BUILTIN: &str = "hostlib_host_lease_acquire";
28const UPDATE_METADATA_BUILTIN: &str = "hostlib_host_lease_update_metadata";
29const RELEASE_BUILTIN: &str = "hostlib_host_lease_release";
30const MAX_WAIT_SLICE_MS: i64 = 5_000;
31
32#[derive(Default)]
39pub struct HostLeaseCapability;
40
41impl HostlibCapability for HostLeaseCapability {
42 fn module_name(&self) -> &'static str {
43 "host_lease"
44 }
45
46 fn register_builtins(&self, registry: &mut BuiltinRegistry) {
47 registry.register_fn("host_lease", STATUS_BUILTIN, "status", handle_status);
48 registry.register_fn("host_lease", ACQUIRE_BUILTIN, "acquire", handle_acquire);
49 registry.register_fn(
50 "host_lease",
51 UPDATE_METADATA_BUILTIN,
52 "update_metadata",
53 handle_update_metadata,
54 );
55 registry.register_fn("host_lease", RELEASE_BUILTIN, "release", handle_release);
56 }
57}
58
59fn handle_update_metadata(args: &[VmValue]) -> Result<VmValue, HostlibError> {
60 let dict = dict_arg(UPDATE_METADATA_BUILTIN, args)?;
61 let host = require_nonempty_string(UPDATE_METADATA_BUILTIN, &dict, "host")?;
62 let resource_class = resource_class(UPDATE_METADATA_BUILTIN, dict.get("resource_class"))?;
63 let domain = optional_string(UPDATE_METADATA_BUILTIN, &dict, "domain")?
64 .unwrap_or_else(|| DEFAULT_HOST_LEASE_DOMAIN.to_string());
65 let lease_id = require_nonempty_string(UPDATE_METADATA_BUILTIN, &dict, "lease_id")?;
66 let metadata = string_map(UPDATE_METADATA_BUILTIN, dict.get("metadata"))?;
67 let receipt = HostLeaseStore::from_env()
68 .and_then(|store| {
69 store.update_metadata_for_domain(&host, resource_class, &domain, &lease_id, metadata)
70 })
71 .map_err(|error| backend(UPDATE_METADATA_BUILTIN, error))?;
72 metadata_update_to_value(&receipt)
73}
74
75fn resolve_host(
86 builtin: &'static str,
87 dict: &harn_vm::value::DictMap,
88) -> Result<String, HostlibError> {
89 match optional_string(builtin, dict, "host")? {
90 Some(host) if !host.trim().is_empty() => Ok(host),
91 Some(_) => Err(HostlibError::InvalidParameter {
92 builtin,
93 param: "host",
94 message: "must be a non-empty string when provided".to_string(),
95 }),
96 None => Ok(HostLeaseStore::default_host()),
97 }
98}
99
100fn handle_status(args: &[VmValue]) -> Result<VmValue, HostlibError> {
101 let dict = dict_arg(STATUS_BUILTIN, args)?;
102 let host = resolve_host(STATUS_BUILTIN, &dict)?;
103 let resource_class = resource_class(STATUS_BUILTIN, dict.get("resource_class"))?;
104 let domain = optional_string(STATUS_BUILTIN, &dict, "domain")?
105 .unwrap_or_else(|| DEFAULT_HOST_LEASE_DOMAIN.to_string());
106 let state = HostLeaseStore::from_env()
107 .and_then(|store| store.status_for_domain(&host, resource_class, &domain))
108 .map_err(|error| HostlibError::Backend {
109 builtin: STATUS_BUILTIN,
110 message: error.to_string(),
111 })?;
112 state_to_value(&state)
113}
114
115fn handle_acquire(args: &[VmValue]) -> Result<VmValue, HostlibError> {
116 let dict = dict_arg(ACQUIRE_BUILTIN, args)?;
117 let host = resolve_host(ACQUIRE_BUILTIN, &dict)?;
118 let owner = require_nonempty_string(ACQUIRE_BUILTIN, &dict, "owner")?;
119 let resource_class = resource_class(ACQUIRE_BUILTIN, dict.get("resource_class"))?;
120 let domain = optional_string(ACQUIRE_BUILTIN, &dict, "domain")?
121 .unwrap_or_else(|| DEFAULT_HOST_LEASE_DOMAIN.to_string());
122 let priority_class = priority_class(ACQUIRE_BUILTIN, dict.get("priority_class"))?;
123 let ttl_ms = optional_positive_u64(ACQUIRE_BUILTIN, &dict, "ttl_ms")?;
124 let wait_slice_ms = optional_int(ACQUIRE_BUILTIN, &dict, "wait_slice_ms", 0)?;
125 if !(0..=MAX_WAIT_SLICE_MS).contains(&wait_slice_ms) {
126 return Err(HostlibError::InvalidParameter {
127 builtin: ACQUIRE_BUILTIN,
128 param: "wait_slice_ms",
129 message: format!("must be between 0 and {MAX_WAIT_SLICE_MS}"),
130 });
131 }
132 let reason = optional_string(ACQUIRE_BUILTIN, &dict, "reason")?;
133 let metadata = string_map(ACQUIRE_BUILTIN, dict.get("metadata"))?;
134 let request = HostLeaseRequest {
135 host,
136 resource_class,
137 domain,
138 execution_context: None,
139 owner,
140 priority_class,
141 ttl_ms,
142 owner_pid: Some(std::process::id()),
143 reason,
144 metadata,
145 };
146 let store = HostLeaseStore::from_env().map_err(|error| backend(ACQUIRE_BUILTIN, error))?;
147 let receipt = if wait_slice_ms == 0 {
148 store.try_acquire(request)
149 } else {
150 store.acquire_wait(request, Duration::from_millis(wait_slice_ms as u64))
151 }
152 .map_err(|error| backend(ACQUIRE_BUILTIN, error))?;
153 acquire_to_value(store, receipt)
154}
155
156fn handle_release(args: &[VmValue]) -> Result<VmValue, HostlibError> {
157 let dict = dict_arg(RELEASE_BUILTIN, args)?;
158 let Some(VmValue::ResourceGuard(guard)) = dict.get("guard") else {
159 return Err(HostlibError::InvalidParameter {
160 builtin: RELEASE_BUILTIN,
161 param: "guard",
162 message: "must be a resource_guard returned by host_lease.acquire".to_string(),
163 });
164 };
165 guard
166 .release()
167 .map_err(|error| backend(RELEASE_BUILTIN, error))
168}
169
170fn backend(builtin: &'static str, error: impl std::fmt::Display) -> HostlibError {
171 HostlibError::Backend {
172 builtin,
173 message: error.to_string(),
174 }
175}
176
177fn resource_class(
178 builtin: &'static str,
179 value: Option<&VmValue>,
180) -> Result<HostLeaseResourceClass, HostlibError> {
181 match value {
182 None | Some(VmValue::Nil) => Ok(HostLeaseResourceClass::WholeMachine),
183 Some(VmValue::String(value)) if value.as_str() == "whole-machine" => {
184 Ok(HostLeaseResourceClass::WholeMachine)
185 }
186 Some(VmValue::String(value)) if value.as_str() == "rust-heavy" => {
187 Ok(HostLeaseResourceClass::RustHeavy)
188 }
189 _ => Err(HostlibError::InvalidParameter {
190 builtin,
191 param: "resource_class",
192 message: "must be whole-machine or rust-heavy".to_string(),
193 }),
194 }
195}
196
197fn priority_class(
198 builtin: &'static str,
199 value: Option<&VmValue>,
200) -> Result<HostLeasePriorityClass, HostlibError> {
201 match value {
202 None | Some(VmValue::Nil) => Ok(HostLeasePriorityClass::Deferrable),
203 Some(VmValue::String(value)) => match value.as_str() {
204 "interactive" => Ok(HostLeasePriorityClass::Interactive),
205 "measurement" => Ok(HostLeasePriorityClass::Measurement),
206 "ci-verify" => Ok(HostLeasePriorityClass::CiVerify),
207 "deferrable" => Ok(HostLeasePriorityClass::Deferrable),
208 _ => Err(HostlibError::InvalidParameter {
209 builtin,
210 param: "priority_class",
211 message: "must be interactive, measurement, ci-verify, or deferrable".to_string(),
212 }),
213 },
214 _ => Err(HostlibError::InvalidParameter {
215 builtin,
216 param: "priority_class",
217 message: "must be a string".to_string(),
218 }),
219 }
220}
221
222fn optional_positive_u64(
223 builtin: &'static str,
224 dict: &harn_vm::value::DictMap,
225 key: &'static str,
226) -> Result<Option<u64>, HostlibError> {
227 match dict.get(key) {
228 None | Some(VmValue::Nil) => Ok(None),
229 Some(VmValue::Int(value)) if *value > 0 => Ok(Some(*value as u64)),
230 _ => Err(HostlibError::InvalidParameter {
231 builtin,
232 param: key,
233 message: "must be a positive integer or nil".to_string(),
234 }),
235 }
236}
237
238fn string_map(
239 builtin: &'static str,
240 value: Option<&VmValue>,
241) -> Result<BTreeMap<String, String>, HostlibError> {
242 let Some(value) = value else {
243 return Ok(BTreeMap::new());
244 };
245 if matches!(value, VmValue::Nil) {
246 return Ok(BTreeMap::new());
247 }
248 let VmValue::Dict(entries) = value else {
249 return Err(HostlibError::InvalidParameter {
250 builtin,
251 param: "metadata",
252 message: "must be a dict of string values".to_string(),
253 });
254 };
255 entries
256 .iter()
257 .map(|(key, value)| match value {
258 VmValue::String(value) => Ok((key.to_string(), value.to_string())),
259 _ => Err(HostlibError::InvalidParameter {
260 builtin,
261 param: "metadata",
262 message: "must contain only string values".to_string(),
263 }),
264 })
265 .collect()
266}
267
268fn acquire_to_value(
269 store: HostLeaseStore,
270 receipt: HostLeaseAcquireReceipt,
271) -> Result<VmValue, HostlibError> {
272 let handle = match receipt.handle.as_ref() {
273 Some(handle) => handle_to_value(ACQUIRE_BUILTIN, handle)?,
274 None => VmValue::Nil,
275 };
276 let defer = receipt
277 .defer
278 .as_ref()
279 .map(defer_to_value)
280 .transpose()?
281 .unwrap_or(VmValue::Nil);
282 let guard = if receipt.status == HostLeaseAcquireStatus::Acquired {
283 let handle = receipt
284 .handle
285 .as_ref()
286 .ok_or_else(|| HostlibError::Backend {
287 builtin: ACQUIRE_BUILTIN,
288 message: "acquired receipt omitted its lease handle".to_string(),
289 })?;
290 let host = handle.host.clone();
291 let resource_class = handle.resource_class;
292 let domain = handle.domain.clone();
293 let lease_id = handle.lease_id.clone();
294 VmValue::resource_guard(VmResourceGuardHandle::new("host_lease", move || {
295 store
296 .release_for_domain(&host, resource_class, &domain, &lease_id)
297 .map(|receipt| release_to_value(&receipt))
298 .map_err(|error| error.to_string())
299 }))
300 } else {
301 VmValue::Nil
302 };
303 Ok(build_dict([
304 (
305 "schema_version",
306 VmValue::Int(i64::from(receipt.schema_version)),
307 ),
308 (
309 "status",
310 str_value(match receipt.status {
311 HostLeaseAcquireStatus::Acquired => "acquired",
312 HostLeaseAcquireStatus::Deferred => "deferred",
313 }),
314 ),
315 ("observed_at_ms", VmValue::Int(receipt.observed_at_ms)),
316 ("waited_ms", VmValue::Int(receipt.waited_ms as i64)),
317 ("handle", handle),
318 ("defer", defer),
319 ("guard", guard),
320 (
321 "recovered_stale_lease",
322 VmValue::Bool(receipt.recovered_stale_lease),
323 ),
324 (
325 "recovered",
326 receipt
327 .recovered
328 .as_ref()
329 .map(|handle| handle_to_value(ACQUIRE_BUILTIN, handle))
330 .transpose()?
331 .unwrap_or(VmValue::Nil),
332 ),
333 ]))
334}
335
336fn defer_to_value(defer: &HostLeaseDeferReceipt) -> Result<VmValue, HostlibError> {
337 Ok(build_dict([
338 ("host", str_value(&defer.host)),
339 ("resource_class", str_value(defer.resource_class.as_str())),
340 ("domain", str_value(&defer.domain)),
341 ("deferred_reason", str_value(defer.deferred_reason.as_str())),
342 ("observed_at_ms", VmValue::Int(defer.observed_at_ms)),
343 (
344 "next_wake_at_ms",
345 defer
346 .next_wake_at_ms
347 .map(VmValue::Int)
348 .unwrap_or(VmValue::Nil),
349 ),
350 (
351 "deadline_at_ms",
352 defer
353 .deadline_at_ms
354 .map(VmValue::Int)
355 .unwrap_or(VmValue::Nil),
356 ),
357 (
358 "active",
359 defer
360 .active
361 .as_ref()
362 .map(|handle| handle_to_value(ACQUIRE_BUILTIN, handle))
363 .transpose()?
364 .unwrap_or(VmValue::Nil),
365 ),
366 ]))
367}
368
369fn release_to_value(receipt: &HostLeaseReleaseReceipt) -> VmValue {
370 build_dict([
371 (
372 "schema_version",
373 VmValue::Int(i64::from(receipt.schema_version)),
374 ),
375 ("released", VmValue::Bool(receipt.released)),
376 ("host", str_value(&receipt.host)),
377 ("resource_class", str_value(receipt.resource_class.as_str())),
378 ("domain", str_value(&receipt.domain)),
379 ("lease_id", str_value(&receipt.lease_id)),
380 ("observed_at_ms", VmValue::Int(receipt.observed_at_ms)),
381 ])
382}
383
384fn metadata_update_to_value(
385 receipt: &HostLeaseMetadataUpdateReceipt,
386) -> Result<VmValue, HostlibError> {
387 Ok(build_dict([
388 (
389 "schema_version",
390 VmValue::Int(i64::from(receipt.schema_version)),
391 ),
392 ("updated", VmValue::Bool(receipt.updated)),
393 ("observed_at_ms", VmValue::Int(receipt.observed_at_ms)),
394 (
395 "handle",
396 receipt
397 .handle
398 .as_ref()
399 .map(|handle| handle_to_value(UPDATE_METADATA_BUILTIN, handle))
400 .transpose()?
401 .unwrap_or(VmValue::Nil),
402 ),
403 ]))
404}
405
406fn require_nonempty_string(
407 builtin: &'static str,
408 dict: &harn_vm::value::DictMap,
409 key: &'static str,
410) -> Result<String, HostlibError> {
411 let value = require_string(builtin, dict, key)?;
412 if value.trim().is_empty() {
413 return Err(HostlibError::InvalidParameter {
414 builtin,
415 param: key,
416 message: "must be a non-empty string".to_string(),
417 });
418 }
419 Ok(value)
420}
421
422fn state_to_value(state: &HostLeaseState) -> Result<VmValue, HostlibError> {
423 let active = match state.active.as_ref() {
424 Some(handle) => handle_to_value(STATUS_BUILTIN, handle)?,
425 None => VmValue::Nil,
426 };
427 Ok(build_dict([
428 (
429 "schema_version",
430 VmValue::Int(i64::from(state.schema_version)),
431 ),
432 ("host", str_value(&state.host)),
433 ("resource_class", str_value(state.resource_class.as_str())),
434 ("domain", str_value(&state.domain)),
435 ("observed_at_ms", VmValue::Int(state.observed_at_ms)),
436 ("active", active),
437 (
438 "recovered_stale_lease",
439 VmValue::Bool(state.recovered_stale_lease),
440 ),
441 (
442 "recovered",
443 state
444 .recovered
445 .as_ref()
446 .map(|handle| handle_to_value(STATUS_BUILTIN, handle))
447 .transpose()?
448 .unwrap_or(VmValue::Nil),
449 ),
450 ]))
451}
452
453fn handle_to_value(
454 builtin: &'static str,
455 handle: &HostLeaseHandle,
456) -> Result<VmValue, HostlibError> {
457 let owner_process_identity = match handle.owner_process_identity {
458 Some(identity) => {
459 VmValue::Int(i64::try_from(identity).map_err(|_| HostlibError::Backend {
460 builtin,
461 message: "owner process identity exceeds the Harn integer range".to_string(),
462 })?)
463 }
464 None => VmValue::Nil,
465 };
466 let metadata = build_dict(
467 handle
468 .metadata
469 .iter()
470 .map(|(key, value)| (key.clone(), str_value(value))),
471 );
472 Ok(build_dict([
473 (
474 "schema_version",
475 VmValue::Int(i64::from(handle.schema_version)),
476 ),
477 ("host", str_value(&handle.host)),
478 ("resource_class", str_value(handle.resource_class.as_str())),
479 ("domain", str_value(&handle.domain)),
480 ("lease_id", str_value(&handle.lease_id)),
481 ("owner", str_value(&handle.owner)),
482 ("priority_class", str_value(handle.priority_class.as_str())),
483 ("acquired_at_ms", VmValue::Int(handle.acquired_at_ms)),
484 ("updated_at_ms", VmValue::Int(handle.updated_at_ms)),
485 (
486 "expires_at_ms",
487 handle
488 .expires_at_ms
489 .map(VmValue::Int)
490 .unwrap_or(VmValue::Nil),
491 ),
492 (
493 "owner_pid",
494 handle
495 .owner_pid
496 .map(i64::from)
497 .map(VmValue::Int)
498 .unwrap_or(VmValue::Nil),
499 ),
500 ("owner_process_identity", owner_process_identity),
501 (
502 "reason",
503 handle
504 .reason
505 .as_deref()
506 .map(str_value)
507 .unwrap_or(VmValue::Nil),
508 ),
509 ("metadata", metadata),
510 ]))
511}
512
513#[cfg(test)]
514mod tests {
515 use std::collections::BTreeMap;
516
517 use super::*;
518 use crate::host_lease::{HostLeasePriorityClass, HostLeaseResourceClass};
519
520 #[test]
521 fn status_value_preserves_active_and_recovery_evidence() {
522 let state = HostLeaseState {
523 schema_version: 1,
524 host: "mac-local".to_string(),
525 resource_class: HostLeaseResourceClass::WholeMachine,
526 domain: DEFAULT_HOST_LEASE_DOMAIN.to_string(),
527 observed_at_ms: 42,
528 active: None,
529 recovered_stale_lease: true,
530 recovered: Some(HostLeaseHandle {
531 schema_version: 1,
532 host: "mac-local".to_string(),
533 resource_class: HostLeaseResourceClass::WholeMachine,
534 domain: DEFAULT_HOST_LEASE_DOMAIN.to_string(),
535 execution_context: None,
536 lease_id: "lease-1".to_string(),
537 owner: "owner".to_string(),
538 priority_class: HostLeasePriorityClass::Measurement,
539 acquired_at_ms: 10,
540 updated_at_ms: 20,
541 expires_at_ms: Some(30),
542 owner_pid: Some(123),
543 owner_process_identity: Some(456),
544 reason: Some("measurement".to_string()),
545 metadata: BTreeMap::from([("lane".to_string(), "meter".to_string())]),
546 }),
547 };
548
549 let value = state_to_value(&state).expect("state converts");
550 let VmValue::Dict(state) = value else {
551 panic!("expected state dict");
552 };
553 assert!(matches!(
554 state.get("recovered_stale_lease"),
555 Some(VmValue::Bool(true))
556 ));
557 assert!(matches!(state.get("active"), Some(VmValue::Nil)));
558 let Some(VmValue::Dict(recovered)) = state.get("recovered") else {
559 panic!("expected recovered lease");
560 };
561 assert_eq!(
562 recovered.get("priority_class").map(VmValue::display),
563 Some("measurement".to_string())
564 );
565 assert_eq!(
566 recovered
567 .get("owner_process_identity")
568 .map(VmValue::display),
569 Some("456".to_string())
570 );
571 }
572}