Skip to main content

hydro_deploy/
deployment.rs

1#![expect(
2    mismatched_lifetime_syntaxes,
3    reason = "https://github.com/BrynCooke/buildstructor/issues/200"
4)]
5#![expect(
6    elided_lifetimes_in_paths,
7    reason = "https://github.com/BrynCooke/buildstructor/issues/200#issuecomment-5195747044"
8)]
9
10use std::collections::HashMap;
11use std::future::Future;
12use std::sync::{Arc, Mutex, Weak};
13
14use anyhow::Result;
15use futures::{FutureExt, StreamExt, TryStreamExt};
16
17use crate::aws::{AwsCloudwatchLogGroup, AwsEc2IamInstanceProfile, AwsNetwork};
18use crate::gcp::GcpNetwork;
19use crate::{
20    AwsEc2Host, AzureHost, CustomService, GcpComputeEngineHost, Host, HostTargetType,
21    LocalhostHost, ResourcePool, ResourceResult, Service, ServiceBuilder, progress,
22};
23
24pub struct Deployment {
25    pub hosts: Vec<Weak<dyn Host>>,
26    pub services: Vec<Weak<dyn Service>>,
27    pub resource_pool: ResourcePool,
28    localhost_host: Option<Arc<LocalhostHost>>,
29    last_resource_result: Option<Arc<ResourceResult>>,
30    next_host_id: usize,
31    next_service_id: usize,
32}
33
34impl Default for Deployment {
35    fn default() -> Self {
36        Self::new()
37    }
38}
39
40impl Deployment {
41    pub fn new() -> Self {
42        let mut ret = Self {
43            hosts: Vec::new(),
44            services: Vec::new(),
45            resource_pool: ResourcePool::default(),
46            localhost_host: None,
47            last_resource_result: None,
48            next_host_id: 0,
49            next_service_id: 0,
50        };
51
52        ret.localhost_host = Some(ret.add_host(LocalhostHost::new));
53        ret
54    }
55
56    #[expect(non_snake_case, reason = "constructor-esque")]
57    pub fn Localhost(&self) -> Arc<LocalhostHost> {
58        self.localhost_host.clone().unwrap()
59    }
60
61    #[expect(non_snake_case, reason = "constructor-esque")]
62    pub fn CustomService(
63        &mut self,
64        on: Arc<dyn Host>,
65        external_ports: Vec<u16>,
66    ) -> Arc<CustomService> {
67        self.add_service(|id, on| CustomService::new(id, on, external_ports), on)
68    }
69
70    /// Runs `deploy()`, and `start()`, waits for the trigger future, then runs `stop()`.
71    pub async fn run_until(&mut self, trigger: impl Future<Output = ()>) -> Result<()> {
72        // TODO(mingwei): should `trigger` interrupt `deploy()` and `start()`? If so make sure shutdown works as expected.
73        self.deploy().await?;
74        self.start().await?;
75        trigger.await;
76        self.stop().await?;
77        Ok(())
78    }
79
80    /// Runs `start()`, waits for the trigger future, then runs `stop()`.
81    /// This is useful if you need to initiate external network connections between
82    /// `deploy()` and `start()`.
83    pub async fn start_until(&mut self, trigger: impl Future<Output = ()>) -> Result<()> {
84        // TODO(mingwei): should `trigger` interrupt `deploy()` and `start()`? If so make sure shutdown works as expected.
85        self.start().await?;
86        trigger.await;
87        self.stop().await?;
88        Ok(())
89    }
90
91    /// Runs `deploy()`, and `start()`, waits for CTRL+C, then runs `stop()`.
92    pub async fn run_ctrl_c(&mut self) -> Result<()> {
93        self.run_until(tokio::signal::ctrl_c().map(|_| ())).await
94    }
95
96    /// Runs `start()`, waits for CTRL+C, then runs `stop()`.
97    /// This is useful if you need to initiate external network connections between
98    /// `deploy()` and `start()`.
99    pub async fn start_ctrl_c(&mut self) -> Result<()> {
100        self.start_until(tokio::signal::ctrl_c().map(|_| ())).await
101    }
102
103    pub async fn deploy(&mut self) -> Result<()> {
104        self.services.retain(|weak| weak.strong_count() > 0);
105
106        progress::ProgressTracker::with_group("deploy", Some(3), || async {
107            let mut resource_batch = super::ResourceBatch::new();
108
109            for service in self.services.iter().filter_map(Weak::upgrade) {
110                service.collect_resources(&mut resource_batch);
111            }
112
113            for host in self.hosts.iter().filter_map(Weak::upgrade) {
114                host.collect_resources(&mut resource_batch);
115            }
116
117            let resource_result = Arc::new(
118                progress::ProgressTracker::with_group("provision", Some(1), || async {
119                    resource_batch
120                        .provision(&mut self.resource_pool, self.last_resource_result.clone())
121                        .await
122                })
123                .await?,
124            );
125            self.last_resource_result = Some(resource_result.clone());
126
127            for host in self.hosts.iter().filter_map(Weak::upgrade) {
128                host.provision(&resource_result);
129            }
130
131            let upgraded_services = self
132                .services
133                .iter()
134                .filter_map(Weak::upgrade)
135                .collect::<Vec<_>>();
136
137            progress::ProgressTracker::with_group("prepare", Some(upgraded_services.len()), || {
138                let services_future = upgraded_services
139                    .iter()
140                    .map(|service: &Arc<dyn Service>| {
141                        let resource_result = &resource_result;
142                        async move { service.deploy(resource_result).await }
143                    })
144                    .collect::<Vec<_>>();
145
146                futures::stream::iter(services_future)
147                    .buffer_unordered(16)
148                    .try_fold((), |_, _| async { Ok(()) })
149            })
150            .await?;
151
152            progress::ProgressTracker::with_group("ready", Some(upgraded_services.len()), || {
153                let all_services_ready =
154                    upgraded_services
155                        .iter()
156                        .map(|service: &Arc<dyn Service>| async move {
157                            service.ready().await?;
158                            Ok(()) as Result<()>
159                        });
160
161                futures::future::try_join_all(all_services_ready)
162            })
163            .await?;
164
165            Ok(())
166        })
167        .await
168    }
169
170    pub async fn start(&mut self) -> Result<()> {
171        self.services.retain(|weak| weak.strong_count() > 0);
172
173        progress::ProgressTracker::with_group("start", None, || {
174            let all_services_start = self.services.iter().filter_map(Weak::upgrade).map(
175                |service: Arc<dyn Service>| async move {
176                    service.start().await?;
177                    Ok(()) as Result<()>
178                },
179            );
180
181            futures::future::try_join_all(all_services_start)
182        })
183        .await?;
184        Ok(())
185    }
186
187    pub async fn stop(&mut self) -> Result<()> {
188        self.services.retain(|weak| weak.strong_count() > 0);
189
190        progress::ProgressTracker::with_group("stop", None, || {
191            let all_services_stop = self.services.iter().filter_map(Weak::upgrade).map(
192                |service: Arc<dyn Service>| async move {
193                    service.stop().await?;
194                    Ok(()) as Result<()>
195                },
196            );
197
198            futures::future::try_join_all(all_services_stop)
199        })
200        .await?;
201        Ok(())
202    }
203}
204
205impl Deployment {
206    pub fn add_host<T: Host + 'static, F: FnOnce(usize) -> T>(&mut self, host: F) -> Arc<T> {
207        let arc = Arc::new(host(self.next_host_id));
208        self.next_host_id += 1;
209
210        self.hosts.push(Arc::downgrade(&arc) as Weak<dyn Host>);
211        arc
212    }
213
214    pub fn add_service<T: Service + 'static>(
215        &mut self,
216        service: impl ServiceBuilder<Service = T>,
217        on: Arc<dyn Host>,
218    ) -> Arc<T> {
219        let arc = Arc::new(service.build(self.next_service_id, on));
220        self.next_service_id += 1;
221
222        self.services
223            .push(Arc::downgrade(&arc) as Weak<dyn Service>);
224        arc
225    }
226}
227
228/// Buildstructor methods.
229#[buildstructor::buildstructor]
230impl Deployment {
231    #[builder(entry = "GcpComputeEngineHost", exit = "add", visibility = "pub")]
232    fn add_gcp_compute_engine_host(
233        &mut self,
234        project: String,
235        machine_type: String,
236        image: String,
237        target_type: Option<HostTargetType>,
238        region: String,
239        network: Arc<GcpNetwork>,
240        user: Option<String>,
241        display_name: Option<String>,
242    ) -> Arc<GcpComputeEngineHost> {
243        self.add_host(|id| {
244            GcpComputeEngineHost::new(
245                id,
246                project,
247                machine_type,
248                image,
249                target_type.unwrap_or(HostTargetType::Linux(crate::LinuxCompileType::Musl)),
250                region,
251                network,
252                user,
253                display_name,
254            )
255        })
256    }
257
258    #[builder(entry = "AzureHost", exit = "add", visibility = "pub")]
259    fn add_azure_host(
260        &mut self,
261        project: String,
262        os_type: String, // linux or windows
263        machine_size: String,
264        image: Option<HashMap<String, String>>,
265        target_type: Option<HostTargetType>,
266        region: String,
267        user: Option<String>,
268    ) -> Arc<AzureHost> {
269        self.add_host(|id| {
270            AzureHost::new(
271                id,
272                project,
273                os_type,
274                machine_size,
275                image,
276                target_type.unwrap_or(HostTargetType::Linux(crate::LinuxCompileType::Musl)),
277                region,
278                user,
279            )
280        })
281    }
282
283    #[builder(entry = "AwsEc2Host", exit = "add", visibility = "pub")]
284    fn add_aws_ec2_host(
285        &mut self,
286        region: String,
287        instance_type: String,
288        target_type: Option<HostTargetType>,
289        ami: String,
290        network: Arc<AwsNetwork>,
291        iam_instance_profile: Option<Arc<Mutex<AwsEc2IamInstanceProfile>>>,
292        cloudwatch_log_group: Option<Arc<Mutex<AwsCloudwatchLogGroup>>>,
293        // `metrics_collected`: https://docs.aws.amazon.com/AmazonCloudWatch/latest/monitoring/CloudWatch-Agent-Configuration-File-Details.html#CloudWatch-Agent-Configuration-File-Metricssection
294        cwa_metrics_collected: Option<serde_json::Value>,
295        user: Option<String>,
296        display_name: Option<String>,
297    ) -> Arc<AwsEc2Host> {
298        self.add_host(|id| {
299            AwsEc2Host::new(
300                id,
301                region,
302                instance_type,
303                target_type.unwrap_or(HostTargetType::Linux(crate::LinuxCompileType::Musl)),
304                ami,
305                network,
306                iam_instance_profile,
307                cloudwatch_log_group,
308                cwa_metrics_collected,
309                user,
310                display_name,
311            )
312        })
313    }
314}