distributed.rs 30.7 KB
Newer Older
1
// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2
// SPDX-License-Identifier: Apache-2.0
Ryan Olson's avatar
Ryan Olson committed
3

4
use crate::component::{Component, Instance};
5
use crate::pipeline::PipelineError;
6
use crate::pipeline::network::manager::NetworkManager;
7
use crate::service::{ServiceClient, ServiceSet};
8
use crate::storage::kv;
Ryan Olson's avatar
Ryan Olson committed
9
use crate::{
10
    component::{self, ComponentBuilder, Endpoint, Namespace},
11
    discovery::Discovery,
12
13
    metrics::PrometheusUpdateCallback,
    metrics::{MetricsHierarchy, MetricsRegistry},
Ryan Olson's avatar
Ryan Olson committed
14
15
    transports::{etcd, nats, tcp},
};
16
use crate::{discovery, system_status_server, transports};
Ryan Olson's avatar
Ryan Olson committed
17

18
use super::utils::GracefulShutdownTracker;
19
20
use crate::SystemHealth;
use crate::runtime::Runtime;
Ryan Olson's avatar
Ryan Olson committed
21

22
// Used instead of std::cell::OnceCell because get_or_try_init there is nightly
23
use async_once_cell::OnceCell;
24

25
use std::fmt;
26
use std::sync::{Arc, OnceLock, Weak};
27
use std::time::Duration;
28
use tokio::sync::watch::Receiver;
29
30

use anyhow::Result;
Ryan Olson's avatar
Ryan Olson committed
31
32
use derive_getters::Dissolve;
use figment::error;
33
34
use std::collections::HashMap;
use tokio::sync::Mutex;
35
use tokio_util::sync::CancellationToken;
Ryan Olson's avatar
Ryan Olson committed
36

37
38
type InstanceMap = HashMap<Endpoint, Weak<Receiver<Vec<Instance>>>>;

39
40
41
42
43
44
45
46
/// Distributed [Runtime] which provides access to shared resources across the cluster, this includes
/// communication protocols and transports.
#[derive(Clone)]
pub struct DistributedRuntime {
    // local runtime
    runtime: Runtime,

    nats_client: Option<transports::nats::Client>,
47
    network_manager: Arc<NetworkManager>,
48
49
    tcp_server: Arc<OnceCell<Arc<transports::tcp::server::TcpStreamServer>>>,
    system_status_server: Arc<OnceLock<Arc<system_status_server::SystemStatusServerInfo>>>,
50
    request_plane: RequestPlaneMode,
51
52
53
54

    // Service discovery client
    discovery_client: Arc<dyn discovery::Discovery>,

55
56
57
58
    // Discovery metadata (only used for Kubernetes backend)
    // Shared with system status server to expose via /metadata endpoint
    discovery_metadata: Option<Arc<tokio::sync::RwLock<discovery::DiscoveryMetadata>>>,

59
60
61
62
63
64
65
    // local registry for components
    // the registry allows us to use share runtime resources across instances of the same component object.
    // take for example two instances of a client to the same remote component. The registry allows us to use
    // a single endpoint watcher for both clients, this keeps the number background tasking watching specific
    // paths in etcd to a minimum.
    component_registry: component::Registry,

66
    instance_sources: Arc<tokio::sync::Mutex<InstanceMap>>,
67
68
69
70

    // Health Status
    system_health: Arc<parking_lot::Mutex<SystemHealth>>,

71
72
73
    // Local endpoint registry for in-process calls
    local_endpoint_registry: crate::local_endpoint_registry::LocalEndpointRegistry,

74
75
    // This hierarchy's own metrics registry
    metrics_registry: MetricsRegistry,
76
77
78

    // Registry for /engine/* route callbacks
    engine_routes: crate::engine_routes::EngineRouteRegistry,
79
80
}

81
impl MetricsHierarchy for DistributedRuntime {
82
83
84
85
    fn basename(&self) -> String {
        "".to_string() // drt has no basename. Basename only begins with the Namespace.
    }

86
87
88
89
90
91
    fn parent_hierarchies(&self) -> Vec<&dyn MetricsHierarchy> {
        vec![] // drt is the root, so no parent hierarchies
    }

    fn get_metrics_registry(&self) -> &MetricsRegistry {
        &self.metrics_registry
92
93
94
    }
}

Ryan Olson's avatar
Ryan Olson committed
95
96
97
98
99
100
impl std::fmt::Debug for DistributedRuntime {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "DistributedRuntime")
    }
}

Ryan Olson's avatar
Ryan Olson committed
101
102
impl DistributedRuntime {
    pub async fn new(runtime: Runtime, config: DistributedConfig) -> Result<Self> {
103
        let (discovery_backend, nats_config, request_plane) = config.dissolve();
Ryan Olson's avatar
Ryan Olson committed
104

105
106
107
108
        let nats_client = match nats_config {
            Some(nc) => Some(nc.connect().await?),
            None => None,
        };
Ryan Olson's avatar
Ryan Olson committed
109

110
        // Start system status server for health and metrics if enabled in configuration
111
112
113
114
115
116
117
118
        let config = crate::config::RuntimeConfig::from_settings().unwrap_or_default();
        // IMPORTANT: We must extract cancel_token from runtime BEFORE moving runtime into the struct below.
        // This is because after moving, runtime is no longer accessible in this scope (ownership rules).
        let cancel_token = if config.system_server_enabled() {
            Some(runtime.clone().child_token())
        } else {
            None
        };
119
120
        let starting_health_status = config.starting_health_status.clone();
        let use_endpoint_health_status = config.use_endpoint_health_status.clone();
121
122
        let health_endpoint_path = config.system_health_path.clone();
        let live_endpoint_path = config.system_live_path.clone();
123
        let system_health = Arc::new(parking_lot::Mutex::new(SystemHealth::new(
124
125
            starting_health_status,
            use_endpoint_health_status,
126
127
            health_endpoint_path,
            live_endpoint_path,
128
        )));
129

130
        // Initialize discovery client based on backend configuration
131
132
        let (discovery_client, discovery_metadata) = match discovery_backend {
            DiscoveryBackend::Kubernetes => {
133
134
135
136
137
138
139
140
141
142
143
144
145
146
                tracing::info!("Initializing Kubernetes discovery backend");
                let metadata = Arc::new(tokio::sync::RwLock::new(
                    crate::discovery::DiscoveryMetadata::new(),
                ));
                let client = crate::discovery::KubeDiscoveryClient::new(
                    metadata.clone(),
                    runtime.primary_token(),
                )
                .await
                .inspect_err(
                    |err| tracing::error!(%err, "Failed to initialize Kubernetes discovery client"),
                )?;
                (Arc::new(client) as Arc<dyn Discovery>, Some(metadata))
            }
147
            DiscoveryBackend::KvStore(kv_selector) => {
148
                tracing::info!("Initializing KV store discovery backend");
149
150
151
152
153
154
155
156
157
158
                let runtime_clone = runtime.clone();
                let store = match kv_selector {
                    kv::Selector::Etcd(etcd_config) => {
                        let etcd_client = etcd::Client::new(*etcd_config, runtime_clone).await.inspect_err(|err|
                            tracing::error!(%err, "Could not connect to etcd. Pass `--discovery-backend ..` to use a different backend or start etcd."))?;
                        kv::Manager::etcd(etcd_client)
                    }
                    kv::Selector::File(root) => kv::Manager::file(runtime.primary_token(), root),
                    kv::Selector::Memory => kv::Manager::memory(),
                };
159
160
                use crate::discovery::KVStoreDiscovery;
                (
161
162
                    Arc::new(KVStoreDiscovery::new(store, runtime.primary_token()))
                        as Arc<dyn Discovery>,
163
164
165
                    None,
                )
            }
166
167
        };

168
        let component_registry = component::Registry::new();
169

170
171
172
173
174
175
176
177
        // NetworkManager for request plane
        let network_manager = NetworkManager::new(
            runtime.child_token(),
            nats_client.clone().map(|c| c.client().clone()),
            component_registry.clone(),
            request_plane,
        );

178
        let distributed_runtime = Self {
Ryan Olson's avatar
Ryan Olson committed
179
            runtime,
180
            network_manager: Arc::new(network_manager),
Ryan Olson's avatar
Ryan Olson committed
181
182
            nats_client,
            tcp_server: Arc::new(OnceCell::new()),
183
            system_status_server: Arc::new(OnceLock::new()),
184
            discovery_client,
185
            discovery_metadata,
186
            component_registry,
187
            instance_sources: Arc::new(Mutex::new(HashMap::new())),
188
            metrics_registry: crate::MetricsRegistry::new(),
189
            system_health,
190
            request_plane,
191
            local_endpoint_registry: crate::local_endpoint_registry::LocalEndpointRegistry::new(),
192
            engine_routes: crate::engine_routes::EngineRouteRegistry::new(),
193
194
        };

195
196
197
198
199
200
        // Initialize the uptime gauge in SystemHealth
        distributed_runtime
            .system_health
            .lock()
            .initialize_uptime_gauge(&distributed_runtime)?;

201
202
203
204
205
206
207
208
209
210
211
212
        // Register an update callback so the uptime gauge is refreshed before
        // every Prometheus scrape (both system status server and frontend).
        {
            let system_health = distributed_runtime.system_health.clone();
            distributed_runtime
                .metrics_registry
                .add_update_callback(std::sync::Arc::new(move || {
                    system_health.lock().update_uptime_gauge();
                    Ok(())
                }));
        }

213
        // Handle system status server initialization
214
        if let Some(cancel_token) = cancel_token {
215
            // System server is enabled - start both the state and HTTP server
216
            let host = config.system_host.clone();
217
            let port = config.system_port as u16;
218

219
            // Start system status server (it creates SystemStatusState internally)
220
            match crate::system_status_server::spawn_system_status_server(
221
222
223
224
                &host,
                port,
                cancel_token,
                Arc::new(distributed_runtime.clone()),
225
                distributed_runtime.discovery_metadata.clone(),
226
227
228
            )
            .await
            {
229
                Ok((addr, handle)) => {
230
                    tracing::info!("System status server started successfully on {}", addr);
231

232
233
234
235
236
237
                    // Store system status server information
                    let system_status_server_info =
                        crate::system_status_server::SystemStatusServerInfo::new(
                            addr,
                            Some(handle),
                        );
238

239
                    // Initialize the system_status_server field
240
                    distributed_runtime
241
242
243
                        .system_status_server
                        .set(Arc::new(system_status_server_info))
                        .expect("System status server info should only be set once");
244
245
                }
                Err(e) => {
246
                    tracing::error!("System status server startup failed: {}", e);
247
                }
248
            }
249
        } else {
250
            // System server HTTP is disabled, but uptime metrics are still being tracked via SystemHealth
251
252
253
            tracing::debug!(
                "System status server HTTP endpoints disabled, but uptime metrics are being tracked"
            );
254
255
        }

256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
        // Start health check manager if enabled
        if config.health_check_enabled {
            let health_check_config = crate::health_check::HealthCheckConfig {
                canary_wait_time: std::time::Duration::from_secs(config.canary_wait_time_secs),
                request_timeout: std::time::Duration::from_secs(
                    config.health_check_request_timeout_secs,
                ),
            };

            // Start the health check manager (spawns per-endpoint monitoring tasks)
            match crate::health_check::start_health_check_manager(
                distributed_runtime.clone(),
                Some(health_check_config),
            )
            .await
            {
                Ok(()) => tracing::info!(
                    "Health check manager started (canary_wait_time: {}s, request_timeout: {}s)",
                    config.canary_wait_time_secs,
                    config.health_check_request_timeout_secs
                ),
                Err(e) => tracing::error!("Health check manager failed to start: {}", e),
            }
        }

281
        Ok(distributed_runtime)
Ryan Olson's avatar
Ryan Olson committed
282
283
284
    }

    pub async fn from_settings(runtime: Runtime) -> Result<Self> {
285
        let config = DistributedConfig::from_settings();
Ryan Olson's avatar
Ryan Olson committed
286
287
288
289
290
291
292
        Self::new(runtime, config).await
    }

    pub fn runtime(&self) -> &Runtime {
        &self.runtime
    }

293
294
295
296
    pub fn primary_token(&self) -> CancellationToken {
        self.runtime.primary_token()
    }

297
298
299
300
301
302
303
304
305
306
307
    // TODO: Don't hand out pointers, instead have methods to use the registry in friendly ways
    // (without being aware of async locks and so on)
    pub fn component_registry(&self) -> &component::Registry {
        &self.component_registry
    }

    // TODO: Don't hand out pointers, instead provide system health related services.
    pub fn system_health(&self) -> Arc<parking_lot::Mutex<SystemHealth>> {
        self.system_health.clone()
    }

308
309
310
311
312
313
314
    /// Get the local endpoint registry for in-process endpoint calls
    pub fn local_endpoint_registry(
        &self,
    ) -> &crate::local_endpoint_registry::LocalEndpointRegistry {
        &self.local_endpoint_registry
    }

315
316
317
318
319
    /// Get the engine route registry for registering custom /engine/* routes
    pub fn engine_routes(&self) -> &crate::engine_routes::EngineRouteRegistry {
        &self.engine_routes
    }

320
    pub fn connection_id(&self) -> u64 {
321
        self.discovery_client.instance_id()
Ryan Olson's avatar
Ryan Olson committed
322
323
324
325
    }

    pub fn shutdown(&self) {
        self.runtime.shutdown();
326
        self.discovery_client.shutdown();
Ryan Olson's avatar
Ryan Olson committed
327
328
329
330
    }

    /// Create a [`Namespace`]
    pub fn namespace(&self, name: impl Into<String>) -> Result<Namespace> {
331
        Namespace::new(self.clone(), name.into())
Ryan Olson's avatar
Ryan Olson committed
332
333
    }

334
335
336
    /// Returns the discovery interface for service registration and discovery
    pub fn discovery(&self) -> Arc<dyn Discovery> {
        self.discovery_client.clone()
337
338
    }

339
    pub async fn tcp_server(&self) -> Result<Arc<tcp::server::TcpStreamServer>> {
Ryan Olson's avatar
Ryan Olson committed
340
341
342
343
344
        Ok(self
            .tcp_server
            .get_or_try_init(async move {
                let options = tcp::server::ServerOptions::default();
                let server = tcp::server::TcpStreamServer::new(options).await?;
345
                Ok::<_, PipelineError>(server)
Ryan Olson's avatar
Ryan Olson committed
346
347
348
349
350
            })
            .await?
            .clone())
    }

351
    /// Get the network manager
352
353
354
    ///
    /// The network manager consolidates all network configuration and provides
    /// unified access to request plane servers and clients.
355
356
    pub fn network_manager(&self) -> Arc<NetworkManager> {
        self.network_manager.clone()
357
358
359
360
361
362
363
364
365
    }

    /// Get the request plane server (convenience method)
    ///
    /// This is a shortcut for `network_manager().await?.server().await`.
    pub async fn request_plane_server(
        &self,
    ) -> Result<Arc<dyn crate::pipeline::network::ingress::unified_server::RequestPlaneServer>>
    {
366
        self.network_manager().server().await
Ryan Olson's avatar
Ryan Olson committed
367
368
    }

369
370
371
372
373
    /// Get system status server information if available
    pub fn system_status_server_info(
        &self,
    ) -> Option<Arc<crate::system_status_server::SystemStatusServerInfo>> {
        self.system_status_server.get().cloned()
374
375
    }

376
377
378
379
380
    /// How the frontend should talk to the backend.
    pub fn request_plane(&self) -> RequestPlaneMode {
        self.request_plane
    }

381
382
383
    pub fn child_token(&self) -> CancellationToken {
        self.runtime.child_token()
    }
384

385
386
387
388
    pub(crate) fn graceful_shutdown_tracker(&self) -> Arc<GracefulShutdownTracker> {
        self.runtime.graceful_shutdown_tracker()
    }

389
    pub fn instance_sources(&self) -> Arc<Mutex<InstanceMap>> {
390
391
        self.instance_sources.clone()
    }
392
393
394

    /// TODO: This is a temporary KV router measure for component/component.rs EventPublisher impl for
    /// Component, to allow it to publish to NATS. KV Router is the only user.
395
396
397
    ///
    /// When NATS is not available (e.g., running in approximate mode with --no-kv-events),
    /// this function returns Ok(()) silently since publishing is optional in that mode.
398
    pub async fn kv_router_nats_publish(
399
400
401
402
403
        &self,
        subject: String,
        payload: bytes::Bytes,
    ) -> anyhow::Result<()> {
        let Some(nats_client) = self.nats_client.as_ref() else {
404
405
406
            // NATS not available - this is expected in approximate mode (--no-kv-events)
            tracing::trace!("Skipping NATS publish (NATS not configured): {}", subject);
            return Ok(());
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
        };
        Ok(nats_client.client().publish(subject, payload).await?)
    }

    /// TODO: This is a temporary KV router measure for component/component.rs EventSubscriber impl for
    /// Component, to allow it to subscribe to NATS. KV Router is the only user.
    pub(crate) async fn kv_router_nats_subscribe(
        &self,
        subject: String,
    ) -> Result<async_nats::Subscriber> {
        let Some(nats_client) = self.nats_client.as_ref() else {
            anyhow::bail!("KV router's EventSubscriber requires NATS");
        };
        Ok(nats_client.client().subscribe(subject).await?)
    }

423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
    /// TODO (karenc): This is a temporary KV router measure for worker query requests.
    /// Allows KV Router to perform request/reply with workers. (versus the pub/sub pattern above)
    /// KV Router is the only user, made public for use in dynamo-llm crate
    pub async fn kv_router_nats_request(
        &self,
        subject: String,
        payload: bytes::Bytes,
        timeout: std::time::Duration,
    ) -> anyhow::Result<async_nats::Message> {
        let Some(nats_client) = self.nats_client.as_ref() else {
            anyhow::bail!("KV router's request requires NATS");
        };
        let response =
            tokio::time::timeout(timeout, nats_client.client().request(subject, payload))
                .await
                .map_err(|_| anyhow::anyhow!("Request timed out after {:?}", timeout))??;
        Ok(response)
    }

442
443
444
    /// DEPRECATED: This method exists only for NATS request plane support.
    /// Once everything uses the TCP request plane, this can be removed along with
    /// the NATS service registration infrastructure.
445
446
447
448
449
450
451
452
453
454
    ///
    /// Returns a receiver that signals when the NATS service registration is complete.
    /// The caller should use `blocking_recv()` to wait for completion.
    pub fn register_nats_service(
        &self,
        component: Component,
    ) -> tokio::sync::mpsc::Receiver<Result<(), String>> {
        // Create a oneshot-style channel (capacity 1) to signal completion
        let (tx, rx) = tokio::sync::mpsc::channel::<Result<(), String>>(1);

455
456
457
458
        let drt = self.clone();
        self.runtime().secondary().spawn(async move {
            let service_name = component.service_name();

459
460
461
462
463
464
465
466
467
468
469
470
            // Pre-check to save cost of creating the service, but don't hold the lock
            if drt
                .component_registry()
                .inner
                .lock()
                .await
                .services
                .contains_key(&service_name)
            {
                // The NATS service is per component, but it is called from `serve_endpoint`, and there
                // are often multiple endpoints for a component (e.g. `clear_kv_blocks` and `generate`).
                tracing::trace!("Service {service_name} already exists");
471
472
                // Signal success - service already exists
                let _ = tx.send(Ok(())).await;
473
474
                return;
            }
475

476
477
            let Some(nats_client) = drt.nats_client.as_ref() else {
                tracing::error!("Cannot create NATS service without NATS.");
478
479
480
                let _ = tx
                    .send(Err("Cannot create NATS service without NATS".to_string()))
                    .await;
481
482
483
484
485
486
487
488
                return;
            };
            let description = None;
            let nats_service = match crate::component::service::build_nats_service(
                nats_client,
                &component,
                description,
            )
489
            .await
490
491
492
493
            {
                Ok(service) => service,
                Err(err) => {
                    tracing::error!(error = %err, component = service_name, "Failed to build NATS service");
494
                    let _ = tx.send(Err(format!("Failed to build NATS service: {err}"))).await;
495
496
497
                    return;
                }
            };
498

499
500
501
502
            let mut guard = drt.component_registry().inner.lock().await;
            if !guard.services.contains_key(&service_name) {
                // Normal case
                guard.services.insert(service_name.clone(), nats_service);
503

504
                tracing::info!("Added NATS service {service_name}");
505

506
507
508
509
510
511
512
                drop(guard);
            } else {
                drop(guard);
                let _ = nats_service.stop().await;
                // The NATS service is per component, but it is called from `serve_endpoint`, and there
                // are often multiple endpoints for a component (e.g. `clear_kv_blocks` and `generate`).
                // TODO: Is this still true?
513
            }
514
515
516

            // Signal completion - service registered successfully
            let _ = tx.send(Ok(())).await;
517
        });
518
519

        rx
520
    }
Ryan Olson's avatar
Ryan Olson committed
521
522
}

523
524
525
526
527
528
529
530
531
/// Selects which discovery backend to use and, for KV store backends, which KV store.
#[derive(Clone, Debug)]
pub enum DiscoveryBackend {
    /// Use Kubernetes API for service discovery (no KV store needed)
    Kubernetes,
    /// Use a KV store (etcd, file, or memory) for service discovery
    KvStore(kv::Selector),
}

Ryan Olson's avatar
Ryan Olson committed
532
533
#[derive(Dissolve)]
pub struct DistributedConfig {
534
    pub discovery_backend: DiscoveryBackend,
535
    pub nats_config: Option<nats::ClientOptions>,
536
    pub request_plane: RequestPlaneMode,
Ryan Olson's avatar
Ryan Olson committed
537
538
539
}

impl DistributedConfig {
540
    pub fn from_settings() -> DistributedConfig {
541
        let request_plane = RequestPlaneMode::from_env();
542
543
544
545
546
547
548
549
550
        // NATS is used for more than just NATS request-plane RPC:
        // - KV router events (JetStream or NATS core + local indexer)
        // - inter-router replica sync (NATS core)
        //
        // Historically we only connected to NATS when the request plane was NATS, which made
        // `DYN_REQUEST_PLANE=tcp|http` incompatible with KV routing modes that rely on NATS.
        // If a NATS server is configured via env, enable the client regardless of request plane.
        let nats_enabled = request_plane.is_nats()
            || std::env::var(crate::config::environment_names::nats::NATS_SERVER).is_ok();
551

552
553
554
555
        // DYN_DISCOVERY_BACKEND selects the discovery mechanism
        // Valid values: "kubernetes", "etcd" (default), "file", "mem"
        let backend_str =
            std::env::var("DYN_DISCOVERY_BACKEND").unwrap_or_else(|_| "etcd".to_string());
556

557
558
559
560
561
562
563
564
565
566
567
568
569
570
        let discovery_backend = match backend_str.as_str() {
            "kubernetes" => {
                tracing::info!("Using Kubernetes discovery backend");
                DiscoveryBackend::Kubernetes
            }
            other => {
                let selector: kv::Selector = other.parse().unwrap_or_else(|_| {
                    panic!(
                        "Unknown DYN_DISCOVERY_BACKEND value: '{other}'. \
                         Valid options: kubernetes, etcd, file, mem"
                    )
                });
                DiscoveryBackend::KvStore(selector)
            }
571
572
        };

Ryan Olson's avatar
Ryan Olson committed
573
        DistributedConfig {
574
            discovery_backend,
575
            nats_config: if nats_enabled {
576
577
578
579
580
                Some(nats::ClientOptions::default())
            } else {
                None
            },
            request_plane,
Ryan Olson's avatar
Ryan Olson committed
581
582
        }
    }
Ryan Olson's avatar
Ryan Olson committed
583
584

    pub fn for_cli() -> DistributedConfig {
585
586
587
588
        let etcd_config = etcd::ClientOptions {
            attach_lease: false,
            ..Default::default()
        };
589
        let request_plane = RequestPlaneMode::from_env();
590
591
        let nats_enabled = request_plane.is_nats()
            || std::env::var(crate::config::environment_names::nats::NATS_SERVER).is_ok();
592
        DistributedConfig {
593
            discovery_backend: DiscoveryBackend::KvStore(kv::Selector::Etcd(Box::new(etcd_config))),
594
            nats_config: if nats_enabled {
595
596
597
598
599
                Some(nats::ClientOptions::default())
            } else {
                None
            },
            request_plane,
600
601
        }
    }
602
603
604
605
606

    /// A DistributedConfig that isn't distributed, for when the frontend and backend are in the
    /// same process.
    pub fn process_local() -> DistributedConfig {
        DistributedConfig {
607
            discovery_backend: DiscoveryBackend::KvStore(kv::Selector::Memory),
608
609
610
611
612
613
            nats_config: None,
            // This won't be used in process local, so we likely need a "none" option to
            // communicate that and avoid opening the ports.
            request_plane: RequestPlaneMode::Tcp,
        }
    }
614
615
616
617
618
}

/// Request plane transport mode configuration
///
/// This determines how requests are distributed from routers to workers:
619
/// - `Nats`: Use NATS for request distribution (legacy)
620
/// - `Http`: Use HTTP/2 for request distribution
621
/// - `Tcp`: Use raw TCP for request distribution with msgpack support (default)
622
623
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum RequestPlaneMode {
624
    /// Use NATS for request plane
625
626
627
628
629
630
631
632
633
    Nats,
    /// Use HTTP/2 for request plane
    Http,
    /// Use raw TCP for request plane with msgpack support
    Tcp,
}

impl Default for RequestPlaneMode {
    fn default() -> Self {
634
        Self::Tcp
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
    }
}

impl fmt::Display for RequestPlaneMode {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::Nats => write!(f, "nats"),
            Self::Http => write!(f, "http"),
            Self::Tcp => write!(f, "tcp"),
        }
    }
}

impl std::str::FromStr for RequestPlaneMode {
    type Err = anyhow::Error;

    fn from_str(s: &str) -> std::result::Result<Self, Self::Err> {
        match s.to_lowercase().as_str() {
            "nats" => Ok(Self::Nats),
            "http" => Ok(Self::Http),
            "tcp" => Ok(Self::Tcp),
            _ => Err(anyhow::anyhow!(
                "Invalid request plane mode: '{}'. Valid options are: 'nats', 'http', 'tcp'",
                s
            )),
660
        }
Ryan Olson's avatar
Ryan Olson committed
661
    }
Ryan Olson's avatar
Ryan Olson committed
662
}
663

664
665
666
667
668
669
670
671
672
impl RequestPlaneMode {
    /// Get the request plane mode from environment variable (uncached)
    /// Reads from `DYN_REQUEST_PLANE` environment variable.
    fn from_env() -> Self {
        std::env::var("DYN_REQUEST_PLANE")
            .ok()
            .and_then(|s| s.parse().ok())
            .unwrap_or_default()
    }
673
674
675
676

    pub fn is_nats(&self) -> bool {
        matches!(self, RequestPlaneMode::Nats)
    }
677
678
}

679
pub mod distributed_test_utils {
680
681
    //! Common test helper functions for DistributedRuntime tests

682
    /// Helper function to create a DRT instance for integration-only tests.
683
    /// Uses from_current to leverage existing tokio runtime
684
    /// Note: Settings are read from environment variables inside DistributedRuntime::from_settings
685
    #[cfg(feature = "integration")]
686
    pub async fn create_test_drt_async() -> super::DistributedRuntime {
687
        use crate::transports::nats;
688

689
        let rt = crate::Runtime::from_current().unwrap();
690
        let config = super::DistributedConfig {
691
692
693
            discovery_backend: super::DiscoveryBackend::KvStore(
                crate::storage::kv::Selector::Memory,
            ),
694
            nats_config: Some(nats::ClientOptions::default()),
695
            request_plane: crate::distributed::RequestPlaneMode::default(),
696
697
        };
        super::DistributedRuntime::new(rt, config).await.unwrap()
698
    }
699
700
701
702
703
704
705
706
707
708

    /// Helper function to create a DRT instance which points at
    /// a (shared) file-backed KV store and ephemeral NATS transport so that
    /// multiple DRT instances may observe the same registration state.
    /// NOTE: This gets around the fact that create_test_drt_async() is
    /// hardcoded to spin up a memory-backed discovery store
    /// which means we can't share discovery state across runtimes.
    pub async fn create_test_shared_drt_async(
        store_path: &std::path::Path,
    ) -> super::DistributedRuntime {
709
        use crate::transports::nats;
710
711
712

        let rt = crate::Runtime::from_current().unwrap();
        let config = super::DistributedConfig {
713
714
715
            discovery_backend: super::DiscoveryBackend::KvStore(
                crate::storage::kv::Selector::File(store_path.to_path_buf()),
            ),
716
717
718
719
720
            nats_config: Some(nats::ClientOptions::default()),
            request_plane: crate::distributed::RequestPlaneMode::default(),
        };
        super::DistributedRuntime::new(rt, config).await.unwrap()
    }
721
}
722

723
#[cfg(all(test, feature = "integration"))]
724
mod tests {
725
    use super::RequestPlaneMode;
726
727
728
729
    use super::distributed_test_utils::create_test_drt_async;

    #[tokio::test]
    async fn test_drt_uptime_after_delay_system_disabled() {
730
        use crate::config::environment_names::runtime::system as env_system;
731
        // Test uptime with system status server disabled
732
        temp_env::async_with_vars([(env_system::DYN_SYSTEM_PORT, None::<&str>)], async {
733
734
735
736
737
738
739
            // Start a DRT
            let drt = create_test_drt_async().await;

            // Wait 50ms
            tokio::time::sleep(tokio::time::Duration::from_millis(50)).await;

            // Check that uptime is 50+ ms
740
            let uptime = drt.system_health.lock().uptime();
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
            assert!(
                uptime >= std::time::Duration::from_millis(50),
                "Expected uptime to be at least 50ms, but got {:?}",
                uptime
            );

            println!(
                "✓ DRT uptime test passed (system disabled): uptime = {:?}",
                uptime
            );
        })
        .await;
    }

    #[tokio::test]
    async fn test_drt_uptime_after_delay_system_enabled() {
757
        use crate::config::environment_names::runtime::system as env_system;
758
        // Test uptime with system status server enabled
759
        temp_env::async_with_vars([(env_system::DYN_SYSTEM_PORT, Some("8081"))], async {
760
761
762
763
764
765
766
            // Start a DRT
            let drt = create_test_drt_async().await;

            // Wait 50ms
            tokio::time::sleep(tokio::time::Duration::from_millis(50)).await;

            // Check that uptime is 50+ ms
767
            let uptime = drt.system_health.lock().uptime();
768
769
770
771
772
773
774
775
776
777
778
779
780
            assert!(
                uptime >= std::time::Duration::from_millis(50),
                "Expected uptime to be at least 50ms, but got {:?}",
                uptime
            );

            println!(
                "✓ DRT uptime test passed (system enabled): uptime = {:?}",
                uptime
            );
        })
        .await;
    }
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816

    #[test]
    fn test_request_plane_mode_from_str() {
        assert_eq!(
            "nats".parse::<RequestPlaneMode>().unwrap(),
            RequestPlaneMode::Nats
        );
        assert_eq!(
            "http".parse::<RequestPlaneMode>().unwrap(),
            RequestPlaneMode::Http
        );
        assert_eq!(
            "tcp".parse::<RequestPlaneMode>().unwrap(),
            RequestPlaneMode::Tcp
        );
        assert_eq!(
            "NATS".parse::<RequestPlaneMode>().unwrap(),
            RequestPlaneMode::Nats
        );
        assert_eq!(
            "HTTP".parse::<RequestPlaneMode>().unwrap(),
            RequestPlaneMode::Http
        );
        assert_eq!(
            "TCP".parse::<RequestPlaneMode>().unwrap(),
            RequestPlaneMode::Tcp
        );
        assert!("invalid".parse::<RequestPlaneMode>().is_err());
    }

    #[test]
    fn test_request_plane_mode_display() {
        assert_eq!(RequestPlaneMode::Nats.to_string(), "nats");
        assert_eq!(RequestPlaneMode::Http.to_string(), "http");
        assert_eq!(RequestPlaneMode::Tcp.to_string(), "tcp");
    }
817
}