1use crate::component::{
5 self, Component, ComponentBuilder, Endpoint, EndpointDiscoverySource, Instance, Namespace,
6};
7use crate::config::environment_names::tcp_response_stream;
8use crate::pipeline::PipelineError;
9use crate::pipeline::network::ResponsePlaneMode;
10use crate::pipeline::network::manager::NetworkManager;
11use crate::service::{ServiceClient, ServiceSet};
12use crate::storage::kv;
13use crate::{discovery, system_status_server, transports};
14use crate::{
15 discovery::{Discovery, DiscoverySpec, EndpointRegistrationLease, EndpointRegistrationManager},
16 metrics::PrometheusUpdateCallback,
17 metrics::{MetricsHierarchy, MetricsRegistry},
18 transports::{etcd, nats, tcp},
19};
20
21use super::utils::GracefulShutdownTracker;
22use crate::SystemHealth;
23use crate::routing_policy::RoutingOccupancyState;
24use crate::runtime::Runtime;
25
26use async_once_cell::OnceCell;
28
29use std::fmt;
30use std::sync::{Arc, OnceLock, Weak};
31use std::time::Duration;
32use tokio::sync::watch::Receiver;
33
34use anyhow::Result;
35use derive_getters::Dissolve;
36use figment::error;
37use std::collections::HashMap;
38use tokio::sync::Mutex;
39use tokio_util::sync::CancellationToken;
40
41type EndpointDiscoverySourceMap = HashMap<Endpoint, Weak<EndpointDiscoverySource>>;
42type RoutingOccupancyMap = HashMap<Endpoint, Weak<RoutingOccupancyState>>;
43
44fn parse_tcp_response_stream_port(value: Option<&str>) -> Result<u16, PipelineError> {
45 let Some(port) = value.map(str::trim).filter(|value| !value.is_empty()) else {
46 return Ok(0);
47 };
48
49 port.parse::<u16>().map_err(|_| {
50 PipelineError::Generic(format!(
51 "invalid {}: '{}' is not a valid port number",
52 tcp_response_stream::DYN_TCP_RESPONSE_STREAM_PORT,
53 port
54 ))
55 })
56}
57
58#[cfg(test)]
59mod unit_tests {
60 use super::parse_tcp_response_stream_port;
61 use crate::pipeline::PipelineError;
62
63 #[tokio::test]
64 async fn completed_runtime_initialization_rejects_shutdown_before_http_bind() {
65 if crate::test_utils::run_isolated(
66 concat!(
67 module_path!(),
68 "::completed_runtime_initialization_rejects_shutdown_before_http_bind"
69 ),
70 &[],
71 ) {
72 return;
73 }
74
75 use super::{DistributedConfig, DistributedRuntime, Runtime};
76 use crate::system_status_server::SystemProbePolicy;
77
78 temp_env::async_with_vars(
79 [
80 ("DYN_SYSTEM_HOST", Some("127.0.0.1")),
81 ("DYN_SYSTEM_PORT", Some("0")),
82 ],
83 async {
84 let runtime = Runtime::from_current().unwrap();
85 runtime.mark_shutting_down();
86 let result = DistributedRuntime::build(
89 runtime.clone(),
90 DistributedConfig::process_local(),
91 SystemProbePolicy::RuntimeOnly,
92 )
93 .await;
94 assert!(
95 result
96 .unwrap_err()
97 .to_string()
98 .contains("runtime shut down during initialization")
99 );
100 runtime.shutdown();
101 },
102 )
103 .await;
104 }
105
106 #[test]
107 fn response_stream_port_trims_and_treats_empty_as_unset() {
108 for value in [None, Some(""), Some(" \t ")] {
109 assert_eq!(parse_tcp_response_stream_port(value).unwrap(), 0);
110 }
111 assert_eq!(
112 parse_tcp_response_stream_port(Some(" 8080 ")).unwrap(),
113 8080
114 );
115 }
116
117 #[test]
118 fn response_stream_port_rejects_invalid_values() {
119 let error = parse_tcp_response_stream_port(Some(" 65536 ")).unwrap_err();
120 assert!(matches!(
121 error,
122 PipelineError::Generic(message)
123 if message
124 == "invalid DYN_TCP_RESPONSE_STREAM_PORT: '65536' is not a valid port number"
125 ));
126 }
127}
128
129#[derive(Clone)]
139pub struct DistributedRuntime {
140 runtime: Runtime,
142
143 nats_client: Option<transports::nats::Client>,
144 network_manager: Arc<NetworkManager>,
145 tcp_server: Arc<OnceCell<Arc<transports::tcp::server::TcpStreamServer>>>,
146 quic_response_server:
147 Arc<OnceCell<Arc<crate::pipeline::network::quic_response::QuicResponseServer>>>,
148 system_status_server: Arc<OnceLock<Arc<system_status_server::SystemStatusServerInfo>>>,
149 request_plane: RequestPlaneMode,
150 response_plane: ResponsePlaneMode,
151
152 discovery_client: Arc<dyn discovery::Discovery>,
154 endpoint_registrations: Arc<EndpointRegistrationManager>,
155
156 discovery_metadata: Option<Arc<tokio::sync::RwLock<discovery::DiscoveryMetadata>>>,
159
160 component_registry: component::Registry,
166
167 endpoint_discovery_sources: Arc<tokio::sync::Mutex<EndpointDiscoverySourceMap>>,
168 routing_occupancy_states: Arc<tokio::sync::Mutex<RoutingOccupancyMap>>,
169
170 system_health: Arc<parking_lot::Mutex<SystemHealth>>,
172
173 local_endpoint_registry: crate::local_endpoint_registry::LocalEndpointRegistry,
175
176 metrics_registry: MetricsRegistry,
178
179 engine_routes: crate::engine_routes::EngineRouteRegistry,
181
182 metadata_artifacts: crate::metadata_registry::MetadataArtifactRegistry,
184
185 event_transport_kind: crate::discovery::EventTransportKind,
188}
189
190impl MetricsHierarchy for DistributedRuntime {
191 fn basename(&self) -> String {
192 "".to_string() }
194
195 fn parent_hierarchies(&self) -> Vec<&dyn MetricsHierarchy> {
196 vec![] }
198
199 fn get_metrics_registry(&self) -> &MetricsRegistry {
200 &self.metrics_registry
201 }
202
203 fn connection_id(&self) -> Option<u64> {
204 Some(self.discovery_client.instance_id())
205 }
206}
207
208impl std::fmt::Debug for DistributedRuntime {
209 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
210 write!(f, "DistributedRuntime")
211 }
212}
213
214impl DistributedRuntime {
215 pub async fn new(runtime: Runtime, config: DistributedConfig) -> Result<Self> {
216 Self::new_with_probe_policy(
217 runtime,
218 config,
219 system_status_server::SystemProbePolicy::Worker,
220 )
221 .await
222 }
223
224 pub async fn new_with_probe_policy(
227 runtime: Runtime,
228 config: DistributedConfig,
229 policy: system_status_server::SystemProbePolicy,
230 ) -> Result<Self> {
231 let shutdown = runtime.shutdown_started_token();
232 tokio::select! {
233 biased;
234 _ = shutdown.cancelled() => anyhow::bail!("runtime shut down during initialization"),
235 result = Self::build(runtime, config, policy) => result,
236 }
237 }
238
239 async fn build(
240 runtime: Runtime,
241 config: DistributedConfig,
242 policy: system_status_server::SystemProbePolicy,
243 ) -> Result<Self> {
244 let (discovery_backend, nats_config, request_plane, response_plane, event_transport_kind) =
245 config.dissolve();
246 let response_plane = match response_plane {
247 Some(mode) => mode,
248 None => ResponsePlaneMode::configured()?,
249 };
250
251 let config = crate::config::RuntimeConfig::from_settings()?;
252 let nats_client = match nats_config {
253 Some(nc) => Some(
254 nc.connect_with_cancellation(runtime.primary_token())
255 .await?,
256 ),
257 None => None,
258 };
259
260 let (discovery_client, discovery_metadata) = match discovery_backend {
262 DiscoveryBackend::Kubernetes => {
263 tracing::info!("Initializing Kubernetes discovery backend");
264 let metadata = Arc::new(tokio::sync::RwLock::new(
265 crate::discovery::DiscoveryMetadata::new(),
266 ));
267 let client = crate::discovery::KubeDiscoveryClient::new(
268 metadata.clone(),
269 runtime.primary_token(),
270 )
271 .await
272 .inspect_err(
273 |err| tracing::error!(%err, "Failed to initialize Kubernetes discovery client"),
274 )?;
275 (Arc::new(client) as Arc<dyn Discovery>, Some(metadata))
276 }
277 DiscoveryBackend::KvStore(kv_selector) => {
278 tracing::info!("Initializing KV store discovery backend: {kv_selector}");
279 let runtime_clone = runtime.clone();
280 let store = match kv_selector {
281 kv::Selector::Etcd(etcd_config) => {
282 let etcd_client = etcd::Client::new(*etcd_config, runtime_clone).await.inspect_err(|err|
283 tracing::error!(%err, "Could not connect to etcd. Pass `--discovery-backend ..` to use a different backend or start etcd."))?;
284 kv::Manager::etcd(etcd_client)
285 }
286 kv::Selector::File(root) => kv::Manager::file(runtime.primary_token(), root),
287 kv::Selector::Memory => kv::Manager::memory(),
288 };
289 use crate::discovery::KVStoreDiscovery;
290 (
291 Arc::new(KVStoreDiscovery::new(store, runtime.primary_token()))
292 as Arc<dyn Discovery>,
293 None,
294 )
295 }
296 };
297
298 let starting_health_status = config.starting_health_status.clone();
299 let use_endpoint_health_status = config.use_endpoint_health_status.clone();
300 let health_endpoint_path = config.system_health_path.clone();
301 let live_endpoint_path = config.system_live_path.clone();
302 let system_health = Arc::new(parking_lot::Mutex::new(SystemHealth::new(
303 starting_health_status,
304 use_endpoint_health_status,
305 config.health_check_enabled,
306 health_endpoint_path,
307 live_endpoint_path,
308 )));
309
310 let component_registry = component::Registry::new();
311
312 let network_manager = NetworkManager::new(
314 runtime.child_token(),
315 nats_client.clone().map(|c| c.client().clone()),
316 component_registry.clone(),
317 request_plane,
318 );
319
320 let endpoint_registrations = EndpointRegistrationManager::new(
321 discovery_client.clone(),
322 runtime.secondary(),
323 runtime.primary_token(),
324 );
325 let distributed_runtime = Self {
326 runtime,
327 network_manager: Arc::new(network_manager),
328 nats_client,
329 tcp_server: Arc::new(OnceCell::new()),
330 quic_response_server: Arc::new(OnceCell::new()),
331 system_status_server: Arc::new(OnceLock::new()),
332 discovery_client,
333 endpoint_registrations,
334 discovery_metadata,
335 component_registry,
336 endpoint_discovery_sources: Arc::new(Mutex::new(HashMap::new())),
337 routing_occupancy_states: Arc::new(Mutex::new(HashMap::new())),
338 metrics_registry: crate::MetricsRegistry::new(),
339 system_health,
340 request_plane,
341 response_plane,
342 local_endpoint_registry: crate::local_endpoint_registry::LocalEndpointRegistry::new(),
343 engine_routes: crate::engine_routes::EngineRouteRegistry::new(),
344 metadata_artifacts: crate::metadata_registry::MetadataArtifactRegistry::new(),
345 event_transport_kind,
346 };
347
348 if response_plane == ResponsePlaneMode::Quic {
349 crate::metrics::quic_response::ensure_registered(
350 distributed_runtime.get_metrics_registry(),
351 );
352 }
353
354 distributed_runtime
356 .system_health
357 .lock()
358 .initialize_uptime_gauge(&distributed_runtime)?;
359
360 {
363 let system_health = distributed_runtime.system_health.clone();
364 distributed_runtime
365 .metrics_registry
366 .add_update_callback(std::sync::Arc::new(move || {
367 system_health.lock().update_uptime_gauge();
368 Ok(())
369 }));
370 }
371
372 match crate::metrics::otlp_export::ExportConfig::from_env() {
381 Ok(Some(export_config)) => {
382 tracing::info!(
383 endpoint = %export_config.endpoint,
384 interval_ms = export_config.interval.as_millis(),
385 "exporting metrics over OTLP"
386 );
387 let shutdown_guard = distributed_runtime
394 .runtime
395 .graceful_shutdown_tracker()
396 .register_task();
397 let registry = distributed_runtime.metrics_registry.clone();
398 let cancel = distributed_runtime.runtime.child_token();
399 tokio::spawn(async move {
400 crate::metrics::otlp_export::run(registry, export_config, cancel).await;
401 drop(shutdown_guard);
402 });
403 }
404 Ok(None) => {}
405 Err(error) => {
406 tracing::error!(%error, "OTLP metrics export is misconfigured; not exporting");
407 }
408 }
409
410 if config.health_check_enabled {
412 let health_check_config = crate::health_check::HealthCheckConfig {
413 canary_wait_time: std::time::Duration::from_secs(config.canary_wait_time_secs),
414 request_timeout: std::time::Duration::from_secs(
415 config.health_check_request_timeout_secs,
416 ),
417 };
418
419 match crate::health_check::start_health_check_manager(
421 distributed_runtime.clone(),
422 Some(health_check_config),
423 )
424 .await
425 {
426 Ok(()) => tracing::info!(
427 "Health check manager started (canary_wait_time: {}s, request_timeout: {}s)",
428 config.canary_wait_time_secs,
429 config.health_check_request_timeout_secs
430 ),
431 Err(e) => tracing::error!("Health check manager failed to start: {e}"),
432 }
433 }
434
435 anyhow::ensure!(
436 !distributed_runtime.runtime.is_shutting_down(),
437 "runtime shut down during initialization"
438 );
439 if config.system_server_enabled() {
440 let stop = match policy {
443 system_status_server::SystemProbePolicy::Worker => {
444 distributed_runtime.runtime.child_token()
445 }
446 system_status_server::SystemProbePolicy::RuntimeOnly => {
447 distributed_runtime.runtime.primary_token().child_token()
448 }
449 };
450 match system_status_server::spawn_system_status_server(
451 &config.system_host,
452 config.system_port as u16,
453 stop,
454 Arc::new(distributed_runtime.clone()),
455 distributed_runtime.discovery_metadata.clone(),
456 policy,
457 )
458 .await
459 {
460 Ok((address, handle)) => {
461 distributed_runtime
462 .system_status_server
463 .set(Arc::new(system_status_server::SystemStatusServerInfo::new(
464 address,
465 Some(handle),
466 )))
467 .expect("System status server info should only be set once");
468 tracing::info!(%address, ?policy, "System HTTP listener started");
469 }
470 Err(error) if policy == system_status_server::SystemProbePolicy::Worker => {
472 tracing::error!(%error, "System status server startup failed");
473 }
474 Err(error) => return Err(error),
475 }
476 }
477 Ok(distributed_runtime)
478 }
479
480 pub async fn from_settings(runtime: Runtime) -> Result<Self> {
481 let config = DistributedConfig::try_from_settings()?;
482 Self::new(runtime, config).await
483 }
484
485 pub(crate) async fn check_dependencies(&self) -> Result<()> {
488 anyhow::ensure!(!self.runtime.is_shutting_down(), "runtime is shutting down");
489 if let Some(client) = &self.nats_client {
490 anyhow::ensure!(
491 client.client().connection_state() == async_nats::connection::State::Connected,
492 "NATS is disconnected"
493 );
494 }
495 self.discovery_client.check_connection().await?;
496 anyhow::ensure!(!self.runtime.is_shutting_down(), "runtime is shutting down");
497 Ok(())
498 }
499
500 pub fn runtime(&self) -> &Runtime {
501 &self.runtime
502 }
503
504 pub fn primary_token(&self) -> CancellationToken {
505 self.runtime.primary_token()
506 }
507
508 pub fn component_registry(&self) -> &component::Registry {
511 &self.component_registry
512 }
513
514 pub fn system_health(&self) -> Arc<parking_lot::Mutex<SystemHealth>> {
516 self.system_health.clone()
517 }
518
519 pub fn local_endpoint_registry(
521 &self,
522 ) -> &crate::local_endpoint_registry::LocalEndpointRegistry {
523 &self.local_endpoint_registry
524 }
525
526 pub fn engine_routes(&self) -> &crate::engine_routes::EngineRouteRegistry {
528 &self.engine_routes
529 }
530
531 pub fn metadata_artifacts(&self) -> &crate::metadata_registry::MetadataArtifactRegistry {
532 &self.metadata_artifacts
533 }
534
535 pub fn connection_id(&self) -> u64 {
540 self.discovery_client.instance_id()
541 }
542
543 pub fn shutdown(&self) {
544 self.runtime.shutdown();
545 self.discovery_client.shutdown();
546 }
547
548 pub fn namespace(&self, name: impl Into<String>) -> Result<Namespace> {
550 Namespace::new(self.clone(), name.into())
551 }
552
553 pub fn discovery(&self) -> Arc<dyn Discovery> {
555 self.discovery_client.clone()
556 }
557
558 pub async fn register_endpoint_lease(
560 &self,
561 spec: DiscoverySpec,
562 ) -> Result<EndpointRegistrationLease> {
563 self.endpoint_registrations.register(spec).await
564 }
565
566 pub async fn tcp_server(&self) -> Result<Arc<tcp::server::TcpStreamServer>> {
567 Ok(self
568 .tcp_server
569 .get_or_try_init(async move {
570 let port_value =
571 std::env::var(tcp_response_stream::DYN_TCP_RESPONSE_STREAM_PORT).ok();
572 let port = parse_tcp_response_stream_port(port_value.as_deref())?;
573 let host = crate::utils::ip_resolver::host_override_from_env(
574 tcp_response_stream::DYN_TCP_RESPONSE_STREAM_HOST,
575 )
576 .map_err(|error| PipelineError::Generic(error.to_string()))?;
577
578 let host_suffix = host
579 .as_ref()
580 .map_or(String::new(), |h| format!(" on host {h}"));
581 if port == 0 {
582 tracing::info!(
583 "TCP request callback server using OS-assigned port{host_suffix}"
584 );
585 } else {
586 tracing::info!(
587 "TCP request callback server using fixed port {port}{host_suffix}"
588 );
589 }
590
591 let options = tcp::server::ServerOptions {
592 port,
593 interface: host,
594 };
595 let server = tcp::server::TcpStreamServer::new(options).await?;
596 Ok::<_, PipelineError>(server)
597 })
598 .await?
599 .clone())
600 }
601
602 pub async fn quic_response_server(
603 &self,
604 ) -> Result<Arc<crate::pipeline::network::quic_response::QuicResponseServer>> {
605 anyhow::ensure!(
606 self.response_plane == ResponsePlaneMode::Quic,
607 "QUIC response server requested while response plane is {}",
608 self.response_plane.name()
609 );
610 Ok(self
611 .quic_response_server
612 .get_or_try_init(async {
613 let tcp_server = self.tcp_server().await?;
614 let tcp_address = tcp_server.local_address()?;
615 let address = std::net::SocketAddr::new(tcp_address.ip(), 0);
620 crate::pipeline::network::quic_response::QuicResponseServer::new(
621 address,
622 address,
623 self.runtime.child_token(),
624 )
625 .map_err(anyhow::Error::from)
626 })
627 .await?
628 .clone())
629 }
630
631 pub fn quic_response_client_pool(
632 &self,
633 ) -> Result<Arc<crate::pipeline::network::quic_response::QuicResponseClientPool>> {
634 anyhow::ensure!(
635 self.response_plane == ResponsePlaneMode::Quic,
636 "QUIC response client pool requested while response plane is {}",
637 self.response_plane.name()
638 );
639 crate::pipeline::network::quic_response::process_client_pool_from_env()
640 .map_err(anyhow::Error::from)
641 }
642
643 pub fn response_plane(&self) -> ResponsePlaneMode {
644 self.response_plane
645 }
646
647 pub fn network_manager(&self) -> Arc<NetworkManager> {
652 self.network_manager.clone()
653 }
654
655 pub async fn request_plane_server(
659 &self,
660 ) -> Result<Arc<dyn crate::pipeline::network::ingress::unified_server::RequestPlaneServer>>
661 {
662 self.network_manager().server().await
663 }
664
665 pub fn system_status_server_info(
667 &self,
668 ) -> Option<Arc<crate::system_status_server::SystemStatusServerInfo>> {
669 self.system_status_server.get().cloned()
670 }
671
672 pub fn request_plane(&self) -> RequestPlaneMode {
674 self.request_plane
675 }
676
677 pub fn default_event_transport_kind(&self) -> crate::discovery::EventTransportKind {
685 self.event_transport_kind
686 }
687
688 pub fn child_token(&self) -> CancellationToken {
689 self.runtime.child_token()
690 }
691
692 pub(crate) fn graceful_shutdown_tracker(&self) -> Arc<GracefulShutdownTracker> {
693 self.runtime.graceful_shutdown_tracker()
694 }
695
696 pub(crate) fn endpoint_discovery_sources(&self) -> Arc<Mutex<EndpointDiscoverySourceMap>> {
697 self.endpoint_discovery_sources.clone()
698 }
699
700 pub fn register_graceful_task(&self) -> crate::utils::GracefulTaskGuard {
706 self.runtime.graceful_shutdown_tracker().register_task()
707 }
708
709 pub(crate) fn routing_occupancy_states(&self) -> Arc<Mutex<RoutingOccupancyMap>> {
710 self.routing_occupancy_states.clone()
711 }
712
713 pub async fn kv_router_nats_publish(
719 &self,
720 subject: String,
721 payload: bytes::Bytes,
722 ) -> anyhow::Result<()> {
723 self.kv_router_nats_publish_subject(subject.into(), payload)
724 .await
725 }
726
727 pub(crate) async fn kv_router_nats_publish_subject(
728 &self,
729 subject: async_nats::Subject,
730 payload: bytes::Bytes,
731 ) -> anyhow::Result<()> {
732 let Some(nats_client) = self.nats_client.as_ref() else {
733 tracing::trace!("Skipping NATS publish (NATS not configured): {subject}");
735 return Ok(());
736 };
737 Ok(nats_client.client().publish(subject, payload).await?)
738 }
739
740 pub(crate) async fn kv_router_nats_subscribe(
743 &self,
744 subject: String,
745 ) -> Result<async_nats::Subscriber> {
746 let Some(nats_client) = self.nats_client.as_ref() else {
747 anyhow::bail!("KV router's EventSubscriber requires NATS");
748 };
749 Ok(nats_client.client().subscribe(subject).await?)
750 }
751
752 pub async fn kv_router_nats_request(
756 &self,
757 subject: String,
758 payload: bytes::Bytes,
759 timeout: std::time::Duration,
760 ) -> anyhow::Result<async_nats::Message> {
761 let Some(nats_client) = self.nats_client.as_ref() else {
762 anyhow::bail!("KV router's request requires NATS");
763 };
764 let response =
765 tokio::time::timeout(timeout, nats_client.client().request(subject, payload))
766 .await
767 .map_err(|_| anyhow::anyhow!("Request timed out after {:?}", timeout))??;
768 Ok(response)
769 }
770
771 pub fn register_nats_service(
778 &self,
779 component: Component,
780 ) -> tokio::sync::mpsc::Receiver<Result<(), String>> {
781 let (tx, rx) = tokio::sync::mpsc::channel::<Result<(), String>>(1);
783
784 let drt = self.clone();
785 self.runtime().secondary().spawn(async move {
786 let service_name = component.service_name();
787
788 if drt
790 .component_registry()
791 .inner
792 .lock()
793 .await
794 .services
795 .contains_key(&service_name)
796 {
797 tracing::trace!("Service {service_name} already exists");
800 let _ = tx.send(Ok(())).await;
802 return;
803 }
804
805 let Some(nats_client) = drt.nats_client.as_ref() else {
806 tracing::error!("Cannot create NATS service without NATS.");
807 let _ = tx
808 .send(Err("Cannot create NATS service without NATS".to_string()))
809 .await;
810 return;
811 };
812 let description = None;
813 let nats_service = match crate::component::service::build_nats_service(
814 nats_client,
815 &component,
816 description,
817 )
818 .await
819 {
820 Ok(service) => service,
821 Err(err) => {
822 tracing::error!(error = %err, component = service_name, "Failed to build NATS service");
823 let _ = tx.send(Err(format!("Failed to build NATS service: {err}"))).await;
824 return;
825 }
826 };
827
828 let mut guard = drt.component_registry().inner.lock().await;
829 if !guard.services.contains_key(&service_name) {
830 guard.services.insert(service_name.clone(), nats_service);
832
833 tracing::info!("Added NATS service {service_name}");
834
835 drop(guard);
836 } else {
837 drop(guard);
838 let _ = nats_service.stop().await;
839 }
843
844 let _ = tx.send(Ok(())).await;
846 });
847
848 rx
849 }
850}
851
852#[cfg(test)]
853mod parser_env_tests {
854 use super::*;
855
856 #[tokio::test]
857 async fn distributed_runtime_rejects_invalid_parser_version() {
858 if crate::test_utils::run_isolated(
859 concat!(
860 module_path!(),
861 "::distributed_runtime_rejects_invalid_parser_version"
862 ),
863 &[(
864 crate::config::environment_names::llm::DYN_PARSER_VERSION,
865 "v3",
866 )],
867 ) {
868 return;
869 }
870
871 let runtime = Runtime::from_current().unwrap();
872 let error = DistributedRuntime::new(runtime, DistributedConfig::process_local())
873 .await
874 .expect_err("invalid parser configuration must fail runtime construction");
875 assert!(error.to_string().contains("DYN_PARSER_VERSION"));
876 }
877}
878
879#[derive(Clone, Debug)]
881pub enum DiscoveryBackend {
882 Kubernetes,
884 KvStore(kv::Selector),
886}
887
888impl DiscoveryBackend {
889 pub fn is_local(&self) -> bool {
893 matches!(
894 self,
895 DiscoveryBackend::KvStore(kv::Selector::File(_))
896 | DiscoveryBackend::KvStore(kv::Selector::Memory)
897 )
898 }
899
900 pub fn resolve_event_transport_kind(&self) -> Result<crate::discovery::EventTransportKind> {
913 use crate::config::environment_names::event_plane::DYN_EVENT_PLANE;
914 Self::event_transport_kind_from_env_result(std::env::var(DYN_EVENT_PLANE))
915 }
916
917 fn event_transport_kind_from_env_result(
918 value: std::result::Result<String, std::env::VarError>,
919 ) -> Result<crate::discovery::EventTransportKind> {
920 use crate::discovery::EventTransportKind;
921 match value {
922 Err(std::env::VarError::NotPresent) => Ok(EventTransportKind::Zmq),
924 Ok(raw) => match raw.trim().to_lowercase().as_str() {
925 "nats" => Ok(EventTransportKind::Nats),
926 "zmq" | "" => Ok(EventTransportKind::Zmq),
927 _ => Err(anyhow::anyhow!(
928 "Invalid DYN_EVENT_PLANE value '{raw}'. Valid values: 'nats', 'zmq'"
929 )),
930 },
931 Err(std::env::VarError::NotUnicode(raw)) => Err(anyhow::anyhow!(
932 "Invalid DYN_EVENT_PLANE value '{}': not valid Unicode. \
933 Valid values: 'nats', 'zmq'",
934 raw.to_string_lossy()
935 )),
936 }
937 }
938}
939
940#[derive(Dissolve)]
941pub struct DistributedConfig {
942 pub discovery_backend: DiscoveryBackend,
943 pub nats_config: Option<nats::ClientOptions>,
944 pub request_plane: RequestPlaneMode,
945 pub response_plane: Option<ResponsePlaneMode>,
948 pub event_transport_kind: crate::discovery::EventTransportKind,
953}
954
955impl DistributedConfig {
956 pub fn from_settings() -> DistributedConfig {
961 Self::try_from_settings().unwrap_or_else(|err| panic!("{err}"))
962 }
963
964 pub fn try_from_settings() -> Result<DistributedConfig> {
965 Self::from_settings_with_overrides(None, None, None)
966 }
967
968 pub fn from_settings_with_overrides(
971 discovery_backend: Option<&str>,
972 request_plane: Option<&str>,
973 event_plane: Option<&str>,
974 ) -> Result<DistributedConfig> {
975 let request_plane = match request_plane {
976 Some(value) => value.parse()?,
977 None => RequestPlaneMode::from_env()?,
978 };
979
980 let backend_str = discovery_backend
983 .map(str::to_owned)
984 .or_else(|| std::env::var("DYN_DISCOVERY_BACKEND").ok())
985 .unwrap_or_else(|| "etcd".to_string());
986
987 let discovery_backend = match backend_str.as_str() {
988 "kubernetes" => {
989 tracing::info!("Using Kubernetes discovery backend");
990 DiscoveryBackend::Kubernetes
991 }
992 other => {
993 let selector: kv::Selector = other.parse().map_err(|_| {
994 anyhow::anyhow!(
995 "Unknown DYN_DISCOVERY_BACKEND value: '{other}'. \
996 Valid options: kubernetes, etcd, file, mem"
997 )
998 })?;
999 DiscoveryBackend::KvStore(selector)
1000 }
1001 };
1002
1003 let event_transport_kind = match event_plane {
1007 Some("nats") => crate::discovery::EventTransportKind::Nats,
1008 Some("zmq" | "") => crate::discovery::EventTransportKind::Zmq,
1009 Some(other) => {
1010 anyhow::bail!("Invalid event plane '{other}'. Valid options are: 'nats', 'zmq'")
1011 }
1012 None => discovery_backend.resolve_event_transport_kind()?,
1013 };
1014
1015 let nats_enabled = request_plane.is_nats()
1024 || std::env::var(crate::config::environment_names::nats::NATS_SERVER).is_ok()
1025 || matches!(
1026 event_transport_kind,
1027 crate::discovery::EventTransportKind::Nats
1028 );
1029
1030 Ok(DistributedConfig {
1031 discovery_backend,
1032 nats_config: if nats_enabled {
1033 Some(nats::ClientOptions::default())
1034 } else {
1035 None
1036 },
1037 request_plane,
1038 response_plane: None,
1039 event_transport_kind,
1040 })
1041 }
1042
1043 pub fn for_cli() -> DistributedConfig {
1044 Self::try_for_cli().unwrap_or_else(|err| panic!("{err}"))
1045 }
1046
1047 pub fn try_for_cli() -> Result<DistributedConfig> {
1048 let etcd_config = etcd::ClientOptions {
1049 attach_lease: false,
1050 ..Default::default()
1051 };
1052 let request_plane = RequestPlaneMode::from_env()?;
1053 let discovery_backend =
1054 DiscoveryBackend::KvStore(kv::Selector::Etcd(Box::new(etcd_config)));
1055 let event_transport_kind = discovery_backend.resolve_event_transport_kind()?;
1056 let nats_enabled = request_plane.is_nats()
1057 || std::env::var(crate::config::environment_names::nats::NATS_SERVER).is_ok()
1058 || matches!(
1059 event_transport_kind,
1060 crate::discovery::EventTransportKind::Nats
1061 );
1062 Ok(DistributedConfig {
1063 discovery_backend,
1064 nats_config: if nats_enabled {
1065 Some(nats::ClientOptions::default())
1066 } else {
1067 None
1068 },
1069 request_plane,
1070 response_plane: None,
1071 event_transport_kind,
1072 })
1073 }
1074
1075 pub fn process_local() -> DistributedConfig {
1078 DistributedConfig {
1079 discovery_backend: DiscoveryBackend::KvStore(kv::Selector::Memory),
1080 nats_config: None,
1081 request_plane: RequestPlaneMode::Tcp,
1084 response_plane: None,
1085 event_transport_kind: crate::discovery::EventTransportKind::Zmq,
1086 }
1087 }
1088}
1089
1090#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
1096pub enum RequestPlaneMode {
1097 Nats,
1099 #[default]
1101 Tcp,
1102}
1103
1104impl fmt::Display for RequestPlaneMode {
1105 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1106 match self {
1107 Self::Nats => write!(f, "nats"),
1108 Self::Tcp => write!(f, "tcp"),
1109 }
1110 }
1111}
1112
1113impl std::str::FromStr for RequestPlaneMode {
1114 type Err = anyhow::Error;
1115
1116 fn from_str(s: &str) -> std::result::Result<Self, Self::Err> {
1117 match s.to_lowercase().as_str() {
1118 "nats" => Ok(Self::Nats),
1119 "tcp" => Ok(Self::Tcp),
1120 _ => Err(anyhow::anyhow!(
1121 "Invalid request plane mode: '{}'. Valid options are: 'nats', 'tcp'",
1122 s
1123 )),
1124 }
1125 }
1126}
1127
1128impl RequestPlaneMode {
1129 fn from_env() -> Result<Self> {
1130 Self::from_env_result(std::env::var(
1131 crate::config::environment_names::request_plane::DYN_REQUEST_PLANE,
1132 ))
1133 }
1134
1135 fn from_env_result(value: std::result::Result<String, std::env::VarError>) -> Result<Self> {
1136 match value {
1137 Err(std::env::VarError::NotPresent) => Ok(Self::default()),
1138 Ok(s) => s.parse(),
1139 Err(std::env::VarError::NotUnicode(raw)) => Err(anyhow::anyhow!(
1140 "Invalid request plane mode: '{}' is not valid Unicode. \
1141 Valid options are: 'nats', 'tcp'",
1142 raw.to_string_lossy()
1143 )),
1144 }
1145 }
1146
1147 pub fn is_nats(&self) -> bool {
1148 matches!(self, RequestPlaneMode::Nats)
1149 }
1150}
1151
1152#[cfg(test)]
1153mod request_plane_env_tests {
1154 use super::RequestPlaneMode;
1155
1156 #[test]
1157 fn absent_request_plane_defaults_to_tcp() {
1158 let mode = RequestPlaneMode::from_env_result(Err(std::env::VarError::NotPresent))
1159 .expect("an absent DYN_REQUEST_PLANE must not be an error");
1160 assert_eq!(mode, RequestPlaneMode::Tcp);
1161 }
1162
1163 #[test]
1164 fn empty_request_plane_is_an_error() {
1165 RequestPlaneMode::from_env_result(Ok(String::new()))
1166 .expect_err("an empty DYN_REQUEST_PLANE must not silently fall back to TCP");
1167 }
1168
1169 #[test]
1170 fn valid_request_plane_value_resolves() {
1171 let mode = RequestPlaneMode::from_env_result(Ok("nats".to_string()))
1172 .expect("DYN_REQUEST_PLANE=nats should resolve");
1173 assert_eq!(mode, RequestPlaneMode::Nats);
1174 }
1175
1176 #[test]
1177 fn invalid_request_plane_is_an_error_naming_value_and_options() {
1178 let err = RequestPlaneMode::from_env_result(Ok("nat".to_string()))
1179 .expect_err("a misspelled DYN_REQUEST_PLANE must not silently fall back to TCP");
1180 let message = err.to_string();
1181 assert!(
1182 message.contains("nat"),
1183 "error should name the offending value, got: {message}"
1184 );
1185 assert!(
1186 message.contains("'nats'") && message.contains("'tcp'"),
1187 "error should list the valid options, got: {message}"
1188 );
1189 }
1190
1191 #[cfg(unix)]
1192 #[test]
1193 fn non_unicode_request_plane_is_an_error_not_a_default() {
1194 use std::ffi::OsString;
1195 use std::os::unix::ffi::OsStringExt;
1196
1197 let raw = OsString::from_vec(b"nat\xff".to_vec());
1198 let err = RequestPlaneMode::from_env_result(Err(std::env::VarError::NotUnicode(raw)))
1199 .expect_err("a non-Unicode DYN_REQUEST_PLANE must not silently fall back to TCP");
1200 let message = err.to_string();
1201 assert!(
1202 message.contains("not valid Unicode"),
1203 "error should say the value was not valid Unicode, got: {message}"
1204 );
1205 assert!(
1206 message.contains("'nats'") && message.contains("'tcp'"),
1207 "error should list the valid options, got: {message}"
1208 );
1209 }
1210}
1211
1212#[cfg(test)]
1213mod event_plane_env_tests {
1214 use super::DiscoveryBackend;
1215 use crate::discovery::EventTransportKind;
1216
1217 fn resolve(
1218 value: std::result::Result<String, std::env::VarError>,
1219 ) -> anyhow::Result<EventTransportKind> {
1220 DiscoveryBackend::event_transport_kind_from_env_result(value)
1221 }
1222
1223 #[test]
1224 fn valid_event_plane_values_resolve_case_insensitively() {
1225 for (value, expected) in [
1226 ("nats", EventTransportKind::Nats),
1227 ("NATS", EventTransportKind::Nats),
1228 (" nats ", EventTransportKind::Nats),
1229 ("zmq", EventTransportKind::Zmq),
1230 ("Zmq", EventTransportKind::Zmq),
1231 ] {
1232 let kind = resolve(Ok(value.to_string()))
1233 .unwrap_or_else(|err| panic!("DYN_EVENT_PLANE={value:?} should resolve: {err}"));
1234 assert_eq!(kind, expected, "value {value:?}");
1235 }
1236 }
1237
1238 #[test]
1239 fn invalid_event_plane_is_an_error_naming_value_and_options() {
1240 for value in ["nat", "x"] {
1241 let err = resolve(Ok(value.to_string()))
1242 .expect_err("a misspelled DYN_EVENT_PLANE must not silently fall back to ZMQ");
1243 let message = err.to_string();
1244 assert!(
1245 message.contains(&format!("'{value}'")),
1246 "error should name the offending value, got: {message}"
1247 );
1248 assert!(
1249 message.contains("Valid values: 'nats', 'zmq'"),
1250 "error should list the valid values, got: {message}"
1251 );
1252 }
1253 }
1254
1255 #[cfg(unix)]
1256 #[test]
1257 fn non_unicode_event_plane_is_an_error_not_a_default() {
1258 use std::ffi::OsString;
1259 use std::os::unix::ffi::OsStringExt;
1260
1261 let raw = OsString::from_vec(b"nats\xff".to_vec());
1262 let err = resolve(Err(std::env::VarError::NotUnicode(raw)))
1263 .expect_err("a non-Unicode DYN_EVENT_PLANE must not silently fall back to ZMQ");
1264 assert!(
1265 err.to_string().contains("not valid Unicode"),
1266 "error should say the value was not valid Unicode, got: {err}"
1267 );
1268 }
1269}
1270
1271pub mod distributed_test_utils {
1272 #[cfg(feature = "integration")]
1278 pub async fn create_test_drt_async() -> super::DistributedRuntime {
1279 use crate::transports::nats;
1280
1281 let rt = crate::Runtime::from_current().unwrap();
1282 let config = super::DistributedConfig {
1283 discovery_backend: super::DiscoveryBackend::KvStore(
1284 crate::storage::kv::Selector::Memory,
1285 ),
1286 nats_config: Some(nats::ClientOptions::default()),
1287 request_plane: crate::distributed::RequestPlaneMode::default(),
1288 response_plane: None,
1289 event_transport_kind: crate::discovery::EventTransportKind::Nats,
1290 };
1291 super::DistributedRuntime::new(rt, config).await.unwrap()
1292 }
1293
1294 pub async fn create_test_shared_drt_async(
1301 store_path: &std::path::Path,
1302 ) -> super::DistributedRuntime {
1303 use crate::transports::nats;
1304
1305 let rt = crate::Runtime::from_current().unwrap();
1306 let config = super::DistributedConfig {
1307 discovery_backend: super::DiscoveryBackend::KvStore(
1308 crate::storage::kv::Selector::File(store_path.to_path_buf()),
1309 ),
1310 nats_config: Some(nats::ClientOptions::default()),
1311 request_plane: crate::distributed::RequestPlaneMode::default(),
1312 response_plane: None,
1313 event_transport_kind: crate::discovery::EventTransportKind::Nats,
1314 };
1315 super::DistributedRuntime::new(rt, config).await.unwrap()
1316 }
1317}
1318
1319#[cfg(all(test, feature = "integration"))]
1320mod tests {
1321 use super::RequestPlaneMode;
1322 use super::distributed_test_utils::create_test_drt_async;
1323
1324 #[tokio::test]
1325 async fn test_drt_uptime_after_delay_system_disabled() {
1326 if crate::test_utils::run_isolated(
1327 concat!(
1328 module_path!(),
1329 "::test_drt_uptime_after_delay_system_disabled"
1330 ),
1331 &[],
1332 ) {
1333 return;
1334 }
1335
1336 use crate::config::environment_names::runtime::system as env_system;
1337 temp_env::async_with_vars([(env_system::DYN_SYSTEM_PORT, None::<&str>)], async {
1339 let drt = create_test_drt_async().await;
1341
1342 tokio::time::sleep(tokio::time::Duration::from_millis(50)).await;
1344
1345 let uptime = drt.system_health.lock().uptime();
1347 assert!(
1348 uptime >= std::time::Duration::from_millis(50),
1349 "Expected uptime to be at least 50ms, but got {:?}",
1350 uptime
1351 );
1352
1353 println!(
1354 "✓ DRT uptime test passed (system disabled): uptime = {:?}",
1355 uptime
1356 );
1357 })
1358 .await;
1359 }
1360
1361 #[tokio::test]
1362 async fn test_drt_uptime_after_delay_system_enabled() {
1363 if crate::test_utils::run_isolated(
1364 concat!(
1365 module_path!(),
1366 "::test_drt_uptime_after_delay_system_enabled"
1367 ),
1368 &[],
1369 ) {
1370 return;
1371 }
1372
1373 use crate::config::environment_names::runtime::system as env_system;
1374 temp_env::async_with_vars([(env_system::DYN_SYSTEM_PORT, Some("0"))], async {
1376 let drt = create_test_drt_async().await;
1378
1379 tokio::time::sleep(tokio::time::Duration::from_millis(50)).await;
1381
1382 let uptime = drt.system_health.lock().uptime();
1384 assert!(
1385 uptime >= std::time::Duration::from_millis(50),
1386 "Expected uptime to be at least 50ms, but got {:?}",
1387 uptime
1388 );
1389
1390 println!(
1391 "✓ DRT uptime test passed (system enabled): uptime = {:?}",
1392 uptime
1393 );
1394 })
1395 .await;
1396 }
1397
1398 #[test]
1399 fn test_request_plane_mode_from_str() {
1400 assert_eq!(
1401 "nats".parse::<RequestPlaneMode>().unwrap(),
1402 RequestPlaneMode::Nats
1403 );
1404 assert_eq!(
1405 "tcp".parse::<RequestPlaneMode>().unwrap(),
1406 RequestPlaneMode::Tcp
1407 );
1408 assert_eq!(
1409 "NATS".parse::<RequestPlaneMode>().unwrap(),
1410 RequestPlaneMode::Nats
1411 );
1412 assert_eq!(
1413 "TCP".parse::<RequestPlaneMode>().unwrap(),
1414 RequestPlaneMode::Tcp
1415 );
1416 assert!("invalid".parse::<RequestPlaneMode>().is_err());
1417 }
1418}