1use crate::contract::{
2 RuntimeActionRequest, RuntimeApplyRequest, RuntimeCapabilities, RuntimeExecRequest,
3 RuntimeExecResult, RuntimeInspection, RuntimeLogChunk, RuntimeLogQuery, RuntimeObservation,
4 RuntimeRemoval, RuntimeUnitState,
5};
6use crate::{
7 RuntimeActionKind, RuntimeClient, RuntimeClock, RuntimeDriver, RuntimeError,
8 RuntimeRequestKind, RuntimeRequestReceipt, RuntimeRequestState, RuntimeResult,
9 RuntimeStateStore, SystemRuntimeClock,
10};
11use async_trait::async_trait;
12use std::future::Future;
13use std::sync::Arc;
14use std::time::Duration;
15
16pub struct ManagedRuntimeClient {
18 state: Arc<dyn RuntimeStateStore>,
19 driver: Arc<dyn RuntimeDriver>,
20 clock: Arc<dyn RuntimeClock>,
21}
22
23impl ManagedRuntimeClient {
24 pub fn new(state: Arc<dyn RuntimeStateStore>, driver: Arc<dyn RuntimeDriver>) -> Self {
25 Self::with_clock(state, driver, Arc::new(SystemRuntimeClock))
26 }
27
28 pub fn with_clock(
29 state: Arc<dyn RuntimeStateStore>,
30 driver: Arc<dyn RuntimeDriver>,
31 clock: Arc<dyn RuntimeClock>,
32 ) -> Self {
33 Self {
34 state,
35 driver,
36 clock,
37 }
38 }
39
40 async fn checked_capabilities(&self) -> RuntimeResult<RuntimeCapabilities> {
41 let capabilities = self.driver.capabilities().await?;
42 capabilities.validate().map_err(RuntimeError::Protocol)?;
43 if &capabilities.provider_id != self.driver.provider_id() {
44 return Err(RuntimeError::Protocol(format!(
45 "Runtime driver {:?} reported capabilities for {:?}",
46 self.driver.provider_id().as_str(),
47 capabilities.provider_id.as_str()
48 )));
49 }
50 Ok(capabilities)
51 }
52
53 fn check_deadline(&self, deadline_at_ms: Option<u64>) -> RuntimeResult<()> {
54 if deadline_at_ms.is_some_and(|deadline| deadline <= self.clock.now_ms()) {
55 return Err(RuntimeError::DeadlineExceeded(
56 "request expired before provider dispatch".into(),
57 ));
58 }
59 Ok(())
60 }
61
62 async fn bounded<T, F>(
63 &self,
64 deadline_at_ms: Option<u64>,
65 stage: &'static str,
66 future: F,
67 ) -> RuntimeResult<T>
68 where
69 T: Send,
70 F: Future<Output = RuntimeResult<T>> + Send,
71 {
72 let Some(deadline_at_ms) = deadline_at_ms else {
73 return future.await;
74 };
75 let now_ms = self.clock.now_ms();
76 let Some(remaining_ms) = deadline_at_ms.checked_sub(now_ms) else {
77 return Err(RuntimeError::DeadlineExceeded(format!(
78 "request expired before {stage}"
79 )));
80 };
81 if remaining_ms == 0 {
82 return Err(RuntimeError::DeadlineExceeded(format!(
83 "request expired before {stage}"
84 )));
85 }
86 tokio::time::timeout(Duration::from_millis(remaining_ms), future)
87 .await
88 .map_err(|_| {
89 RuntimeError::DeadlineExceeded(format!("request deadline elapsed during {stage}"))
90 })?
91 }
92
93 fn exec_deadline(&self, request: &RuntimeExecRequest, started_at_ms: u64) -> u64 {
94 let relative = started_at_ms.saturating_add(request.timeout_ms);
95 request
96 .deadline_at_ms
97 .map_or(relative, |absolute| absolute.min(relative))
98 }
99
100 async fn matching_receipt(
101 &self,
102 unit_id: &str,
103 generation: u64,
104 request_id: &str,
105 kind: RuntimeRequestKind,
106 request_digest: &str,
107 ) -> RuntimeResult<Option<RuntimeRequestReceipt>> {
108 let receipt = match self.state.load_request(unit_id, request_id).await {
109 Ok(receipt) => receipt,
110 Err(RuntimeError::RequestNotFound { .. }) => return Ok(None),
111 Err(error) => return Err(error),
112 };
113 receipt.validate().map_err(RuntimeError::Protocol)?;
114 if receipt.unit_id != unit_id || receipt.request_id != request_id {
115 return Err(RuntimeError::Protocol(
116 "Runtime request receipt storage key mismatch".into(),
117 ));
118 }
119 if receipt.generation != generation
120 || receipt.kind != kind
121 || receipt.request_digest != request_digest
122 {
123 return Err(RuntimeError::RequestConflict {
124 request_id: request_id.into(),
125 });
126 }
127 Ok(Some(receipt))
128 }
129
130 async fn completed_replay(
131 &self,
132 unit_id: &str,
133 generation: u64,
134 request_id: &str,
135 kind: RuntimeRequestKind,
136 request_digest: &str,
137 ) -> RuntimeResult<Option<RuntimeRequestReceipt>> {
138 Ok(self
139 .matching_receipt(unit_id, generation, request_id, kind, request_digest)
140 .await?
141 .filter(|receipt| receipt.state == RuntimeRequestState::Completed))
142 }
143}
144
145#[async_trait]
146impl RuntimeClient for ManagedRuntimeClient {
147 async fn capabilities(&self) -> RuntimeResult<RuntimeCapabilities> {
148 self.checked_capabilities().await
149 }
150
151 async fn apply(&self, request: &RuntimeApplyRequest) -> RuntimeResult<RuntimeObservation> {
152 request.validate().map_err(RuntimeError::InvalidRequest)?;
153 let request_digest = request.digest().map_err(RuntimeError::InvalidRequest)?;
154 if self
155 .completed_replay(
156 &request.spec.unit_id,
157 request.spec.generation,
158 &request.request_id,
159 RuntimeRequestKind::Apply,
160 &request_digest,
161 )
162 .await?
163 .is_some()
164 {
165 let _lease = self
166 .state
167 .acquire_operation_lease(&request.spec.unit_id)
168 .await?;
169 let reservation = self
170 .state
171 .reserve_apply(request, self.clock.now_ms())
172 .await?;
173 if reservation.dispatch {
174 return Err(RuntimeError::Protocol(
175 "completed apply receipt regressed to pending".into(),
176 ));
177 }
178 return reservation.receipt.observation.ok_or_else(|| {
179 RuntimeError::Protocol("completed apply receipt has no observation".into())
180 });
181 }
182 self.check_deadline(request.deadline_at_ms)?;
183 let capabilities = self
184 .bounded(
185 request.deadline_at_ms,
186 "capability query",
187 self.checked_capabilities(),
188 )
189 .await?;
190 let missing = capabilities
191 .missing_for(&request.spec)
192 .map_err(RuntimeError::InvalidRequest)?;
193 if !missing.is_empty() {
194 return Err(RuntimeError::UnsupportedCapabilities(missing));
195 }
196
197 let _lease = self
198 .bounded(
199 request.deadline_at_ms,
200 "operation lease wait",
201 self.state.acquire_operation_lease(&request.spec.unit_id),
202 )
203 .await?;
204 self.check_deadline(request.deadline_at_ms)?;
205
206 let reservation = self
207 .state
208 .reserve_apply(request, self.clock.now_ms())
209 .await?;
210 if !reservation.dispatch {
211 return reservation.receipt.observation.ok_or_else(|| {
212 RuntimeError::Protocol("completed apply receipt has no observation".into())
213 });
214 }
215
216 let observation = self
217 .bounded(
218 request.deadline_at_ms,
219 "provider apply",
220 self.driver
221 .apply(&request.spec, &reservation.record.observation),
222 )
223 .await?;
224 observation
225 .validate_against(&request.spec)
226 .map_err(RuntimeError::Protocol)?;
227 ensure_apply_result(&request.spec, &observation)?;
228 Ok(self
229 .state
230 .update_observation(Some(&request.request_id), &observation)
231 .await?
232 .observation)
233 }
234
235 async fn inspect(&self, unit_id: &str) -> RuntimeResult<RuntimeInspection> {
236 let _lease = self.state.acquire_operation_lease(unit_id).await?;
237 let record = match self.state.load(unit_id).await {
238 Ok(record) => record,
239 Err(RuntimeError::NotFound { .. }) => {
240 return Ok(RuntimeInspection::NotFound {
241 schema: RuntimeInspection::SCHEMA.into(),
242 unit_id: unit_id.into(),
243 last_generation: None,
244 });
245 }
246 Err(error) => return Err(error),
247 };
248 if record.removed_at_ms.is_some() {
249 return Ok(RuntimeInspection::NotFound {
250 schema: RuntimeInspection::SCHEMA.into(),
251 unit_id: unit_id.into(),
252 last_generation: Some(record.spec.generation),
253 });
254 }
255 if record.observation.state.is_terminal() {
256 return Ok(RuntimeInspection::Found {
257 schema: RuntimeInspection::SCHEMA.into(),
258 observation: Box::new(record.observation),
259 });
260 }
261
262 let inspection = self.driver.inspect(&record).await?;
263 inspection.validate().map_err(RuntimeError::Protocol)?;
264 match inspection {
265 RuntimeInspection::Found { observation, .. } => {
266 observation
267 .validate_against(&record.spec)
268 .map_err(RuntimeError::Protocol)?;
269 let record = self
270 .state
271 .update_observation(None, observation.as_ref())
272 .await?;
273 Ok(RuntimeInspection::Found {
274 schema: RuntimeInspection::SCHEMA.into(),
275 observation: Box::new(record.observation),
276 })
277 }
278 RuntimeInspection::NotFound { .. } => {
279 let mut unknown = record.observation;
280 unknown.state = RuntimeUnitState::Unknown;
281 unknown.observed_at_ms = unknown.observed_at_ms.max(self.clock.now_ms());
282 unknown.finished_at_ms = None;
283 unknown.health = None;
284 unknown.liveness = None;
285 unknown.outputs.clear();
286 unknown.clear_service_endpoints();
287 unknown.failure = None;
288 let record = self.state.update_observation(None, &unknown).await?;
289 Ok(RuntimeInspection::Found {
290 schema: RuntimeInspection::SCHEMA.into(),
291 observation: Box::new(record.observation),
292 })
293 }
294 }
295 }
296
297 async fn stop(&self, request: &RuntimeActionRequest) -> RuntimeResult<RuntimeInspection> {
298 request.validate().map_err(RuntimeError::InvalidRequest)?;
299 let request_digest = request.digest().map_err(RuntimeError::InvalidRequest)?;
300 if self
301 .completed_replay(
302 &request.unit_id,
303 request.generation,
304 &request.request_id,
305 RuntimeRequestKind::Stop,
306 &request_digest,
307 )
308 .await?
309 .is_some()
310 {
311 let _lease = self.state.acquire_operation_lease(&request.unit_id).await?;
312 let reservation = self
313 .state
314 .reserve_action(RuntimeActionKind::Stop, request, self.clock.now_ms())
315 .await?;
316 if reservation.dispatch {
317 return Err(RuntimeError::Protocol(
318 "completed stop receipt regressed to pending".into(),
319 ));
320 }
321 return Ok(RuntimeInspection::Found {
322 schema: RuntimeInspection::SCHEMA.into(),
323 observation: Box::new(reservation.receipt.observation.ok_or_else(|| {
324 RuntimeError::Protocol("completed stop receipt has no observation".into())
325 })?),
326 });
327 }
328 self.check_deadline(request.deadline_at_ms)?;
329 let capabilities = self
330 .bounded(
331 request.deadline_at_ms,
332 "capability query",
333 self.checked_capabilities(),
334 )
335 .await?;
336 if !capabilities.supports_feature(crate::contract::RuntimeFeature::Stop) {
337 return Err(RuntimeError::UnsupportedCapabilities(vec![
338 "feature:Stop".into()
339 ]));
340 }
341 let _lease = self
342 .bounded(
343 request.deadline_at_ms,
344 "operation lease wait",
345 self.state.acquire_operation_lease(&request.unit_id),
346 )
347 .await?;
348 self.check_deadline(request.deadline_at_ms)?;
349 let reservation = self
350 .state
351 .reserve_action(RuntimeActionKind::Stop, request, self.clock.now_ms())
352 .await?;
353 if !reservation.dispatch {
354 return Ok(RuntimeInspection::Found {
355 schema: RuntimeInspection::SCHEMA.into(),
356 observation: Box::new(reservation.receipt.observation.ok_or_else(|| {
357 RuntimeError::Protocol("completed stop receipt has no observation".into())
358 })?),
359 });
360 }
361 let observation = self
362 .bounded(
363 request.deadline_at_ms,
364 "provider stop",
365 self.driver.stop(&reservation.record, request),
366 )
367 .await?;
368 observation
369 .validate_against(&reservation.record.spec)
370 .map_err(RuntimeError::Protocol)?;
371 ensure_stop_result(&reservation.record.observation, &observation)?;
372 let record = self
373 .state
374 .update_observation(Some(&request.request_id), &observation)
375 .await?;
376 Ok(RuntimeInspection::Found {
377 schema: RuntimeInspection::SCHEMA.into(),
378 observation: Box::new(record.observation),
379 })
380 }
381
382 async fn remove(&self, request: &RuntimeActionRequest) -> RuntimeResult<RuntimeRemoval> {
383 request.validate().map_err(RuntimeError::InvalidRequest)?;
384 let request_digest = request.digest().map_err(RuntimeError::InvalidRequest)?;
385 if self
386 .completed_replay(
387 &request.unit_id,
388 request.generation,
389 &request.request_id,
390 RuntimeRequestKind::Remove,
391 &request_digest,
392 )
393 .await?
394 .is_some()
395 {
396 let _lease = self.state.acquire_operation_lease(&request.unit_id).await?;
397 let reservation = self
398 .state
399 .reserve_action(RuntimeActionKind::Remove, request, self.clock.now_ms())
400 .await?;
401 if reservation.dispatch {
402 return Err(RuntimeError::Protocol(
403 "completed remove receipt regressed to pending".into(),
404 ));
405 }
406 return reservation.receipt.removal.ok_or_else(|| {
407 RuntimeError::Protocol("completed remove receipt has no removal".into())
408 });
409 }
410 self.check_deadline(request.deadline_at_ms)?;
411 let capabilities = self
412 .bounded(
413 request.deadline_at_ms,
414 "capability query",
415 self.checked_capabilities(),
416 )
417 .await?;
418 if !capabilities.supports_feature(crate::contract::RuntimeFeature::Remove) {
419 return Err(RuntimeError::UnsupportedCapabilities(vec![
420 "feature:Remove".into(),
421 ]));
422 }
423 let _lease = self
424 .bounded(
425 request.deadline_at_ms,
426 "operation lease wait",
427 self.state.acquire_operation_lease(&request.unit_id),
428 )
429 .await?;
430 self.check_deadline(request.deadline_at_ms)?;
431 let reservation = self
432 .state
433 .reserve_action(RuntimeActionKind::Remove, request, self.clock.now_ms())
434 .await?;
435 if !reservation.dispatch {
436 return reservation.receipt.removal.ok_or_else(|| {
437 RuntimeError::Protocol("completed remove receipt has no removal".into())
438 });
439 }
440 let removal = self
441 .bounded(
442 request.deadline_at_ms,
443 "provider remove",
444 self.driver.remove(&reservation.record, request),
445 )
446 .await?;
447 removal.validate().map_err(RuntimeError::Protocol)?;
448 if removal.request_id != request.request_id
449 || removal.unit_id != request.unit_id
450 || removal.generation != request.generation
451 {
452 return Err(RuntimeError::Protocol(
453 "provider removal changed immutable request identity".into(),
454 ));
455 }
456 self.state.complete_removal(&removal).await?;
457 Ok(removal)
458 }
459
460 async fn logs(&self, query: &RuntimeLogQuery) -> RuntimeResult<Vec<RuntimeLogChunk>> {
461 query.validate().map_err(RuntimeError::InvalidRequest)?;
462 let capabilities = self.checked_capabilities().await?;
463 if !capabilities.supports_feature(crate::contract::RuntimeFeature::Logs) {
464 return Err(RuntimeError::UnsupportedCapabilities(vec![
465 "feature:Logs".into()
466 ]));
467 }
468 let _lease = self.state.acquire_operation_lease(&query.unit_id).await?;
469 let record = self.state.load(&query.unit_id).await?;
470 ensure_current_generation(&record, query.generation)?;
471 let chunks = self.driver.logs(&record, query).await?;
472 for chunk in &chunks {
473 chunk.validate().map_err(RuntimeError::Protocol)?;
474 }
475 if chunks
476 .windows(2)
477 .any(|pair| pair[0].sequence >= pair[1].sequence)
478 {
479 return Err(RuntimeError::Protocol(
480 "provider returned unordered log chunks".into(),
481 ));
482 }
483 Ok(chunks)
484 }
485
486 async fn exec(&self, request: &RuntimeExecRequest) -> RuntimeResult<RuntimeExecResult> {
487 request.validate().map_err(RuntimeError::InvalidRequest)?;
488 let request_digest = request.digest().map_err(RuntimeError::InvalidRequest)?;
489 let operation_started_at_ms = self.clock.now_ms();
490 let existing = self
491 .matching_receipt(
492 &request.unit_id,
493 request.generation,
494 &request.request_id,
495 RuntimeRequestKind::Exec,
496 &request_digest,
497 )
498 .await?;
499 if existing
500 .as_ref()
501 .is_some_and(|receipt| receipt.state == RuntimeRequestState::Completed)
502 {
503 let _lease = self.state.acquire_operation_lease(&request.unit_id).await?;
504 let reservation = self
505 .state
506 .reserve_exec(request, operation_started_at_ms)
507 .await?;
508 if reservation.dispatch {
509 return Err(RuntimeError::Protocol(
510 "completed exec receipt regressed to pending".into(),
511 ));
512 }
513 return reservation.receipt.exec_result.ok_or_else(|| {
514 RuntimeError::Protocol("completed exec receipt has no result".into())
515 });
516 }
517 let deadline_at_ms = existing
518 .as_ref()
519 .and_then(|receipt| receipt.deadline_at_ms)
520 .unwrap_or_else(|| self.exec_deadline(request, operation_started_at_ms));
521 let deadline_at_ms = Some(deadline_at_ms);
522 self.check_deadline(deadline_at_ms)?;
523 let capabilities = self
524 .bounded(
525 deadline_at_ms,
526 "capability query",
527 self.checked_capabilities(),
528 )
529 .await?;
530 if !capabilities.supports_feature(crate::contract::RuntimeFeature::Exec) {
531 return Err(RuntimeError::UnsupportedCapabilities(vec![
532 "feature:Exec".into()
533 ]));
534 }
535 let _lease = self
536 .bounded(
537 deadline_at_ms,
538 "operation lease wait",
539 self.state.acquire_operation_lease(&request.unit_id),
540 )
541 .await?;
542 self.check_deadline(deadline_at_ms)?;
543 let reservation = self
544 .state
545 .reserve_exec(request, operation_started_at_ms)
546 .await?;
547 if !reservation.dispatch {
548 return reservation.receipt.exec_result.ok_or_else(|| {
549 RuntimeError::Protocol("completed exec receipt has no result".into())
550 });
551 }
552 let deadline_at_ms = reservation.receipt.deadline_at_ms.ok_or_else(|| {
553 RuntimeError::Protocol("pending exec receipt has no effective deadline".into())
554 })?;
555 let deadline_at_ms = Some(deadline_at_ms);
556 self.check_deadline(deadline_at_ms)?;
557 let mut provider_request = request.clone();
558 provider_request.deadline_at_ms = deadline_at_ms;
559 let result = self
560 .bounded(
561 deadline_at_ms,
562 "provider exec",
563 self.driver.exec(&reservation.record, &provider_request),
564 )
565 .await?;
566 result.validate().map_err(RuntimeError::Protocol)?;
567 result
568 .observation
569 .validate_against(&reservation.record.spec)
570 .map_err(RuntimeError::Protocol)?;
571 if result.request_id != request.request_id
572 || result.observation.unit_id != request.unit_id
573 || result.observation.generation != request.generation
574 {
575 return Err(RuntimeError::Protocol(
576 "provider exec changed immutable request identity".into(),
577 ));
578 }
579 self.state.complete_exec(&result).await?;
580 Ok(result)
581 }
582}
583
584fn ensure_current_generation(
585 record: &crate::RuntimeUnitRecord,
586 requested: u64,
587) -> RuntimeResult<()> {
588 if record.removed_at_ms.is_some() {
589 return Err(RuntimeError::NotFound {
590 unit_id: record.spec.unit_id.clone(),
591 });
592 }
593 if requested < record.spec.generation {
594 return Err(RuntimeError::StaleGeneration {
595 unit_id: record.spec.unit_id.clone(),
596 requested,
597 current: record.spec.generation,
598 });
599 }
600 if requested != record.spec.generation {
601 return Err(RuntimeError::GenerationConflict {
602 unit_id: record.spec.unit_id.clone(),
603 generation: requested,
604 });
605 }
606 Ok(())
607}
608
609fn ensure_apply_result(
610 spec: &crate::contract::RuntimeUnitSpec,
611 observation: &RuntimeObservation,
612) -> RuntimeResult<()> {
613 let allowed = match spec.class {
614 crate::contract::RuntimeUnitClass::Task => matches!(
615 observation.state,
616 RuntimeUnitState::Succeeded | RuntimeUnitState::Failed
617 ),
618 crate::contract::RuntimeUnitClass::Service => matches!(
619 observation.state,
620 RuntimeUnitState::Running
621 | RuntimeUnitState::Stopped
622 | RuntimeUnitState::Failed
623 | RuntimeUnitState::Unknown
624 ),
625 };
626 if !allowed {
627 return Err(RuntimeError::Protocol(format!(
628 "provider apply returned invalid {:?} result {:?}",
629 spec.class, observation.state
630 )));
631 }
632 if observation.provider_resource_id.is_none() || observation.provider_build.is_none() {
633 return Err(RuntimeError::Protocol(
634 "provider apply returned an observation without provider identity".into(),
635 ));
636 }
637 Ok(())
638}
639
640fn ensure_stop_result(
641 current: &RuntimeObservation,
642 observation: &RuntimeObservation,
643) -> RuntimeResult<()> {
644 if observation.state == RuntimeUnitState::Stopped
645 || observation.state == RuntimeUnitState::Unknown
646 || current.state.is_terminal() && observation == current
647 {
648 return Ok(());
649 }
650 Err(RuntimeError::Protocol(format!(
651 "provider stop returned nonterminal state {:?}",
652 observation.state
653 )))
654}