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