hydro_deploy/rust_crate/
ports.rs

1use std::any::Any;
2use std::collections::{BTreeMap, HashMap};
3use std::fmt::Debug;
4use std::ops::Deref;
5use std::sync::{Arc, Weak};
6
7use anyhow::{Result, bail};
8use append_only_vec::AppendOnlyVec;
9use async_recursion::async_recursion;
10use hydro_deploy_integration::ServerPort;
11use tokio::sync::RwLock;
12
13use super::RustCrateService;
14use crate::{ClientStrategy, Host, LaunchedHost, PortNetworkHint, ServerStrategy};
15
16pub trait RustCrateSource: Send + Sync {
17    fn source_path(&self) -> SourcePath;
18    fn record_server_config(&self, config: ServerConfig);
19
20    fn host(&self) -> Arc<dyn Host>;
21    fn server(&self) -> Arc<dyn RustCrateServer>;
22    fn record_server_strategy(&self, config: ServerStrategy);
23
24    fn wrap_reverse_server_config(&self, config: ServerConfig) -> ServerConfig {
25        config
26    }
27
28    fn send_to(&self, sink: &dyn RustCrateSink) {
29        let forward_res = sink.instantiate(&self.source_path());
30        if let Ok(instantiated) = forward_res {
31            self.record_server_config(instantiated());
32        } else {
33            drop(forward_res);
34            let instantiated = sink
35                .instantiate_reverse(&self.host(), self.server(), &|p| {
36                    self.wrap_reverse_server_config(p)
37                })
38                .unwrap();
39            self.record_server_strategy(instantiated(sink));
40        }
41    }
42}
43
44pub trait RustCrateServer: Debug + Send + Sync {
45    fn get_port(&self) -> ServerPort;
46    fn launched_host(&self) -> Arc<dyn LaunchedHost>;
47}
48
49pub type ReverseSinkInstantiator = Box<dyn FnOnce(&dyn Any) -> ServerStrategy>;
50
51pub trait RustCrateSink: Any + Send + Sync {
52    /// Instantiate the sink as the source host connecting to the sink host.
53    /// Returns a thunk that can be called to perform mutations that instantiate the sink.
54    fn instantiate(&self, client_path: &SourcePath) -> Result<Box<dyn FnOnce() -> ServerConfig>>;
55
56    /// Instantiate the sink, but as the sink host connecting to the source host.
57    /// Returns a thunk that can be called to perform mutations that instantiate the sink, taking a mutable reference to this sink.
58    fn instantiate_reverse(
59        &self,
60        server_host: &Arc<dyn Host>,
61        server_sink: Arc<dyn RustCrateServer>,
62        wrap_client_port: &dyn Fn(ServerConfig) -> ServerConfig,
63    ) -> Result<ReverseSinkInstantiator>;
64}
65
66pub struct TaggedSource {
67    pub source: Arc<dyn RustCrateSource>,
68    pub tag: u32,
69}
70
71impl RustCrateSource for TaggedSource {
72    fn source_path(&self) -> SourcePath {
73        SourcePath::Tagged(Box::new(self.source.source_path()), self.tag)
74    }
75
76    fn record_server_config(&self, config: ServerConfig) {
77        self.source.record_server_config(config);
78    }
79
80    fn host(&self) -> Arc<dyn Host> {
81        self.source.host()
82    }
83
84    fn server(&self) -> Arc<dyn RustCrateServer> {
85        self.source.server()
86    }
87
88    fn wrap_reverse_server_config(&self, config: ServerConfig) -> ServerConfig {
89        ServerConfig::Tagged(Box::new(config), self.tag)
90    }
91
92    fn record_server_strategy(&self, config: ServerStrategy) {
93        self.source.record_server_strategy(config);
94    }
95}
96
97pub struct NullSourceSink;
98
99impl RustCrateSource for NullSourceSink {
100    fn source_path(&self) -> SourcePath {
101        SourcePath::Null
102    }
103
104    fn host(&self) -> Arc<dyn Host> {
105        panic!("null source has no host")
106    }
107
108    fn server(&self) -> Arc<dyn RustCrateServer> {
109        panic!("null source has no server")
110    }
111
112    fn record_server_config(&self, _config: ServerConfig) {}
113    fn record_server_strategy(&self, _config: ServerStrategy) {}
114}
115
116impl RustCrateSink for NullSourceSink {
117    fn instantiate(&self, _client_path: &SourcePath) -> Result<Box<dyn FnOnce() -> ServerConfig>> {
118        Ok(Box::new(|| ServerConfig::Null))
119    }
120
121    fn instantiate_reverse(
122        &self,
123        _server_host: &Arc<dyn Host>,
124        _server_sink: Arc<dyn RustCrateServer>,
125        _wrap_client_port: &dyn Fn(ServerConfig) -> ServerConfig,
126    ) -> Result<ReverseSinkInstantiator> {
127        Ok(Box::new(|_| ServerStrategy::Null))
128    }
129}
130
131pub struct DemuxSink {
132    pub demux: HashMap<u32, Arc<dyn RustCrateSink>>,
133}
134
135impl RustCrateSink for DemuxSink {
136    fn instantiate(&self, client_host: &SourcePath) -> Result<Box<dyn FnOnce() -> ServerConfig>> {
137        let mut thunk_map = HashMap::new();
138        for (key, target) in &self.demux {
139            thunk_map.insert(*key, target.instantiate(client_host)?);
140        }
141
142        Ok(Box::new(move || {
143            let instantiated_map = thunk_map
144                .into_iter()
145                .map(|(key, thunk)| (key, (thunk)()))
146                .collect();
147
148            ServerConfig::Demux(instantiated_map)
149        }))
150    }
151
152    fn instantiate_reverse(
153        &self,
154        server_host: &Arc<dyn Host>,
155        server_sink: Arc<dyn RustCrateServer>,
156        wrap_client_port: &dyn Fn(ServerConfig) -> ServerConfig,
157    ) -> Result<ReverseSinkInstantiator> {
158        let mut thunk_map = HashMap::new();
159        for (key, target) in &self.demux {
160            thunk_map.insert(
161                *key,
162                target.instantiate_reverse(
163                    server_host,
164                    server_sink.clone(),
165                    // the parent wrapper selects the demux port for the parent defn, so do that first
166                    &|p| ServerConfig::DemuxSelect(Box::new(wrap_client_port(p)), *key),
167                )?,
168            );
169        }
170
171        Ok(Box::new(move |me| {
172            let me = me.downcast_ref::<DemuxSink>().unwrap();
173            let instantiated_map = thunk_map
174                .into_iter()
175                .map(|(key, thunk)| (key, (thunk)(me.demux.get(&key).unwrap())))
176                .collect();
177
178            ServerStrategy::Demux(instantiated_map)
179        }))
180    }
181}
182
183#[derive(Clone, Debug)]
184pub struct RustCratePortConfig {
185    pub service: Weak<RustCrateService>,
186    pub service_host: Arc<dyn Host>,
187    pub service_server_defns: Arc<RwLock<HashMap<String, ServerPort>>>,
188    pub network_hint: PortNetworkHint,
189    pub port: String,
190    pub merge: bool,
191}
192
193impl RustCratePortConfig {
194    #[must_use]
195    pub fn merge(mut self) -> Self {
196        self.merge = true;
197        self
198    }
199}
200
201impl RustCrateSource for RustCratePortConfig {
202    fn source_path(&self) -> SourcePath {
203        SourcePath::Direct(self.service.upgrade().unwrap().on.clone())
204    }
205
206    fn host(&self) -> Arc<dyn Host> {
207        self.service_host.clone()
208    }
209
210    fn server(&self) -> Arc<dyn RustCrateServer> {
211        let from = self.service.upgrade().unwrap();
212
213        Arc::new(RustCratePortConfig {
214            service: Arc::downgrade(&from),
215            service_host: from.on.clone(),
216            service_server_defns: from.server_defns.clone(),
217            network_hint: self.network_hint,
218            port: self.port.clone(),
219            merge: false,
220        })
221    }
222
223    fn record_server_config(&self, config: ServerConfig) {
224        let from = self.service.upgrade().unwrap();
225        // TODO(shadaj): if already in this map, we want to broadcast
226        assert!(
227            from.port_to_server.insert(self.port.clone(), config),
228            "The port configuration is incorrect, for example, are you using a ConnectedDirect instead of a ConnectedDemux?"
229        );
230    }
231
232    fn record_server_strategy(&self, config: ServerStrategy) {
233        let from = self.service.upgrade().unwrap();
234        assert!(
235            from.port_to_bind.insert(self.port.clone(), config),
236            "port already set!"
237        );
238    }
239}
240
241impl RustCrateServer for RustCratePortConfig {
242    fn get_port(&self) -> ServerPort {
243        // we are in `deployment.start()`, so no one should be writing
244        let server_defns = self.service_server_defns.try_read().unwrap();
245        server_defns.get(&self.port).unwrap().clone()
246    }
247
248    fn launched_host(&self) -> Arc<dyn LaunchedHost> {
249        self.service_host.launched().unwrap()
250    }
251}
252
253pub enum SourcePath {
254    Null,
255    Direct(Arc<dyn Host>),
256    Many(Arc<dyn Host>),
257    Tagged(Box<SourcePath>, u32),
258}
259
260impl SourcePath {
261    #[expect(
262        clippy::type_complexity,
263        reason = "internals (dyn Fn to defer instantiation)"
264    )]
265    fn plan<T: RustCrateServer + Clone + 'static>(
266        &self,
267        server: &T,
268        server_host: &dyn Host,
269        network_hint: PortNetworkHint,
270    ) -> Result<(Box<dyn FnOnce(&dyn Any) -> ServerStrategy>, ServerConfig)> {
271        match self {
272            SourcePath::Direct(client_host) => {
273                let (conn_type, bind_type) =
274                    server_host.strategy_as_server(client_host.deref(), network_hint)?;
275                let base_config = ServerConfig::from_strategy(&conn_type, Arc::new(server.clone()));
276                Ok((
277                    Box::new(|host| ServerStrategy::Direct(bind_type(host))),
278                    base_config,
279                ))
280            }
281
282            SourcePath::Many(client_host) => {
283                let (conn_type, bind_type) =
284                    server_host.strategy_as_server(client_host.deref(), network_hint)?;
285                let base_config = ServerConfig::from_strategy(&conn_type, Arc::new(server.clone()));
286                Ok((
287                    Box::new(|host| ServerStrategy::Many(bind_type(host))),
288                    base_config,
289                ))
290            }
291
292            SourcePath::Tagged(underlying, tag) => {
293                let (bind_type, base_config) =
294                    underlying.plan(server, server_host, network_hint)?;
295                let tag = *tag;
296                Ok((
297                    Box::new(move |host| ServerStrategy::Tagged(Box::new(bind_type(host)), tag)),
298                    ServerConfig::TaggedUnwrap(Box::new(base_config)),
299                ))
300            }
301
302            SourcePath::Null => Ok((Box::new(|_| ServerStrategy::Null), ServerConfig::Null)),
303        }
304    }
305}
306
307impl RustCrateSink for RustCratePortConfig {
308    fn instantiate(&self, client_path: &SourcePath) -> Result<Box<dyn FnOnce() -> ServerConfig>> {
309        let server = self.service.upgrade().unwrap();
310
311        let server_host = server.on.clone();
312
313        let (bind_type, base_config) =
314            client_path.plan(self, server_host.deref(), self.network_hint)?;
315
316        let server = server.clone();
317        let merge = self.merge;
318        let port = self.port.clone();
319        Ok(Box::new(move || {
320            let bind_type = (bind_type)(&*server.on);
321
322            if merge {
323                let merge_config = server
324                    .port_to_bind
325                    .get_or_insert_owned(port, || ServerStrategy::Merge(Default::default()));
326                let ServerStrategy::Merge(merge) = merge_config else {
327                    panic!("Expected a merge connection definition")
328                };
329                merge.push(bind_type);
330                ServerConfig::MergeSelect(Box::new(base_config), merge.len() - 1)
331            } else {
332                assert!(
333                    server.port_to_bind.insert(port.clone(), bind_type),
334                    "port already set!"
335                );
336                base_config
337            }
338        }))
339    }
340
341    fn instantiate_reverse(
342        &self,
343        server_host: &Arc<dyn Host>,
344        server_sink: Arc<dyn RustCrateServer>,
345        wrap_client_port: &dyn Fn(ServerConfig) -> ServerConfig,
346    ) -> Result<ReverseSinkInstantiator> {
347        if !matches!(self.network_hint, PortNetworkHint::Auto) {
348            bail!("Trying to form collection where I am the client, but I have server hint")
349        }
350
351        let client = self.service.upgrade().unwrap();
352
353        let server_host = server_host.clone();
354
355        let (conn_type, bind_type) =
356            server_host.strategy_as_server(&*client.on, PortNetworkHint::Auto)?;
357        let client_port = wrap_client_port(ServerConfig::from_strategy(&conn_type, server_sink));
358
359        let client = client.clone();
360        let merge = self.merge;
361        let port = self.port.clone();
362        Ok(Box::new(move |_| {
363            if merge {
364                let merge_config = client
365                    .port_to_server
366                    .get_or_insert_owned(port, || ServerConfig::Merge(Default::default()));
367                let ServerConfig::Merge(merge) = merge_config else {
368                    panic!()
369                };
370                merge.push(client_port);
371            } else {
372                assert!(
373                    client.port_to_server.insert(port.clone(), client_port),
374                    "port already set!"
375                );
376            };
377
378            ServerStrategy::Direct((bind_type)(&*client.on))
379        }))
380    }
381}
382
383#[derive(Clone, Debug)]
384pub enum ServerConfig {
385    Direct(Arc<dyn RustCrateServer>),
386    Forwarded(Arc<dyn RustCrateServer>),
387    /// A demux that will be used at runtime to listen to many connections.
388    Demux(HashMap<u32, ServerConfig>),
389    /// The other side of a demux, with a port to extract the appropriate connection.
390    DemuxSelect(Box<ServerConfig>, u32),
391    /// A merge that will be used at runtime to combine many connections.
392    /// `AppendOnlyVec` has a quite large inline array, so we box it.
393    Merge(Box<AppendOnlyVec<ServerConfig>>),
394    /// The other side of a merge, with a port to extract the appropriate connection.
395    MergeSelect(Box<ServerConfig>, usize),
396    Tagged(Box<ServerConfig>, u32),
397    TaggedUnwrap(Box<ServerConfig>),
398    Null,
399}
400
401impl ServerConfig {
402    pub fn from_strategy(
403        strategy: &ClientStrategy<'_>,
404        server: Arc<dyn RustCrateServer>,
405    ) -> ServerConfig {
406        match strategy {
407            ClientStrategy::UnixSocket(_) | ClientStrategy::InternalTcpPort(_) => {
408                ServerConfig::Direct(server)
409            }
410            ClientStrategy::ForwardedTcpPort(_) => ServerConfig::Forwarded(server),
411        }
412    }
413}
414
415#[async_recursion]
416async fn forward_connection(conn: &ServerPort, target: &dyn LaunchedHost) -> ServerPort {
417    match conn {
418        ServerPort::UnixSocket(_) => panic!("Expected a TCP port to be forwarded"),
419        ServerPort::TcpPort(addr) => ServerPort::TcpPort(target.forward_port(addr).await.unwrap()),
420        ServerPort::Demux(demux) => {
421            let mut forwarded_map = BTreeMap::new();
422            for (key, conn) in demux {
423                forwarded_map.insert(*key, forward_connection(conn, target).await); // TODO(mingwei): do in parallel?
424            }
425            ServerPort::Demux(forwarded_map)
426        }
427        ServerPort::Merge(merge) => {
428            let mut forwarded_vec = Vec::new();
429            for conn in merge {
430                forwarded_vec.push(forward_connection(conn, target).await); // TODO(mingwei): do in parallel?
431            }
432            ServerPort::Merge(forwarded_vec)
433        }
434        ServerPort::Tagged(underlying, id) => {
435            ServerPort::Tagged(Box::new(forward_connection(underlying, target).await), *id)
436        }
437        ServerPort::Null => ServerPort::Null,
438    }
439}
440
441impl ServerConfig {
442    #[async_recursion]
443    pub async fn load_instantiated(
444        &self,
445        select: &(dyn Fn(ServerPort) -> ServerPort + Send + Sync),
446    ) -> ServerPort {
447        match self {
448            ServerConfig::Direct(server) => select(server.get_port()),
449
450            ServerConfig::Forwarded(server) => {
451                let selected = select(server.get_port());
452                forward_connection(&selected, server.launched_host().as_ref()).await
453            }
454
455            ServerConfig::Demux(demux) => {
456                let mut demux_map = BTreeMap::new();
457                for (key, conn) in demux {
458                    demux_map.insert(*key, conn.load_instantiated(select).await); // TODO(mingwei): do in parallel?
459                }
460                ServerPort::Demux(demux_map)
461            }
462
463            ServerConfig::DemuxSelect(underlying, key) => {
464                let key = *key;
465                underlying
466                    .load_instantiated(
467                        &(move |p| {
468                            if let ServerPort::Demux(mut mapping) = p {
469                                select(mapping.remove(&key).unwrap())
470                            } else {
471                                panic!("Expected a demux connection definition")
472                            }
473                        }),
474                    )
475                    .await
476            }
477
478            ServerConfig::Merge(merge) => {
479                let mut merge_vec = Vec::new();
480                for conn in merge.iter() {
481                    merge_vec.push(conn.load_instantiated(select).await); // TODO(mingwei): do in parallel?
482                }
483                ServerPort::Merge(merge_vec)
484            }
485
486            ServerConfig::MergeSelect(underlying, key) => {
487                let key = *key;
488                underlying
489                    .load_instantiated(
490                        &(move |p| {
491                            if let ServerPort::Merge(mut mapping) = p {
492                                select(mapping.remove(key))
493                            } else {
494                                panic!("Expected a merge connection definition")
495                            }
496                        }),
497                    )
498                    .await
499            }
500
501            ServerConfig::Tagged(underlying, id) => {
502                ServerPort::Tagged(Box::new(underlying.load_instantiated(select).await), *id)
503            }
504
505            ServerConfig::TaggedUnwrap(underlying) => {
506                let loaded = underlying.load_instantiated(select).await;
507                if let ServerPort::Tagged(underlying, _) = loaded {
508                    *underlying
509                } else {
510                    panic!("Expected a tagged connection definition")
511                }
512            }
513
514            ServerConfig::Null => ServerPort::Null,
515        }
516    }
517}
hydro_deploy/rust_crate/ports.rs

hydro_deploy/rust_crate/
ports.rs
