hydro_lang/deploy/
deploy_runtime_containerized_ecs.rs

1#![allow(
2    unused,
3    reason = "unused in trybuild but the __staged version is needed"
4)]
5#![allow(missing_docs, reason = "used internally")]
6
7use std::time::Duration;
8
9use futures::{SinkExt, Stream, StreamExt};
10use sinktools::lazy::{LazySink, LazySource};
11use sinktools::lazy_sink_source::LazySinkSource;
12use stageleft::{QuotedWithContext, q};
13use tokio::net::TcpStream;
14use tokio::net::tcp::{OwnedReadHalf, OwnedWriteHalf};
15use tokio_util::codec::{FramedRead, FramedWrite, LengthDelimitedCodec};
16use tracing::{Instrument, debug, instrument, span, trace, trace_span};
17
18pub use super::deploy_runtime_containerized::{
19    CHANNEL_MAGIC, CHANNEL_MUX_PORT, CHANNEL_PROTOCOL_VERSION, ChannelHandshake, ChannelMagic,
20    ChannelMux, ChannelProtocolVersion, SocketIdent, cluster_ids, connect_channel,
21    get_or_init_channel_mux, send_handshake,
22};
23use crate::location::dynamic::LocationId;
24use crate::location::member_id::TaglessMemberId;
25use crate::location::{LocationKey, MembershipEvent};
26
27pub fn deploy_containerized_o2o(
28    target_task_family: &str,
29    channel_name: &str,
30) -> (syn::Expr, syn::Expr) {
31    (
32        q!(LazySink::<_, _, _, bytes::Bytes>::new(move || Box::pin(
33            async move {
34                let channel_name = channel_name;
35                let target_task_family = target_task_family;
36                let task_id = self::resolve_task_family_to_task_id(target_task_family).await;
37                let ip = self::resolve_task_ip(&task_id).await;
38                let target = format!("{}:{}", ip, self::CHANNEL_MUX_PORT);
39                debug!(name: "connecting", %target, %target_task_family, %task_id, %channel_name);
40
41                let stream = self::connect_channel(&target).await?;
42                let mut sink = FramedWrite::new(stream, LengthDelimitedCodec::new());
43
44                self::send_handshake(&mut sink, channel_name, None).await?;
45
46                Result::<_, std::io::Error>::Ok(sink)
47            }
48        )))
49        .splice_untyped_ctx(&()),
50        q!(LazySource::new(move || Box::pin(async move {
51            let channel_name = channel_name;
52            let mux = self::get_or_init_channel_mux();
53            let mut rx = mux.register(channel_name.to_owned());
54
55            let (_sender_id, source) = rx.recv().await.ok_or_else(|| {
56                std::io::Error::new(std::io::ErrorKind::ConnectionReset, "channel mux closed")
57            })?;
58
59            debug!(name: "o2o_channel_connected", %channel_name);
60
61            Result::<_, std::io::Error>::Ok(source)
62        })))
63        .splice_untyped_ctx(&()),
64    )
65}
66
67pub fn deploy_containerized_o2m(channel_name: &str) -> (syn::Expr, syn::Expr) {
68    (
69        q!(sinktools::demux_map_lazy::<_, _, _, _>(
70            move |key: &TaglessMemberId| {
71                let key = key.clone();
72                let channel_name = channel_name.to_owned();
73
74                LazySink::<_, _, _, bytes::Bytes>::new(move || {
75                    Box::pin(async move {
76                        let task_id = key.get_container_name();
77                        let ip = self::resolve_task_ip(task_id).await;
78                        let target = format!("{}:{}", ip, self::CHANNEL_MUX_PORT);
79                        debug!(name: "connecting", %target, %task_id, channel_name = %channel_name);
80
81                        let stream = self::connect_channel(&target).await?;
82                        let mut sink = FramedWrite::new(stream, LengthDelimitedCodec::new());
83
84                        self::send_handshake(&mut sink, &channel_name, None).await?;
85
86                        Result::<_, std::io::Error>::Ok(sink)
87                    })
88                })
89            }
90        ))
91        .splice_untyped_ctx(&()),
92        q!(LazySource::new(move || Box::pin(async move {
93            let channel_name = channel_name;
94            let mux = self::get_or_init_channel_mux();
95            let mut rx = mux.register(channel_name.to_owned());
96
97            let (_sender_id, source) = rx.recv().await.ok_or_else(|| {
98                std::io::Error::new(std::io::ErrorKind::ConnectionReset, "channel mux closed")
99            })?;
100
101            debug!(name: "o2m_channel_connected", %channel_name);
102
103            Result::<_, std::io::Error>::Ok(source)
104        })))
105        .splice_untyped_ctx(&()),
106    )
107}
108
109pub fn deploy_containerized_m2o(
110    target_task_family: &str,
111    channel_name: &str,
112) -> (syn::Expr, syn::Expr) {
113    (
114        q!(LazySink::<_, _, _, bytes::Bytes>::new(move || {
115            Box::pin(async move {
116                let channel_name = channel_name;
117                let target_task_family = target_task_family;
118                let target_task_id = self::resolve_task_family_to_task_id(target_task_family).await;
119                let ip = self::resolve_task_ip(&target_task_id).await;
120                let target = format!("{}:{}", ip, self::CHANNEL_MUX_PORT);
121                debug!(name: "connecting", %target, %target_task_family, %target_task_id, %channel_name);
122
123                let stream = self::connect_channel(&target).await?;
124                let mut sink = FramedWrite::new(stream, LengthDelimitedCodec::new());
125
126                let self_task_id = self::get_self_task_id();
127                self::send_handshake(&mut sink, channel_name, Some(&self_task_id)).await?;
128
129                Result::<_, std::io::Error>::Ok(sink)
130            })
131        }))
132        .splice_untyped_ctx(&()),
133        q!(LazySource::new(move || Box::pin(async move {
134            let channel_name = channel_name;
135            let mux = self::get_or_init_channel_mux();
136            let mut rx = mux.register(channel_name.to_owned());
137
138            Result::<_, std::io::Error>::Ok(
139                futures::stream::unfold(rx, |mut rx| {
140                    Box::pin(async move {
141                        let (sender_id, source) = rx.recv().await?;
142                        let from_task_id = sender_id
143                            .expect("m2o sender must provide task ID");
144
145                        debug!(name: "m2o_channel_connected", %from_task_id);
146
147                        Some((
148                            source.map(move |v| {
149                                v.map(|v| (TaglessMemberId::from_container_name(from_task_id.clone()), v))
150                            }),
151                            rx,
152                        ))
153                    })
154                })
155                .flatten_unordered(None),
156            )
157        })))
158        .splice_untyped_ctx(&()),
159    )
160}
161
162pub fn deploy_containerized_m2m(channel_name: &str) -> (syn::Expr, syn::Expr) {
163    (
164        q!(sinktools::demux_map_lazy::<_, _, _, _>(
165            move |key: &TaglessMemberId| {
166                let key = key.clone();
167                let channel_name = channel_name.to_owned();
168
169                LazySink::<_, _, _, bytes::Bytes>::new(move || {
170                    Box::pin(async move {
171                        let task_id = key.get_container_name();
172                        let ip = self::resolve_task_ip(task_id).await;
173                        let target = format!("{}:{}", ip, self::CHANNEL_MUX_PORT);
174                        debug!(name: "connecting", %target, %task_id, channel_name = %channel_name);
175
176                        let stream = self::connect_channel(&target).await?;
177                        let mut sink = FramedWrite::new(stream, LengthDelimitedCodec::new());
178
179                        let self_task_id = self::get_self_task_id();
180                        self::send_handshake(&mut sink, &channel_name, Some(&self_task_id)).await?;
181
182                        Result::<_, std::io::Error>::Ok(sink)
183                    })
184                })
185            }
186        ))
187        .splice_untyped_ctx(&()),
188        q!(LazySource::new(move || Box::pin(async move {
189            let channel_name = channel_name;
190            let mux = self::get_or_init_channel_mux();
191            let mut rx = mux.register(channel_name.to_owned());
192
193            Result::<_, std::io::Error>::Ok(
194                futures::stream::unfold(rx, |mut rx| {
195                    Box::pin(async move {
196                        let (sender_id, source) = rx.recv().await?;
197                        let from_task_id = sender_id.expect("m2m sender must provide task ID");
198
199                        debug!(name: "m2m_channel_connected", %from_task_id);
200
201                        Some((
202                            source.map(move |v| {
203                                v.map(|v| {
204                                    (
205                                        TaglessMemberId::from_container_name(from_task_id.clone()),
206                                        v,
207                                    )
208                                })
209                            }),
210                            rx,
211                        ))
212                    })
213                })
214                .flatten_unordered(None),
215            )
216        })))
217        .splice_untyped_ctx(&()),
218    )
219}
220
221pub fn deploy_containerized_external_sink_source_ident(
222    bind_addr: String,
223    socket_ident: syn::Ident,
224) -> syn::Expr {
225    let socket_ident = SocketIdent { socket_ident };
226
227    q!(LazySinkSource::<
228        _,
229        FramedRead<OwnedReadHalf, LengthDelimitedCodec>,
230        FramedWrite<OwnedWriteHalf, LengthDelimitedCodec>,
231        bytes::Bytes,
232        std::io::Error,
233    >::new(async move {
234        let span = span!(tracing::Level::TRACE, "lazy_sink_source");
235        let guard = span.enter();
236        let bind_addr = bind_addr;
237        trace!(name: "attempting to accept from external", %bind_addr);
238        std::mem::drop(guard);
239        let (stream, peer) = socket_ident.accept().instrument(span.clone()).await?;
240        let guard = span.enter();
241
242        debug!(name: "external accepting", ?peer);
243        let (rx, tx) = stream.into_split();
244
245        let fr = FramedRead::new(rx, LengthDelimitedCodec::new());
246        let fw = FramedWrite::new(tx, LengthDelimitedCodec::new());
247
248        Result::<_, std::io::Error>::Ok((fr, fw))
249    },))
250    .splice_untyped_ctx(&())
251}
252
253pub fn cluster_self_id<'a>() -> impl QuotedWithContext<'a, TaglessMemberId, ()> + Clone + 'a {
254    q!(TaglessMemberId::from_container_name(
255        self::get_self_task_id()
256    ))
257}
258
259pub fn cluster_membership_stream<'a>(
260    location_id: &LocationId,
261) -> impl QuotedWithContext<'a, Box<dyn Stream<Item = (TaglessMemberId, MembershipEvent)> + Unpin>, ()>
262{
263    let location_key = location_id.key();
264
265    q!(Box::new(self::ecs_membership_stream(
266        std::env::var("CLUSTER_NAME").unwrap(),
267        location_key
268    ))
269        as Box<
270            dyn Stream<Item = (TaglessMemberId, MembershipEvent)> + Unpin,
271        >)
272}
273
274#[instrument(skip_all, fields(%cluster_name, %location_key))]
275fn ecs_membership_stream(
276    cluster_name: String,
277    location_key: LocationKey,
278) -> impl Stream<Item = (TaglessMemberId, MembershipEvent)> + Unpin {
279    use std::collections::HashSet;
280
281    use futures::stream::{StreamExt, once};
282
283    trace!(name: "ecs_membership_stream_created", %cluster_name, %location_key);
284
285    let ecs_poller_span = trace_span!("ecs_poller");
286
287    // Task family format: hy-{name_hint}-loc{idx}v{version}
288    // Example: hy-p1-loc2v1
289    let task_definition_arn_parser =
290        regex::Regex::new(r#"arn:aws:ecs:(?<region>.*):(?<account_id>.*):task-definition\/(?<container_id>hy-(?<type>[^-]+)-loc(?<location_idx>[0-9]+)v(?<location_version>[0-9]+)(?:-(?<instance_id>.*))?):.*"#).unwrap();
291
292    let poll_stream = futures::stream::unfold(
293        (HashSet::<String>::new(), cluster_name, location_key),
294        move |(known_tasks, cluster_name, location_key)| {
295            let task_definition_arn_parser = task_definition_arn_parser.clone();
296
297            async move {
298                let config = aws_config::load_defaults(aws_config::BehaviorVersion::latest()).await;
299                let ecs_client = aws_sdk_ecs::Client::new(&config);
300
301                let tasks = match ecs_client.list_tasks().cluster(&cluster_name).send().await {
302                    Ok(tasks) => tasks,
303                    Err(e) => {
304                        trace!(name: "list_tasks_error", error = %e);
305                        tokio::time::sleep(Duration::from_secs(2)).await;
306                        return Some((Vec::new(), (known_tasks, cluster_name, location_key)));
307                    }
308                };
309
310                let task_arns = tasks.task_arns().to_vec();
311
312                let mut events = Vec::new();
313                let mut current_tasks = HashSet::<String>::new();
314
315                if !task_arns.is_empty() {
316                    let task_details = match ecs_client
317                        .describe_tasks()
318                        .cluster(&cluster_name)
319                        .set_tasks(Some(task_arns.clone()))
320                        .send()
321                        .await
322                    {
323                        Ok(details) => details,
324                        Err(e) => {
325                            trace!(name: "describe_tasks_error", error = %e);
326                            tokio::time::sleep(Duration::from_secs(2)).await;
327                            return Some((Vec::new(), (known_tasks, cluster_name, location_key)));
328                        }
329                    };
330
331                    for task in task_details.tasks() {
332                        let Some(last_status) = task.last_status() else {
333                            continue;
334                        };
335
336                        if last_status != "RUNNING" {
337                            continue;
338                        }
339
340                        let Some(task_def_arn) = task.task_definition_arn() else {
341                            continue;
342                        };
343
344                        let Some(captures) = task_definition_arn_parser.captures(task_def_arn)
345                        else {
346                            continue;
347                        };
348
349                        let Some(location_idx) = captures.name("location_idx") else {
350                            continue;
351                        };
352                        let Some(location_version) = captures.name("location_version") else {
353                            continue;
354                        };
355                        // Reconstruct the location key string and parse it
356                        let location_key_str =
357                            format!("loc{}v{}", location_idx.as_str(), location_version.as_str());
358                        let task_location_key: LocationKey = match location_key_str.parse() {
359                            Ok(key) => key,
360                            Err(_) => {
361                                continue;
362                            }
363                        };
364
365                        // Filter by location_id - only include tasks for this specific cluster
366                        if task_location_key != location_key {
367                            continue;
368                        }
369
370                        // Extract task ID from task ARN (last segment after final /)
371                        // Task ARN format: arn:aws:ecs:region:account:task/cluster-name/task-id
372                        let Some(task_arn) = task.task_arn() else {
373                            continue;
374                        };
375                        let Some(task_id) = task_arn.rsplit('/').next() else {
376                            continue;
377                        };
378
379                        // Use task_id as the member identifier
380                        current_tasks.insert(task_id.to_owned());
381                        if !known_tasks.contains(task_id) {
382                            trace!(name: "task_joined", %task_id);
383                            events.push((task_id.to_owned(), MembershipEvent::Joined));
384                        }
385                    }
386                }
387
388                #[expect(
389                    clippy::disallowed_methods,
390                    reason = "nondeterministic iteration order, container events are not deterministically ordered"
391                )]
392                for task_id in known_tasks.iter() {
393                    if !current_tasks.contains(task_id) {
394                        trace!(name: "task_left", %task_id);
395                        events.push((task_id.to_owned(), MembershipEvent::Left));
396                    }
397                }
398
399                tokio::time::sleep(Duration::from_secs(2)).await;
400
401                Some((events, (current_tasks, cluster_name, location_key)))
402            }
403            .instrument(ecs_poller_span.clone())
404        },
405    )
406    .flat_map(futures::stream::iter);
407
408    Box::pin(
409        poll_stream
410            .map(|(k, v)| (TaglessMemberId::from_container_name(k), v))
411            .inspect(|(member_id, event)| trace!(name: "membership_event", ?member_id, ?event)),
412    )
413}
414
415/// Resolve a task ID to its private IP address via ECS API.
416async fn resolve_task_ip(task_id: &str) -> String {
417    let cluster_name = std::env::var("CLUSTER_NAME").unwrap();
418
419    let config = aws_config::load_defaults(aws_config::BehaviorVersion::latest()).await;
420    let ecs_client = aws_sdk_ecs::Client::new(&config);
421
422    loop {
423        let tasks = match ecs_client.list_tasks().cluster(&cluster_name).send().await {
424            Ok(t) => t,
425            Err(e) => {
426                trace!(name: "resolve_ip_list_error", %task_id, error = %e);
427                tokio::time::sleep(Duration::from_secs(1)).await;
428                continue;
429            }
430        };
431
432        let task_arns: Vec<_> = tasks.task_arns().to_vec();
433        if task_arns.is_empty() {
434            trace!(name: "resolve_ip_no_tasks", %task_id);
435            tokio::time::sleep(Duration::from_secs(1)).await;
436            continue;
437        }
438
439        let task_details = match ecs_client
440            .describe_tasks()
441            .cluster(&cluster_name)
442            .set_tasks(Some(task_arns))
443            .send()
444            .await
445        {
446            Ok(d) => d,
447            Err(e) => {
448                trace!(name: "resolve_ip_describe_error", %task_id, error = %e);
449                tokio::time::sleep(Duration::from_secs(1)).await;
450                continue;
451            }
452        };
453
454        // Find the task with matching task ID
455        for task in task_details.tasks() {
456            let Some(task_arn) = task.task_arn() else {
457                continue;
458            };
459            let current_task_id = task_arn.rsplit('/').next().unwrap_or_default();
460
461            if current_task_id == task_id
462                && let Some(ip) = task
463                    .attachments()
464                    .iter()
465                    .flat_map(|a| a.details())
466                    .find(|d| d.name() == Some("privateIPv4Address"))
467                    .and_then(|d| d.value())
468            {
469                trace!(name: "resolved_ip", %task_id, %ip);
470                return ip.to_owned();
471            }
472        }
473
474        trace!(name: "resolve_ip_not_found", %task_id);
475        tokio::time::sleep(Duration::from_secs(1)).await;
476    }
477}
478
479/// Resolve a task family name to its task ID via ECS API.
480/// Used for process-to-process connections where the target is known by task family at compile time.
481async fn resolve_task_family_to_task_id(task_family: &str) -> String {
482    let cluster_name = std::env::var("CLUSTER_NAME").unwrap();
483
484    let config = aws_config::load_defaults(aws_config::BehaviorVersion::latest()).await;
485    let ecs_client = aws_sdk_ecs::Client::new(&config);
486
487    loop {
488        let tasks = match ecs_client
489            .list_tasks()
490            .cluster(&cluster_name)
491            .family(task_family)
492            .send()
493            .await
494        {
495            Ok(t) => t,
496            Err(e) => {
497                trace!(name: "resolve_family_list_error", %task_family, error = %e);
498                tokio::time::sleep(Duration::from_secs(1)).await;
499                continue;
500            }
501        };
502
503        let Some(task_arn) = tasks.task_arns().first() else {
504            trace!(name: "resolve_family_no_task", %task_family);
505            tokio::time::sleep(Duration::from_secs(1)).await;
506            continue;
507        };
508
509        // Extract task ID from ARN
510        let task_id = task_arn.rsplit('/').next().unwrap_or_default();
511        if !task_id.is_empty() {
512            trace!(name: "resolved_task_id", %task_family, %task_id);
513            return task_id.to_owned();
514        }
515
516        trace!(name: "resolve_family_invalid_arn", %task_family, %task_arn);
517        tokio::time::sleep(Duration::from_secs(1)).await;
518    }
519}
520
521/// Parse the ECS task ID from a metadata URI.
522///
523/// URI format: `http://169.254.170.2/v4/{task_id}-{runtime_id}`
524/// where task_id is 32 hex chars and runtime_id is a numeric suffix.
525fn parse_task_id_from_metadata_uri(metadata_uri: &str) -> Option<String> {
526    let re = regex::Regex::new(r"/v4/(?P<task_id>[0-9a-f]{32})-\d+$").unwrap();
527    re.captures(metadata_uri)
528        .and_then(|c| c.name("task_id"))
529        .map(|m| m.as_str().to_owned())
530}
531
532fn get_self_task_id() -> String {
533    let metadata_uri = std::env::var("ECS_CONTAINER_METADATA_URI_V4")
534        .expect("ECS_CONTAINER_METADATA_URI_V4 not set - are we running in ECS?");
535    parse_task_id_from_metadata_uri(&metadata_uri).unwrap_or_else(|| {
536        panic!("ECS_CONTAINER_METADATA_URI_V4 does not match expected format /v4/{{task_id}}-{{runtime_id}}: {metadata_uri}")
537    })
538}
539
540#[cfg(test)]
541mod tests {
542    use super::*;
543
544    #[test]
545    fn test_parse_task_id_from_metadata_uri() {
546        let uri = "http://169.254.170.2/v4/01234567890abcdef01234567890abcd-123456789";
547        assert_eq!(
548            parse_task_id_from_metadata_uri(uri),
549            Some("01234567890abcdef01234567890abcd".to_owned())
550        );
551    }
552
553    #[test]
554    fn test_parse_task_id_invalid_uri() {
555        assert_eq!(
556            parse_task_id_from_metadata_uri("http://169.254.170.2/v4/tooshort-123"),
557            None
558        );
559        assert_eq!(
560            parse_task_id_from_metadata_uri(
561                "http://169.254.170.2/v3/01234567890abcdef01234567890abcd-123"
562            ),
563            None
564        );
565        assert_eq!(parse_task_id_from_metadata_uri(""), None);
566    }
567}
hydro_lang/deploy/deploy_runtime_containerized_ecs.rs

hydro_lang/deploy/
deploy_runtime_containerized_ecs.rs
