1use std::{collections::BTreeMap, num::NonZeroU16, sync::Arc};
4
5use a3s_box_core::{
6 scale::ScaleEndpoint, CreateExecutionRequest, ExecutionGeneration, ExecutionId, ExecutionLease,
7 ExecutionManager, ExecutionPortConnector, ExecutionState, OperationId, ReconcileOutcome,
8};
9use async_trait::async_trait;
10use sha2::{Digest, Sha256};
11use thiserror::Error;
12use tokio::sync::Mutex;
13
14use crate::{LocalExecutionManager, ManagedExecutionState};
15
16use super::catalog::{
17 ScaleServiceCatalog, SCALE_GUEST_PORT_LABEL, SCALE_MANAGED_LABEL, SCALE_SERVICE_LABEL,
18 SCALE_SLOT_LABEL, SCALE_TEMPLATE_DIGEST_LABEL,
19};
20use super::endpoints::{ScaleEndpointConfig, ScaleEndpointOwner, ScaleEndpointTarget};
21
22#[derive(Debug, Clone, Copy, PartialEq, Eq)]
23enum InstancePhase {
24 Active,
25 Ready,
26 Terminal,
27 Removing,
28}
29
30#[derive(Debug, Clone, PartialEq, Eq)]
31struct ScaleExecution {
32 execution_id: ExecutionId,
33 generation: ExecutionGeneration,
34 service: String,
35 slot: u32,
36 template_digest: String,
37 guest_port: Option<NonZeroU16>,
38 phase: InstancePhase,
39}
40
41#[derive(Debug, Clone, PartialEq, Eq)]
42pub struct ScaleReconcileObservation {
43 pub ready_replicas: u32,
44 pub endpoints: Vec<ScaleEndpoint>,
45}
46
47#[derive(Debug, Clone, PartialEq, Eq)]
48pub struct ScaleReconcileReport {
49 pub service: String,
50 pub desired_replicas: u32,
51 pub ready_replicas: u32,
52 pub created: u32,
53 pub removed: u32,
54 pub endpoints: Vec<ScaleEndpoint>,
55}
56
57#[derive(Debug, Error)]
58pub enum ScaleReconcileError {
59 #[error("unknown scale service {0:?}")]
60 UnknownService(String),
61 #[error("invalid scale template for {service:?}: {message}")]
62 Template { service: String, message: String },
63 #[error("scale lifecycle error: {0}")]
64 Lifecycle(String),
65}
66
67#[async_trait]
68trait ScaleExecutionLifecycle: Send + Sync {
69 async fn inventory(&self) -> Result<Vec<ScaleExecution>, ScaleReconcileError>;
70
71 async fn ensure_running(
72 &self,
73 request: CreateExecutionRequest,
74 operation_id: OperationId,
75 ) -> Result<(), ScaleReconcileError>;
76
77 async fn ensure_removed(&self, execution: &ScaleExecution) -> Result<(), ScaleReconcileError>;
78}
79
80pub struct LocalScaleReconciler {
82 catalog: ScaleServiceCatalog,
83 lifecycle: Arc<dyn ScaleExecutionLifecycle>,
84 endpoint_owner: Option<ScaleEndpointOwner>,
85 reconcile_lock: Mutex<()>,
86}
87
88impl LocalScaleReconciler {
89 pub fn new(manager: LocalExecutionManager, catalog: ScaleServiceCatalog) -> Self {
90 Self::with_endpoint_config(manager, catalog, ScaleEndpointConfig::default())
91 }
92
93 pub fn with_endpoint_config(
94 manager: LocalExecutionManager,
95 catalog: ScaleServiceCatalog,
96 endpoint_config: ScaleEndpointConfig,
97 ) -> Self {
98 let connector: Arc<dyn ExecutionPortConnector> = Arc::new(manager.clone());
99 Self {
100 catalog,
101 lifecycle: Arc::new(LocalScaleExecutionLifecycle { manager }),
102 endpoint_owner: Some(ScaleEndpointOwner::new(endpoint_config, connector)),
103 reconcile_lock: Mutex::new(()),
104 }
105 }
106
107 #[cfg(test)]
108 fn with_lifecycle(
109 catalog: ScaleServiceCatalog,
110 lifecycle: Arc<dyn ScaleExecutionLifecycle>,
111 ) -> Self {
112 Self {
113 catalog,
114 lifecycle,
115 endpoint_owner: None,
116 reconcile_lock: Mutex::new(()),
117 }
118 }
119
120 #[cfg(test)]
121 fn with_lifecycle_and_endpoint(
122 catalog: ScaleServiceCatalog,
123 lifecycle: Arc<dyn ScaleExecutionLifecycle>,
124 connector: Arc<dyn ExecutionPortConnector>,
125 endpoint_config: ScaleEndpointConfig,
126 ) -> Self {
127 Self {
128 catalog,
129 lifecycle,
130 endpoint_owner: Some(ScaleEndpointOwner::new(endpoint_config, connector)),
131 reconcile_lock: Mutex::new(()),
132 }
133 }
134
135 pub fn knows_service(&self, service: &str) -> bool {
136 self.catalog.contains(service)
137 }
138
139 pub fn services(&self) -> Vec<String> {
140 self.catalog.services()
141 }
142
143 pub async fn observation(
144 &self,
145 service: &str,
146 desired_replicas: u32,
147 ) -> Result<ScaleReconcileObservation, ScaleReconcileError> {
148 let _guard = self.reconcile_lock.lock().await;
149 if !self.catalog.contains(service) {
150 return Err(ScaleReconcileError::UnknownService(service.to_string()));
151 }
152 let desired_templates = self.desired_templates(service, desired_replicas)?;
153 self.observation_unlocked(service, &desired_templates).await
154 }
155
156 pub async fn reconcile(
157 &self,
158 service: &str,
159 desired_replicas: u32,
160 ) -> Result<ScaleReconcileReport, ScaleReconcileError> {
161 let _guard = self.reconcile_lock.lock().await;
162 if !self.catalog.contains(service) {
163 return Err(ScaleReconcileError::UnknownService(service.to_string()));
164 }
165
166 let desired_templates = self.desired_templates(service, desired_replicas)?;
167 let mut inventory = self.service_inventory(service).await?;
168 inventory.sort_by(|left, right| {
169 left.slot
170 .cmp(&right.slot)
171 .then_with(|| left.execution_id.as_str().cmp(right.execution_id.as_str()))
172 });
173
174 let mut retained = BTreeMap::<u32, ScaleExecution>::new();
175 let mut removed = 0_u32;
176 for execution in inventory {
177 let expected = desired_templates.get(&execution.slot);
178 let should_retain = expected
179 .is_some_and(|request| execution_matches_template(&execution, request))
180 && matches!(
181 execution.phase,
182 InstancePhase::Active | InstancePhase::Ready
183 )
184 && !retained.contains_key(&execution.slot);
185 if should_retain {
186 retained.insert(execution.slot, execution);
187 } else {
188 if let Some(endpoint_owner) = &self.endpoint_owner {
189 endpoint_owner.remove(&execution.execution_id).await;
190 }
191 self.lifecycle.ensure_removed(&execution).await?;
192 removed = removed.saturating_add(1);
193 }
194 }
195
196 let mut created = 0_u32;
197 for (slot, request) in &desired_templates {
198 let retained_phase = retained.get(slot).map(|execution| execution.phase);
199 match retained_phase {
200 Some(InstancePhase::Ready) => continue,
201 Some(InstancePhase::Active) => {}
202 Some(InstancePhase::Terminal | InstancePhase::Removing) => {
203 unreachable!("only active or ready executions are retained")
204 }
205 None => created = created.saturating_add(1),
206 }
207 let operation_id = create_operation_id(
208 service,
209 *slot,
210 template_digest(request).ok_or_else(|| ScaleReconcileError::Template {
211 service: service.to_string(),
212 message: "generated template has no digest label".to_string(),
213 })?,
214 )?;
215 self.lifecycle
216 .ensure_running(request.clone(), operation_id)
217 .await?;
218 }
219
220 let observation = self
221 .observation_unlocked(service, &desired_templates)
222 .await?;
223
224 Ok(ScaleReconcileReport {
225 service: service.to_string(),
226 desired_replicas,
227 ready_replicas: observation.ready_replicas,
228 created,
229 removed,
230 endpoints: observation.endpoints,
231 })
232 }
233
234 fn desired_templates(
235 &self,
236 service: &str,
237 desired_replicas: u32,
238 ) -> Result<BTreeMap<u32, CreateExecutionRequest>, ScaleReconcileError> {
239 (0..desired_replicas)
240 .map(|slot| {
241 self.catalog
242 .create_request(service, slot)
243 .map(|request| (slot, request))
244 .map_err(|error| ScaleReconcileError::Template {
245 service: service.to_string(),
246 message: error.to_string(),
247 })
248 })
249 .collect()
250 }
251
252 async fn observation_unlocked(
253 &self,
254 service: &str,
255 desired_templates: &BTreeMap<u32, CreateExecutionRequest>,
256 ) -> Result<ScaleReconcileObservation, ScaleReconcileError> {
257 let mut inventory = self.service_inventory(service).await?;
258 inventory.sort_by(|left, right| {
259 left.slot
260 .cmp(&right.slot)
261 .then_with(|| left.execution_id.as_str().cmp(right.execution_id.as_str()))
262 });
263 let mut ready = BTreeMap::new();
264 for execution in inventory {
265 let is_desired = desired_templates
266 .get(&execution.slot)
267 .is_some_and(|request| execution_matches_template(&execution, request));
268 if is_desired
269 && execution.phase == InstancePhase::Ready
270 && !ready.contains_key(&execution.slot)
271 {
272 ready.insert(execution.slot, execution);
273 }
274 }
275
276 let targets = ready
277 .values()
278 .filter_map(|execution| {
279 execution.guest_port.map(|guest_port| ScaleEndpointTarget {
280 execution_id: execution.execution_id.clone(),
281 generation: execution.generation,
282 service: execution.service.clone(),
283 slot: execution.slot,
284 guest_port,
285 })
286 })
287 .collect::<Vec<_>>();
288 let endpoints = match (&self.endpoint_owner, targets.is_empty()) {
289 (Some(owner), _) => owner.reconcile_service(service, &targets).await?,
290 (None, true) => Vec::new(),
291 (None, false) => {
292 return Err(ScaleReconcileError::Lifecycle(
293 "scale endpoint publication is not configured".to_string(),
294 ))
295 }
296 };
297 Ok(ScaleReconcileObservation {
298 ready_replicas: ready.len() as u32,
299 endpoints,
300 })
301 }
302
303 async fn service_inventory(
304 &self,
305 service: &str,
306 ) -> Result<Vec<ScaleExecution>, ScaleReconcileError> {
307 Ok(self
308 .lifecycle
309 .inventory()
310 .await?
311 .into_iter()
312 .filter(|execution| execution.service == service)
313 .collect())
314 }
315}
316
317fn template_digest(request: &CreateExecutionRequest) -> Option<&str> {
318 request
319 .labels
320 .get(SCALE_TEMPLATE_DIGEST_LABEL)
321 .map(String::as_str)
322}
323
324fn template_guest_port(request: &CreateExecutionRequest) -> Option<NonZeroU16> {
325 request
326 .labels
327 .get(SCALE_GUEST_PORT_LABEL)
328 .and_then(|value| value.parse::<u16>().ok())
329 .and_then(NonZeroU16::new)
330}
331
332fn execution_matches_template(
333 execution: &ScaleExecution,
334 request: &CreateExecutionRequest,
335) -> bool {
336 template_digest(request).is_some_and(|digest| digest == execution.template_digest)
337 && template_guest_port(request) == execution.guest_port
338}
339
340fn create_operation_id(
341 service: &str,
342 slot: u32,
343 template_digest: &str,
344) -> Result<OperationId, ScaleReconcileError> {
345 let identity = format!("{service}\0{slot}\0{template_digest}");
346 OperationId::new(format!(
347 "scale-create-v1-{:x}",
348 Sha256::digest(identity.as_bytes())
349 ))
350 .map_err(|error| ScaleReconcileError::Lifecycle(error.to_string()))
351}
352
353struct LocalScaleExecutionLifecycle {
354 manager: LocalExecutionManager,
355}
356
357impl LocalScaleExecutionLifecycle {
358 async fn ensure_lease_running(&self, lease: ExecutionLease) -> Result<(), ScaleReconcileError> {
359 let status = self
360 .manager
361 .inspect(&lease.execution_id)
362 .await
363 .map_err(lifecycle_error)?;
364 if status.generation != lease.generation {
365 return Err(ScaleReconcileError::Lifecycle(format!(
366 "scale execution {} changed from generation {} to {} while ensuring it is running",
367 lease.execution_id,
368 lease.generation.get(),
369 status.generation.get()
370 )));
371 }
372 match status.state {
373 ExecutionState::Running => Ok(()),
374 ExecutionState::Paused => self
375 .manager
376 .resume(&lease.execution_id, lease.generation)
377 .await
378 .map(|_| ())
379 .map_err(lifecycle_error),
380 ExecutionState::Created | ExecutionState::Creating => {
381 Err(ScaleReconcileError::Lifecycle(format!(
382 "scale execution {} is still converging",
383 lease.execution_id
384 )))
385 }
386 ExecutionState::Stopped | ExecutionState::Failed => {
387 Err(ScaleReconcileError::Lifecycle(format!(
388 "scale execution {} is terminal",
389 lease.execution_id
390 )))
391 }
392 }
393 }
394}
395
396#[async_trait]
397impl ScaleExecutionLifecycle for LocalScaleExecutionLifecycle {
398 async fn inventory(&self) -> Result<Vec<ScaleExecution>, ScaleReconcileError> {
399 let records = self
400 .manager
401 .managed_records()
402 .await
403 .map_err(lifecycle_error)?;
404 let mut inventory = Vec::new();
405 for record in records {
406 if record.labels.get(SCALE_MANAGED_LABEL).map(String::as_str) != Some("true") {
407 continue;
408 }
409 let service = required_label(&record.labels, SCALE_SERVICE_LABEL, &record.id)?;
410 let slot = required_label(&record.labels, SCALE_SLOT_LABEL, &record.id)?
411 .parse::<u32>()
412 .map_err(|error| {
413 ScaleReconcileError::Lifecycle(format!(
414 "scale execution {} has invalid slot label: {error}",
415 record.id
416 ))
417 })?;
418 let template_digest =
419 required_label(&record.labels, SCALE_TEMPLATE_DIGEST_LABEL, &record.id)?;
420 let guest_port = record
421 .labels
422 .get(SCALE_GUEST_PORT_LABEL)
423 .map(|value| {
424 value
425 .parse::<u16>()
426 .ok()
427 .and_then(NonZeroU16::new)
428 .ok_or_else(|| {
429 ScaleReconcileError::Lifecycle(format!(
430 "scale execution {} has invalid guest port label",
431 record.id
432 ))
433 })
434 })
435 .transpose()?;
436 let metadata = record.managed_execution.as_ref().ok_or_else(|| {
437 ScaleReconcileError::Lifecycle(format!(
438 "scale execution {} has no managed lifecycle metadata",
439 record.id
440 ))
441 })?;
442 let execution_id = ExecutionId::new(record.id.clone()).map_err(lifecycle_error)?;
443 let internal = ManagedExecutionState::from_status(&record.status)
444 .map_err(|error| ScaleReconcileError::Lifecycle(error.to_string()))?;
445 let phase = if internal == ManagedExecutionState::Removing {
446 InstancePhase::Removing
447 } else if internal.is_terminal() {
448 InstancePhase::Terminal
449 } else {
450 let status = self
451 .manager
452 .inspect(&execution_id)
453 .await
454 .map_err(lifecycle_error)?;
455 match status.state {
456 ExecutionState::Running => InstancePhase::Ready,
457 ExecutionState::Stopped | ExecutionState::Failed => InstancePhase::Terminal,
458 ExecutionState::Created | ExecutionState::Creating | ExecutionState::Paused => {
459 InstancePhase::Active
460 }
461 }
462 };
463 inventory.push(ScaleExecution {
464 execution_id,
465 generation: metadata.generation,
466 service,
467 slot,
468 template_digest,
469 guest_port,
470 phase,
471 });
472 }
473 Ok(inventory)
474 }
475
476 async fn ensure_running(
477 &self,
478 request: CreateExecutionRequest,
479 operation_id: OperationId,
480 ) -> Result<(), ScaleReconcileError> {
481 let outcome = self
482 .manager
483 .reconcile(&operation_id)
484 .await
485 .map_err(lifecycle_error)?;
486 let reservation = match outcome {
487 ReconcileOutcome::Ready(lease) => return self.ensure_lease_running(lease).await,
488 ReconcileOutcome::Created(reservation) => reservation,
489 ReconcileOutcome::Absent => match self.manager.create(request, &operation_id).await {
490 Ok(reservation) => reservation,
491 Err(create_error) => match self.manager.reconcile(&operation_id).await {
492 Ok(ReconcileOutcome::Ready(lease)) => {
493 return self.ensure_lease_running(lease).await;
494 }
495 Ok(ReconcileOutcome::Created(reservation)) => reservation,
496 _ => return Err(lifecycle_error(create_error)),
497 },
498 },
499 ReconcileOutcome::Creating => {
500 return Err(ScaleReconcileError::Lifecycle(format!(
501 "scale create operation {operation_id} is still converging"
502 )))
503 }
504 ReconcileOutcome::Failed => {
505 return Err(ScaleReconcileError::Lifecycle(format!(
506 "scale create operation {operation_id} is terminal"
507 )))
508 }
509 };
510
511 match self
512 .manager
513 .start(&reservation.execution_id, reservation.generation)
514 .await
515 {
516 Ok(_) => Ok(()),
517 Err(start_error) => match self.manager.reconcile(&operation_id).await {
518 Ok(ReconcileOutcome::Ready(lease)) => self.ensure_lease_running(lease).await,
519 _ => Err(lifecycle_error(start_error)),
520 },
521 }
522 }
523
524 async fn ensure_removed(&self, execution: &ScaleExecution) -> Result<(), ScaleReconcileError> {
525 if !matches!(
526 execution.phase,
527 InstancePhase::Terminal | InstancePhase::Removing
528 ) {
529 if let Err(kill_error) = self
530 .manager
531 .kill(&execution.execution_id, execution.generation)
532 .await
533 {
534 let status = self.manager.inspect(&execution.execution_id).await;
535 if !matches!(
536 status,
537 Ok(status)
538 if matches!(status.state, ExecutionState::Stopped | ExecutionState::Failed)
539 ) {
540 return Err(lifecycle_error(kill_error));
541 }
542 }
543 }
544 self.manager
545 .remove(&execution.execution_id, execution.generation)
546 .await
547 .map(|_| ())
548 .map_err(lifecycle_error)
549 }
550}
551
552fn required_label(
553 labels: &std::collections::HashMap<String, String>,
554 name: &str,
555 execution_id: &str,
556) -> Result<String, ScaleReconcileError> {
557 labels.get(name).cloned().ok_or_else(|| {
558 ScaleReconcileError::Lifecycle(format!(
559 "scale execution {execution_id} is missing label {name}"
560 ))
561 })
562}
563
564fn lifecycle_error(error: impl std::fmt::Display) -> ScaleReconcileError {
565 ScaleReconcileError::Lifecycle(error.to_string())
566}
567
568#[cfg(test)]
569mod tests {
570 use std::{
571 collections::{HashMap, HashSet},
572 sync::{
573 atomic::{AtomicBool, AtomicUsize, Ordering},
574 Mutex as StdMutex,
575 },
576 };
577
578 use a3s_box_core::{
579 ExecutionIsolation, ExecutionManagerError, ExecutionManagerResult, KillOutcome,
580 };
581 use chrono::Utc;
582 use tokio::{
583 io::{AsyncReadExt, AsyncWriteExt},
584 net::TcpStream,
585 };
586
587 use crate::{
588 BoxRecord, LocalExecutionBackend, LocalExecutionHandle, LocalExecutionObservation,
589 };
590
591 use super::*;
592
593 const CATALOG: &str = r#"service "api" { image = "api:v1" }"#;
594 const ENDPOINT_CATALOG: &str = r#"service "api" { image = "api:v1"; ports = ["0:8080"] }"#;
595
596 struct EchoConnector;
597
598 #[async_trait]
599 impl ExecutionPortConnector for EchoConnector {
600 async fn connect_port(
601 &self,
602 _execution_id: &ExecutionId,
603 _generation: ExecutionGeneration,
604 _port: NonZeroU16,
605 _timeout: std::time::Duration,
606 ) -> ExecutionManagerResult<a3s_box_core::ExecutionPortStream> {
607 let (client, mut server) = tokio::io::duplex(1_024);
608 tokio::spawn(async move {
609 let mut buffer = [0_u8; 1_024];
610 loop {
611 let count = match server.read(&mut buffer).await {
612 Ok(0) | Err(_) => return,
613 Ok(count) => count,
614 };
615 if server.write_all(&buffer[..count]).await.is_err() {
616 return;
617 }
618 }
619 });
620 Ok(Box::pin(client))
621 }
622 }
623
624 struct FakeLifecycle {
625 executions: Mutex<HashMap<String, ScaleExecution>>,
626 lose_next_create_response: AtomicBool,
627 }
628
629 impl FakeLifecycle {
630 fn new() -> Self {
631 Self {
632 executions: Mutex::new(HashMap::new()),
633 lose_next_create_response: AtomicBool::new(false),
634 }
635 }
636 }
637
638 #[async_trait]
639 impl ScaleExecutionLifecycle for FakeLifecycle {
640 async fn inventory(&self) -> Result<Vec<ScaleExecution>, ScaleReconcileError> {
641 Ok(self.executions.lock().await.values().cloned().collect())
642 }
643
644 async fn ensure_running(
645 &self,
646 request: CreateExecutionRequest,
647 _operation_id: OperationId,
648 ) -> Result<(), ScaleReconcileError> {
649 let service = request.labels[SCALE_SERVICE_LABEL].clone();
650 let slot = request.labels[SCALE_SLOT_LABEL].parse().unwrap();
651 let key = format!("{service}-{slot}");
652 self.executions
653 .lock()
654 .await
655 .entry(key.clone())
656 .or_insert(ScaleExecution {
657 execution_id: ExecutionId::new(key).unwrap(),
658 generation: ExecutionGeneration::INITIAL,
659 service,
660 slot,
661 template_digest: request.labels[SCALE_TEMPLATE_DIGEST_LABEL].clone(),
662 guest_port: template_guest_port(&request),
663 phase: InstancePhase::Ready,
664 });
665 if self.lose_next_create_response.swap(false, Ordering::SeqCst) {
666 return Err(ScaleReconcileError::Lifecycle("response lost".to_string()));
667 }
668 Ok(())
669 }
670
671 async fn ensure_removed(
672 &self,
673 execution: &ScaleExecution,
674 ) -> Result<(), ScaleReconcileError> {
675 self.executions
676 .lock()
677 .await
678 .retain(|_, current| current.execution_id != execution.execution_id);
679 Ok(())
680 }
681 }
682
683 fn reconciler(lifecycle: Arc<FakeLifecycle>) -> LocalScaleReconciler {
684 let catalog = ScaleServiceCatalog::from_acl_str(
685 CATALOG,
686 "gateway-scale",
687 ExecutionIsolation::Sandbox,
688 )
689 .unwrap();
690 LocalScaleReconciler::with_lifecycle(catalog, lifecycle)
691 }
692
693 #[tokio::test]
694 async fn converges_up_and_down_using_stable_slots() {
695 let lifecycle = Arc::new(FakeLifecycle::new());
696 let reconciler = reconciler(lifecycle.clone());
697 let up = reconciler.reconcile("api", 3).await.unwrap();
698 assert_eq!(up.ready_replicas, 3);
699 assert_eq!(up.created, 3);
700
701 let replay = reconciler.reconcile("api", 3).await.unwrap();
702 assert_eq!(replay.created, 0);
703 assert_eq!(replay.removed, 0);
704
705 let down = reconciler.reconcile("api", 1).await.unwrap();
706 assert_eq!(down.ready_replicas, 1);
707 assert_eq!(down.removed, 2);
708 let inventory = lifecycle.inventory().await.unwrap();
709 assert_eq!(inventory[0].slot, 0);
710 }
711
712 #[tokio::test]
713 async fn downscale_waits_for_endpoint_relays_before_removing_the_execution() {
714 let lifecycle = Arc::new(FakeLifecycle::new());
715 let catalog = ScaleServiceCatalog::from_acl_str(
716 ENDPOINT_CATALOG,
717 "gateway-scale",
718 ExecutionIsolation::Sandbox,
719 )
720 .unwrap();
721 let reconciler = Arc::new(LocalScaleReconciler::with_lifecycle_and_endpoint(
722 catalog,
723 lifecycle.clone(),
724 Arc::new(EchoConnector),
725 ScaleEndpointConfig::loopback().with_drain_timeout(std::time::Duration::from_secs(5)),
726 ));
727 let up = reconciler.reconcile("api", 1).await.unwrap();
728 let address = up.endpoints[0]
729 .url
730 .strip_prefix("http://")
731 .unwrap()
732 .parse::<std::net::SocketAddr>()
733 .unwrap();
734 let mut stream = TcpStream::connect(address).await.unwrap();
735 stream.write_all(b"before").await.unwrap();
736 let mut reply = [0_u8; 6];
737 stream.read_exact(&mut reply).await.unwrap();
738
739 let downscale_reconciler = Arc::clone(&reconciler);
740 let downscale = tokio::spawn(async move { downscale_reconciler.reconcile("api", 0).await });
741 tokio::time::sleep(std::time::Duration::from_millis(20)).await;
742 assert_eq!(lifecycle.inventory().await.unwrap().len(), 1);
743 stream.write_all(b"during").await.unwrap();
744 stream.read_exact(&mut reply).await.unwrap();
745 assert_eq!(&reply, b"during");
746 drop(stream);
747
748 let report = tokio::time::timeout(std::time::Duration::from_secs(1), downscale)
749 .await
750 .expect("downscale did not finish after endpoint drain")
751 .unwrap()
752 .unwrap();
753 assert_eq!(report.removed, 1);
754 assert!(lifecycle.inventory().await.unwrap().is_empty());
755 }
756
757 #[tokio::test]
758 async fn ambiguous_create_is_adopted_without_a_duplicate_on_retry() {
759 let lifecycle = Arc::new(FakeLifecycle::new());
760 lifecycle
761 .lose_next_create_response
762 .store(true, Ordering::SeqCst);
763 let reconciler = reconciler(lifecycle.clone());
764 assert!(reconciler.reconcile("api", 1).await.is_err());
765
766 let recovered = reconciler.reconcile("api", 1).await.unwrap();
767 assert_eq!(recovered.ready_replicas, 1);
768 assert_eq!(recovered.created, 0);
769 assert_eq!(lifecycle.inventory().await.unwrap().len(), 1);
770 }
771
772 struct ConvergingLifecycle {
773 execution: Mutex<Option<ScaleExecution>>,
774 ensure_calls: AtomicUsize,
775 }
776
777 #[async_trait]
778 impl ScaleExecutionLifecycle for ConvergingLifecycle {
779 async fn inventory(&self) -> Result<Vec<ScaleExecution>, ScaleReconcileError> {
780 Ok(self.execution.lock().await.clone().into_iter().collect())
781 }
782
783 async fn ensure_running(
784 &self,
785 request: CreateExecutionRequest,
786 _operation_id: OperationId,
787 ) -> Result<(), ScaleReconcileError> {
788 let phase = if self.ensure_calls.fetch_add(1, Ordering::SeqCst) == 0 {
789 InstancePhase::Active
790 } else {
791 InstancePhase::Ready
792 };
793 *self.execution.lock().await = Some(ScaleExecution {
794 execution_id: ExecutionId::new("api-0").unwrap(),
795 generation: ExecutionGeneration::INITIAL,
796 service: request.labels[SCALE_SERVICE_LABEL].clone(),
797 slot: request.labels[SCALE_SLOT_LABEL].parse().unwrap(),
798 template_digest: request.labels[SCALE_TEMPLATE_DIGEST_LABEL].clone(),
799 guest_port: template_guest_port(&request),
800 phase,
801 });
802 Ok(())
803 }
804
805 async fn ensure_removed(
806 &self,
807 _execution: &ScaleExecution,
808 ) -> Result<(), ScaleReconcileError> {
809 *self.execution.lock().await = None;
810 Ok(())
811 }
812 }
813
814 #[tokio::test]
815 async fn retained_active_slot_continues_startup_convergence() {
816 let lifecycle = Arc::new(ConvergingLifecycle {
817 execution: Mutex::new(None),
818 ensure_calls: AtomicUsize::new(0),
819 });
820 let catalog = ScaleServiceCatalog::from_acl_str(
821 CATALOG,
822 "gateway-scale",
823 ExecutionIsolation::Sandbox,
824 )
825 .unwrap();
826 let reconciler = LocalScaleReconciler::with_lifecycle(catalog, lifecycle.clone());
827
828 let first = reconciler.reconcile("api", 1).await.unwrap();
829 assert_eq!(first.ready_replicas, 0);
830 assert_eq!(first.created, 1);
831
832 let second = reconciler.reconcile("api", 1).await.unwrap();
833 assert_eq!(second.ready_replicas, 1);
834 assert_eq!(second.created, 0);
835 assert_eq!(lifecycle.ensure_calls.load(Ordering::SeqCst), 2);
836
837 let replay = reconciler.reconcile("api", 1).await.unwrap();
838 assert_eq!(replay.created, 0);
839 assert_eq!(lifecycle.ensure_calls.load(Ordering::SeqCst), 2);
840 }
841
842 #[tokio::test]
843 async fn unknown_service_fails_before_any_lifecycle_side_effect() {
844 let lifecycle = Arc::new(FakeLifecycle::new());
845 let reconciler = reconciler(lifecycle.clone());
846 assert!(matches!(
847 reconciler.reconcile("missing", 1).await,
848 Err(ScaleReconcileError::UnknownService(_))
849 ));
850 assert!(lifecycle.inventory().await.unwrap().is_empty());
851 }
852
853 struct RecordingBackend {
854 running: StdMutex<HashSet<String>>,
855 paused: StdMutex<HashSet<String>>,
856 starts: AtomicUsize,
857 lose_next_start_response: AtomicBool,
858 }
859
860 impl RecordingBackend {
861 fn new() -> Self {
862 Self {
863 running: StdMutex::new(HashSet::new()),
864 paused: StdMutex::new(HashSet::new()),
865 starts: AtomicUsize::new(0),
866 lose_next_start_response: AtomicBool::new(false),
867 }
868 }
869
870 fn handle(record: &BoxRecord) -> LocalExecutionHandle {
871 LocalExecutionHandle {
872 started_at: Utc::now(),
873 pid: None,
874 pid_start_time: None,
875 exec_socket_path: record.box_dir.join("sockets/exec.sock"),
876 console_log: record.box_dir.join("logs/console.log"),
877 anonymous_volumes: Vec::new(),
878 oci_runtime: None,
879 }
880 }
881 }
882
883 #[async_trait]
884 impl LocalExecutionBackend for RecordingBackend {
885 async fn start(&self, record: &BoxRecord) -> ExecutionManagerResult<LocalExecutionHandle> {
886 self.starts.fetch_add(1, Ordering::SeqCst);
887 self.running.lock().unwrap().insert(record.id.clone());
888 if self.lose_next_start_response.swap(false, Ordering::SeqCst) {
889 return Err(ExecutionManagerError::Unavailable(
890 "start response lost".to_string(),
891 ));
892 }
893 Ok(Self::handle(record))
894 }
895
896 async fn inspect(
897 &self,
898 record: &BoxRecord,
899 ) -> ExecutionManagerResult<LocalExecutionObservation> {
900 let state = if self.running.lock().unwrap().contains(&record.id) {
901 Some(ExecutionState::Running)
902 } else if self.paused.lock().unwrap().contains(&record.id) {
903 Some(ExecutionState::Paused)
904 } else {
905 None
906 };
907 if let Some(state) = state {
908 Ok(LocalExecutionObservation {
909 state,
910 handle: Some(Self::handle(record)),
911 exit_code: None,
912 })
913 } else {
914 Ok(LocalExecutionObservation {
915 state: ExecutionState::Stopped,
916 handle: None,
917 exit_code: Some(0),
918 })
919 }
920 }
921
922 async fn pause(
923 &self,
924 record: &BoxRecord,
925 _keep_memory: bool,
926 ) -> ExecutionManagerResult<LocalExecutionHandle> {
927 self.running.lock().unwrap().remove(&record.id);
928 self.paused.lock().unwrap().insert(record.id.clone());
929 Ok(Self::handle(record))
930 }
931
932 async fn resume(&self, record: &BoxRecord) -> ExecutionManagerResult<LocalExecutionHandle> {
933 self.paused.lock().unwrap().remove(&record.id);
934 self.running.lock().unwrap().insert(record.id.clone());
935 Ok(Self::handle(record))
936 }
937
938 async fn kill(&self, record: &BoxRecord) -> ExecutionManagerResult<KillOutcome> {
939 let was_running = self.running.lock().unwrap().remove(&record.id);
940 let was_paused = self.paused.lock().unwrap().remove(&record.id);
941 let removed = was_running || was_paused;
942 Ok(if removed {
943 KillOutcome::Killed
944 } else {
945 KillOutcome::AlreadyStopped
946 })
947 }
948 }
949
950 #[tokio::test]
951 async fn local_execution_facade_converges_and_recovers_lost_start_response() {
952 let directory = tempfile::tempdir().unwrap();
953 let home = directory.path().join("home");
954 let state = directory.path().join("boxes.json");
955 let backend = Arc::new(RecordingBackend::new());
956 backend
957 .lose_next_start_response
958 .store(true, Ordering::SeqCst);
959 let manager = LocalExecutionManager::new(&state, &home, backend.clone());
960 let catalog = ScaleServiceCatalog::from_acl_str(
961 CATALOG,
962 "gateway-scale",
963 ExecutionIsolation::Sandbox,
964 )
965 .unwrap();
966 let reconciler = LocalScaleReconciler::new(manager, catalog.clone());
967
968 let up = reconciler.reconcile("api", 2).await.unwrap();
969 assert_eq!(up.ready_replicas, 2);
970 assert_eq!(backend.starts.load(Ordering::SeqCst), 2);
971
972 let reopened = LocalExecutionManager::new(&state, &home, backend.clone());
973 let restarted = LocalScaleReconciler::new(reopened.clone(), catalog);
974 let replay = restarted.reconcile("api", 2).await.unwrap();
975 assert_eq!(replay.created, 0);
976 assert_eq!(replay.ready_replicas, 2);
977 assert_eq!(backend.starts.load(Ordering::SeqCst), 2);
978
979 let down = restarted.reconcile("api", 0).await.unwrap();
980 assert_eq!(down.removed, 2);
981 assert_eq!(down.ready_replicas, 0);
982 assert!(reopened.managed_records().await.unwrap().is_empty());
983 }
984
985 #[tokio::test]
986 async fn local_execution_facade_resumes_a_retained_paused_slot() {
987 let directory = tempfile::tempdir().unwrap();
988 let home = directory.path().join("home");
989 let state = directory.path().join("boxes.json");
990 let backend = Arc::new(RecordingBackend::new());
991 let manager = LocalExecutionManager::new(&state, &home, backend);
992 let catalog = ScaleServiceCatalog::from_acl_str(
993 CATALOG,
994 "gateway-scale",
995 ExecutionIsolation::Sandbox,
996 )
997 .unwrap();
998 let reconciler = LocalScaleReconciler::new(manager.clone(), catalog);
999
1000 let up = reconciler.reconcile("api", 1).await.unwrap();
1001 assert_eq!(up.ready_replicas, 1);
1002 let record = manager.managed_records().await.unwrap().pop().unwrap();
1003 let execution_id = ExecutionId::new(record.id).unwrap();
1004 let generation = record.managed_execution.unwrap().generation;
1005 manager
1006 .pause(&execution_id, generation, true)
1007 .await
1008 .unwrap();
1009
1010 let resumed = reconciler.reconcile("api", 1).await.unwrap();
1011 assert_eq!(resumed.ready_replicas, 1);
1012 assert_eq!(resumed.created, 0);
1013 assert_eq!(
1014 manager.inspect(&execution_id).await.unwrap().state,
1015 ExecutionState::Running
1016 );
1017 }
1018}