diff --git a/crates/client/Cargo.toml b/crates/client/Cargo.toml index 7b538a67d..22a6f400f 100644 --- a/crates/client/Cargo.toml +++ b/crates/client/Cargo.toml @@ -35,7 +35,7 @@ hyper-util = "0.1.16" opentelemetry = { workspace = true, features = ["metrics"], optional = true } parking_lot = "0.12" thiserror = { workspace = true } -tokio = "1.47" +tokio = { version = "1.47", features = ["net", "time"] } tonic = { workspace = true, features = ["tls-ring", "tls-native-roots"] } tower = { version = "0.5", features = ["util"] } tracing = "0.1" diff --git a/crates/client/src/dns.rs b/crates/client/src/dns.rs new file mode 100644 index 000000000..e90e58283 --- /dev/null +++ b/crates/client/src/dns.rs @@ -0,0 +1,306 @@ +use crate::{ + add_tls_to_channel, + errors::ClientConnectError, + options_structs::{ + ClientKeepAliveOptions, ConnectionOptions, DnsLoadBalancingOptions, TlsOptions, + }, +}; +use http::Uri; +use std::{collections::HashSet, net::SocketAddr, sync::Arc, time::Duration}; +use tokio::sync::mpsc; +use tonic::transport::{Channel, Endpoint, channel::Change}; +use url::Url; + +/// Validates DNS load balancing configuration and returns the options if DNS LB should be used. +/// +/// Returns `Err` if `dns_load_balancing` is set alongside `service_override` or +/// `http_connect_proxy`. Returns `Ok(None)` if DNS LB is disabled or the target is an IP literal. +pub(crate) fn validate_and_get_dns_lb( + options: &ConnectionOptions, +) -> Result, ClientConnectError> { + let Some(dns_opts) = options.dns_load_balancing.as_ref() else { + return Ok(None); + }; + + if options.service_override.is_some() { + return Err(ClientConnectError::InvalidConfig( + "dns_load_balancing cannot be used with service_override".to_owned(), + )); + } + if options.http_connect_proxy.is_some() { + return Err(ClientConnectError::InvalidConfig( + "dns_load_balancing cannot be used with http_connect_proxy".to_owned(), + )); + } + + let host = options + .target + .host() + .ok_or_else(|| ClientConnectError::InvalidConfig("target URL has no host".to_owned()))?; + + match host { + url::Host::Domain("localhost") => Ok(None), + url::Host::Domain(_) => Ok(Some(dns_opts)), + url::Host::Ipv4(_) | url::Host::Ipv6(_) => Ok(None), + } +} + +async fn resolve_host(host: &str, port: u16) -> Result, std::io::Error> { + tokio::net::lookup_host(format!("{host}:{port}")) + .await + .map(|addrs| addrs.collect()) +} + +fn endpoint_uri(addr: SocketAddr, scheme: &str) -> String { + match addr { + SocketAddr::V4(v4) => format!("{scheme}://{v4}"), + SocketAddr::V6(v6) => format!("{scheme}://[{}]:{}", v6.ip(), v6.port()), + } +} + +async fn build_endpoint( + addr: SocketAddr, + original_host: &str, + scheme: &str, + tls_options: Option<&TlsOptions>, + keep_alive: Option<&ClientKeepAliveOptions>, + override_origin: Option<&Uri>, +) -> Result { + let uri = endpoint_uri(addr, scheme); + let channel = Channel::from_shared(uri)?; + + // When connecting to an IP with TLS, SNI must use the original hostname. + let tls_for_ip = tls_options.map(|tls| { + if tls.domain.is_some() { + tls.clone() + } else { + let mut patched = tls.clone(); + patched.domain = Some(original_host.to_owned()); + patched + } + }); + let channel = add_tls_to_channel(tls_for_ip.as_ref().or(tls_options), channel).await?; + + let channel = if let Some(keep_alive) = keep_alive { + channel + .keep_alive_while_idle(true) + .http2_keep_alive_interval(keep_alive.interval) + .keep_alive_timeout(keep_alive.timeout) + } else { + channel + }; + + let channel = if let Some(origin) = override_origin.cloned() { + channel.origin(origin) + } else { + channel + }; + + Ok(channel) +} + +/// Creates a balanced channel backed by all DNS-resolved addresses for the target. +pub(crate) async fn create_balanced_channel( + options: &ConnectionOptions, +) -> Result<(Channel, mpsc::Sender>), ClientConnectError> { + let host = options + .target + .host_str() + .ok_or_else(|| ClientConnectError::InvalidConfig("target URL has no host".to_owned()))?; + let port = options.target.port_or_known_default().unwrap_or(7233); + let scheme = options.target.scheme(); + + let addrs = resolve_host(host, port).await.map_err(|source| { + ClientConnectError::DnsResolutionError { + host: host.to_owned(), + source, + } + })?; + if addrs.is_empty() { + return Err(ClientConnectError::DnsResolutionError { + host: host.to_owned(), + source: std::io::Error::new( + std::io::ErrorKind::NotFound, + "DNS resolution returned no addresses", + ), + }); + } + + let (channel, sender) = Channel::balance_channel(addrs.len()); + + for addr in addrs { + let endpoint = build_endpoint( + addr, + host, + scheme, + options.tls_options.as_ref(), + options.keep_alive.as_ref(), + options.override_origin.as_ref(), + ) + .await?; + // Unbounded-ish send into the freshly-created channel; can't realistically fail. + let _ = sender.send(Change::Insert(addr, endpoint)).await; + } + + Ok((channel, sender)) +} + +/// Handle that aborts the DNS re-resolution task when dropped. +pub(crate) struct DnsReresolutionHandle { + abort_handle: tokio::task::AbortHandle, +} + +impl Drop for DnsReresolutionHandle { + fn drop(&mut self) { + self.abort_handle.abort(); + } +} + +/// Spawns a background task that periodically re-resolves DNS and updates the balanced channel. +pub(crate) fn spawn_dns_reresolution( + sender: mpsc::Sender>, + target: Url, + tls_options: Option, + keep_alive: Option, + override_origin: Option, + resolution_interval: Duration, +) -> Arc { + let host = target.host_str().unwrap_or("").to_owned(); + let port = target.port_or_known_default().unwrap_or(7233); + let scheme = target.scheme().to_owned(); + + let handle = tokio::spawn(async move { + let mut current_addrs: HashSet = HashSet::new(); + // Populate initial set from the channel we already seeded + if let Ok(initial) = resolve_host(&host, port).await { + current_addrs.extend(initial); + } + + loop { + tokio::time::sleep(resolution_interval).await; + + let new_addrs = match resolve_host(&host, port).await { + Ok(addrs) => addrs.into_iter().collect::>(), + Err(e) => { + warn!( + host = %host, + error = %e, + "DNS re-resolution failed, keeping existing endpoints" + ); + continue; + } + }; + + if new_addrs.is_empty() { + warn!( + host = %host, + "DNS re-resolution returned no addresses, keeping existing endpoints" + ); + continue; + } + + // Remove stale endpoints + for addr in current_addrs.difference(&new_addrs) { + if sender.send(Change::Remove(*addr)).await.is_err() { + return; + } + } + + // Add new endpoints + for addr in new_addrs.difference(¤t_addrs) { + match build_endpoint( + *addr, + &host, + &scheme, + tls_options.as_ref(), + keep_alive.as_ref(), + override_origin.as_ref(), + ) + .await + { + Ok(endpoint) => { + if sender.send(Change::Insert(*addr, endpoint)).await.is_err() { + return; + } + } + Err(e) => { + warn!( + addr = %addr, + error = %e, + "Failed to build endpoint for resolved address" + ); + } + } + } + + current_addrs = new_addrs; + } + }); + + Arc::new(DnsReresolutionHandle { + abort_handle: handle.abort_handle(), + }) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn ip_v4_target_returns_none() { + let opts = ConnectionOptions::new(Url::parse("http://1.2.3.4:7233").unwrap()).build(); + assert!(validate_and_get_dns_lb(&opts).unwrap().is_none()); + } + + #[test] + fn ip_v6_target_returns_none() { + let opts = ConnectionOptions::new(Url::parse("http://[::1]:7233").unwrap()).build(); + assert!(validate_and_get_dns_lb(&opts).unwrap().is_none()); + } + + #[test] + fn domain_target_returns_some() { + let opts = + ConnectionOptions::new(Url::parse("http://temporal.example.com:7233").unwrap()).build(); + assert!(validate_and_get_dns_lb(&opts).unwrap().is_some()); + } + + #[test] + fn disabled_returns_none() { + let opts = ConnectionOptions::new(Url::parse("http://temporal.example.com:7233").unwrap()) + .dns_load_balancing(None) + .build(); + assert!(validate_and_get_dns_lb(&opts).unwrap().is_none()); + } + + #[test] + fn service_override_with_dns_lb_is_error() { + use crate::callback_based::CallbackBasedGrpcService; + + let svc = CallbackBasedGrpcService { + callback: Arc::new(|_| Box::pin(async { unreachable!() })), + }; + let opts = ConnectionOptions::new(Url::parse("http://temporal.example.com:7233").unwrap()) + .service_override(svc) + .build(); + assert!(validate_and_get_dns_lb(&opts).is_err()); + } + + #[test] + fn localhost_returns_none() { + let opts = ConnectionOptions::new(Url::parse("http://localhost:7233").unwrap()).build(); + assert!(validate_and_get_dns_lb(&opts).unwrap().is_none()); + } + + #[test] + fn endpoint_uri_v4() { + let addr: SocketAddr = "1.2.3.4:7233".parse().unwrap(); + assert_eq!(endpoint_uri(addr, "https"), "https://1.2.3.4:7233"); + } + + #[test] + fn endpoint_uri_v6() { + let addr: SocketAddr = "[::1]:7233".parse().unwrap(); + assert_eq!(endpoint_uri(addr, "https"), "https://[::1]:7233"); + } +} diff --git a/crates/client/src/errors.rs b/crates/client/src/errors.rs index 02b297bfa..51c52564f 100644 --- a/crates/client/src/errors.rs +++ b/crates/client/src/errors.rs @@ -24,6 +24,18 @@ pub enum ClientConnectError { /// server capabilities / verify server is responding. #[error("`get_system_info` call error after connection: {0:?}")] SystemInfoCallError(tonic::Status), + /// DNS resolution failed when attempting load-balanced connection. + #[error("DNS resolution error for '{host}': {source}")] + DnsResolutionError { + /// The host that failed to resolve. + host: String, + /// The underlying IO error. + #[source] + source: std::io::Error, + }, + /// Invalid client configuration. + #[error("Invalid client configuration: {0}")] + InvalidConfig(String), } /// Errors thrown when a gRPC metadata header is invalid. diff --git a/crates/client/src/grpc.rs b/crates/client/src/grpc.rs index f67524563..bea678741 100644 --- a/crates/client/src/grpc.rs +++ b/crates/client/src/grpc.rs @@ -2094,6 +2094,7 @@ mod tests { let opts = ConnectionOptions::new(url::Url::parse("http://localhost:7233").unwrap()) .skip_get_system_info(true) .service_override(service_override) + .dns_load_balancing(None) .build(); let mut connection = crate::Connection::connect(opts).await.unwrap(); diff --git a/crates/client/src/lib.rs b/crates/client/src/lib.rs index 75d6dbfde..c99f7b06d 100644 --- a/crates/client/src/lib.rs +++ b/crates/client/src/lib.rs @@ -9,6 +9,7 @@ extern crate tracing; mod async_activity_handle; pub mod callback_based; +mod dns; /// Configuration loading from environment variables and TOML files. #[cfg(feature = "envconfig")] pub mod envconfig; @@ -144,17 +145,42 @@ struct ConnectionInner { /// Capabilities as read from the `get_system_info` RPC call made on client connection capabilities: Option, workers: Arc, + _dns_task: Option>, } impl Connection { /// Connect to a Temporal service. pub async fn connect(options: ConnectionOptions) -> Result { - let service = if let Some(service_override) = options.service_override { - GrpcMetricSvc { - inner: ChannelOrGrpcOverride::GrpcOverride(service_override), - metrics: options.metrics_meter.clone().map(MetricsContext::new), - disable_errcode_label: options.disable_error_code_metric_tags, - } + let dns_lb_opts = dns::validate_and_get_dns_lb(&options)?.cloned(); + let (service, dns_task) = if let Some(service_override) = options.service_override { + ( + GrpcMetricSvc { + inner: ChannelOrGrpcOverride::GrpcOverride(service_override), + metrics: options.metrics_meter.clone().map(MetricsContext::new), + disable_errcode_label: options.disable_error_code_metric_tags, + }, + None, + ) + } else if let Some(dns_opts) = &dns_lb_opts { + let (channel, sender) = dns::create_balanced_channel(&options).await?; + let handle = dns::spawn_dns_reresolution( + sender, + options.target.clone(), + options.tls_options.clone(), + options.keep_alive.clone(), + options.override_origin.clone(), + dns_opts.resolution_interval, + ); + ( + ServiceBuilder::new() + .layer_fn(move |channel| GrpcMetricSvc { + inner: ChannelOrGrpcOverride::Channel(channel), + metrics: options.metrics_meter.clone().map(MetricsContext::new), + disable_errcode_label: options.disable_error_code_metric_tags, + }) + .service(channel), + Some(handle), + ) } else { let channel = Channel::from_shared(options.target.to_string())?; let channel = add_tls_to_channel(options.tls_options.as_ref(), channel).await?; @@ -177,13 +203,16 @@ impl Connection { } else { channel.connect().await? }; - ServiceBuilder::new() - .layer_fn(move |channel| GrpcMetricSvc { - inner: ChannelOrGrpcOverride::Channel(channel), - metrics: options.metrics_meter.clone().map(MetricsContext::new), - disable_errcode_label: options.disable_error_code_metric_tags, - }) - .service(channel) + ( + ServiceBuilder::new() + .layer_fn(move |channel| GrpcMetricSvc { + inner: ChannelOrGrpcOverride::Channel(channel), + metrics: options.metrics_meter.clone().map(MetricsContext::new), + disable_errcode_label: options.disable_error_code_metric_tags, + }) + .service(channel), + None, + ) }; let headers = Arc::new(RwLock::new(ClientHeaders { @@ -225,6 +254,7 @@ impl Connection { client_version: options.client_version, capabilities, workers: Arc::new(ClientWorkerSet::new()), + _dns_task: dns_task, }), }) } diff --git a/crates/client/src/options_structs.rs b/crates/client/src/options_structs.rs index 6153522db..a05923c98 100644 --- a/crates/client/src/options_structs.rs +++ b/crates/client/src/options_structs.rs @@ -66,6 +66,12 @@ pub struct ConnectionOptions { pub binary_headers: Option>>, /// HTTP CONNECT proxy to use for this client. pub http_connect_proxy: Option, + /// If set, DNS-based load balancing is enabled. When the target is a hostname (not an IP + /// literal), DNS is resolved to all addresses and requests are distributed across them. + /// Incompatible with `service_override` and `http_connect_proxy`. Setting either in addition + /// to this field is an error. Set to `None` to disable. + #[builder(required, default = Some(DnsLoadBalancingOptions::default()))] + pub dns_load_balancing: Option, /// If set true, error code labels will not be included on request failure metrics. #[builder(default)] pub disable_error_code_metric_tags: bool, @@ -166,6 +172,22 @@ impl Default for ClientKeepAliveOptions { } } +/// Options for DNS-based load balancing. +#[derive(Clone, Debug)] +#[non_exhaustive] +pub struct DnsLoadBalancingOptions { + /// How often to re-resolve DNS. Defaults to 30 seconds. + pub resolution_interval: Duration, +} + +impl Default for DnsLoadBalancingOptions { + fn default() -> Self { + Self { + resolution_interval: Duration::from_secs(30), + } + } +} + impl std::fmt::Debug for ClientTlsOptions { // Intentionally omit details here since they could leak a key if ever printed fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { diff --git a/crates/sdk-core-c-bridge/src/client.rs b/crates/sdk-core-c-bridge/src/client.rs index b5b988c88..3073fb06a 100644 --- a/crates/sdk-core-c-bridge/src/client.rs +++ b/crates/sdk-core-c-bridge/src/client.rs @@ -130,6 +130,7 @@ pub extern "C" fn temporal_core_client_connect( cb, options.grpc_override_callback_user_data, )); + connection_options.dns_load_balancing = None; } // Spawn async call let user_data = UserDataHandle(user_data); @@ -1272,7 +1273,12 @@ impl TryFrom<&ConnectionOptions> for temporalio_client::ConnectionOptions { .maybe_headers(headers) .maybe_binary_headers(binary_headers) .maybe_api_key(api_key) - .maybe_http_connect_proxy(http_connect_proxy) + .maybe_http_connect_proxy(http_connect_proxy.clone()) + .dns_load_balancing(if http_connect_proxy.is_some() { + None + } else { + Some(temporalio_client::DnsLoadBalancingOptions::default()) + }) .maybe_tls_options(tls_cfg) .build(), ) diff --git a/crates/sdk-core/tests/common/fake_grpc_server.rs b/crates/sdk-core/tests/common/fake_grpc_server.rs index 86bc1c297..da0da139a 100644 --- a/crates/sdk-core/tests/common/fake_grpc_server.rs +++ b/crates/sdk-core/tests/common/fake_grpc_server.rs @@ -70,7 +70,7 @@ where let (shutdown_tx, shutdown_rx) = oneshot::channel::<()>(); let (header_tx, header_rx) = tokio::sync::mpsc::unbounded_channel(); - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let listener = TcpListener::bind("[::]:0").await.unwrap(); let addr = listener.local_addr().unwrap(); let server_handle = tokio::spawn(async move { diff --git a/crates/sdk-core/tests/integ_tests/client_tests.rs b/crates/sdk-core/tests/integ_tests/client_tests.rs index c0b8c1260..ff0cf92b7 100644 --- a/crates/sdk-core/tests/integ_tests/client_tests.rs +++ b/crates/sdk-core/tests/integ_tests/client_tests.rs @@ -255,7 +255,7 @@ async fn namespace_header_attached_to_relevant_calls() { let (shutdown_tx, shutdown_rx) = oneshot::channel::<()>(); let (header_tx, mut header_rx) = tokio::sync::mpsc::unbounded_channel(); - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let listener = TcpListener::bind("[::]:0").await.unwrap(); let addr = listener.local_addr().unwrap(); let server_handle = tokio::spawn(async move { @@ -394,6 +394,7 @@ async fn http_proxy() { target_addr: tcp_proxy_addr.to_string(), basic_auth: None, }); + opts.dns_load_balancing = None; let connection = Connection::connect(opts.clone()).await.unwrap(); let client_opts = temporalio_client::ClientOptions::new("my-namespace").build(); let proxied_client = temporalio_client::Client::new(connection, client_opts).unwrap(); @@ -422,6 +423,7 @@ async fn http_proxy() { target_addr: format!("unix:{}", sock_path.to_str().unwrap()), basic_auth: None, }); + opts.dns_load_balancing = None; let connection = Connection::connect(opts.clone()).await.unwrap(); let client_opts = temporalio_client::ClientOptions::new("my-namespace").build(); let proxied_client = temporalio_client::Client::new(connection, client_opts).unwrap();