Skip to main content

dynamo_runtime/
distributed.rs

1// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4use 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
26// Used instead of std::cell::OnceCell because get_or_try_init there is nightly
27use 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                // Exercise the final bind guard directly: the public constructor
87                // would reject shutdown before polling build at all.
88                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/// Distributed [Runtime] providing cluster-wide communication, transport, and discovery resources.
130///
131/// `DistributedRuntime` is not a process singleton. Calling [`DistributedRuntime::new`] more than
132/// once creates independent DRT instances with distinct discovery connection IDs, even when they
133/// share a process. Cloning a DRT continues to share the original instance and connection ID.
134///
135/// Production services should normally treat one DRT per service replica/process as a soft
136/// invariant. Multiple DRTs in one process are primarily supported for single-process test
137/// topologies and for the mocker, which models multiple isolated workers in one process.
138#[derive(Clone)]
139pub struct DistributedRuntime {
140    // local runtime
141    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    // Service discovery client
153    discovery_client: Arc<dyn discovery::Discovery>,
154    endpoint_registrations: Arc<EndpointRegistrationManager>,
155
156    // Discovery metadata (only used for Kubernetes backend)
157    // Shared with system status server to expose via /metadata endpoint
158    discovery_metadata: Option<Arc<tokio::sync::RwLock<discovery::DiscoveryMetadata>>>,
159
160    // local registry for components
161    // the registry allows us to use share runtime resources across instances of the same component object.
162    // take for example two instances of a client to the same remote component. The registry allows us to use
163    // a single endpoint watcher for both clients, this keeps the number background tasking watching specific
164    // paths in etcd to a minimum.
165    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    // Health Status
171    system_health: Arc<parking_lot::Mutex<SystemHealth>>,
172
173    // Local endpoint registry for in-process calls
174    local_endpoint_registry: crate::local_endpoint_registry::LocalEndpointRegistry,
175
176    // This hierarchy's own metrics registry
177    metrics_registry: MetricsRegistry,
178
179    // Registry for /engine/* route callbacks
180    engine_routes: crate::engine_routes::EngineRouteRegistry,
181
182    // Backs `/v1/metadata/{model_slug}/{model_suffix}/{filename}`.
183    metadata_artifacts: crate::metadata_registry::MetadataArtifactRegistry,
184
185    // Resolved event transport kind — set once at construction time from
186    // DYN_EVENT_PLANE + discovery backend; returned by default_event_transport_kind().
187    event_transport_kind: crate::discovery::EventTransportKind,
188}
189
190impl MetricsHierarchy for DistributedRuntime {
191    fn basename(&self) -> String {
192        "".to_string() // drt has no basename. Basename only begins with the Namespace.
193    }
194
195    fn parent_hierarchies(&self) -> Vec<&dyn MetricsHierarchy> {
196        vec![] // drt is the root, so no parent hierarchies
197    }
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    /// Initialize runtime dependencies, then bind HTTP with the selected probe policy.
225    /// Shutdown cancels pending initialization; no listener exists until it completes.
226    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        // Initialize discovery client based on backend configuration
261        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        // NetworkManager for request plane
313        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        // Initialize the uptime gauge in SystemHealth
355        distributed_runtime
356            .system_health
357            .lock()
358            .initialize_uptime_gauge(&distributed_runtime)?;
359
360        // Register an update callback so the uptime gauge is refreshed before
361        // every Prometheus scrape (both system status server and frontend).
362        {
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        // Opt-in OTLP metrics export. Deliberately not tied to the system
373        // status server: that server is disabled by default
374        // (DYN_SYSTEM_PORT=-1), and gating export on it would make
375        // OTEL_METRICS_EXPORTER=otlp a silent no-op in the default
376        // configuration. Traces and logs are set up in logging::init() for the
377        // same reason -- an OTEL_* variable should mean the same thing for
378        // every signal. Metrics cannot join them there because the exporter
379        // needs the registry, which only exists once the runtime does.
380        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                // Hold a graceful-shutdown guard for the task's life so the
388                // final export is not abandoned mid-RPC. `child_token()`
389                // derives from the endpoint shutdown token, which Phase 1
390                // cancels *before* the Phase 2 wait, so the exporter is told to
391                // stop and then waited for -- it cannot deadlock the wait on a
392                // token that only fires in Phase 3.
393                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        // Start health check manager if enabled
411        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            // Start the health check manager (spawns per-endpoint monitoring tasks)
420            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            // Keep sidecar probes alive through unregister/drain. Ordinary workers
441            // retain their existing endpoint-shutdown lifetime.
442            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                // Preserve ordinary workers' optional-HTTP failure behavior.
471                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    /// Check configured runtime dependencies, independently of model registration.
486    /// The HTTP caller bounds this operation; it never performs inference.
487    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    // TODO: Don't hand out pointers, instead have methods to use the registry in friendly ways
509    // (without being aware of async locks and so on)
510    pub fn component_registry(&self) -> &component::Registry {
511        &self.component_registry
512    }
513
514    // TODO: Don't hand out pointers, instead provide system health related services.
515    pub fn system_health(&self) -> Arc<parking_lot::Mutex<SystemHealth>> {
516        self.system_health.clone()
517    }
518
519    /// Get the local endpoint registry for in-process endpoint calls
520    pub fn local_endpoint_registry(
521        &self,
522    ) -> &crate::local_endpoint_registry::LocalEndpointRegistry {
523        &self.local_endpoint_registry
524    }
525
526    /// Get the engine route registry for registering custom /engine/* routes
527    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    /// Returns this DRT instance's discovery identity.
536    ///
537    /// This identifies the DRT, not the operating-system process. Multiple DRTs in one process
538    /// receive distinct connection IDs.
539    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    /// Create a [`Namespace`]
549    pub fn namespace(&self, name: impl Into<String>) -> Result<Namespace> {
550        Namespace::new(self.clone(), name.into())
551    }
552
553    /// Returns the discovery interface for service registration and discovery
554    pub fn discovery(&self) -> Arc<dyn Discovery> {
555        self.discovery_client.clone()
556    }
557
558    /// Register an endpoint until the last runtime-wide owner drops its lease.
559    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                // Keep the selected interface, but let the UDP stack choose a
616                // free port. A TCP ephemeral port can already be in use by an
617                // unrelated UDP socket because the two protocols allocate
618                // ports independently.
619                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    /// Get the network manager
648    ///
649    /// The network manager consolidates all network configuration and provides
650    /// unified access to request plane servers and clients.
651    pub fn network_manager(&self) -> Arc<NetworkManager> {
652        self.network_manager.clone()
653    }
654
655    /// Get the request plane server (convenience method)
656    ///
657    /// This is a shortcut for `network_manager().await?.server().await`.
658    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    /// Get system status server information if available
666    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    /// How the frontend should talk to the backend.
673    pub fn request_plane(&self) -> RequestPlaneMode {
674        self.request_plane
675    }
676
677    /// Returns the event transport kind this runtime was configured with.
678    ///
679    /// The value is resolved once at construction time by `DiscoveryBackend::resolve_event_transport_kind`:
680    /// if `DYN_EVENT_PLANE` is set explicitly that value wins; otherwise the default is ZMQ.
681    ///
682    /// Use this instead of `EventTransportKind::from_env_or_default` wherever you have
683    /// access to a `DistributedRuntime`.
684    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    /// Register an external long-running shutdown task with this runtime's
701    /// graceful-shutdown tracker. While the returned guard is alive,
702    /// `Runtime::shutdown` will keep waiting in Phase 2 (rather than
703    /// advancing to Phase 3 / NATS+etcd teardown). Drop the guard once
704    /// the task has finished.
705    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    /// TODO: This is a temporary KV router measure for component/component.rs EventPublisher impl for
714    /// Component, to allow it to publish to NATS. KV Router is the only user.
715    ///
716    /// When NATS is not available (e.g., running in approximate mode with --no-kv-events),
717    /// this function returns Ok(()) silently since publishing is optional in that mode.
718    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            // NATS not available - this is expected in approximate mode (--no-kv-events)
734            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    /// TODO: This is a temporary KV router measure for component/component.rs EventSubscriber impl for
741    /// Component, to allow it to subscribe to NATS. KV Router is the only user.
742    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    /// TODO (karenc): This is a temporary KV router measure for worker query requests.
753    /// Allows KV Router to perform request/reply with workers. (versus the pub/sub pattern above)
754    /// KV Router is the only user, made public for use in dynamo-llm crate
755    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    /// DEPRECATED: This method exists only for NATS request plane support.
772    /// Once everything uses the TCP request plane, this can be removed along with
773    /// the NATS service registration infrastructure.
774    ///
775    /// Returns a receiver that signals when the NATS service registration is complete.
776    /// The caller should use `blocking_recv()` to wait for completion.
777    pub fn register_nats_service(
778        &self,
779        component: Component,
780    ) -> tokio::sync::mpsc::Receiver<Result<(), String>> {
781        // Create a oneshot-style channel (capacity 1) to signal completion
782        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            // Pre-check to save cost of creating the service, but don't hold the lock
789            if drt
790                .component_registry()
791                .inner
792                .lock()
793                .await
794                .services
795                .contains_key(&service_name)
796            {
797                // The NATS service is per component, but it is called from `serve_endpoint`, and there
798                // are often multiple endpoints for a component (e.g. `clear_kv_blocks` and `generate`).
799                tracing::trace!("Service {service_name} already exists");
800                // Signal success - service already exists
801                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                // Normal case
831                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                // The NATS service is per component, but it is called from `serve_endpoint`, and there
840                // are often multiple endpoints for a component (e.g. `clear_kv_blocks` and `generate`).
841                // TODO: Is this still true?
842            }
843
844            // Signal completion - service registered successfully
845            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/// Selects which discovery backend to use and, for KV store backends, which KV store.
880#[derive(Clone, Debug)]
881pub enum DiscoveryBackend {
882    /// Use Kubernetes API for service discovery (no KV store needed)
883    Kubernetes,
884    /// Use a KV store (etcd, file, or memory) for service discovery
885    KvStore(kv::Selector),
886}
887
888impl DiscoveryBackend {
889    /// Returns true if this backend requires no external services (file or in-memory).
890    ///
891    /// Local backends do not need etcd, NATS, or any other infrastructure daemon.
892    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    /// Resolve the event transport kind for this backend.
901    ///
902    /// This is the single authoritative mapping of `DYN_EVENT_PLANE` →
903    /// `EventTransportKind`. ZMQ is the default event plane for all backends
904    /// (`file`/`mem`/`etcd`/`kubernetes`); NATS is an explicit opt-in via
905    /// `DYN_EVENT_PLANE=nats`. The value is trimmed and matched case-insensitively.
906    ///
907    /// Call this once at startup and store the result; do not call it repeatedly.
908    ///
909    /// # Errors
910    /// Returns an error if `DYN_EVENT_PLANE` is set to anything other than
911    /// `nats`, `zmq`, or an empty value, so a typo cannot silently select ZMQ.
912    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            // Unset or empty: ZMQ is the default for every backend.
923            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    /// Explicit response transport. `None` reads `DYN_RESPONSE_PLANE` for
946    /// standalone Rust entry points.
947    pub response_plane: Option<ResponsePlaneMode>,
948    /// Resolved event transport kind — computed once at config time from
949    /// `DYN_EVENT_PLANE` and the discovery backend, then stored on the runtime
950    /// so callers always get the same answer regardless of which other services
951    /// happen to be reachable.
952    pub event_transport_kind: crate::discovery::EventTransportKind,
953}
954
955impl DistributedConfig {
956    /// Build distributed runtime configuration from environment defaults.
957    ///
958    /// # Panics
959    /// Panics if a discovery or transport setting is invalid.
960    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    /// Resolve per-worker options before environment defaults, without mutating
969    /// the process environment. Safe to use after the Tokio runtime has started.
970    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        // Determine the discovery backend first — we need it to compute the NATS default below.
981        // Valid values for DYN_DISCOVERY_BACKEND: "kubernetes", "etcd" (default), "file", "mem"
982        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        // Resolve event transport kind once — the single source of truth used both to
1004        // decide whether to open a NATS connection and to answer
1005        // `DistributedRuntime::default_event_transport_kind()` later.
1006        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        // NATS is used for more than just NATS request-plane RPC:
1016        // - KV router events (NATS core event plane)
1017        // - inter-router replica sync (NATS core)
1018        //
1019        // Enable the NATS client when any of these hold:
1020        // 1. Request plane is NATS
1021        // 2. NATS_SERVER is explicitly configured by the user
1022        // 3. The resolved event transport kind is NATS
1023        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    /// A DistributedConfig that isn't distributed, for when the frontend and backend are in the
1076    /// same process.
1077    pub fn process_local() -> DistributedConfig {
1078        DistributedConfig {
1079            discovery_backend: DiscoveryBackend::KvStore(kv::Selector::Memory),
1080            nats_config: None,
1081            // This won't be used in process local, so we likely need a "none" option to
1082            // communicate that and avoid opening the ports.
1083            request_plane: RequestPlaneMode::Tcp,
1084            response_plane: None,
1085            event_transport_kind: crate::discovery::EventTransportKind::Zmq,
1086        }
1087    }
1088}
1089
1090/// Request plane transport mode configuration
1091///
1092/// This determines how requests are distributed from routers to workers:
1093/// - `Nats`: Use NATS for request distribution (legacy)
1094/// - `Tcp`: Use raw TCP for request distribution with msgpack support (default)
1095#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
1096pub enum RequestPlaneMode {
1097    /// Use NATS for request plane
1098    Nats,
1099    /// Use raw TCP for request plane with msgpack support
1100    #[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    //! Common test helper functions for DistributedRuntime tests
1273
1274    /// Helper function to create a DRT instance for integration-only tests.
1275    /// Uses from_current to leverage existing tokio runtime
1276    /// Note: Settings are read from environment variables inside DistributedRuntime::from_settings
1277    #[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    /// Helper function to create a DRT instance which points at
1295    /// a (shared) file-backed KV store and ephemeral NATS transport so that
1296    /// multiple DRT instances may observe the same registration state.
1297    /// NOTE: This gets around the fact that create_test_drt_async() is
1298    /// hardcoded to spin up a memory-backed discovery store
1299    /// which means we can't share discovery state across runtimes.
1300    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        // Test uptime with system status server disabled
1338        temp_env::async_with_vars([(env_system::DYN_SYSTEM_PORT, None::<&str>)], async {
1339            // Start a DRT
1340            let drt = create_test_drt_async().await;
1341
1342            // Wait 50ms
1343            tokio::time::sleep(tokio::time::Duration::from_millis(50)).await;
1344
1345            // Check that uptime is 50+ ms
1346            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        // Test uptime with system status server enabled
1375        temp_env::async_with_vars([(env_system::DYN_SYSTEM_PORT, Some("0"))], async {
1376            // Start a DRT
1377            let drt = create_test_drt_async().await;
1378
1379            // Wait 50ms
1380            tokio::time::sleep(tokio::time::Duration::from_millis(50)).await;
1381
1382            // Check that uptime is 50+ ms
1383            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}