"lib/llm/vscode:/vscode.git/clone" did not exist on "f8213242d768dd4dda051a150785eb4824e50212"
Unverified Commit 06b0ebef authored by Biswa Panda's avatar Biswa Panda Committed by GitHub
Browse files

feat: transport agnostic request plane for dynamo - natless (#4246)

parent 381c428c
This diff is collapsed.
......@@ -97,6 +97,7 @@ reqwest = { version = "0.12.22", default-features = false, features = [
"stream",
"rustls-tls",
] }
rmp-serde = { version = "1" }
serde = { version = "1", features = ["derive"] }
serde_json = { version = "1" }
strum = { version = "0.27", features = ["derive"] }
......
#!/bin/bash
# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
set -e
trap 'echo Cleaning up...; kill 0' EXIT
# Parse command-line arguments for request plane mode
REQUEST_PLANE="tcp" # Default to TCP
while [[ $# -gt 0 ]]; do
case $1 in
--tcp)
REQUEST_PLANE="tcp"
shift
;;
--http)
REQUEST_PLANE="http"
shift
;;
--nats)
REQUEST_PLANE="nats"
shift
;;
-h|--help)
echo "Usage: $0 [--tcp|--http|--nats]"
echo " --tcp Use TCP request plane (default)"
echo " --http Use HTTP/2 request plane"
echo " --nats Use NATS request plane"
exit 0
;;
*)
echo "Unknown option: $1"
echo "Use --help for usage information"
exit 1
;;
esac
done
# Set the request plane mode
export DYN_REQUEST_PLANE=$REQUEST_PLANE
echo "Using request plane mode: $REQUEST_PLANE"
# Frontend
python -m dynamo.frontend --http-port=8000 &
DYN_SYSTEM_PORT=8081 \
python -m dynamo.vllm --model Qwen/Qwen3-0.6B --enforce-eager --connector none
......@@ -39,7 +39,7 @@ backoff = { version = "0.4.0", features = ["tokio"] }
base64 = "0.22.1"
futures = "0.3.31"
rand = "0.9.0"
reqwest = { version = "0.12.12", features = [
reqwest = { version = "0.12.24", features = [
"json",
"stream",
"multipart",
......
This diff is collapsed.
......@@ -169,7 +169,7 @@ where
let message = message.replace(['\r', '\n'], "");
return Err(http_error::HttpError { code, message })?;
}
Err(error!("Python Error: {}", py_err.to_string()))
Err(error!("Python Error: {}", py_err))
})
} else {
Err(e)
......
......@@ -106,8 +106,11 @@ unicode-segmentation = "1.12"
# http-service
axum = { workspace = true }
axum-server = { version = "0.7", features = ["tls-rustls"] }
hyper = { version = "1.5", features = ["http2", "server"] }
hyper-util = { version = "0.1", features = ["tokio", "server", "server-auto"] }
tower-http = { workspace = true }
rustls = { version = "0.23" }
tower = { version = "0.5", features = ["util", "make"] }
utoipa = { version = "5.3", features = ["axum_extras"] }
utoipa-swagger-ui = { version = "9.0", features = ["axum"] }
......
......@@ -82,10 +82,6 @@ impl CheckedFile {
}
}
pub fn is_nats_url(&self) -> bool {
matches!(self.path.as_ref(), Either::Right(u) if u.scheme() == "nats")
}
pub fn checksum(&self) -> &Checksum {
&self.checksum
}
......
......@@ -18,6 +18,7 @@ integration = []
testing-etcd = [] # Tests that require an active ETCD server
tokio-console = ["dep:console-subscriber", "tokio/tracing"]
compute-validation = [] # Enable validation and timing for compute macros
tcp-low-latency = [] # Enable Linux-specific TCP optimizations (TCP_QUICKACK, SO_BUSY_POLL)
[dependencies]
# Use workspace dependencies where available
......@@ -30,6 +31,7 @@ axum = { workspace = true }
blake3 = { workspace = true }
bytes = { workspace = true }
chrono = { workspace = true }
dashmap = { workspace = true }
derive_builder = { workspace = true }
derive-getters = { workspace = true }
either = { workspace = true }
......@@ -40,13 +42,17 @@ parking_lot = { workspace = true }
prometheus = { workspace = true }
rand = { workspace = true }
reqwest = { workspace = true }
rmp-serde = { workspace = true }
serde = { workspace = true }
serde_json = { workspace = true }
tokio = { workspace = true }
tokio-stream = { workspace = true }
tokio-util = { workspace = true }
tower = { version = "0.5" }
tower-http = { workspace = true }
tracing = { workspace = true }
hyper = { version = "1.5", features = ["http2", "server"] }
hyper-util = { version = "0.1", features = ["tokio", "server", "server-auto"] }
tracing-subscriber = { workspace = true }
tracing-opentelemetry = { workspace = true }
opentelemetry = { workspace = true }
......@@ -65,6 +71,8 @@ console-subscriber = { version = "0.4", optional = true }
educe = { version = "0.6.0" }
figment = { version = "0.10.19", features = ["env", "json", "toml", "test"] }
notify = { version = "6.1", default-features = false, features = ["macos_fsevent"] }
inotify = { version = "0.11" }
libc = { version = "0.2" }
local-ip-address = { version = "0.6.3" }
log = { version = "0.4" }
nid = { version = "3.0.0", features = ["serde"] }
......@@ -84,7 +92,6 @@ k8s-openapi = { version = "0.26.0", features = ["v1_32"] }
assert_matches = { version = "1.5.0" }
criterion = { version = "0.5", features = ["async_tokio"] }
env_logger = { version = "0.11" }
reqwest = { workspace = true }
rstest = { version = "0.23.0" }
temp-env = { version = "0.3.6" , features=["async_closure"] }
stdio-override = {version= "0.2.0"}
......@@ -94,3 +101,7 @@ tempfile = { workspace = true }
[[bench]]
name = "compute_pool_overhead"
harness = false
[[bench]]
name = "tcp_codec_perf"
harness = false
// SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0
//! Micro-benchmarks for TCP codec performance
//!
//! Run with: cargo bench --bench tcp_codec_perf
use bytes::Bytes;
use criterion::{BenchmarkId, Criterion, Throughput, black_box, criterion_group, criterion_main};
use dynamo_runtime::pipeline::network::codec::{TcpRequestMessage, TcpResponseMessage};
/// Benchmark request encoding (hot path operation)
fn bench_request_encoding(c: &mut Criterion) {
let mut group = c.benchmark_group("tcp_request_encoding");
// Test different payload sizes
for size in [102_400, 1_024_000, 31_000_000].iter() {
let payload = Bytes::from(vec![0u8; *size]);
let endpoint = "api.endpoint.test".to_string();
group.throughput(Throughput::Bytes(*size as u64));
group.bench_with_input(BenchmarkId::from_parameter(size), size, |b, _| {
b.iter(|| {
let msg = TcpRequestMessage::new(endpoint.clone(), payload.clone());
let encoded = msg.encode().unwrap();
black_box(encoded);
});
});
}
group.finish();
}
/// Benchmark response encoding
fn bench_response_encoding(c: &mut Criterion) {
let mut group = c.benchmark_group("tcp_response_encoding");
for size in [100, 1024, 10_240, 102_400].iter() {
let data = Bytes::from(vec![0u8; *size]);
group.throughput(Throughput::Bytes(*size as u64));
group.bench_with_input(BenchmarkId::from_parameter(size), size, |b, _| {
b.iter(|| {
let msg = TcpResponseMessage::new(data.clone());
let encoded = msg.encode().unwrap();
black_box(encoded);
});
});
}
group.finish();
}
/// Benchmark request decoding
fn bench_request_decoding(c: &mut Criterion) {
let mut group = c.benchmark_group("tcp_request_decoding");
for size in [100, 1024, 10_240, 102_400].iter() {
let payload = Bytes::from(vec![0u8; *size]);
let msg = TcpRequestMessage::new("api.endpoint.test".to_string(), payload);
let encoded = msg.encode().unwrap();
group.throughput(Throughput::Bytes(*size as u64));
group.bench_with_input(BenchmarkId::from_parameter(size), size, |b, _| {
b.iter(|| {
let decoded = TcpRequestMessage::decode(&encoded).unwrap();
black_box(decoded);
});
});
}
group.finish();
}
/// Benchmark response decoding
fn bench_response_decoding(c: &mut Criterion) {
let mut group = c.benchmark_group("tcp_response_decoding");
for size in [100, 1024, 10_240, 102_400].iter() {
let data = Bytes::from(vec![0u8; *size]);
let msg = TcpResponseMessage::new(data);
let encoded = msg.encode().unwrap();
group.throughput(Throughput::Bytes(*size as u64));
group.bench_with_input(BenchmarkId::from_parameter(size), size, |b, _| {
b.iter(|| {
let decoded = TcpResponseMessage::decode(&encoded).unwrap();
black_box(decoded);
});
});
}
group.finish();
}
/// Benchmark full encode-decode cycle for requests
fn bench_request_roundtrip(c: &mut Criterion) {
let mut group = c.benchmark_group("tcp_request_roundtrip");
for size in [100, 1024, 10_240].iter() {
let payload = Bytes::from(vec![0u8; *size]);
let endpoint = "api.endpoint.test".to_string();
group.throughput(Throughput::Bytes(*size as u64));
group.bench_with_input(BenchmarkId::from_parameter(size), size, |b, _| {
b.iter(|| {
let msg = TcpRequestMessage::new(endpoint.clone(), payload.clone());
let encoded = msg.encode().unwrap();
let decoded = TcpRequestMessage::decode(&encoded).unwrap();
black_box(decoded);
});
});
}
group.finish();
}
criterion_group!(
benches,
bench_request_encoding,
bench_response_encoding,
bench_request_decoding,
bench_response_decoding,
bench_request_roundtrip,
);
criterion_main!(benches);
......@@ -32,7 +32,7 @@
use std::fmt;
use crate::{
config::HealthStatus,
config::{HealthStatus, RequestPlaneMode},
metrics::{MetricsHierarchy, MetricsRegistry, prometheus_names},
service::ServiceSet,
transports::etcd::{ETCD_ROOT_PATH, EtcdPath},
......@@ -78,18 +78,22 @@ pub const INSTANCE_ROOT_PATH: &str = "v1/instances";
#[derive(Debug, Clone, Serialize, Deserialize, Eq, PartialEq, Hash)]
#[serde(rename_all = "snake_case")]
pub enum TransportType {
NatsTcp(String),
#[serde(rename = "nats_tcp")]
Nats(String),
Http(String),
Tcp(String),
}
#[derive(Default)]
pub struct RegistryInner {
services: HashMap<String, Service>,
stats_handlers: HashMap<String, Arc<parking_lot::Mutex<HashMap<String, EndpointStatsHandler>>>>,
pub(crate) services: HashMap<String, Service>,
pub(crate) stats_handlers:
HashMap<String, Arc<parking_lot::Mutex<HashMap<String, EndpointStatsHandler>>>>,
}
#[derive(Clone)]
pub struct Registry {
inner: Arc<tokio::sync::Mutex<RegistryInner>>,
pub(crate) inner: Arc<tokio::sync::Mutex<RegistryInner>>,
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
......@@ -407,8 +411,25 @@ impl Component {
}
// Register metrics callback. CRITICAL: Never fail service creation for metrics issues.
// Only enable NATS service metrics collection when using NATS request plane mode
let request_plane_mode = RequestPlaneMode::get();
match request_plane_mode {
RequestPlaneMode::Nats => {
if let Err(err) = self.start_scraping_nats_service_component_metrics() {
tracing::debug!(service_name, error = %err, "Metrics registration failed");
tracing::debug!(
"Metrics registration failed for '{}': {}",
self.service_name(),
err
);
}
}
_ => {
tracing::info!(
"Skipping NATS service metrics collection for '{}' - request plane mode is '{}'",
self.service_name(),
request_plane_mode
);
}
}
Ok(())
}
......
......@@ -12,6 +12,7 @@ use tokio_util::sync::CancellationToken;
use crate::{
component::{Endpoint, Instance, TransportType, service::EndpointStatsHandler},
config::RequestPlaneMode,
pipeline::network::{PushWorkHandler, ingress::push_endpoint::PushEndpoint},
storage::key_value_store,
traits::DistributedRuntimeProvider,
......@@ -87,8 +88,9 @@ impl EndpointConfigBuilder {
let registry = endpoint.drt().component_registry().inner.lock().await;
// get the group
let group = registry
// Note: NATS service group is no longer needed here as the NetworkManager
// handles all transport-specific initialization internally
let _group = registry
.services
.get(&service_name)
.map(|service| service.group(endpoint.component.service_name()))
......@@ -110,11 +112,12 @@ impl EndpointConfigBuilder {
.insert(endpoint.subject_to(connection_id), stats_handler);
}
// creates an endpoint for the service
let service_endpoint = group
.endpoint(&endpoint.name_with_id(connection_id))
.await
.map_err(|e| anyhow::anyhow!("Failed to start endpoint: {e}"))?;
// Determine request plane mode
let request_plane_mode = RequestPlaneMode::get();
tracing::info!(
"Endpoint starting with request plane mode: {:?}",
request_plane_mode
);
// This creates a child token of the runtime's endpoint_shutdown_token. That token is
// cancelled first as part of graceful shutdown. See Runtime::shutdown.
......@@ -129,12 +132,20 @@ impl EndpointConfigBuilder {
// Register health check target in SystemHealth if provided
if let Some(health_check_payload) = &health_check_payload {
// Build transport based on request plane mode
let transport = build_transport_type(
request_plane_mode,
&endpoint_name,
&subject,
TransportContext::HealthCheck,
);
let instance = Instance {
component: component_name.clone(),
endpoint: endpoint_name.clone(),
namespace: namespace_name.clone(),
instance_id: connection_id,
transport: TransportType::NatsTcp(subject.clone()),
transport,
};
tracing::debug!(endpoint_name = %endpoint_name, "Registering endpoint health check target");
let guard = system_health.lock();
......@@ -160,13 +171,7 @@ impl EndpointConfigBuilder {
tracing::debug!("Endpoint '{}' has graceful_shutdown=false", endpoint.name);
}
let push_endpoint = PushEndpoint::builder()
.service_handler(handler)
.cancellation_token(endpoint_shutdown_token.clone())
.graceful_shutdown(graceful_shutdown)
.build()
.map_err(|e| anyhow::anyhow!("Failed to build push endpoint: {e}"))?;
// Launch endpoint based on request plane mode
let tracker_clone = if graceful_shutdown {
Some(endpoint.drt().graceful_shutdown_tracker())
} else {
......@@ -178,25 +183,59 @@ impl EndpointConfigBuilder {
let component_name_for_task = component_name.clone();
let endpoint_name_for_task = endpoint_name.clone();
let task = tokio::spawn(async move {
let result = push_endpoint
.start(
service_endpoint,
namespace_name_for_task,
component_name_for_task,
endpoint_name_for_task,
// Get the unified request plane server (works for all transport types)
let server = endpoint.drt().request_plane_server().await?;
tracing::info!(
endpoint = %endpoint_name_for_task,
transport = server.transport_name(),
"Registering endpoint with request plane server"
);
// Register endpoint with the server (unified interface)
server
.register_endpoint(
endpoint_name_for_task.clone(),
handler,
connection_id,
system_health,
namespace_name_for_task.clone(),
component_name_for_task.clone(),
system_health.clone(),
)
.await;
.await?;
// Create cleanup task that unregisters on cancellation
let endpoint_name_for_cleanup = endpoint_name_for_task.clone();
let server_for_cleanup = server.clone();
let cancel_token_for_cleanup = endpoint_shutdown_token.clone();
let task: tokio::task::JoinHandle<anyhow::Result<()>> = tokio::spawn(async move {
cancel_token_for_cleanup.cancelled().await;
tracing::debug!(
endpoint = %endpoint_name_for_cleanup,
"Unregistering endpoint from request plane server"
);
// Unregister from server
if let Err(e) = server_for_cleanup
.unregister_endpoint(&endpoint_name_for_cleanup)
.await
{
tracing::warn!(
endpoint = %endpoint_name_for_cleanup,
error = %e,
"Failed to unregister endpoint"
);
}
// Unregister from graceful shutdown tracker
if let Some(tracker) = tracker_clone {
tracing::debug!("Unregistering endpoint from graceful shutdown tracker");
tracing::debug!("Unregister endpoint from graceful shutdown tracker");
tracker.unregister_endpoint();
}
result
anyhow::Ok(())
});
// Register this endpoint instance in the discovery plane
......@@ -204,11 +243,19 @@ impl EndpointConfigBuilder {
// consistent registration/discovery across the system.
let discovery = endpoint.drt().discovery();
// Build transport for discovery service based on request plane mode
let transport = build_transport_type(
request_plane_mode,
&endpoint_name,
&subject,
TransportContext::Discovery,
);
let discovery_spec = crate::discovery::DiscoverySpec::Endpoint {
namespace: namespace_name.clone(),
component: component_name.clone(),
endpoint: endpoint_name.clone(),
transport: TransportType::NatsTcp(subject.clone()),
transport,
};
if let Err(e) = discovery.register(discovery_spec).await {
......@@ -229,3 +276,66 @@ impl EndpointConfigBuilder {
Ok(())
}
}
/// Context for building transport type - determines port and formatting differences
enum TransportContext {
/// For health check targets
HealthCheck,
/// For discovery service registration
Discovery,
}
/// Build transport type based on request plane mode and context
///
/// This unified function handles both health check and discovery transport building,
/// with context-specific differences:
/// - HTTP: Both use the same port (default 8888, configurable via DYN_HTTP_RPC_PORT)
/// - TCP: Health check omits endpoint suffix, discovery includes it for routing
/// - NATS: Identical for both contexts
fn build_transport_type(
mode: RequestPlaneMode,
endpoint_name: &str,
subject: &str,
context: TransportContext,
) -> TransportType {
match mode {
RequestPlaneMode::Http => {
let http_host = crate::utils::get_http_rpc_host_from_env();
// Both health check and discovery use the same port (8888) where the HTTP server binds
let http_port = std::env::var("DYN_HTTP_RPC_PORT")
.ok()
.and_then(|p| p.parse::<u16>().ok())
.unwrap_or(8888);
let rpc_root =
std::env::var("DYN_HTTP_RPC_ROOT_PATH").unwrap_or_else(|_| "/v1/rpc".to_string());
let http_endpoint = format!(
"http://{}:{}{}/{}",
http_host, http_port, rpc_root, endpoint_name
);
TransportType::Http(http_endpoint)
}
RequestPlaneMode::Tcp => {
let tcp_host = crate::utils::get_tcp_rpc_host_from_env();
let tcp_port = std::env::var("DYN_TCP_RPC_PORT")
.ok()
.and_then(|p| p.parse::<u16>().ok())
.unwrap_or(9999);
let tcp_endpoint = match context {
TransportContext::HealthCheck => {
// Health check uses simple host:port format
format!("{}:{}", tcp_host, tcp_port)
}
TransportContext::Discovery => {
// Discovery includes endpoint name for routing
format!("{}:{}/{}", tcp_host, tcp_port, endpoint_name)
}
};
TransportType::Tcp(tcp_endpoint)
}
RequestPlaneMode::Nats => TransportType::Nats(subject.to_string()),
}
}
// SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0
use super::*;
use crate::component::Component;
use crate::config::RequestPlaneMode;
use async_nats::service::Service as NatsService;
use async_nats::service::ServiceExt as _;
use derive_builder::Builder;
......@@ -9,8 +12,6 @@ use parking_lot::Mutex;
use std::collections::HashMap;
use std::sync::Arc;
use crate::component::Component;
pub use super::endpoint::EndpointStats;
type StatsHandlerRegistry = Arc<Mutex<HashMap<String, EndpointStatsHandler>>>;
......
......@@ -9,6 +9,7 @@ use figment::{
};
use serde::{Deserialize, Serialize};
use std::fmt;
use std::sync::OnceLock;
use validator::Validate;
/// Default system host for health and metrics endpoints
......@@ -468,6 +469,75 @@ pub fn use_local_timezone() -> bool {
env_is_truthy("DYN_LOG_USE_LOCAL_TZ")
}
/// Request plane transport mode configuration
///
/// This determines how requests are distributed from routers to workers:
/// - `Nats`: Use NATS for request distribution (default, legacy)
/// - `Http`: Use HTTP/2 for request distribution
/// - `Tcp`: Use raw TCP for request distribution with msgpack support
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum RequestPlaneMode {
/// Use NATS for request plane (default for backward compatibility)
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 {
Self::Nats
}
}
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
)),
}
}
}
/// Global cached request plane mode
static REQUEST_PLANE_MODE: OnceLock<RequestPlaneMode> = OnceLock::new();
impl RequestPlaneMode {
/// The cached request plane mode, initialized from `DYN_REQUEST_PLANE` environment variable
/// or defaulting to NATS if not set or invalid.
pub fn get() -> Self {
*REQUEST_PLANE_MODE.get_or_init(Self::from_env)
}
/// Get the request plane mode from environment variable (uncached)
/// Reads from `DYN_REQUEST_PLANE` environment variable.
pub fn from_env() -> Self {
std::env::var("DYN_REQUEST_PLANE")
.ok()
.and_then(|s| s.parse().ok())
.unwrap_or_default()
}
}
#[cfg(test)]
mod tests {
use super::*;
......@@ -671,4 +741,62 @@ mod tests {
assert!(!env_is_falsey("TEST_MISSING"));
});
}
#[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");
}
#[test]
fn test_request_plane_mode_default() {
assert_eq!(RequestPlaneMode::default(), RequestPlaneMode::Nats);
}
#[test]
fn test_request_plane_mode_get_cached() {
// Test that get() returns a consistent value
let mode1 = RequestPlaneMode::get();
let mode2 = RequestPlaneMode::get();
assert_eq!(mode1, mode2, "Cached mode should be consistent");
// Verify it's one of the valid modes
assert!(
matches!(
mode1,
RequestPlaneMode::Nats | RequestPlaneMode::Http | RequestPlaneMode::Tcp
),
"Mode should be a valid variant"
);
}
}
......@@ -403,7 +403,7 @@ mod tests {
namespace: "test".to_string(),
component: "comp1".to_string(),
endpoint: "ep1".to_string(),
transport: TransportType::NatsTcp("nats://localhost:4222".to_string()),
transport: TransportType::Nats("nats://localhost:4222".to_string()),
};
let instance = client.register(spec).await.unwrap();
......@@ -429,7 +429,7 @@ mod tests {
namespace: "ns1".to_string(),
component: "comp1".to_string(),
endpoint: "ep1".to_string(),
transport: TransportType::NatsTcp("nats://localhost:4222".to_string()),
transport: TransportType::Nats("nats://localhost:4222".to_string()),
};
client.register(spec1).await.unwrap();
......@@ -437,7 +437,7 @@ mod tests {
namespace: "ns1".to_string(),
component: "comp1".to_string(),
endpoint: "ep2".to_string(),
transport: TransportType::NatsTcp("nats://localhost:4222".to_string()),
transport: TransportType::Nats("nats://localhost:4222".to_string()),
};
client.register(spec2).await.unwrap();
......@@ -445,7 +445,7 @@ mod tests {
namespace: "ns2".to_string(),
component: "comp2".to_string(),
endpoint: "ep1".to_string(),
transport: TransportType::NatsTcp("nats://localhost:4222".to_string()),
transport: TransportType::Nats("nats://localhost:4222".to_string()),
};
client.register(spec3).await.unwrap();
......@@ -493,7 +493,7 @@ mod tests {
namespace: "test".to_string(),
component: "comp1".to_string(),
endpoint: "ep1".to_string(),
transport: TransportType::NatsTcp("nats://localhost:4222".to_string()),
transport: TransportType::Nats("nats://localhost:4222".to_string()),
};
client_clone.register(spec).await.unwrap();
});
......
......@@ -234,7 +234,7 @@ mod tests {
component: "comp1".to_string(),
endpoint: "ep1".to_string(),
instance_id: 123,
transport: TransportType::NatsTcp("nats://localhost:4222".to_string()),
transport: TransportType::Nats("nats://localhost:4222".to_string()),
});
metadata.register_endpoint(instance).unwrap();
......@@ -266,7 +266,7 @@ mod tests {
component: "comp1".to_string(),
endpoint: format!("ep{}", i),
instance_id: i,
transport: TransportType::NatsTcp("nats://localhost:4222".to_string()),
transport: TransportType::Nats("nats://localhost:4222".to_string()),
});
meta.register_endpoint(instance).unwrap();
})
......@@ -294,7 +294,7 @@ mod tests {
component: "comp1".to_string(),
endpoint: format!("ep{}", i),
instance_id: i,
transport: TransportType::NatsTcp("nats://localhost:4222".to_string()),
transport: TransportType::Nats("nats://localhost:4222".to_string()),
});
metadata.register_endpoint(instance).unwrap();
}
......
......@@ -218,7 +218,7 @@ mod tests {
namespace: "test-ns".to_string(),
component: "test-comp".to_string(),
endpoint: "test-ep".to_string(),
transport: crate::component::TransportType::NatsTcp("test-subject".to_string()),
transport: crate::component::TransportType::Nats("test-subject".to_string()),
};
let query = DiscoveryQuery::Endpoint {
......
......@@ -42,11 +42,12 @@ pub struct DistributedRuntime {
// local runtime
runtime: Runtime,
// we might consider a unifed transport manager here
// Unified transport manager
etcd_client: Option<transports::etcd::Client>,
nats_client: Option<transports::nats::Client>,
store: KeyValueStoreManager,
tcp_server: Arc<OnceCell<Arc<transports::tcp::server::TcpStreamServer>>>,
network_manager: Arc<OnceCell<Arc<crate::pipeline::network::manager::NetworkManager>>>,
system_status_server: Arc<OnceLock<Arc<system_status_server::SystemStatusServerInfo>>>,
// Service discovery client
......@@ -174,6 +175,7 @@ impl DistributedRuntime {
store,
nats_client,
tcp_server: Arc::new(OnceCell::new()),
network_manager: Arc::new(OnceCell::new()),
system_status_server: Arc::new(OnceLock::new()),
discovery_client,
discovery_metadata,
......@@ -337,6 +339,72 @@ impl DistributedRuntime {
.clone())
}
/// Get the network manager (lazy initialization)
///
/// The network manager consolidates all network configuration and provides
/// unified access to request plane servers and clients.
pub async fn network_manager(
&self,
) -> Result<Arc<crate::pipeline::network::manager::NetworkManager>> {
use crate::pipeline::network::manager::NetworkManager;
let manager = self
.network_manager
.get_or_try_init(async {
// Get NATS client if available
let nats_client = self.nats_client().map(|c| c.client().clone());
// NetworkManager handles all config reading and mode selection
anyhow::Ok(NetworkManager::new(
self.child_token(),
nats_client,
self.component_registry.clone(),
))
})
.await?;
Ok(manager.clone())
}
/// 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>>
{
let manager = self.network_manager().await?;
manager.server().await
}
/// DEPRECATED: Use network_manager().server() instead
#[deprecated(note = "Use request_plane_server() or network_manager().server() instead")]
pub async fn http_server(
&self,
) -> Result<Arc<crate::pipeline::network::ingress::http_endpoint::SharedHttpServer>> {
// For backward compatibility, try to downcast
let _server = self.request_plane_server().await?;
// This will only work if we're actually in HTTP mode
// For now, just return an error suggesting the new API
anyhow::bail!(
"http_server() is deprecated. Use request_plane_server() instead, which returns a trait object that works with all transport types."
)
}
/// DEPRECATED: Use network_manager().server() instead
#[deprecated(note = "Use request_plane_server() or network_manager().server() instead")]
pub async fn shared_tcp_server(
&self,
) -> Result<Arc<crate::pipeline::network::ingress::shared_tcp_endpoint::SharedTcpServer>> {
// For backward compatibility, try to downcast
let _server = self.request_plane_server().await?;
// This will only work if we're actually in TCP mode
// For now, just return an error suggesting the new API
anyhow::bail!(
"shared_tcp_server() is deprecated. Use request_plane_server() instead, which returns a trait object that works with all transport types."
)
}
pub fn nats_client(&self) -> Option<&nats::Client> {
self.nats_client.as_ref()
}
......
This diff is collapsed.
This diff is collapsed.
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment