Skip to main content

wfe_core/models/
service.rs

1use std::collections::HashMap;
2
3use serde::{Deserialize, Serialize};
4
5/// An infrastructure service that runs alongside workflow steps.
6#[derive(Debug, Clone, Serialize, Deserialize)]
7pub struct ServiceDefinition {
8    /// Service name -- used as DNS hostname (K8s) or env prefix (containerd).
9    pub name: String,
10    /// Container image to run.
11    pub image: String,
12    /// Ports exposed by the service.
13    #[serde(default)]
14    pub ports: Vec<ServicePort>,
15    /// Environment variables for the service container.
16    #[serde(default)]
17    pub env: HashMap<String, String>,
18    /// How to check if the service is ready to accept connections.
19    #[serde(default)]
20    pub readiness: Option<ReadinessProbe>,
21    /// Override the container entrypoint.
22    #[serde(default)]
23    pub command: Vec<String>,
24    /// Override the container command/args.
25    #[serde(default)]
26    pub args: Vec<String>,
27    /// Memory limit (e.g., "512Mi").
28    #[serde(default)]
29    pub memory: Option<String>,
30    /// CPU limit (e.g., "500m").
31    #[serde(default)]
32    pub cpu: Option<String>,
33}
34
35/// A port exposed by a service.
36#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
37pub struct ServicePort {
38    pub container_port: u16,
39    #[serde(default)]
40    pub name: Option<String>,
41    /// Protocol: "TCP" (default) or "UDP".
42    #[serde(default = "default_protocol")]
43    pub protocol: String,
44}
45
46impl ServicePort {
47    pub fn tcp(port: u16) -> Self {
48        Self {
49            container_port: port,
50            name: None,
51            protocol: "TCP".into(),
52        }
53    }
54}
55
56fn default_protocol() -> String {
57    "TCP".into()
58}
59
60/// How to determine if a service is ready.
61#[derive(Debug, Clone, Serialize, Deserialize)]
62pub struct ReadinessProbe {
63    pub check: ReadinessCheck,
64    /// Poll interval in milliseconds.
65    #[serde(default = "default_5000")]
66    pub interval_ms: u64,
67    /// Total timeout in milliseconds.
68    #[serde(default = "default_60000")]
69    pub timeout_ms: u64,
70    /// Maximum number of retries before giving up.
71    #[serde(default = "default_12")]
72    pub retries: u32,
73}
74
75/// The type of readiness check to perform.
76#[derive(Debug, Clone, Serialize, Deserialize)]
77#[serde(rename_all = "snake_case")]
78pub enum ReadinessCheck {
79    /// Run a command inside the container.
80    Exec(Vec<String>),
81    /// Check if a TCP port is accepting connections.
82    TcpSocket(u16),
83    /// Make an HTTP GET request.
84    HttpGet { port: u16, path: String },
85}
86
87/// Runtime endpoint info for a provisioned service.
88#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
89pub struct ServiceEndpoint {
90    pub name: String,
91    pub host: String,
92    pub ports: Vec<ServicePort>,
93}
94
95fn default_5000() -> u64 {
96    5000
97}
98fn default_60000() -> u64 {
99    60000
100}
101fn default_12() -> u32 {
102    12
103}
104
105#[cfg(test)]
106mod tests {
107    use super::*;
108    use pretty_assertions::assert_eq;
109
110    #[test]
111    fn service_definition_minimal_serde() {
112        let json = r#"{"name":"postgres","image":"postgres:15"}"#;
113        let svc: ServiceDefinition = serde_json::from_str(json).unwrap();
114        assert_eq!(svc.name, "postgres");
115        assert_eq!(svc.image, "postgres:15");
116        assert!(svc.ports.is_empty());
117        assert!(svc.env.is_empty());
118        assert!(svc.readiness.is_none());
119        assert!(svc.command.is_empty());
120        assert!(svc.memory.is_none());
121    }
122
123    #[test]
124    fn service_definition_full_round_trip() {
125        let svc = ServiceDefinition {
126            name: "redis".into(),
127            image: "redis:7-alpine".into(),
128            ports: vec![ServicePort::tcp(6379)],
129            env: [("REDIS_PASSWORD".into(), "secret".into())].into(),
130            readiness: Some(ReadinessProbe {
131                check: ReadinessCheck::TcpSocket(6379),
132                interval_ms: 2000,
133                timeout_ms: 30000,
134                retries: 15,
135            }),
136            command: vec![],
137            args: vec!["--requirepass".into(), "secret".into()],
138            memory: Some("256Mi".into()),
139            cpu: Some("250m".into()),
140        };
141        let json = serde_json::to_string(&svc).unwrap();
142        let parsed: ServiceDefinition = serde_json::from_str(&json).unwrap();
143        assert_eq!(parsed.name, "redis");
144        assert_eq!(parsed.ports.len(), 1);
145        assert_eq!(parsed.ports[0].container_port, 6379);
146        assert_eq!(parsed.args, vec!["--requirepass", "secret"]);
147        assert_eq!(parsed.memory, Some("256Mi".into()));
148    }
149
150    #[test]
151    fn service_port_tcp_helper() {
152        let port = ServicePort::tcp(5432);
153        assert_eq!(port.container_port, 5432);
154        assert_eq!(port.protocol, "TCP");
155        assert!(port.name.is_none());
156    }
157
158    #[test]
159    fn service_port_default_protocol() {
160        let json = r#"{"container_port": 8080}"#;
161        let port: ServicePort = serde_json::from_str(json).unwrap();
162        assert_eq!(port.protocol, "TCP");
163    }
164
165    #[test]
166    fn readiness_probe_exec() {
167        let probe = ReadinessProbe {
168            check: ReadinessCheck::Exec(vec!["pg_isready".into(), "-U".into(), "postgres".into()]),
169            interval_ms: 5000,
170            timeout_ms: 60000,
171            retries: 12,
172        };
173        let json = serde_json::to_string(&probe).unwrap();
174        let parsed: ReadinessProbe = serde_json::from_str(&json).unwrap();
175        match parsed.check {
176            ReadinessCheck::Exec(cmd) => assert_eq!(cmd, vec!["pg_isready", "-U", "postgres"]),
177            _ => panic!("expected Exec"),
178        }
179    }
180
181    #[test]
182    fn readiness_probe_tcp_socket() {
183        let probe = ReadinessProbe {
184            check: ReadinessCheck::TcpSocket(6379),
185            interval_ms: 2000,
186            timeout_ms: 30000,
187            retries: 15,
188        };
189        let json = serde_json::to_string(&probe).unwrap();
190        let parsed: ReadinessProbe = serde_json::from_str(&json).unwrap();
191        match parsed.check {
192            ReadinessCheck::TcpSocket(port) => assert_eq!(port, 6379),
193            _ => panic!("expected TcpSocket"),
194        }
195    }
196
197    #[test]
198    fn readiness_probe_http_get() {
199        let probe = ReadinessProbe {
200            check: ReadinessCheck::HttpGet {
201                port: 8080,
202                path: "/health".into(),
203            },
204            interval_ms: 5000,
205            timeout_ms: 60000,
206            retries: 12,
207        };
208        let json = serde_json::to_string(&probe).unwrap();
209        let parsed: ReadinessProbe = serde_json::from_str(&json).unwrap();
210        match parsed.check {
211            ReadinessCheck::HttpGet { port, path } => {
212                assert_eq!(port, 8080);
213                assert_eq!(path, "/health");
214            }
215            _ => panic!("expected HttpGet"),
216        }
217    }
218
219    #[test]
220    fn readiness_probe_defaults() {
221        let json = r#"{"check": {"tcp_socket": 5432}}"#;
222        let probe: ReadinessProbe = serde_json::from_str(json).unwrap();
223        assert_eq!(probe.interval_ms, 5000);
224        assert_eq!(probe.timeout_ms, 60000);
225        assert_eq!(probe.retries, 12);
226    }
227
228    #[test]
229    fn service_endpoint_serde() {
230        let ep = ServiceEndpoint {
231            name: "postgres".into(),
232            host: "postgres.wfe-abc.svc.cluster.local".into(),
233            ports: vec![ServicePort::tcp(5432)],
234        };
235        let json = serde_json::to_string(&ep).unwrap();
236        let parsed: ServiceEndpoint = serde_json::from_str(&json).unwrap();
237        assert_eq!(parsed.name, "postgres");
238        assert_eq!(parsed.host, "postgres.wfe-abc.svc.cluster.local");
239        assert_eq!(parsed.ports.len(), 1);
240    }
241
242    #[test]
243    fn service_definition_with_env() {
244        let svc = ServiceDefinition {
245            name: "postgres".into(),
246            image: "postgres:15".into(),
247            ports: vec![ServicePort::tcp(5432)],
248            env: [
249                ("POSTGRES_PASSWORD".into(), "test".into()),
250                ("POSTGRES_DB".into(), "myapp".into()),
251            ]
252            .into(),
253            readiness: None,
254            command: vec![],
255            args: vec![],
256            memory: None,
257            cpu: None,
258        };
259        let json = serde_json::to_string(&svc).unwrap();
260        let parsed: ServiceDefinition = serde_json::from_str(&json).unwrap();
261        assert_eq!(parsed.env.get("POSTGRES_PASSWORD"), Some(&"test".into()));
262        assert_eq!(parsed.env.get("POSTGRES_DB"), Some(&"myapp".into()));
263    }
264
265    #[test]
266    fn service_definition_with_command() {
267        let svc = ServiceDefinition {
268            name: "custom".into(),
269            image: "myimage:latest".into(),
270            ports: vec![],
271            env: HashMap::new(),
272            readiness: None,
273            command: vec!["/usr/bin/myserver".into()],
274            args: vec!["--config".into(), "/etc/config.yaml".into()],
275            memory: None,
276            cpu: None,
277        };
278        let json = serde_json::to_string(&svc).unwrap();
279        let parsed: ServiceDefinition = serde_json::from_str(&json).unwrap();
280        assert_eq!(parsed.command, vec!["/usr/bin/myserver"]);
281        assert_eq!(parsed.args, vec!["--config", "/etc/config.yaml"]);
282    }
283}