1use crate::machines::{
2 Checks, CpuKind, DnsConfig, FileConfig, GpuKind, GuestConfig, InitConfig, MetricsConfig,
3 MountConfig, ProcessConfig, RestartPolicy, RestartPolicyEnum, ServiceConfig, StaticConfig,
4 StopConfig,
5};
6use serde::{Deserialize, Serialize};
7use std::collections::HashMap;
8
9#[derive(Debug, Serialize, Deserialize)]
10#[serde(rename_all = "lowercase")]
11pub enum MachineState {
12 Started,
13 Stopped,
14 Suspended,
15 Destroyed,
16}
17
18impl std::fmt::Display for MachineState {
19 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
20 let state_str = match self {
21 MachineState::Started => "started",
22 MachineState::Stopped => "stopped",
23 MachineState::Suspended => "suspended",
24 MachineState::Destroyed => "destroyed",
25 };
26 write!(f, "{}", state_str)
27 }
28}
29
30#[derive(Serialize, Deserialize, Debug)]
31pub struct MachineConfig {
32 pub image: String,
33 pub guest: Option<GuestConfig>,
34 pub auto_destroy: Option<bool>,
35 pub init: Option<InitConfig>,
36 pub env: Option<HashMap<String, String>>,
37 pub processes: Option<Vec<ProcessConfig>>,
38 pub mounts: Option<Vec<MountConfig>>,
39 pub restart: Option<RestartPolicy>,
40 pub checks: Option<Checks>,
41 pub dns: Option<DnsConfig>,
42 pub files: Option<Vec<FileConfig>>,
43 pub metadata: Option<HashMap<String, String>>,
44 pub metrics: Option<MetricsConfig>,
45 pub schedule: Option<String>,
46 pub services: Option<Vec<ServiceConfig>>,
47 pub standbys: Option<Vec<String>>,
48 pub statics: Option<Vec<StaticConfig>>,
49 pub stop_config: Option<StopConfig>,
50}
51
52impl Default for MachineConfig {
53 fn default() -> Self {
54 Self {
55 image: "ubuntu:20.04".to_string(),
56 env: None,
57 processes: None,
58 init: None,
59 auto_destroy: Some(false),
60 checks: None,
61 dns: None,
62 files: None,
63 guest: None,
64 metadata: None,
65 metrics: None,
66 mounts: None,
67 restart: None,
68 schedule: None,
69 services: None,
70 standbys: None,
71 statics: None,
72 stop_config: None,
73 }
74 }
75}
76
77impl MachineConfig {
78 pub fn builder() -> MachineConfigBuilder {
79 MachineConfigBuilder::new()
80 }
81
82 pub fn new(
83 image: String,
84 auto_destroy: Option<bool>,
85 guest: Option<GuestConfig>,
86 restart: Option<RestartPolicy>,
87 env: Option<HashMap<String, String>>,
88 processes: Option<Vec<ProcessConfig>>,
89 mounts: Option<Vec<MountConfig>>,
90 checks: Option<Checks>,
91 dns: Option<DnsConfig>,
92 files: Option<Vec<FileConfig>>,
93 init: Option<InitConfig>,
94 metrics: Option<MetricsConfig>,
95 schedule: Option<String>,
96 services: Option<Vec<ServiceConfig>>,
97 standbys: Option<Vec<String>>,
98 statics: Option<Vec<StaticConfig>>,
99 stop_config: Option<StopConfig>,
100 ) -> Self {
101 Self {
102 image,
103 auto_destroy,
104 guest,
105 restart,
106 env,
107 processes,
108 mounts,
109 checks,
110 dns,
111 files,
112 init,
113 metadata: None,
114 metrics,
115 schedule,
116 services,
117 standbys,
118 statics,
119 stop_config,
120 }
121 }
122}
123
124pub struct MachineConfigBuilder {
125 config: MachineConfig,
126}
127
128impl MachineConfigBuilder {
129 pub fn new() -> Self {
130 Self {
131 config: MachineConfig {
132 image: "ubuntu:22.04".to_string(),
133 auto_destroy: Some(false),
134 restart: Some(RestartPolicy::default()),
135 guest: Some(GuestConfig::default()),
136 ..Default::default()
137 },
138 }
139 }
140
141 pub fn image(mut self, image: &str) -> Self {
142 self.config.image = image.to_string();
143 self
144 }
145
146 pub fn auto_destroy(mut self, auto_destroy: bool) -> Self {
147 self.config.auto_destroy = Some(auto_destroy);
148 self
149 }
150
151 pub fn restart_policy(
152 mut self,
153 policy: RestartPolicyEnum,
154 max_retries: Option<u32>,
155 gpu_bid_price: Option<f64>,
156 ) -> Self {
157 self.config.restart = Some(RestartPolicy {
158 policy: policy,
159 max_retries,
160 gpu_bid_price,
161 });
162 self
163 }
164
165 pub fn cpus(mut self, cpus: u64) -> Self {
166 if let Some(guest) = &mut self.config.guest {
167 guest.cpus = Some(cpus);
168 }
169 self
170 }
171
172 pub fn memory(mut self, memory_mb: u64) -> Self {
173 if let Some(guest) = &mut self.config.guest {
174 guest.memory_mb = Some(memory_mb);
175 }
176 self
177 }
178
179 pub fn cpu_kind(mut self, cpu_kind: CpuKind) -> Self {
180 if let Some(guest) = &mut self.config.guest {
181 guest.cpu_kind = Some(cpu_kind);
182 }
183 self
184 }
185
186 pub fn gpus(mut self, gpus: u64) -> Self {
187 if let Some(guest) = &mut self.config.guest {
188 guest.gpus = Some(gpus);
189 }
190 self
191 }
192
193 pub fn gpu_kind(mut self, gpu_kind: GpuKind) -> Self {
194 if let Some(guest) = &mut self.config.guest {
195 guest.gpu_kind = Some(gpu_kind);
196 }
197 self
198 }
199
200 pub fn checks(mut self, config: Checks) -> Self {
201 self.config.checks = Some(config);
202 self
203 }
204
205 pub fn dns(mut self, dns_config: DnsConfig) -> Self {
206 self.config.dns = Some(dns_config);
207 self
208 }
209
210 pub fn add_env(mut self, key: &str, value: &str) -> Self {
211 if let Some(env) = &mut self.config.env {
212 env.insert(key.to_string(), value.to_string());
213 } else {
214 let mut env = HashMap::new();
215 env.insert(key.to_string(), value.to_string());
216 self.config.env = Some(env);
217 }
218 self
219 }
220
221 pub fn add_file(mut self, file_config: FileConfig) -> Self {
222 if let Some(files) = &mut self.config.files {
223 files.push(file_config);
224 } else {
225 self.config.files = Some(vec![file_config]);
226 }
227 self
228 }
229
230 pub fn init(mut self, init_config: InitConfig) -> Self {
231 self.config.init = Some(init_config);
232 self
233 }
234
235 pub fn metrics(mut self, metrics_config: MetricsConfig) -> Self {
236 self.config.metrics = Some(metrics_config);
237 self
238 }
239
240 pub fn add_mount(mut self, mount_config: MountConfig) -> Self {
241 if let Some(mounts) = &mut self.config.mounts {
242 mounts.push(mount_config);
243 } else {
244 self.config.mounts = Some(vec![mount_config]);
245 }
246 self
247 }
248
249 pub fn add_process(mut self, process_config: ProcessConfig) -> Self {
250 if let Some(processes) = &mut self.config.processes {
251 processes.push(process_config);
252 } else {
253 self.config.processes = Some(vec![process_config]);
254 }
255 self
256 }
257
258 pub fn schedule(mut self, schedule: &str) -> Self {
259 self.config.schedule = Some(schedule.to_string());
260 self
261 }
262
263 pub fn add_service(mut self, service_config: ServiceConfig) -> Self {
264 if let Some(services) = &mut self.config.services {
265 services.push(service_config);
266 } else {
267 self.config.services = Some(vec![service_config]);
268 }
269 self
270 }
271
272 pub fn add_standby(mut self, standby: &str) -> Self {
273 if let Some(standbys) = &mut self.config.standbys {
274 standbys.push(standby.to_string());
275 } else {
276 self.config.standbys = Some(vec![standby.to_string()]);
277 }
278 self
279 }
280
281 pub fn add_static(mut self, static_config: StaticConfig) -> Self {
282 if let Some(statics) = &mut self.config.statics {
283 statics.push(static_config);
284 } else {
285 self.config.statics = Some(vec![static_config]);
286 }
287 self
288 }
289
290 pub fn stop_config(mut self, stop_config: StopConfig) -> Self {
291 self.config.stop_config = Some(stop_config);
292 self
293 }
294
295 pub fn build(self) -> MachineConfig {
296 self.config
297 }
298}