1use std::{
4 collections::{BTreeMap, BTreeSet},
5 net::{IpAddr, Ipv4Addr, SocketAddr},
6 num::NonZeroU16,
7 sync::Arc,
8 time::Duration,
9};
10
11use a3s_box_core::{
12 scale::ScaleEndpoint, ExecutionGeneration, ExecutionId, ExecutionPortConnector,
13 ExecutionPortStream,
14};
15use thiserror::Error;
16use tokio::{
17 net::{TcpListener, TcpStream},
18 sync::{oneshot, Mutex, Semaphore},
19 task::{JoinHandle, JoinSet},
20};
21
22use super::reconciler::ScaleReconcileError;
23
24const ENDPOINT_CONNECT_TIMEOUT: Duration = Duration::from_secs(5);
25const DEFAULT_ENDPOINT_DRAIN_TIMEOUT: Duration = Duration::from_secs(3);
26const MAX_ENDPOINT_CONNECTIONS: usize = 1_024;
27
28#[derive(Debug, Error)]
29pub enum ScaleEndpointConfigError {
30 #[error("invalid scale endpoint advertise host {host:?}: {message}")]
31 InvalidAdvertiseHost { host: String, message: String },
32}
33
34#[derive(Debug, Clone, PartialEq, Eq)]
36pub struct ScaleEndpointConfig {
37 bind_address: IpAddr,
38 advertise_host: url::Host<String>,
39 drain_timeout: Duration,
40}
41
42impl ScaleEndpointConfig {
43 pub fn new(
44 bind_address: IpAddr,
45 advertise_host: impl Into<String>,
46 ) -> Result<Self, ScaleEndpointConfigError> {
47 let advertise_host = advertise_host.into();
48 let parsed = match advertise_host.parse::<IpAddr>() {
49 Ok(IpAddr::V4(address)) => url::Host::Ipv4(address),
50 Ok(IpAddr::V6(address)) => url::Host::Ipv6(address),
51 Err(_) => url::Host::parse(&advertise_host).map_err(|error| {
52 ScaleEndpointConfigError::InvalidAdvertiseHost {
53 host: advertise_host.clone(),
54 message: error.to_string(),
55 }
56 })?,
57 };
58 if matches!(&parsed, url::Host::Domain(domain) if domain.trim().is_empty()) {
59 return Err(ScaleEndpointConfigError::InvalidAdvertiseHost {
60 host: advertise_host,
61 message: "host must not be empty".to_string(),
62 });
63 }
64 if matches!(&parsed, url::Host::Ipv4(address) if address.is_unspecified())
65 || matches!(&parsed, url::Host::Ipv6(address) if address.is_unspecified())
66 {
67 return Err(ScaleEndpointConfigError::InvalidAdvertiseHost {
68 host: advertise_host,
69 message: "an unspecified address cannot be advertised to Gateway".to_string(),
70 });
71 }
72 let address_family_mismatch = matches!(
73 (bind_address, &parsed),
74 (IpAddr::V4(_), url::Host::Ipv6(_)) | (IpAddr::V6(_), url::Host::Ipv4(_))
75 );
76 if address_family_mismatch {
77 return Err(ScaleEndpointConfigError::InvalidAdvertiseHost {
78 host: advertise_host,
79 message: format!(
80 "literal address family does not match endpoint bind address {bind_address}"
81 ),
82 });
83 }
84 Ok(Self {
85 bind_address,
86 advertise_host: parsed,
87 drain_timeout: DEFAULT_ENDPOINT_DRAIN_TIMEOUT,
88 })
89 }
90
91 pub fn loopback() -> Self {
92 Self {
93 bind_address: IpAddr::V4(Ipv4Addr::LOCALHOST),
94 advertise_host: url::Host::Ipv4(Ipv4Addr::LOCALHOST),
95 drain_timeout: DEFAULT_ENDPOINT_DRAIN_TIMEOUT,
96 }
97 }
98
99 pub fn with_drain_timeout(mut self, drain_timeout: Duration) -> Self {
100 self.drain_timeout = drain_timeout;
101 self
102 }
103
104 pub fn bind_address(&self) -> IpAddr {
105 self.bind_address
106 }
107
108 pub fn advertise_host(&self) -> String {
109 match &self.advertise_host {
110 url::Host::Domain(host) => host.clone(),
111 url::Host::Ipv4(host) => host.to_string(),
112 url::Host::Ipv6(host) => host.to_string(),
113 }
114 }
115
116 fn endpoint_url(&self, port: u16) -> String {
117 match &self.advertise_host {
118 url::Host::Ipv6(host) => format!("http://[{host}]:{port}"),
119 url::Host::Domain(host) => format!("http://{host}:{port}"),
120 url::Host::Ipv4(host) => format!("http://{host}:{port}"),
121 }
122 }
123}
124
125impl Default for ScaleEndpointConfig {
126 fn default() -> Self {
127 Self::loopback()
128 }
129}
130
131#[derive(Debug, Clone, PartialEq, Eq)]
132pub(super) struct ScaleEndpointTarget {
133 pub execution_id: ExecutionId,
134 pub generation: ExecutionGeneration,
135 pub service: String,
136 pub slot: u32,
137 pub guest_port: NonZeroU16,
138}
139
140struct EndpointLease {
141 target: ScaleEndpointTarget,
142 endpoint: ScaleEndpoint,
143 shutdown: Option<oneshot::Sender<()>>,
144 task: Option<JoinHandle<()>>,
145}
146
147impl EndpointLease {
148 async fn drain(mut self) {
149 if let Some(shutdown) = self.shutdown.take() {
150 let _ = shutdown.send(());
151 }
152 let Some(task) = self.task.take() else {
153 return;
154 };
155 if let Err(error) = task.await {
156 tracing::warn!(
157 execution_id = %self.target.execution_id,
158 service = self.target.service,
159 slot = self.target.slot,
160 %error,
161 "Scale endpoint drain task failed"
162 );
163 }
164 }
165}
166
167impl Drop for EndpointLease {
168 fn drop(&mut self) {
169 if let Some(task) = &self.task {
170 task.abort();
171 }
172 }
173}
174
175pub(super) struct ScaleEndpointOwner {
176 config: ScaleEndpointConfig,
177 connector: Arc<dyn ExecutionPortConnector>,
178 connection_limit: Arc<Semaphore>,
179 leases: Mutex<BTreeMap<String, EndpointLease>>,
180}
181
182impl ScaleEndpointOwner {
183 pub fn new(config: ScaleEndpointConfig, connector: Arc<dyn ExecutionPortConnector>) -> Self {
184 Self {
185 config,
186 connector,
187 connection_limit: Arc::new(Semaphore::new(MAX_ENDPOINT_CONNECTIONS)),
188 leases: Mutex::new(BTreeMap::new()),
189 }
190 }
191
192 pub async fn reconcile_service(
193 &self,
194 service: &str,
195 targets: &[ScaleEndpointTarget],
196 ) -> Result<Vec<ScaleEndpoint>, ScaleReconcileError> {
197 validate_targets(service, targets)?;
198 let desired = targets
199 .iter()
200 .map(|target| (target.execution_id.as_str().to_string(), target))
201 .collect::<BTreeMap<_, _>>();
202 let stale = {
203 let mut leases = self.leases.lock().await;
204 let stale_ids = leases
205 .iter()
206 .filter_map(|(execution_id, lease)| {
207 let retain = lease.target.service != service
208 || (lease.task.as_ref().is_some_and(|task| !task.is_finished())
209 && desired
210 .get(execution_id)
211 .is_some_and(|target| lease.target == **target));
212 (!retain).then(|| execution_id.clone())
213 })
214 .collect::<Vec<_>>();
215 stale_ids
216 .into_iter()
217 .filter_map(|execution_id| leases.remove(&execution_id))
218 .collect::<Vec<_>>()
219 };
220 drain_leases(stale).await;
221
222 let mut leases = self.leases.lock().await;
223
224 for target in targets {
225 if leases.contains_key(target.execution_id.as_str()) {
226 continue;
227 }
228 let listener = TcpListener::bind(SocketAddr::new(self.config.bind_address, 0))
229 .await
230 .map_err(|error| {
231 ScaleReconcileError::Lifecycle(format!(
232 "failed to bind service endpoint for {} slot {}: {error}",
233 target.service, target.slot
234 ))
235 })?;
236 let address = listener.local_addr().map_err(|error| {
237 ScaleReconcileError::Lifecycle(format!(
238 "failed to inspect service endpoint for {} slot {}: {error}",
239 target.service, target.slot
240 ))
241 })?;
242
243 self.connector
246 .connect_port(
247 &target.execution_id,
248 target.generation,
249 target.guest_port,
250 ENDPOINT_CONNECT_TIMEOUT,
251 )
252 .await
253 .map_err(|error| {
254 ScaleReconcileError::Lifecycle(format!(
255 "service endpoint for {} slot {} is not ready: {error}",
256 target.service, target.slot
257 ))
258 })?;
259
260 let endpoint = ScaleEndpoint {
261 instance_id: target.execution_id.as_str().to_string(),
262 slot: target.slot,
263 url: self.config.endpoint_url(address.port()),
264 };
265 let (shutdown, shutdown_receiver) = oneshot::channel();
266 let task = tokio::spawn(serve_endpoint(
267 listener,
268 Arc::clone(&self.connector),
269 Arc::clone(&self.connection_limit),
270 target.clone(),
271 shutdown_receiver,
272 self.config.drain_timeout,
273 ));
274 leases.insert(
275 target.execution_id.as_str().to_string(),
276 EndpointLease {
277 target: target.clone(),
278 endpoint,
279 shutdown: Some(shutdown),
280 task: Some(task),
281 },
282 );
283 }
284
285 let mut endpoints = leases
286 .values()
287 .filter(|lease| lease.target.service == service)
288 .map(|lease| lease.endpoint.clone())
289 .collect::<Vec<_>>();
290 endpoints.sort_by(|left, right| {
291 left.slot
292 .cmp(&right.slot)
293 .then_with(|| left.instance_id.cmp(&right.instance_id))
294 });
295 Ok(endpoints)
296 }
297
298 pub async fn remove(&self, execution_id: &ExecutionId) {
299 let lease = self.leases.lock().await.remove(execution_id.as_str());
300 if let Some(lease) = lease {
301 lease.drain().await;
302 }
303 }
304}
305
306async fn drain_leases(leases: Vec<EndpointLease>) {
307 let mut drains = JoinSet::new();
308 for lease in leases {
309 drains.spawn(lease.drain());
310 }
311 while let Some(result) = drains.join_next().await {
312 if let Err(error) = result {
313 tracing::warn!(%error, "Scale endpoint lease drain failed");
314 }
315 }
316}
317
318fn validate_targets(
319 service: &str,
320 targets: &[ScaleEndpointTarget],
321) -> Result<(), ScaleReconcileError> {
322 let mut executions = BTreeSet::new();
323 let mut service_slots = BTreeSet::new();
324 for target in targets {
325 if target.service != service {
326 return Err(ScaleReconcileError::Lifecycle(format!(
327 "endpoint target for service {:?} was reconciled in scope {service:?}",
328 target.service
329 )));
330 }
331 if !executions.insert(target.execution_id.as_str()) {
332 return Err(ScaleReconcileError::Lifecycle(format!(
333 "duplicate endpoint target for execution {}",
334 target.execution_id
335 )));
336 }
337 if !service_slots.insert((target.service.as_str(), target.slot)) {
338 return Err(ScaleReconcileError::Lifecycle(format!(
339 "duplicate endpoint target for service {:?} slot {}",
340 target.service, target.slot
341 )));
342 }
343 }
344 Ok(())
345}
346
347async fn serve_endpoint(
348 listener: TcpListener,
349 connector: Arc<dyn ExecutionPortConnector>,
350 connection_limit: Arc<Semaphore>,
351 target: ScaleEndpointTarget,
352 mut shutdown: oneshot::Receiver<()>,
353 drain_timeout: Duration,
354) {
355 let mut relays = JoinSet::new();
356 loop {
357 let accepted = tokio::select! {
358 biased;
359 _ = &mut shutdown => break,
360 accepted = listener.accept() => accepted,
361 };
362 let (host_stream, peer) = match accepted {
363 Ok(connection) => connection,
364 Err(error) => {
365 tracing::warn!(
366 execution_id = %target.execution_id,
367 service = target.service,
368 slot = target.slot,
369 %error,
370 "Scale endpoint accept failed"
371 );
372 tokio::time::sleep(Duration::from_millis(100)).await;
373 continue;
374 }
375 };
376 let permit = tokio::select! {
377 biased;
378 _ = &mut shutdown => break,
379 permit = Arc::clone(&connection_limit).acquire_owned() => match permit {
380 Ok(permit) => permit,
381 Err(_) => return,
382 },
383 };
384 let connector = Arc::clone(&connector);
385 let relay_target = target.clone();
386 relays.spawn(async move {
387 let _permit = permit;
388 if let Err(error) =
389 relay_connection(connector.as_ref(), &relay_target, host_stream).await
390 {
391 tracing::warn!(
392 execution_id = %relay_target.execution_id,
393 service = relay_target.service,
394 slot = relay_target.slot,
395 %peer,
396 %error,
397 "Scale endpoint relay failed"
398 );
399 }
400 });
401 while let Some(result) = relays.try_join_next() {
402 if let Err(error) = result {
403 tracing::warn!(
404 execution_id = %target.execution_id,
405 service = target.service,
406 slot = target.slot,
407 %error,
408 "Scale endpoint relay task failed"
409 );
410 }
411 }
412 }
413 drop(listener);
414
415 let relay_count = relays.len();
416 let drained = tokio::time::timeout(drain_timeout, async {
417 while let Some(result) = relays.join_next().await {
418 if let Err(error) = result {
419 tracing::warn!(
420 execution_id = %target.execution_id,
421 service = target.service,
422 slot = target.slot,
423 %error,
424 "Scale endpoint relay task failed during drain"
425 );
426 }
427 }
428 })
429 .await
430 .is_ok();
431 if drained {
432 tracing::debug!(
433 execution_id = %target.execution_id,
434 service = target.service,
435 slot = target.slot,
436 relay_count,
437 "Scale endpoint relays drained"
438 );
439 } else {
440 let forced_relays = relays.len();
441 tracing::warn!(
442 execution_id = %target.execution_id,
443 service = target.service,
444 slot = target.slot,
445 forced_relays,
446 drain_timeout_ms = drain_timeout.as_millis(),
447 "Scale endpoint drain timed out; aborting remaining relays"
448 );
449 relays.abort_all();
450 while relays.join_next().await.is_some() {}
451 }
452}
453
454async fn relay_connection(
455 connector: &dyn ExecutionPortConnector,
456 target: &ScaleEndpointTarget,
457 mut host_stream: TcpStream,
458) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
459 let mut guest_stream: ExecutionPortStream = connector
460 .connect_port(
461 &target.execution_id,
462 target.generation,
463 target.guest_port,
464 ENDPOINT_CONNECT_TIMEOUT,
465 )
466 .await?;
467 tokio::io::copy_bidirectional(&mut host_stream, &mut guest_stream)
468 .await
469 .map_err(|error| format!("failed to relay scale endpoint traffic: {error}"))?;
470 Ok(())
471}
472
473#[cfg(test)]
474mod tests {
475 use super::*;
476 use a3s_box_core::ExecutionManagerResult;
477 use async_trait::async_trait;
478 use tokio::io::{AsyncReadExt, AsyncWriteExt};
479
480 struct EchoConnector;
481
482 #[async_trait]
483 impl ExecutionPortConnector for EchoConnector {
484 async fn connect_port(
485 &self,
486 _execution_id: &ExecutionId,
487 _generation: ExecutionGeneration,
488 _port: NonZeroU16,
489 _timeout: Duration,
490 ) -> ExecutionManagerResult<ExecutionPortStream> {
491 let (client, mut server) = tokio::io::duplex(1_024);
492 tokio::spawn(async move {
493 let mut buffer = [0_u8; 1_024];
494 loop {
495 let count = match server.read(&mut buffer).await {
496 Ok(0) | Err(_) => return,
497 Ok(count) => count,
498 };
499 if server.write_all(&buffer[..count]).await.is_err() {
500 return;
501 }
502 }
503 });
504 Ok(Box::pin(client))
505 }
506 }
507
508 fn target() -> ScaleEndpointTarget {
509 ScaleEndpointTarget {
510 execution_id: ExecutionId::new("scale-api-0").unwrap(),
511 generation: ExecutionGeneration::INITIAL,
512 service: "api".to_string(),
513 slot: 0,
514 guest_port: NonZeroU16::new(8080).unwrap(),
515 }
516 }
517
518 #[tokio::test]
519 async fn endpoint_is_stable_and_relays_to_the_exact_generation() {
520 let owner =
521 ScaleEndpointOwner::new(ScaleEndpointConfig::loopback(), Arc::new(EchoConnector));
522 let first = owner.reconcile_service("api", &[target()]).await.unwrap();
523 let replay = owner.reconcile_service("api", &[target()]).await.unwrap();
524 assert_eq!(replay, first);
525 assert_eq!(first[0].slot, 0);
526
527 let worker = ScaleEndpointTarget {
528 execution_id: ExecutionId::new("scale-worker-0").unwrap(),
529 service: "worker".to_string(),
530 ..target()
531 };
532 assert_eq!(
533 owner
534 .reconcile_service("worker", &[worker])
535 .await
536 .unwrap()
537 .len(),
538 1
539 );
540 assert_eq!(
541 owner.reconcile_service("api", &[target()]).await.unwrap(),
542 first
543 );
544
545 let address = first[0]
546 .url
547 .strip_prefix("http://")
548 .unwrap()
549 .parse::<SocketAddr>()
550 .unwrap();
551 let mut stream = TcpStream::connect(address).await.unwrap();
552 stream.write_all(b"ready").await.unwrap();
553 let mut reply = [0_u8; 5];
554 stream.read_exact(&mut reply).await.unwrap();
555 assert_eq!(&reply, b"ready");
556 drop(stream);
557
558 assert!(owner
559 .reconcile_service("api", &[])
560 .await
561 .unwrap()
562 .is_empty());
563 }
564
565 #[tokio::test]
566 async fn endpoint_removal_stops_accepting_before_existing_relays_drain() {
567 let owner = Arc::new(ScaleEndpointOwner::new(
568 ScaleEndpointConfig::loopback().with_drain_timeout(Duration::from_secs(5)),
569 Arc::new(EchoConnector),
570 ));
571 let endpoint = owner
572 .reconcile_service("api", &[target()])
573 .await
574 .unwrap()
575 .remove(0);
576 let address = endpoint
577 .url
578 .strip_prefix("http://")
579 .unwrap()
580 .parse::<SocketAddr>()
581 .unwrap();
582 let mut stream = TcpStream::connect(address).await.unwrap();
583 stream.write_all(b"before").await.unwrap();
584 let mut reply = [0_u8; 6];
585 stream.read_exact(&mut reply).await.unwrap();
586 assert_eq!(&reply, b"before");
587
588 let removal_owner = Arc::clone(&owner);
589 let execution_id = target().execution_id;
590 let removal = tokio::spawn(async move {
591 removal_owner.remove(&execution_id).await;
592 });
593 tokio::time::sleep(Duration::from_millis(20)).await;
594 assert!(!removal.is_finished());
595 if let Ok(Ok(mut rejected)) =
596 tokio::time::timeout(Duration::from_millis(100), TcpStream::connect(address)).await
597 {
598 let _ = rejected.write_all(b"new").await;
599 let mut rejected_reply = [0_u8; 3];
600 assert!(
601 tokio::time::timeout(
602 Duration::from_millis(50),
603 rejected.read_exact(&mut rejected_reply),
604 )
605 .await
606 .map_or(true, |result| result.is_err()),
607 "retiring endpoint relayed a newly opened connection"
608 );
609 }
610 stream.write_all(b"during").await.unwrap();
611 stream.read_exact(&mut reply).await.unwrap();
612 assert_eq!(&reply, b"during");
613 drop(stream);
614
615 tokio::time::timeout(Duration::from_secs(1), removal)
616 .await
617 .expect("endpoint relay did not drain")
618 .unwrap();
619 }
620
621 #[tokio::test]
622 async fn endpoint_removal_force_closes_relays_after_the_bounded_deadline() {
623 let owner = ScaleEndpointOwner::new(
624 ScaleEndpointConfig::loopback().with_drain_timeout(Duration::from_millis(20)),
625 Arc::new(EchoConnector),
626 );
627 let endpoint = owner
628 .reconcile_service("api", &[target()])
629 .await
630 .unwrap()
631 .remove(0);
632 let address = endpoint
633 .url
634 .strip_prefix("http://")
635 .unwrap()
636 .parse::<SocketAddr>()
637 .unwrap();
638 let mut stream = TcpStream::connect(address).await.unwrap();
639 stream.write_all(b"ready").await.unwrap();
640 let mut reply = [0_u8; 5];
641 stream.read_exact(&mut reply).await.unwrap();
642
643 tokio::time::timeout(Duration::from_secs(1), owner.remove(&target().execution_id))
644 .await
645 .expect("endpoint drain exceeded its bounded deadline");
646
647 let mut byte = [0_u8; 1];
648 let read = tokio::time::timeout(Duration::from_secs(1), stream.read(&mut byte))
649 .await
650 .expect("forced endpoint relay remained open");
651 assert!(matches!(read, Ok(0) | Err(_)));
652 }
653
654 #[test]
655 fn endpoint_config_validates_and_formats_advertised_hosts() {
656 let ipv6 = ScaleEndpointConfig::new(IpAddr::V6("::".parse().unwrap()), "::1").unwrap();
657 assert_eq!(ipv6.endpoint_url(8080), "http://[::1]:8080");
658 assert!(ScaleEndpointConfig::new(IpAddr::V4(Ipv4Addr::LOCALHOST), "bad host").is_err());
659 assert!(ScaleEndpointConfig::new(
660 IpAddr::V4(Ipv4Addr::UNSPECIFIED),
661 Ipv4Addr::UNSPECIFIED.to_string(),
662 )
663 .is_err());
664 assert!(ScaleEndpointConfig::new(IpAddr::V4(Ipv4Addr::UNSPECIFIED), "::1").is_err());
665 }
666}