hydro_deploy/rust_crate/
service.rs

1use std::collections::HashMap;
2use std::sync::{Arc, OnceLock};
3use std::time::Duration;
4
5use anyhow::{Context, Result, bail};
6use async_trait::async_trait;
7use futures::Future;
8use hydro_deploy_integration::{InitConfig, ServerPort};
9use memo_map::MemoMap;
10use serde::Serialize;
11use tokio::sync::{OnceCell, RwLock, mpsc};
12
13use super::build::{BuildError, BuildOutput, BuildParams, build_crate_memoized};
14use super::ports::{self, RustCratePortConfig};
15use super::tracing_options::TracingOptions;
16#[cfg(feature = "profile-folding")]
17use crate::TracingResults;
18use crate::progress::ProgressTracker;
19use crate::{
20    BaseServerStrategy, Host, LaunchedBinary, LaunchedHost, PortNetworkHint, ResourceBatch,
21    ResourceResult, ServerStrategy, Service,
22};
23
24pub struct RustCrateService {
25    id: usize,
26    pub(super) on: Arc<dyn Host>,
27    build_params: BuildParams,
28    tracing: Option<TracingOptions>,
29    args: Option<Vec<String>>,
30    display_id: Option<String>,
31    external_ports: Vec<u16>,
32    env: HashMap<String, String>,
33    pin_to_core: Option<usize>,
34
35    meta: OnceLock<String>,
36
37    /// Configuration for the ports this service will connect to as a client.
38    pub(super) port_to_server: MemoMap<String, ports::ServerConfig>,
39    /// Configuration for the ports that this service will listen on a port for.
40    pub(super) port_to_bind: MemoMap<String, ServerStrategy>,
41
42    launched_host: OnceCell<Arc<dyn LaunchedHost>>,
43
44    /// A map of port names to config for how other services can connect to this one.
45    /// Only valid after `ready` has been called, only contains ports that are configured
46    /// in `server_ports`.
47    pub(super) server_defns: Arc<RwLock<HashMap<String, ServerPort>>>,
48
49    launched_binary: OnceCell<Box<dyn LaunchedBinary>>,
50    started: OnceCell<()>,
51}
52
53impl RustCrateService {
54    #[expect(clippy::too_many_arguments, reason = "internal use")]
55    pub fn new(
56        id: usize,
57        on: Arc<dyn Host>,
58        build_params: BuildParams,
59        tracing: Option<TracingOptions>,
60        args: Option<Vec<String>>,
61        display_id: Option<String>,
62        external_ports: Vec<u16>,
63        env: HashMap<String, String>,
64        pin_to_core: Option<usize>,
65    ) -> Self {
66        Self {
67            id,
68            on,
69            build_params,
70            tracing,
71            args,
72            display_id,
73            external_ports,
74            env,
75            pin_to_core,
76            meta: OnceLock::new(),
77            port_to_server: MemoMap::new(),
78            port_to_bind: MemoMap::new(),
79            launched_host: OnceCell::new(),
80            server_defns: Arc::new(RwLock::new(HashMap::new())),
81            launched_binary: OnceCell::new(),
82            started: OnceCell::new(),
83        }
84    }
85
86    pub fn update_meta<T: Serialize>(&self, meta: T) {
87        assert!(
88            self.launched_binary.get().is_none(),
89            "Cannot update meta after binary has been launched"
90        );
91        self.meta
92            .set(serde_json::to_string(&meta).unwrap())
93            .expect("Cannot set meta twice.");
94    }
95
96    pub fn get_port(self: &Arc<Self>, name: String) -> RustCratePortConfig {
97        RustCratePortConfig {
98            service: Arc::downgrade(self),
99            service_host: self.on.clone(),
100            service_server_defns: self.server_defns.clone(),
101            network_hint: PortNetworkHint::Auto,
102            port: name,
103            merge: false,
104        }
105    }
106
107    pub fn get_port_with_hint(
108        self: &Arc<Self>,
109        name: String,
110        network_hint: PortNetworkHint,
111    ) -> RustCratePortConfig {
112        RustCratePortConfig {
113            service: Arc::downgrade(self),
114            service_host: self.on.clone(),
115            service_server_defns: self.server_defns.clone(),
116            network_hint,
117            port: name,
118            merge: false,
119        }
120    }
121
122    pub fn stdout(&self) -> mpsc::UnboundedReceiver<String> {
123        self.launched_binary.get().unwrap().stdout()
124    }
125
126    pub fn stderr(&self) -> mpsc::UnboundedReceiver<String> {
127        self.launched_binary.get().unwrap().stderr()
128    }
129
130    pub fn stdout_filter(&self, prefix: String) -> mpsc::UnboundedReceiver<String> {
131        self.launched_binary.get().unwrap().stdout_filter(prefix)
132    }
133
134    pub fn stderr_filter(&self, prefix: String) -> mpsc::UnboundedReceiver<String> {
135        self.launched_binary.get().unwrap().stderr_filter(prefix)
136    }
137
138    #[cfg(feature = "profile-folding")]
139    pub fn tracing_results(&self) -> Option<&TracingResults> {
140        self.launched_binary.get().unwrap().tracing_results()
141    }
142
143    pub fn exit_code(&self) -> Option<i32> {
144        self.launched_binary.get().unwrap().exit_code()
145    }
146
147    fn build(
148        &self,
149    ) -> impl use<> + 'static + Future<Output = Result<&'static BuildOutput, BuildError>> {
150        // Memoized, so no caching in `self` is needed.
151        build_crate_memoized(self.build_params.clone())
152    }
153}
154
155#[async_trait]
156impl Service for RustCrateService {
157    fn collect_resources(&self, _resource_batch: &mut ResourceBatch) {
158        if self.launched_host.get().is_some() {
159            return;
160        }
161
162        tokio::task::spawn(self.build());
163
164        let host = &self.on;
165
166        host.request_custom_binary();
167        for (_, bind_type) in self.port_to_bind.iter() {
168            host.request_port(bind_type);
169        }
170
171        for port in self.external_ports.iter() {
172            host.request_port_base(&BaseServerStrategy::ExternalTcpPort(*port));
173        }
174    }
175
176    async fn deploy(&self, resource_result: &Arc<ResourceResult>) -> Result<()> {
177        self.launched_host
178            .get_or_try_init::<anyhow::Error, _, _>(|| {
179                ProgressTracker::with_group(
180                    self.display_id
181                        .clone()
182                        .unwrap_or_else(|| format!("service/{}", self.id)),
183                    None,
184                    || async {
185                        let built = self.build().await?;
186
187                        let host = &self.on;
188                        let launched = host.provision(resource_result);
189
190                        launched.copy_binary(built).await?;
191                        Ok(launched)
192                    },
193                )
194            })
195            .await?;
196        Ok(())
197    }
198
199    async fn ready(&self) -> Result<()> {
200        self.launched_binary
201            .get_or_try_init(|| {
202                ProgressTracker::with_group(
203                    self.display_id
204                        .clone()
205                        .unwrap_or_else(|| format!("service/{}", self.id)),
206                    None,
207                    || async {
208                        let launched_host = self.launched_host.get().unwrap();
209
210                        let built = self.build().await?;
211                        let args = self.args.clone().unwrap_or_default();
212
213                        let binary = launched_host
214                            .launch_binary(
215                                self.display_id
216                                    .clone()
217                                    .unwrap_or_else(|| format!("service/{}", self.id)),
218                                built,
219                                &args,
220                                self.tracing.clone(),
221                                &self.env,
222                                self.pin_to_core,
223                            )
224                            .await?;
225
226                        let bind_config = self
227                            .port_to_bind
228                            .iter()
229                            .map(|(port_name, bind_type)| {
230                                (port_name.clone(), launched_host.server_config(bind_type))
231                            })
232                            .collect::<HashMap<_, _>>();
233
234                        let formatted_bind_config = serde_json::to_string::<InitConfig<'_>>(&(
235                            bind_config,
236                            self.meta.get().map(|s| s.as_str().into()),
237                        ))
238                        .unwrap();
239
240                        // request stdout before sending config so we don't miss the "ready" response
241                        let stdout_receiver = binary.deploy_stdout();
242
243                        binary.stdin().send(format!("{formatted_bind_config}\n"))?;
244
245                        let ready_line = ProgressTracker::leaf(
246                            "waiting for ready",
247                            tokio::time::timeout(Duration::from_secs(60), stdout_receiver),
248                        )
249                        .await
250                        .context("Timed out waiting for ready")?
251                        .context("Program unexpectedly quit")?;
252                        if let Some(line_rest) = ready_line.strip_prefix("ready: ") {
253                            *self.server_defns.try_write().unwrap() =
254                                serde_json::from_str(line_rest).unwrap();
255                        } else {
256                            bail!("expected ready");
257                        }
258                        Ok(binary)
259                    },
260                )
261            })
262            .await?;
263        Ok(())
264    }
265
266    async fn start(&self) -> Result<()> {
267        self.started
268            .get_or_try_init(|| async {
269                let sink_ports_futures =
270                    self.port_to_server
271                        .iter()
272                        .map(|(port_name, outgoing)| async {
273                            (&**port_name, outgoing.load_instantiated(&|p| p).await)
274                        });
275                let sink_ports = futures::future::join_all(sink_ports_futures)
276                    .await
277                    .into_iter()
278                    .collect::<HashMap<_, _>>();
279
280                let formatted_defns = serde_json::to_string(&sink_ports).unwrap();
281
282                let stdout_receiver = self.launched_binary.get().unwrap().deploy_stdout();
283
284                self.launched_binary
285                    .get()
286                    .unwrap()
287                    .stdin()
288                    .send(format!("start: {formatted_defns}\n"))
289                    .unwrap();
290
291                let start_ack_line = ProgressTracker::leaf(
292                    format!(
293                        "{} / waiting for ack start",
294                        self.display_id
295                            .clone()
296                            .unwrap_or_else(|| format!("service/{}", self.id))
297                    ),
298                    tokio::time::timeout(Duration::from_secs(60), stdout_receiver),
299                )
300                .await??;
301                if !start_ack_line.starts_with("ack start") {
302                    bail!("expected ack start");
303                }
304
305                Ok(())
306            })
307            .await?;
308
309        Ok(())
310    }
311
312    async fn stop(&self) -> Result<()> {
313        ProgressTracker::with_group(
314            self.display_id
315                .clone()
316                .unwrap_or_else(|| format!("service/{}", self.id)),
317            None,
318            || async {
319                let launched_binary = self.launched_binary.get().unwrap();
320                launched_binary.stdin().send("stop\n".to_owned())?;
321
322                let timeout_result = ProgressTracker::leaf(
323                    "waiting for exit",
324                    tokio::time::timeout(Duration::from_secs(60), launched_binary.wait()),
325                )
326                .await;
327                match timeout_result {
328                    Err(_timeout) => {} // `wait()` timed out, but stop will force quit.
329                    Ok(Err(unexpected_error)) => return Err(unexpected_error), // `wait()` errored.
330                    Ok(Ok(_exit_status)) => {}
331                }
332                launched_binary.stop().await?;
333
334                Ok(())
335            },
336        )
337        .await
338    }
339}
hydro_deploy/rust_crate/service.rs

hydro_deploy/rust_crate/
service.rs
