Skip to main content

a3s_box_core/
scale.rs

1//! Scale API types for Gateway ↔ Box communication.
2//!
3//! Defines the request/response types for the internal Scale API that
4//! Gateway uses to request instance scale-up/scale-down in standalone mode.
5//!
6//! ## Protocol
7//!
8//! Gateway → Box: `ScaleRequest` (service, desired replicas, config)
9//! Box → Gateway: `ScaleResponse` (actual replicas, instance states)
10//! Box → Gateway: `InstanceEvent` (state transitions, health updates)
11
12use std::collections::HashMap;
13
14use chrono::{DateTime, Utc};
15use serde::{Deserialize, Serialize};
16
17/// Instance lifecycle state for readiness signaling.
18#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
19pub enum InstanceState {
20    /// Instance is being created (image pull, rootfs build)
21    Creating,
22    /// VM is booting (kernel + init)
23    Booting,
24    /// Agent is ready, accepting requests
25    Ready,
26    /// Instance is actively processing a request
27    Busy,
28    /// Instance is draining in-flight requests before shutdown
29    Draining,
30    /// Instance is shutting down
31    Stopping,
32    /// Instance has terminated
33    Stopped,
34    /// Instance failed to start or crashed
35    Failed,
36}
37
38impl std::fmt::Display for InstanceState {
39    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
40        match self {
41            Self::Creating => write!(f, "creating"),
42            Self::Booting => write!(f, "booting"),
43            Self::Ready => write!(f, "ready"),
44            Self::Busy => write!(f, "busy"),
45            Self::Draining => write!(f, "draining"),
46            Self::Stopping => write!(f, "stopping"),
47            Self::Stopped => write!(f, "stopped"),
48            Self::Failed => write!(f, "failed"),
49        }
50    }
51}
52
53/// Request from Gateway to scale a service.
54#[derive(Debug, Clone, Serialize, Deserialize)]
55pub struct ScaleRequest {
56    /// Service identifier
57    pub service: String,
58    /// Desired number of running instances
59    pub replicas: u32,
60    /// Instance configuration overrides
61    #[serde(default)]
62    pub config: ScaleConfig,
63    /// Request ID for correlation
64    #[serde(default)]
65    pub request_id: String,
66}
67
68/// Instance configuration for scale requests.
69#[derive(Debug, Clone, Default, Serialize, Deserialize)]
70pub struct ScaleConfig {
71    /// OCI image to use (overrides service default)
72    #[serde(default)]
73    pub image: Option<String>,
74    /// vCPUs per instance
75    #[serde(default)]
76    pub vcpus: Option<u8>,
77    /// Memory in MiB per instance
78    #[serde(default)]
79    pub memory_mib: Option<u32>,
80    /// Environment variables
81    #[serde(default)]
82    pub env: HashMap<String, String>,
83    /// Port mappings
84    #[serde(default)]
85    pub port_map: Vec<String>,
86}
87
88/// Response from Box after processing a scale request.
89#[derive(Debug, Clone, Serialize, Deserialize)]
90pub struct ScaleResponse {
91    /// Request ID for correlation
92    pub request_id: String,
93    /// Whether the scale operation was accepted
94    pub accepted: bool,
95    /// Current number of running instances
96    pub current_replicas: u32,
97    /// Target number of instances (may differ from requested if at capacity)
98    pub target_replicas: u32,
99    /// Per-instance status
100    pub instances: Vec<InstanceInfo>,
101    /// Error message if not accepted
102    #[serde(default)]
103    pub error: Option<String>,
104}
105
106/// Information about a single instance.
107#[derive(Debug, Clone, Serialize, Deserialize)]
108pub struct InstanceInfo {
109    /// Instance (box) ID
110    pub id: String,
111    /// Current state
112    pub state: InstanceState,
113    /// Service this instance belongs to
114    pub service: String,
115    /// When the instance was created
116    pub created_at: DateTime<Utc>,
117    /// When the instance became ready (None if not yet ready)
118    #[serde(default)]
119    pub ready_at: Option<DateTime<Utc>>,
120    /// Instance endpoint (host:port) for traffic routing
121    #[serde(default)]
122    pub endpoint: Option<String>,
123    /// Health metrics
124    #[serde(default)]
125    pub health: InstanceHealth,
126}
127
128/// Health metrics for an instance.
129#[derive(Debug, Clone, Serialize, Deserialize)]
130pub struct InstanceHealth {
131    /// CPU usage percentage (0-100)
132    #[serde(default)]
133    pub cpu_percent: Option<f32>,
134    /// Memory usage in bytes
135    #[serde(default)]
136    pub memory_bytes: Option<u64>,
137    /// Number of in-flight requests
138    #[serde(default)]
139    pub inflight_requests: u32,
140    /// Whether the instance is healthy
141    #[serde(default = "default_true")]
142    pub healthy: bool,
143}
144
145impl Default for InstanceHealth {
146    fn default() -> Self {
147        Self {
148            cpu_percent: None,
149            memory_bytes: None,
150            inflight_requests: 0,
151            healthy: true,
152        }
153    }
154}
155
156fn default_true() -> bool {
157    true
158}
159
160/// Event emitted by Box when an instance state changes.
161#[derive(Debug, Clone, Serialize, Deserialize)]
162pub struct InstanceEvent {
163    /// Instance (box) ID
164    pub instance_id: String,
165    /// Service this instance belongs to
166    pub service: String,
167    /// Previous state
168    pub from_state: InstanceState,
169    /// New state
170    pub to_state: InstanceState,
171    /// When the transition occurred
172    pub timestamp: DateTime<Utc>,
173    /// Optional message (e.g., error details for Failed state)
174    #[serde(default)]
175    pub message: String,
176}
177
178impl InstanceEvent {
179    /// Create a new state transition event.
180    pub fn transition(
181        instance_id: &str,
182        service: &str,
183        from: InstanceState,
184        to: InstanceState,
185    ) -> Self {
186        Self {
187            instance_id: instance_id.to_string(),
188            service: service.to_string(),
189            from_state: from,
190            to_state: to,
191            timestamp: Utc::now(),
192            message: String::new(),
193        }
194    }
195
196    /// Add a message to the event.
197    pub fn with_message(mut self, msg: &str) -> Self {
198        self.message = msg.to_string();
199        self
200    }
201}
202
203/// Registration payload for instance self-registration with Gateway.
204#[derive(Debug, Clone, Serialize, Deserialize)]
205pub struct InstanceRegistration {
206    /// Instance (box) ID
207    pub instance_id: String,
208    /// Service this instance belongs to
209    pub service: String,
210    /// Endpoint for traffic routing (host:port)
211    pub endpoint: String,
212    /// Instance metadata
213    #[serde(default)]
214    pub metadata: HashMap<String, String>,
215    /// When the instance started
216    pub started_at: DateTime<Utc>,
217}
218
219/// Deregistration payload when an instance is shutting down.
220#[derive(Debug, Clone, Serialize, Deserialize)]
221pub struct InstanceDeregistration {
222    /// Instance (box) ID
223    pub instance_id: String,
224    /// Service this instance belongs to
225    pub service: String,
226    /// Reason for deregistration
227    #[serde(default)]
228    pub reason: String,
229}
230
231#[cfg(test)]
232mod tests {
233    use super::*;
234
235    #[test]
236    fn test_instance_state_display() {
237        assert_eq!(InstanceState::Creating.to_string(), "creating");
238        assert_eq!(InstanceState::Booting.to_string(), "booting");
239        assert_eq!(InstanceState::Ready.to_string(), "ready");
240        assert_eq!(InstanceState::Busy.to_string(), "busy");
241        assert_eq!(InstanceState::Draining.to_string(), "draining");
242        assert_eq!(InstanceState::Stopping.to_string(), "stopping");
243        assert_eq!(InstanceState::Stopped.to_string(), "stopped");
244        assert_eq!(InstanceState::Failed.to_string(), "failed");
245    }
246
247    #[test]
248    fn test_scale_request_serde() {
249        let req = ScaleRequest {
250            service: "my-service".to_string(),
251            replicas: 3,
252            config: ScaleConfig {
253                image: Some("nginx:latest".to_string()),
254                vcpus: Some(2),
255                memory_mib: Some(512),
256                env: HashMap::from([("PORT".to_string(), "8080".to_string())]),
257                port_map: vec!["8080:80".to_string()],
258            },
259            request_id: "req-001".to_string(),
260        };
261        let json = serde_json::to_string(&req).unwrap();
262        let parsed: ScaleRequest = serde_json::from_str(&json).unwrap();
263        assert_eq!(parsed.service, "my-service");
264        assert_eq!(parsed.replicas, 3);
265        assert_eq!(parsed.config.image, Some("nginx:latest".to_string()));
266        assert_eq!(parsed.config.vcpus, Some(2));
267        assert_eq!(parsed.config.env.get("PORT").unwrap(), "8080");
268    }
269
270    #[test]
271    fn test_scale_request_minimal() {
272        let json = r#"{"service":"svc","replicas":1}"#;
273        let req: ScaleRequest = serde_json::from_str(json).unwrap();
274        assert_eq!(req.service, "svc");
275        assert_eq!(req.replicas, 1);
276        assert!(req.config.image.is_none());
277        assert!(req.request_id.is_empty());
278    }
279
280    #[test]
281    fn test_scale_response_accepted() {
282        let resp = ScaleResponse {
283            request_id: "req-001".to_string(),
284            accepted: true,
285            current_replicas: 2,
286            target_replicas: 3,
287            instances: vec![InstanceInfo {
288                id: "box-1".to_string(),
289                state: InstanceState::Ready,
290                service: "svc".to_string(),
291                created_at: Utc::now(),
292                ready_at: Some(Utc::now()),
293                endpoint: Some("10.0.0.2:8080".to_string()),
294                health: InstanceHealth::default(),
295            }],
296            error: None,
297        };
298        let json = serde_json::to_string(&resp).unwrap();
299        let parsed: ScaleResponse = serde_json::from_str(&json).unwrap();
300        assert!(parsed.accepted);
301        assert_eq!(parsed.current_replicas, 2);
302        assert_eq!(parsed.target_replicas, 3);
303        assert_eq!(parsed.instances.len(), 1);
304        assert_eq!(parsed.instances[0].state, InstanceState::Ready);
305    }
306
307    #[test]
308    fn test_scale_response_rejected() {
309        let resp = ScaleResponse {
310            request_id: "req-002".to_string(),
311            accepted: false,
312            current_replicas: 5,
313            target_replicas: 5,
314            instances: vec![],
315            error: Some("At maximum capacity".to_string()),
316        };
317        let json = serde_json::to_string(&resp).unwrap();
318        let parsed: ScaleResponse = serde_json::from_str(&json).unwrap();
319        assert!(!parsed.accepted);
320        assert_eq!(parsed.error, Some("At maximum capacity".to_string()));
321    }
322
323    #[test]
324    fn test_instance_info_serde() {
325        let info = InstanceInfo {
326            id: "box-abc".to_string(),
327            state: InstanceState::Busy,
328            service: "api".to_string(),
329            created_at: Utc::now(),
330            ready_at: Some(Utc::now()),
331            endpoint: Some("10.0.0.5:3000".to_string()),
332            health: InstanceHealth {
333                cpu_percent: Some(45.2),
334                memory_bytes: Some(256 * 1024 * 1024),
335                inflight_requests: 3,
336                healthy: true,
337            },
338        };
339        let json = serde_json::to_string(&info).unwrap();
340        let parsed: InstanceInfo = serde_json::from_str(&json).unwrap();
341        assert_eq!(parsed.id, "box-abc");
342        assert_eq!(parsed.state, InstanceState::Busy);
343        assert_eq!(parsed.health.cpu_percent, Some(45.2));
344        assert_eq!(parsed.health.inflight_requests, 3);
345    }
346
347    #[test]
348    fn test_instance_health_default() {
349        let health = InstanceHealth::default();
350        assert!(health.cpu_percent.is_none());
351        assert!(health.memory_bytes.is_none());
352        assert_eq!(health.inflight_requests, 0);
353        assert!(health.healthy);
354    }
355
356    #[test]
357    fn test_instance_event_transition() {
358        let event = InstanceEvent::transition(
359            "box-123",
360            "my-svc",
361            InstanceState::Booting,
362            InstanceState::Ready,
363        );
364        assert_eq!(event.instance_id, "box-123");
365        assert_eq!(event.service, "my-svc");
366        assert_eq!(event.from_state, InstanceState::Booting);
367        assert_eq!(event.to_state, InstanceState::Ready);
368        assert!(event.message.is_empty());
369    }
370
371    #[test]
372    fn test_instance_event_with_message() {
373        let event = InstanceEvent::transition(
374            "box-456",
375            "svc",
376            InstanceState::Booting,
377            InstanceState::Failed,
378        )
379        .with_message("OOM killed");
380        assert_eq!(event.message, "OOM killed");
381        assert_eq!(event.to_state, InstanceState::Failed);
382    }
383
384    #[test]
385    fn test_instance_event_serde() {
386        let event = InstanceEvent::transition(
387            "box-789",
388            "api",
389            InstanceState::Ready,
390            InstanceState::Draining,
391        );
392        let json = serde_json::to_string(&event).unwrap();
393        let parsed: InstanceEvent = serde_json::from_str(&json).unwrap();
394        assert_eq!(parsed.instance_id, "box-789");
395        assert_eq!(parsed.from_state, InstanceState::Ready);
396        assert_eq!(parsed.to_state, InstanceState::Draining);
397    }
398
399    #[test]
400    fn test_instance_registration_serde() {
401        let reg = InstanceRegistration {
402            instance_id: "box-reg".to_string(),
403            service: "web".to_string(),
404            endpoint: "10.0.0.10:8080".to_string(),
405            metadata: HashMap::from([("version".to_string(), "v1.2".to_string())]),
406            started_at: Utc::now(),
407        };
408        let json = serde_json::to_string(&reg).unwrap();
409        let parsed: InstanceRegistration = serde_json::from_str(&json).unwrap();
410        assert_eq!(parsed.instance_id, "box-reg");
411        assert_eq!(parsed.endpoint, "10.0.0.10:8080");
412        assert_eq!(parsed.metadata.get("version").unwrap(), "v1.2");
413    }
414
415    #[test]
416    fn test_instance_deregistration_serde() {
417        let dereg = InstanceDeregistration {
418            instance_id: "box-dereg".to_string(),
419            service: "web".to_string(),
420            reason: "scale-down".to_string(),
421        };
422        let json = serde_json::to_string(&dereg).unwrap();
423        let parsed: InstanceDeregistration = serde_json::from_str(&json).unwrap();
424        assert_eq!(parsed.instance_id, "box-dereg");
425        assert_eq!(parsed.reason, "scale-down");
426    }
427
428    #[test]
429    fn test_scale_config_default() {
430        let config = ScaleConfig::default();
431        assert!(config.image.is_none());
432        assert!(config.vcpus.is_none());
433        assert!(config.memory_mib.is_none());
434        assert!(config.env.is_empty());
435        assert!(config.port_map.is_empty());
436    }
437
438    #[test]
439    fn test_instance_state_equality() {
440        assert_eq!(InstanceState::Ready, InstanceState::Ready);
441        assert_ne!(InstanceState::Ready, InstanceState::Busy);
442    }
443
444    #[test]
445    fn test_instance_state_hash() {
446        use std::collections::HashSet;
447        let mut set = HashSet::new();
448        set.insert(InstanceState::Ready);
449        set.insert(InstanceState::Busy);
450        set.insert(InstanceState::Ready); // duplicate
451        assert_eq!(set.len(), 2);
452    }
453}