Plumb through option for configuring DNS resolver (#1501)

Author: Sushisource

temporalio/bridge/client.py

@@ -57,6 +57,15 @@ class ClientHttpConnectProxyConfig:
     basic_auth: tuple[str, str] | None
 
 
+@dataclass
+class ClientDnsLoadBalancingConfig:
+    """Python representation of the Rust struct for configuring DNS load
+    balancing.
+    """
+
+    resolution_interval_millis: int
+
+
 @dataclass
 class ClientConfig:
     """Python representation of the Rust struct for configuring the client."""
@@ -71,6 +80,7 @@ class ClientConfig:
     client_name: str
     client_version: str
     http_connect_proxy_config: ClientHttpConnectProxyConfig | None
+    dns_load_balancing_config: ClientDnsLoadBalancingConfig | None
 
 
 @dataclass

temporalio/bridge/src/client.rs

@@ -11,6 +11,7 @@ use temporalio_client::{
     ClientKeepAliveOptions as CoreClientKeepAliveConfig, Connection, ConnectionOptions,
     DnsLoadBalancingOptions, HttpConnectProxyOptions, RetryOptions,
 };
+use tracing::warn;
 use url::Url;
 
 use crate::runtime;
@@ -35,6 +36,7 @@ pub struct ClientConfig {
     retry_config: Option<ClientRetryConfig>,
     keep_alive_config: Option<ClientKeepAliveConfig>,
     http_connect_proxy_config: Option<ClientHttpConnectProxyConfig>,
+    dns_load_balancing_config: Option<ClientDnsLoadBalancingConfig>,
 }
 
 #[derive(FromPyObject)]
@@ -67,6 +69,11 @@ struct ClientHttpConnectProxyConfig {
     pub basic_auth: Option<(String, String)>,
 }
 
+#[derive(FromPyObject)]
+struct ClientDnsLoadBalancingConfig {
+    pub resolution_interval_millis: u64,
+}
+
 #[derive(FromPyObject)]
 pub(crate) struct RpcCall {
     pub(crate) rpc: String,
@@ -236,6 +243,15 @@ impl ClientConfig {
     ) -> PyResult<ConnectionOptions> {
         let (ascii_headers, binary_headers) = partition_headers(self.metadata);
         let has_proxy = self.http_connect_proxy_config.is_some();
+        // Core rejects DNS load balancing alongside an HTTP CONNECT proxy, so
+        // suppress DNS LB whenever a proxy is configured to keep the
+        // pre-existing behavior even if a caller leaves the default.
+        let dns_load_balancing = if has_proxy {
+            warn!("Disabling DNS load balancing because http_connect_proxy_config is set");
+            None
+        } else {
+            self.dns_load_balancing_config.map(Into::into)
+        };
         let conn_opts = ConnectionOptions::new(
             Url::parse(&self.target_url)
                 .map_err(|err| PyValueError::new_err(format!("invalid target URL: {err}")))?,
@@ -249,11 +265,7 @@ impl ClientConfig {
         )
         .keep_alive(self.keep_alive_config.map(Into::into))
         .maybe_http_connect_proxy(self.http_connect_proxy_config.map(Into::into))
-        .dns_load_balancing(if has_proxy {
-            None
-        } else {
-            Some(DnsLoadBalancingOptions::default())
-        })
+        .dns_load_balancing(dns_load_balancing)
         .headers(ascii_headers)
         .binary_headers(binary_headers)
         .maybe_api_key(self.api_key)
@@ -322,3 +334,11 @@ impl From<ClientHttpConnectProxyConfig> for HttpConnectProxyOptions {
         }
     }
 }
+
+impl From<ClientDnsLoadBalancingConfig> for DnsLoadBalancingOptions {
+    fn from(conf: ClientDnsLoadBalancingConfig) -> Self {
+        let mut opts = DnsLoadBalancingOptions::default();
+        opts.resolution_interval = Duration::from_millis(conf.resolution_interval_millis);
+        opts
+    }
+}

temporalio/client.py

@@ -74,6 +74,7 @@
 )
 from temporalio.service import (
     ConnectConfig,
+    DnsLoadBalancingConfig,
     HttpConnectProxyConfig,
     KeepAliveConfig,
     RetryConfig,
@@ -139,6 +140,8 @@ async def connect(
         lazy: bool = False,
         runtime: temporalio.runtime.Runtime | None = None,
         http_connect_proxy_config: HttpConnectProxyConfig | None = None,
+        dns_load_balancing_config: DnsLoadBalancingConfig
+        | None = DnsLoadBalancingConfig.default,
         header_codec_behavior: HeaderCodecBehavior = HeaderCodecBehavior.NO_CODEC,
     ) -> Self:
         """Connect to a Temporal server.
@@ -194,6 +197,11 @@ async def connect(
                 used for workers.
             runtime: The runtime for this client, or the default if unset.
             http_connect_proxy_config: Configuration for HTTP CONNECT proxy.
+            dns_load_balancing_config: DNS load balancing configuration for the
+                client connection. Default is to re-resolve DNS every 30s. Can
+                be set to ``None`` to disable. Silently disabled when
+                ``http_connect_proxy_config`` is set, since the two are mutually
+                exclusive.
             header_codec_behavior: Encoding behavior for headers sent by the client.
         """
         connect_config = temporalio.service.ConnectConfig(
@@ -207,6 +215,7 @@ async def connect(
             lazy=lazy,
             runtime=runtime,
             http_connect_proxy_config=http_connect_proxy_config,
+            dns_load_balancing_config=dns_load_balancing_config,
         )
 
         def make_lambda(
@@ -2873,6 +2882,7 @@ class ClientConnectConfig(TypedDict, total=False):
     lazy: bool
     runtime: temporalio.runtime.Runtime | None
     http_connect_proxy_config: HttpConnectProxyConfig | None
+    dns_load_balancing_config: DnsLoadBalancingConfig | None
     header_codec_behavior: HeaderCodecBehavior
 
 
@@ -9774,6 +9784,7 @@ async def connect(
         lazy: bool = False,
         runtime: temporalio.runtime.Runtime | None = None,
         http_connect_proxy_config: HttpConnectProxyConfig | None = None,
+        dns_load_balancing_config: DnsLoadBalancingConfig | None = None,
     ) -> CloudOperationsClient:
         """Connect to a Temporal Cloud Operations API.
 
@@ -9810,6 +9821,10 @@ async def connect(
                 used for workers.
             runtime: The runtime for this client, or the default if unset.
             http_connect_proxy_config: Configuration for HTTP CONNECT proxy.
+            dns_load_balancing_config: DNS load balancing configuration for the
+                client connection. Default is disabled. Silently disabled when
+                ``http_connect_proxy_config`` is set, since the two are mutually
+                exclusive.
         """
         # Add version if given
         if version:
@@ -9826,6 +9841,7 @@ async def connect(
             lazy=lazy,
             runtime=runtime,
             http_connect_proxy_config=http_connect_proxy_config,
+            dns_load_balancing_config=dns_load_balancing_config,
         )
         return CloudOperationsClient(
             await temporalio.service.ServiceClient.connect(connect_config)

temporalio/service.py

@@ -132,6 +132,32 @@ def _to_bridge_config(
         )
 
 
+@dataclass(frozen=True)
+class DnsLoadBalancingConfig:
+    """DNS load balancing configuration for client connections.
+
+    When enabled, Core periodically re-resolves the target host's DNS records
+    and round-robins requests across the resolved addresses. Cannot be used
+    together with :py:class:`HttpConnectProxyConfig` -- DNS load balancing is
+    silently disabled when an HTTP CONNECT proxy is configured.
+    """
+
+    resolution_interval_millis: int = 30000
+    """How often to re-resolve DNS, in milliseconds."""
+    default: ClassVar[DnsLoadBalancingConfig]
+    """Default DNS load balancing config."""
+
+    def _to_bridge_config(
+        self,
+    ) -> temporalio.bridge.client.ClientDnsLoadBalancingConfig:
+        return temporalio.bridge.client.ClientDnsLoadBalancingConfig(
+            resolution_interval_millis=self.resolution_interval_millis,
+        )
+
+
+DnsLoadBalancingConfig.default = DnsLoadBalancingConfig()
+
+
 @dataclass
 class ConnectConfig:
     """Config for connecting to the server."""
@@ -146,6 +172,9 @@ class ConnectConfig:
     lazy: bool = False
     runtime: temporalio.runtime.Runtime | None = None
     http_connect_proxy_config: HttpConnectProxyConfig | None = None
+    dns_load_balancing_config: DnsLoadBalancingConfig | None = (
+        DnsLoadBalancingConfig.default
+    )
 
     def __post_init__(self) -> None:
         """Set extra defaults on unset properties."""
@@ -203,6 +232,11 @@ def _to_bridge_config(self) -> temporalio.bridge.client.ClientConfig:
                 if self.http_connect_proxy_config
                 else None
             ),
+            dns_load_balancing_config=(
+                self.dns_load_balancing_config._to_bridge_config()
+                if self.dns_load_balancing_config
+                else None
+            ),
         )
 
 

tests/test_service.py

@@ -219,6 +219,29 @@ def test_connect_config_tls_explicit_config_preserved():
     assert config.tls == tls_config
 
 
+def test_connect_config_dns_load_balancing_custom():
+    """Custom DnsLoadBalancingConfig is forwarded to the bridge unchanged."""
+    config = temporalio.service.ConnectConfig(
+        target_host="localhost:7233",
+        dns_load_balancing_config=temporalio.service.DnsLoadBalancingConfig(
+            resolution_interval_millis=5000,
+        ),
+    )
+    bridge_config = config._to_bridge_config()
+    assert bridge_config.dns_load_balancing_config is not None
+    assert bridge_config.dns_load_balancing_config.resolution_interval_millis == 5000
+
+
+def test_connect_config_dns_load_balancing_disabled():
+    """Setting dns_load_balancing_config=None forwards None to the bridge."""
+    config = temporalio.service.ConnectConfig(
+        target_host="localhost:7233",
+        dns_load_balancing_config=None,
+    )
+    bridge_config = config._to_bridge_config()
+    assert bridge_config.dns_load_balancing_config is None
+
+
 async def test_rpc_execution_not_unknown(client: Client):
     """
     Execute each rpc method and expect a failure, but ensure the failure is not that the rpc method is unknown