Skip to main content

a3s_box_runtime/scale/
catalog.rs

1//! Box-owned workload templates for Gateway-driven scaling.
2
3use std::{collections::BTreeMap, path::Path};
4
5use a3s_box_core::{
6    compose::ComposeConfig, CreateExecutionRequest, ExecutionIsolation, ExecutionRecordPolicy,
7};
8use sha2::{Digest, Sha256};
9use thiserror::Error;
10
11use crate::ComposeRuntimePlan;
12
13pub const SCALE_MANAGED_LABEL: &str = "com.a3s.scale.managed";
14pub const SCALE_SERVICE_LABEL: &str = "com.a3s.scale.service";
15pub const SCALE_SLOT_LABEL: &str = "com.a3s.scale.slot";
16pub const SCALE_TEMPLATE_DIGEST_LABEL: &str = "com.a3s.scale.template-digest";
17pub const SCALE_GUEST_PORT_LABEL: &str = "com.a3s.scale.guest-port";
18
19#[derive(Debug, Error)]
20pub enum ScaleCatalogError {
21    #[error("failed to read scale service catalog {path}: {source}")]
22    Read {
23        path: String,
24        source: std::io::Error,
25    },
26    #[error("invalid scale service catalog: {0}")]
27    Invalid(String),
28}
29
30/// Validated Compose ACL templates owned by the Box execution plane.
31#[derive(Debug, Clone)]
32pub struct ScaleServiceCatalog {
33    plan: ComposeRuntimePlan,
34    isolation: ExecutionIsolation,
35}
36
37impl ScaleServiceCatalog {
38    pub fn from_acl_file(
39        path: &Path,
40        project_name: impl Into<String>,
41        isolation: ExecutionIsolation,
42    ) -> Result<Self, ScaleCatalogError> {
43        if path.extension().and_then(|extension| extension.to_str()) != Some("acl") {
44            return Err(ScaleCatalogError::Invalid(format!(
45                "scale service catalog {} must use the .acl format",
46                path.display()
47            )));
48        }
49        let source = std::fs::read_to_string(path).map_err(|source| ScaleCatalogError::Read {
50            path: path.display().to_string(),
51            source,
52        })?;
53        Self::from_acl_str_with_base_dir(
54            &source,
55            project_name,
56            path.parent().unwrap_or_else(|| Path::new(".")),
57            isolation,
58        )
59    }
60
61    pub fn from_acl_str(
62        source: &str,
63        project_name: impl Into<String>,
64        isolation: ExecutionIsolation,
65    ) -> Result<Self, ScaleCatalogError> {
66        Self::from_acl_str_with_base_dir(source, project_name, Path::new("."), isolation)
67    }
68
69    fn from_acl_str_with_base_dir(
70        source: &str,
71        project_name: impl Into<String>,
72        base_dir: &Path,
73        isolation: ExecutionIsolation,
74    ) -> Result<Self, ScaleCatalogError> {
75        let config = ComposeConfig::from_acl_str(source)
76            .map_err(|error| ScaleCatalogError::Invalid(error.to_string()))?;
77        validate_stateless_templates(&config)?;
78        let plan = ComposeRuntimePlan::with_base_dir(project_name, config, base_dir)
79            .map_err(|error| ScaleCatalogError::Invalid(error.to_string()))?;
80        Ok(Self { plan, isolation })
81    }
82
83    pub fn contains(&self, service: &str) -> bool {
84        self.plan.config.services.contains_key(service)
85    }
86
87    pub fn services(&self) -> Vec<String> {
88        let mut services = self
89            .plan
90            .config
91            .services
92            .keys()
93            .cloned()
94            .collect::<Vec<_>>();
95        services.sort();
96        services
97    }
98
99    pub fn guest_port(&self, service: &str) -> Result<Option<u16>, ScaleCatalogError> {
100        let service_config = self.plan.config.services.get(service).ok_or_else(|| {
101            ScaleCatalogError::Invalid(format!("service {service:?} not found in compose config"))
102        })?;
103        service_guest_port(service, service_config)
104    }
105
106    pub fn create_request(
107        &self,
108        service: &str,
109        slot: u32,
110    ) -> Result<CreateExecutionRequest, ScaleCatalogError> {
111        let mut config = self
112            .plan
113            .build_box_config(service, None)
114            .map_err(|error| ScaleCatalogError::Invalid(error.to_string()))?;
115        config.isolation = self.isolation;
116        let guest_port = self.guest_port(service)?;
117        // Scale endpoints are owned by a generation-fenced host relay. Do not
118        // also ask the VM backend to publish the Compose marker port.
119        config.port_map.clear();
120        let encoded = serde_json::to_vec(&(&config, guest_port)).map_err(|error| {
121            ScaleCatalogError::Invalid(format!(
122                "failed to encode service {service:?} template: {error}"
123            ))
124        })?;
125        let digest = format!("sha256:{:x}", Sha256::digest(encoded));
126        let mut labels = BTreeMap::new();
127        labels.insert(SCALE_MANAGED_LABEL.to_string(), "true".to_string());
128        labels.insert(SCALE_SERVICE_LABEL.to_string(), service.to_string());
129        labels.insert(SCALE_SLOT_LABEL.to_string(), slot.to_string());
130        labels.insert(SCALE_TEMPLATE_DIGEST_LABEL.to_string(), digest);
131        if let Some(guest_port) = guest_port {
132            labels.insert(SCALE_GUEST_PORT_LABEL.to_string(), guest_port.to_string());
133        }
134
135        Ok(CreateExecutionRequest {
136            external_sandbox_id: format!("scale-{service}-{slot}"),
137            config,
138            labels,
139            policy: ExecutionRecordPolicy::default(),
140            rootfs_snapshot_id: None,
141        })
142    }
143}
144
145fn validate_stateless_templates(config: &ComposeConfig) -> Result<(), ScaleCatalogError> {
146    for (name, service) in &config.services {
147        if !service.depends_on.services().is_empty() {
148            return Err(ScaleCatalogError::Invalid(format!(
149                "service {name:?} uses depends_on; independently scaled templates cannot own dependency lifecycles"
150            )));
151        }
152        service_guest_port(name, service)?;
153        if !service.volumes.is_empty() {
154            return Err(ScaleCatalogError::Invalid(format!(
155                "service {name:?} mounts shared volumes; Gateway scaling currently accepts only stateless templates"
156            )));
157        }
158        if !service.networks.names().is_empty() {
159            return Err(ScaleCatalogError::Invalid(format!(
160                "service {name:?} selects a Compose network; independently scaled replicas currently use Box's default TSI network"
161            )));
162        }
163    }
164    Ok(())
165}
166
167fn service_guest_port(
168    name: &str,
169    service: &a3s_box_core::compose::ServiceConfig,
170) -> Result<Option<u16>, ScaleCatalogError> {
171    let [entry] = service.ports.as_slice() else {
172        return if service.ports.is_empty() {
173            Ok(None)
174        } else {
175            Err(ScaleCatalogError::Invalid(format!(
176                "service {name:?} declares multiple ports; Gateway scaling currently publishes exactly one HTTP endpoint"
177            )))
178        };
179    };
180    let mapping = a3s_box_core::parse_port_mapping(entry).map_err(ScaleCatalogError::Invalid)?;
181    if mapping.host_port != 0 {
182        return Err(ScaleCatalogError::Invalid(format!(
183            "service {name:?} publishes fixed host port {}; use 0:{} for a runtime-discovered endpoint",
184            mapping.host_port, mapping.guest_port
185        )));
186    }
187    Ok(Some(mapping.guest_port))
188}
189
190#[cfg(test)]
191mod tests {
192    use super::*;
193
194    const CATALOG: &str = r#"
195        service "worker" {
196            image       = "ghcr.io/a3s-lab/worker:v1"
197            command     = ["serve", "--port", "8080"]
198            environment = { MODE = "production" }
199            cpus        = 2
200            mem_limit   = "768m"
201        }
202        service "api" {
203            image = "ghcr.io/a3s-lab/api:v2"
204        }
205    "#;
206
207    #[test]
208    fn catalog_builds_deterministic_labeled_execution_templates() {
209        let catalog = ScaleServiceCatalog::from_acl_str(
210            CATALOG,
211            "gateway-scale",
212            ExecutionIsolation::Sandbox,
213        )
214        .unwrap();
215        assert_eq!(catalog.services(), vec!["api", "worker"]);
216
217        let first = catalog.create_request("worker", 3).unwrap();
218        let replay = catalog.create_request("worker", 3).unwrap();
219        assert_eq!(
220            serde_json::to_value(&first).unwrap(),
221            serde_json::to_value(&replay).unwrap()
222        );
223        assert_eq!(first.external_sandbox_id, "scale-worker-3");
224        assert_eq!(first.config.isolation, ExecutionIsolation::Sandbox);
225        assert_eq!(first.config.resources.vcpus, 2);
226        assert_eq!(first.config.resources.memory_mb, 768);
227        assert_eq!(first.labels[SCALE_SERVICE_LABEL], "worker");
228        assert_eq!(first.labels[SCALE_SLOT_LABEL], "3");
229        assert!(first.labels[SCALE_TEMPLATE_DIGEST_LABEL].starts_with("sha256:"));
230    }
231
232    #[test]
233    fn catalog_rejects_stateful_or_fixed_endpoint_templates() {
234        for (field, expected) in [
235            ("ports = [\"8080:80\"]", "fixed host port"),
236            ("ports = [\"0:8080\", \"0:9090\"]", "multiple ports"),
237            ("volumes = [\"data:/data\"]", "shared volumes"),
238            ("depends_on = [\"db\"]", "depends_on"),
239            ("networks = [\"backend\"]", "Compose network"),
240        ] {
241            let source = format!(
242                "service \"api\" {{ image = \"api:v1\"; {field} }}\nservice \"db\" {{ image = \"db:v1\" }}"
243            );
244            let error = ScaleServiceCatalog::from_acl_str(
245                &source,
246                "gateway-scale",
247                ExecutionIsolation::Microvm,
248            )
249            .unwrap_err();
250            assert!(error.to_string().contains(expected), "{error}");
251        }
252    }
253
254    #[test]
255    fn unknown_service_fails_without_fabricating_a_template() {
256        let catalog = ScaleServiceCatalog::from_acl_str(
257            CATALOG,
258            "gateway-scale",
259            ExecutionIsolation::Microvm,
260        )
261        .unwrap();
262        let error = catalog.create_request("missing", 0).unwrap_err();
263        assert!(error.to_string().contains("not found"));
264    }
265
266    #[test]
267    fn catalog_converts_one_dynamic_port_into_endpoint_metadata() {
268        let catalog = ScaleServiceCatalog::from_acl_str(
269            r#"service "api" { image = "api:v1"; ports = ["0:8080"] }"#,
270            "gateway-scale",
271            ExecutionIsolation::Sandbox,
272        )
273        .unwrap();
274
275        let request = catalog.create_request("api", 2).unwrap();
276        assert!(request.config.port_map.is_empty());
277        assert_eq!(request.labels[SCALE_GUEST_PORT_LABEL], "8080");
278        assert_eq!(catalog.guest_port("api").unwrap(), Some(8080));
279    }
280
281    #[test]
282    fn file_catalog_requires_an_acl_extension() {
283        let directory = tempfile::tempdir().unwrap();
284        let path = directory.path().join("services.yaml");
285        std::fs::write(&path, CATALOG).unwrap();
286
287        let error =
288            ScaleServiceCatalog::from_acl_file(&path, "gateway-scale", ExecutionIsolation::Microvm)
289                .unwrap_err();
290
291        assert!(error.to_string().contains(".acl format"), "{error}");
292    }
293}