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