1use std::{net::SocketAddr, sync::Arc};
4
5use a3s_box_core::scale::{ScaleObservation, ScaleOperationConflict, ScaleOperationRequest};
6use axum::{
7 extract::{DefaultBodyLimit, Path, State},
8 http::StatusCode,
9 response::{IntoResponse, Response},
10 routing::get,
11 Json, Router,
12};
13use tokio::sync::Mutex;
14
15use super::{
16 DurableScaleAuthority, LocalScaleReconciler, ScaleAuthorityError, ScaleReconcileError,
17};
18
19pub type SharedScaleAuthority = Arc<ScaleApiState>;
20
21pub struct ScaleApiState {
22 authority: Mutex<DurableScaleAuthority>,
23 reconciler: Option<Arc<LocalScaleReconciler>>,
24}
25
26impl ScaleApiState {
27 pub fn authority_only(authority: DurableScaleAuthority) -> SharedScaleAuthority {
28 Arc::new(Self {
29 authority: Mutex::new(authority),
30 reconciler: None,
31 })
32 }
33
34 pub fn with_reconciler(
35 authority: DurableScaleAuthority,
36 reconciler: LocalScaleReconciler,
37 ) -> SharedScaleAuthority {
38 Arc::new(Self {
39 authority: Mutex::new(authority),
40 reconciler: Some(Arc::new(reconciler)),
41 })
42 }
43
44 async fn observation(&self, service: &str) -> Result<ScaleObservation, ScaleReconcileError> {
45 let mut observation = self.authority.lock().await.observation(service);
46 if let Some(reconciler) = &self.reconciler {
47 let workloads = reconciler
48 .observation(service, observation.replicas)
49 .await?;
50 observation.ready_replicas = workloads.ready_replicas;
51 observation.endpoints = workloads.endpoints;
52 }
53 Ok(observation)
54 }
55
56 async fn reconcile_desired_services(&self) {
57 let Some(reconciler) = &self.reconciler else {
58 return;
59 };
60 for service in reconciler.services() {
61 let desired = self.authority.lock().await.observation(&service).replicas;
62 if let Err(error) = reconciler.reconcile(&service, desired).await {
63 tracing::warn!(%service, %error, "Scale workload reconciliation failed");
64 }
65 }
66 }
67}
68
69pub fn scale_router(authority: SharedScaleAuthority) -> Router {
70 Router::new()
71 .route("/v1/scale/:service", get(observe).post(apply))
72 .layer(DefaultBodyLimit::max(64 * 1024))
73 .with_state(authority)
74}
75
76pub async fn serve_scale_api(
77 address: SocketAddr,
78 authority: SharedScaleAuthority,
79) -> Result<(), std::io::Error> {
80 let convergence_state = Arc::downgrade(&authority);
81 let convergence = tokio::spawn(async move {
82 let mut interval = tokio::time::interval(std::time::Duration::from_secs(2));
83 loop {
84 interval.tick().await;
85 let Some(state) = convergence_state.upgrade() else {
86 return;
87 };
88 state.reconcile_desired_services().await;
89 }
90 });
91 let result = axum::Server::bind(&address)
92 .serve(scale_router(authority).into_make_service())
93 .await
94 .map_err(std::io::Error::other);
95 convergence.abort();
96 result
97}
98
99async fn observe(
100 Path(service): Path<String>,
101 State(authority): State<SharedScaleAuthority>,
102) -> Response {
103 match authority.observation(&service).await {
104 Ok(observation) => (StatusCode::OK, Json(observation)).into_response(),
105 Err(error) => {
106 let (status, code) = if matches!(&error, ScaleReconcileError::UnknownService(_)) {
107 (StatusCode::NOT_FOUND, "unknown_service")
108 } else {
109 (StatusCode::SERVICE_UNAVAILABLE, "observation_failed")
110 };
111 conflict_response(
112 status,
113 ScaleOperationConflict {
114 code: code.to_string(),
115 message: error.to_string(),
116 observation: authority.authority.lock().await.observation(&service),
117 },
118 )
119 }
120 }
121}
122
123async fn apply(
124 Path(service): Path<String>,
125 State(authority): State<SharedScaleAuthority>,
126 Json(request): Json<ScaleOperationRequest>,
127) -> Response {
128 if request.service != service {
129 return conflict_response(
130 StatusCode::UNPROCESSABLE_ENTITY,
131 ScaleOperationConflict {
132 code: "service_mismatch".to_string(),
133 message: format!(
134 "request service {:?} does not match path service {:?}",
135 request.service, service
136 ),
137 observation: authority.authority.lock().await.observation(&service),
138 },
139 );
140 }
141
142 if let Some(reconciler) = &authority.reconciler {
143 if !reconciler.knows_service(&service) {
144 return conflict_response(
145 StatusCode::NOT_FOUND,
146 ScaleOperationConflict {
147 code: "unknown_service".to_string(),
148 message: format!("service {service:?} has no Box workload template"),
149 observation: authority.authority.lock().await.observation(&service),
150 },
151 );
152 }
153 }
154
155 let applied = {
156 let mut authority = authority.authority.lock().await;
157 authority.apply(&request)
158 };
159 match applied {
160 Ok(mut response) => {
161 if let Some(reconciler) = &authority.reconciler {
162 match reconciler
163 .reconcile(&service, request.desired_replicas)
164 .await
165 {
166 Ok(report) => {
167 response.actual_replicas = report.ready_replicas;
168 response.message = format!(
169 "Box reconciled service '{}' to {} ready replicas",
170 service, report.ready_replicas
171 );
172 match authority
173 .authority
174 .lock()
175 .await
176 .finalize(&request, response.clone())
177 {
178 Ok(response) => (StatusCode::OK, Json(response)).into_response(),
179 Err(error) => {
180 authority_error_response(&authority, &service, error).await
181 }
182 }
183 }
184 Err(error) => {
185 let observation = match authority.observation(&service).await {
186 Ok(observation) => observation,
187 Err(_) => authority.authority.lock().await.observation(&service),
188 };
189 conflict_response(
190 StatusCode::SERVICE_UNAVAILABLE,
191 ScaleOperationConflict {
192 code: "reconcile_failed".to_string(),
193 message: error.to_string(),
194 observation,
195 },
196 )
197 }
198 }
199 } else {
200 (StatusCode::OK, Json(response)).into_response()
201 }
202 }
203 Err(ScaleAuthorityError::Conflict(_, conflict)) => {
204 let status = match conflict.code.as_str() {
205 "stale_revision" | "operation_conflict" => StatusCode::CONFLICT,
206 "capacity_exceeded" => StatusCode::INSUFFICIENT_STORAGE,
207 _ => StatusCode::UNPROCESSABLE_ENTITY,
208 };
209 conflict_response(status, conflict)
210 }
211 Err(error) => authority_error_response(&authority, &service, error).await,
212 }
213}
214
215async fn authority_error_response(
216 authority: &SharedScaleAuthority,
217 service: &str,
218 error: ScaleAuthorityError,
219) -> Response {
220 (
221 StatusCode::INTERNAL_SERVER_ERROR,
222 Json(ScaleOperationConflict {
223 code: "authority_state_error".to_string(),
224 message: error.to_string(),
225 observation: authority.authority.lock().await.observation(service),
226 }),
227 )
228 .into_response()
229}
230
231fn conflict_response(status: StatusCode, conflict: ScaleOperationConflict) -> Response {
232 (status, Json(conflict)).into_response()
233}
234
235#[cfg(test)]
236mod tests {
237 use std::{
238 collections::HashSet,
239 sync::{
240 atomic::{AtomicUsize, Ordering},
241 Mutex as StdMutex,
242 },
243 };
244
245 use super::*;
246 use a3s_box_core::{
247 scale::{ScaleDirection, ScaleOperationResponse, SCALE_OPERATION_SCHEMA_VERSION},
248 ExecutionIsolation, ExecutionManagerError, ExecutionManagerResult, ExecutionState,
249 KillOutcome,
250 };
251 use async_trait::async_trait;
252 use chrono::Utc;
253
254 use crate::{
255 BoxRecord, LocalExecutionBackend, LocalExecutionHandle, LocalExecutionManager,
256 LocalExecutionObservation, ScaleServiceCatalog,
257 };
258
259 fn request(operation_id: &str, revision: &str) -> ScaleOperationRequest {
260 ScaleOperationRequest {
261 schema_version: SCALE_OPERATION_SCHEMA_VERSION,
262 operation_id: operation_id.to_string(),
263 service: "api".to_string(),
264 expected_revision: Some(revision.to_string()),
265 direction: ScaleDirection::Up,
266 current_replicas: 0,
267 desired_replicas: 2,
268 reason: "load".to_string(),
269 }
270 }
271
272 #[tokio::test]
273 async fn tcp_boundary_applies_replays_and_rejects_stale_operations() {
274 let directory = tempfile::tempdir().unwrap();
275 let authority = ScaleApiState::authority_only(
276 DurableScaleAuthority::open(directory.path().join("state.json"), 10).unwrap(),
277 );
278 let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
279 let address = listener.local_addr().unwrap();
280 drop(listener);
281 let server = tokio::spawn(serve_scale_api(address, authority));
282 let client = reqwest::Client::new();
283 let url = format!("http://{address}/v1/scale/api");
284 for _ in 0..50 {
285 if client.get(&url).send().await.is_ok() {
286 break;
287 }
288 tokio::time::sleep(std::time::Duration::from_millis(10)).await;
289 }
290
291 let operation = request("scale-v1-http", "0");
292 let accepted = client.post(&url).json(&operation).send().await.unwrap();
293 assert_eq!(accepted.status(), StatusCode::OK);
294 let accepted: ScaleOperationResponse = accepted.json().await.unwrap();
295
296 let replayed: ScaleOperationResponse = client
297 .post(&url)
298 .json(&operation)
299 .send()
300 .await
301 .unwrap()
302 .json()
303 .await
304 .unwrap();
305 assert_eq!(replayed, accepted);
306
307 let stale = client
308 .post(&url)
309 .json(&request("scale-v1-stale", "0"))
310 .send()
311 .await
312 .unwrap();
313 assert_eq!(stale.status(), StatusCode::CONFLICT);
314 server.abort();
315 }
316
317 struct RecordingBackend {
318 running: StdMutex<HashSet<String>>,
319 starts: AtomicUsize,
320 }
321
322 impl RecordingBackend {
323 fn new() -> Self {
324 Self {
325 running: StdMutex::new(HashSet::new()),
326 starts: AtomicUsize::new(0),
327 }
328 }
329
330 fn handle(record: &BoxRecord) -> LocalExecutionHandle {
331 LocalExecutionHandle {
332 started_at: Utc::now(),
333 pid: None,
334 pid_start_time: None,
335 exec_socket_path: record.box_dir.join("sockets/exec.sock"),
336 console_log: record.box_dir.join("logs/console.log"),
337 anonymous_volumes: Vec::new(),
338 oci_runtime: None,
339 }
340 }
341 }
342
343 #[async_trait]
344 impl LocalExecutionBackend for RecordingBackend {
345 async fn start(&self, record: &BoxRecord) -> ExecutionManagerResult<LocalExecutionHandle> {
346 self.starts.fetch_add(1, Ordering::SeqCst);
347 self.running.lock().unwrap().insert(record.id.clone());
348 Ok(Self::handle(record))
349 }
350
351 async fn inspect(
352 &self,
353 record: &BoxRecord,
354 ) -> ExecutionManagerResult<LocalExecutionObservation> {
355 if self.running.lock().unwrap().contains(&record.id) {
356 Ok(LocalExecutionObservation {
357 state: ExecutionState::Running,
358 handle: Some(Self::handle(record)),
359 exit_code: None,
360 })
361 } else {
362 Ok(LocalExecutionObservation {
363 state: ExecutionState::Stopped,
364 handle: None,
365 exit_code: Some(0),
366 })
367 }
368 }
369
370 async fn pause(
371 &self,
372 _record: &BoxRecord,
373 _keep_memory: bool,
374 ) -> ExecutionManagerResult<LocalExecutionHandle> {
375 Err(ExecutionManagerError::Unavailable(
376 "pause unsupported".to_string(),
377 ))
378 }
379
380 async fn resume(
381 &self,
382 _record: &BoxRecord,
383 ) -> ExecutionManagerResult<LocalExecutionHandle> {
384 Err(ExecutionManagerError::Unavailable(
385 "resume unsupported".to_string(),
386 ))
387 }
388
389 async fn kill(&self, record: &BoxRecord) -> ExecutionManagerResult<KillOutcome> {
390 let removed = self.running.lock().unwrap().remove(&record.id);
391 Ok(if removed {
392 KillOutcome::Killed
393 } else {
394 KillOutcome::AlreadyStopped
395 })
396 }
397 }
398
399 fn reconciled_state(
400 authority_path: &std::path::Path,
401 box_state_path: &std::path::Path,
402 home: &std::path::Path,
403 backend: Arc<RecordingBackend>,
404 ) -> SharedScaleAuthority {
405 let authority = DurableScaleAuthority::open(authority_path, 10).unwrap();
406 let manager = LocalExecutionManager::new(box_state_path, home, backend);
407 let catalog = ScaleServiceCatalog::from_acl_str(
408 r#"service "api" { image = "api:v1" }"#,
409 "gateway-scale",
410 ExecutionIsolation::Sandbox,
411 )
412 .unwrap();
413 ScaleApiState::with_reconciler(authority, LocalScaleReconciler::new(manager, catalog))
414 }
415
416 async fn wait_for_server(client: &reqwest::Client, url: &str) {
417 for _ in 0..100 {
418 if client.get(url).send().await.is_ok() {
419 return;
420 }
421 tokio::time::sleep(std::time::Duration::from_millis(10)).await;
422 }
423 panic!("scale API did not become ready at {url}");
424 }
425
426 #[tokio::test]
427 async fn tcp_boundary_reconciles_real_facade_and_recovers_after_restart() {
428 let directory = tempfile::tempdir().unwrap();
429 let authority_path = directory.path().join("scale-authority.json");
430 let box_state_path = directory.path().join("boxes.json");
431 let home = directory.path().join("home");
432 let backend = Arc::new(RecordingBackend::new());
433 let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
434 let address = listener.local_addr().unwrap();
435 drop(listener);
436 let client = reqwest::Client::new();
437 let url = format!("http://{address}/v1/scale/api");
438
439 let first_state =
440 reconciled_state(&authority_path, &box_state_path, &home, backend.clone());
441 let first_server = tokio::spawn(serve_scale_api(address, first_state));
442 wait_for_server(&client, &url).await;
443 let up = request("scale-v1-facade-up", "0");
444 let accepted: ScaleOperationResponse = client
445 .post(&url)
446 .json(&up)
447 .send()
448 .await
449 .unwrap()
450 .error_for_status()
451 .unwrap()
452 .json()
453 .await
454 .unwrap();
455 assert_eq!(accepted.actual_replicas, 2);
456 assert_eq!(backend.starts.load(Ordering::SeqCst), 2);
457 let observation: ScaleObservation = client
458 .get(&url)
459 .send()
460 .await
461 .unwrap()
462 .error_for_status()
463 .unwrap()
464 .json()
465 .await
466 .unwrap();
467 assert_eq!(observation.replicas, 2);
468 assert_eq!(observation.ready_replicas, 2);
469 assert!(observation.endpoints.is_empty());
470 first_server.abort();
471 let _ = first_server.await;
472
473 let restarted_state =
474 reconciled_state(&authority_path, &box_state_path, &home, backend.clone());
475 let restarted_server = tokio::spawn(serve_scale_api(address, restarted_state));
476 wait_for_server(&client, &url).await;
477 let replayed: ScaleOperationResponse = client
478 .post(&url)
479 .json(&up)
480 .send()
481 .await
482 .unwrap()
483 .error_for_status()
484 .unwrap()
485 .json()
486 .await
487 .unwrap();
488 assert_eq!(replayed, accepted);
489 assert_eq!(backend.starts.load(Ordering::SeqCst), 2);
490
491 let down = ScaleOperationRequest {
492 operation_id: "scale-v1-facade-down".to_string(),
493 expected_revision: Some("1".to_string()),
494 direction: ScaleDirection::Down,
495 current_replicas: 2,
496 desired_replicas: 0,
497 ..request("unused", "unused")
498 };
499 let removed: ScaleOperationResponse = client
500 .post(&url)
501 .json(&down)
502 .send()
503 .await
504 .unwrap()
505 .error_for_status()
506 .unwrap()
507 .json()
508 .await
509 .unwrap();
510 assert_eq!(removed.actual_replicas, 0);
511 assert!(backend.running.lock().unwrap().is_empty());
512 restarted_server.abort();
513 let _ = restarted_server.await;
514 }
515}