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.outputs.clear();
285 unknown.clear_service_endpoints();
286 unknown.failure = None;
287 let record = self.state.update_observation(None, &unknown).await?;
288 Ok(RuntimeInspection::Found {
289 schema: RuntimeInspection::SCHEMA.into(),
290 observation: Box::new(record.observation),
291 })
292 }
293 }
294 }
295
296 async fn stop(&self, request: &RuntimeActionRequest) -> RuntimeResult<RuntimeInspection> {
297 request.validate().map_err(RuntimeError::InvalidRequest)?;
298 let request_digest = request.digest().map_err(RuntimeError::InvalidRequest)?;
299 if self
300 .completed_replay(
301 &request.unit_id,
302 request.generation,
303 &request.request_id,
304 RuntimeRequestKind::Stop,
305 &request_digest,
306 )
307 .await?
308 .is_some()
309 {
310 let _lease = self.state.acquire_operation_lease(&request.unit_id).await?;
311 let reservation = self
312 .state
313 .reserve_action(RuntimeActionKind::Stop, request, self.clock.now_ms())
314 .await?;
315 if reservation.dispatch {
316 return Err(RuntimeError::Protocol(
317 "completed stop receipt regressed to pending".into(),
318 ));
319 }
320 return Ok(RuntimeInspection::Found {
321 schema: RuntimeInspection::SCHEMA.into(),
322 observation: Box::new(reservation.receipt.observation.ok_or_else(|| {
323 RuntimeError::Protocol("completed stop receipt has no observation".into())
324 })?),
325 });
326 }
327 self.check_deadline(request.deadline_at_ms)?;
328 let capabilities = self
329 .bounded(
330 request.deadline_at_ms,
331 "capability query",
332 self.checked_capabilities(),
333 )
334 .await?;
335 if !capabilities.supports_feature(crate::contract::RuntimeFeature::Stop) {
336 return Err(RuntimeError::UnsupportedCapabilities(vec![
337 "feature:Stop".into()
338 ]));
339 }
340 let _lease = self
341 .bounded(
342 request.deadline_at_ms,
343 "operation lease wait",
344 self.state.acquire_operation_lease(&request.unit_id),
345 )
346 .await?;
347 self.check_deadline(request.deadline_at_ms)?;
348 let reservation = self
349 .state
350 .reserve_action(RuntimeActionKind::Stop, request, self.clock.now_ms())
351 .await?;
352 if !reservation.dispatch {
353 return Ok(RuntimeInspection::Found {
354 schema: RuntimeInspection::SCHEMA.into(),
355 observation: Box::new(reservation.receipt.observation.ok_or_else(|| {
356 RuntimeError::Protocol("completed stop receipt has no observation".into())
357 })?),
358 });
359 }
360 let observation = self
361 .bounded(
362 request.deadline_at_ms,
363 "provider stop",
364 self.driver.stop(&reservation.record, request),
365 )
366 .await?;
367 observation
368 .validate_against(&reservation.record.spec)
369 .map_err(RuntimeError::Protocol)?;
370 ensure_stop_result(&reservation.record.observation, &observation)?;
371 let record = self
372 .state
373 .update_observation(Some(&request.request_id), &observation)
374 .await?;
375 Ok(RuntimeInspection::Found {
376 schema: RuntimeInspection::SCHEMA.into(),
377 observation: Box::new(record.observation),
378 })
379 }
380
381 async fn remove(&self, request: &RuntimeActionRequest) -> RuntimeResult<RuntimeRemoval> {
382 request.validate().map_err(RuntimeError::InvalidRequest)?;
383 let request_digest = request.digest().map_err(RuntimeError::InvalidRequest)?;
384 if self
385 .completed_replay(
386 &request.unit_id,
387 request.generation,
388 &request.request_id,
389 RuntimeRequestKind::Remove,
390 &request_digest,
391 )
392 .await?
393 .is_some()
394 {
395 let _lease = self.state.acquire_operation_lease(&request.unit_id).await?;
396 let reservation = self
397 .state
398 .reserve_action(RuntimeActionKind::Remove, request, self.clock.now_ms())
399 .await?;
400 if reservation.dispatch {
401 return Err(RuntimeError::Protocol(
402 "completed remove receipt regressed to pending".into(),
403 ));
404 }
405 return reservation.receipt.removal.ok_or_else(|| {
406 RuntimeError::Protocol("completed remove receipt has no removal".into())
407 });
408 }
409 self.check_deadline(request.deadline_at_ms)?;
410 let capabilities = self
411 .bounded(
412 request.deadline_at_ms,
413 "capability query",
414 self.checked_capabilities(),
415 )
416 .await?;
417 if !capabilities.supports_feature(crate::contract::RuntimeFeature::Remove) {
418 return Err(RuntimeError::UnsupportedCapabilities(vec![
419 "feature:Remove".into(),
420 ]));
421 }
422 let _lease = self
423 .bounded(
424 request.deadline_at_ms,
425 "operation lease wait",
426 self.state.acquire_operation_lease(&request.unit_id),
427 )
428 .await?;
429 self.check_deadline(request.deadline_at_ms)?;
430 let reservation = self
431 .state
432 .reserve_action(RuntimeActionKind::Remove, request, self.clock.now_ms())
433 .await?;
434 if !reservation.dispatch {
435 return reservation.receipt.removal.ok_or_else(|| {
436 RuntimeError::Protocol("completed remove receipt has no removal".into())
437 });
438 }
439 let removal = self
440 .bounded(
441 request.deadline_at_ms,
442 "provider remove",
443 self.driver.remove(&reservation.record, request),
444 )
445 .await?;
446 removal.validate().map_err(RuntimeError::Protocol)?;
447 if removal.request_id != request.request_id
448 || removal.unit_id != request.unit_id
449 || removal.generation != request.generation
450 {
451 return Err(RuntimeError::Protocol(
452 "provider removal changed immutable request identity".into(),
453 ));
454 }
455 self.state.complete_removal(&removal).await?;
456 Ok(removal)
457 }
458
459 async fn logs(&self, query: &RuntimeLogQuery) -> RuntimeResult<Vec<RuntimeLogChunk>> {
460 query.validate().map_err(RuntimeError::InvalidRequest)?;
461 let capabilities = self.checked_capabilities().await?;
462 if !capabilities.supports_feature(crate::contract::RuntimeFeature::Logs) {
463 return Err(RuntimeError::UnsupportedCapabilities(vec![
464 "feature:Logs".into()
465 ]));
466 }
467 let _lease = self.state.acquire_operation_lease(&query.unit_id).await?;
468 let record = self.state.load(&query.unit_id).await?;
469 ensure_current_generation(&record, query.generation)?;
470 let chunks = self.driver.logs(&record, query).await?;
471 for chunk in &chunks {
472 chunk.validate().map_err(RuntimeError::Protocol)?;
473 }
474 if chunks
475 .windows(2)
476 .any(|pair| pair[0].sequence >= pair[1].sequence)
477 {
478 return Err(RuntimeError::Protocol(
479 "provider returned unordered log chunks".into(),
480 ));
481 }
482 Ok(chunks)
483 }
484
485 async fn exec(&self, request: &RuntimeExecRequest) -> RuntimeResult<RuntimeExecResult> {
486 request.validate().map_err(RuntimeError::InvalidRequest)?;
487 let request_digest = request.digest().map_err(RuntimeError::InvalidRequest)?;
488 let operation_started_at_ms = self.clock.now_ms();
489 let existing = self
490 .matching_receipt(
491 &request.unit_id,
492 request.generation,
493 &request.request_id,
494 RuntimeRequestKind::Exec,
495 &request_digest,
496 )
497 .await?;
498 if existing
499 .as_ref()
500 .is_some_and(|receipt| receipt.state == RuntimeRequestState::Completed)
501 {
502 let _lease = self.state.acquire_operation_lease(&request.unit_id).await?;
503 let reservation = self
504 .state
505 .reserve_exec(request, operation_started_at_ms)
506 .await?;
507 if reservation.dispatch {
508 return Err(RuntimeError::Protocol(
509 "completed exec receipt regressed to pending".into(),
510 ));
511 }
512 return reservation.receipt.exec_result.ok_or_else(|| {
513 RuntimeError::Protocol("completed exec receipt has no result".into())
514 });
515 }
516 let deadline_at_ms = existing
517 .as_ref()
518 .and_then(|receipt| receipt.deadline_at_ms)
519 .unwrap_or_else(|| self.exec_deadline(request, operation_started_at_ms));
520 let deadline_at_ms = Some(deadline_at_ms);
521 self.check_deadline(deadline_at_ms)?;
522 let capabilities = self
523 .bounded(
524 deadline_at_ms,
525 "capability query",
526 self.checked_capabilities(),
527 )
528 .await?;
529 if !capabilities.supports_feature(crate::contract::RuntimeFeature::Exec) {
530 return Err(RuntimeError::UnsupportedCapabilities(vec![
531 "feature:Exec".into()
532 ]));
533 }
534 let _lease = self
535 .bounded(
536 deadline_at_ms,
537 "operation lease wait",
538 self.state.acquire_operation_lease(&request.unit_id),
539 )
540 .await?;
541 self.check_deadline(deadline_at_ms)?;
542 let reservation = self
543 .state
544 .reserve_exec(request, operation_started_at_ms)
545 .await?;
546 if !reservation.dispatch {
547 return reservation.receipt.exec_result.ok_or_else(|| {
548 RuntimeError::Protocol("completed exec receipt has no result".into())
549 });
550 }
551 let deadline_at_ms = reservation.receipt.deadline_at_ms.ok_or_else(|| {
552 RuntimeError::Protocol("pending exec receipt has no effective deadline".into())
553 })?;
554 let deadline_at_ms = Some(deadline_at_ms);
555 self.check_deadline(deadline_at_ms)?;
556 let mut provider_request = request.clone();
557 provider_request.deadline_at_ms = deadline_at_ms;
558 let result = self
559 .bounded(
560 deadline_at_ms,
561 "provider exec",
562 self.driver.exec(&reservation.record, &provider_request),
563 )
564 .await?;
565 result.validate().map_err(RuntimeError::Protocol)?;
566 result
567 .observation
568 .validate_against(&reservation.record.spec)
569 .map_err(RuntimeError::Protocol)?;
570 if result.request_id != request.request_id
571 || result.observation.unit_id != request.unit_id
572 || result.observation.generation != request.generation
573 {
574 return Err(RuntimeError::Protocol(
575 "provider exec changed immutable request identity".into(),
576 ));
577 }
578 self.state.complete_exec(&result).await?;
579 Ok(result)
580 }
581}
582
583fn ensure_current_generation(
584 record: &crate::RuntimeUnitRecord,
585 requested: u64,
586) -> RuntimeResult<()> {
587 if record.removed_at_ms.is_some() {
588 return Err(RuntimeError::NotFound {
589 unit_id: record.spec.unit_id.clone(),
590 });
591 }
592 if requested < record.spec.generation {
593 return Err(RuntimeError::StaleGeneration {
594 unit_id: record.spec.unit_id.clone(),
595 requested,
596 current: record.spec.generation,
597 });
598 }
599 if requested != record.spec.generation {
600 return Err(RuntimeError::GenerationConflict {
601 unit_id: record.spec.unit_id.clone(),
602 generation: requested,
603 });
604 }
605 Ok(())
606}
607
608fn ensure_apply_result(
609 spec: &crate::contract::RuntimeUnitSpec,
610 observation: &RuntimeObservation,
611) -> RuntimeResult<()> {
612 let allowed = match spec.class {
613 crate::contract::RuntimeUnitClass::Task => matches!(
614 observation.state,
615 RuntimeUnitState::Succeeded | RuntimeUnitState::Failed
616 ),
617 crate::contract::RuntimeUnitClass::Service => matches!(
618 observation.state,
619 RuntimeUnitState::Running
620 | RuntimeUnitState::Stopped
621 | RuntimeUnitState::Failed
622 | RuntimeUnitState::Unknown
623 ),
624 };
625 if !allowed {
626 return Err(RuntimeError::Protocol(format!(
627 "provider apply returned invalid {:?} result {:?}",
628 spec.class, observation.state
629 )));
630 }
631 if observation.provider_resource_id.is_none() || observation.provider_build.is_none() {
632 return Err(RuntimeError::Protocol(
633 "provider apply returned an observation without provider identity".into(),
634 ));
635 }
636 Ok(())
637}
638
639fn ensure_stop_result(
640 current: &RuntimeObservation,
641 observation: &RuntimeObservation,
642) -> RuntimeResult<()> {
643 if observation.state == RuntimeUnitState::Stopped
644 || observation.state == RuntimeUnitState::Unknown
645 || current.state.is_terminal() && observation == current
646 {
647 return Ok(());
648 }
649 Err(RuntimeError::Protocol(format!(
650 "provider stop returned nonterminal state {:?}",
651 observation.state
652 )))
653}