This is an automated email from the ASF dual-hosted git repository.
alamb pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-rs-object-store.git
The following commit(s) were added to refs/heads/main by this push:
new 699c4fc feat: Support custom DNS resolver (#728)
699c4fc is described below
commit 699c4fc675df8bc8d8fbd756c4cc426186513f97
Author: Koen Denecker <[email protected]>
AuthorDate: Wed Aug 12 17:18:31 2026 +0200
feat: Support custom DNS resolver (#728)
* feat: Support custom DNS resolver
* add local `DnsResolver` trait
* refactor for wasm32 feature and reqwest as optional
---
src/client/dns.rs | 148 +++++++++++++++++++++++++++++++++++++++++++++---------
src/client/mod.rs | 79 +++++++++++++++++++++++++++--
2 files changed, 200 insertions(+), 27 deletions(-)
diff --git a/src/client/dns.rs b/src/client/dns.rs
index b1abe59..7a4cc85 100644
--- a/src/client/dns.rs
+++ b/src/client/dns.rs
@@ -15,36 +15,136 @@
// specific language governing permissions and limitations
// under the License.
-use std::net::ToSocketAddrs;
+//! Customizable DNS resolution for remote object stores
-use rand::prelude::SliceRandom;
-use reqwest::dns::{Addrs, Name, Resolve, Resolving};
-use tokio::task::JoinSet;
+use std::fmt::Debug;
+use std::future::Future;
+use std::net::IpAddr;
+use std::pin::Pin;
-type DynErr = Box<dyn std::error::Error + Send + Sync>;
+/// Error returned by a [`DnsResolver`]
+pub type DnsError = Box<dyn std::error::Error + Send + Sync>;
-#[derive(Debug)]
-pub(crate) struct ShuffleResolver;
+/// Future returned by [`DnsResolver::resolve`]
+// NOTE: the use cases requiring SocketAddr over IpAddr (i.e.,
resolver-supplied ports
+// and IPv6 scope IDs) do not apply to object_store, where the port is always
+// determined by the endpoint URL/scheme and endpoints are never link-local.
+pub type DnsFuture = Pin<Box<dyn Future<Output = Result<Vec<IpAddr>,
DnsError>> + Send>>;
-impl Resolve for ShuffleResolver {
- fn resolve(&self, name: Name) -> Resolving {
- Box::pin(async move {
- // use `JoinSet` to propagate cancelation to tasks that haven't
started running yet.
- let mut tasks = JoinSet::new();
- tasks.spawn_blocking(move || {
- let it = (name.as_str(), 0).to_socket_addrs()?;
- let mut addrs = it.collect::<Vec<_>>();
+/// A custom DNS resolver used when establishing connections to remote object
stores
+///
+/// This can be used to implement custom resolution logic such as caching,
+/// address shuffling, or split-horizon DNS, independent of the underlying
+/// HTTP transport.
+///
+/// Configure via [`ClientOptions::with_dns_resolver`]. The built-in
+/// reqwest-based transport honors this automatically; custom
+/// [`HttpConnector`] implementations should retrieve it via
+/// [`ClientOptions::dns_resolver`] and apply it themselves.
+///
+/// [`ClientOptions::with_dns_resolver`]:
crate::ClientOptions::with_dns_resolver
+/// [`ClientOptions::dns_resolver`]: crate::ClientOptions::dns_resolver
+/// [`HttpConnector`]: crate::client::HttpConnector
+pub trait DnsResolver: Debug + Send + Sync {
+ /// Resolve `host` to one or more IP addresses
+ ///
+ /// The returned addresses are tried in order until a connection succeeds,
+ /// so implementations are responsible for any ordering they require, e.g.
+ /// shuffling or interleaving of address families. The port is determined
+ /// by the transport from the URL, not by the resolver.
+ fn resolve(&self, host: &str) -> DnsFuture;
+}
+
+#[cfg(feature = "reqwest")]
+mod reqwest_impl {
+ use super::{DnsError, DnsFuture, DnsResolver};
+ use rand::prelude::SliceRandom;
+ use std::net::{SocketAddr, ToSocketAddrs};
+ use std::sync::Arc;
+ use tokio::task::JoinSet;
+
+ /// Adapts a [`DnsResolver`] to [`reqwest::dns::Resolve`]
+ ///
+ /// This is deliberately private: it is the only place where [`reqwest`]'s
+ /// resolver API appears, keeping it out of this crate's public interface.
+ pub(crate) struct ReqwestResolver(pub(crate) Arc<dyn DnsResolver>);
+
+ impl reqwest::dns::Resolve for ReqwestResolver {
+ fn resolve(&self, name: reqwest::dns::Name) -> reqwest::dns::Resolving
{
+ let resolver = Arc::clone(&self.0);
+ let host = name.as_str().to_string();
+ Box::pin(async move {
+ let ips = resolver.resolve(&host).await?;
+ // Port 0 is a placeholder: reqwest documents that it is
replaced
+ // by the port from the URL or the scheme's conventional port
+ let addrs: reqwest::dns::Addrs =
+ Box::new(ips.into_iter().map(|ip| SocketAddr::new(ip, 0)));
+ Ok(addrs)
+ })
+ }
+ }
+
+ /// The built-in shuffling [`DnsResolver`], randomizing the order of the
returned
+ /// addresses to spread load across servers, see
[`ClientConfigKey::RandomizeAddresses`]
+ ///
+ /// [`ClientConfigKey::RandomizeAddresses`]:
crate::ClientConfigKey::RandomizeAddresses
+ #[derive(Debug)]
+ pub(crate) struct ShuffleResolver;
+
+ impl DnsResolver for ShuffleResolver {
+ fn resolve(&self, host: &str) -> DnsFuture {
+ let host = host.to_string();
+ Box::pin(async move {
+ // use `JoinSet` to propagate cancellation to tasks that
haven't started running yet.
+ let mut tasks = JoinSet::new();
+ tasks.spawn_blocking(move || {
+ let it = (host.as_str(), 0).to_socket_addrs()?;
+ let mut addrs = it.map(|addr|
addr.ip()).collect::<Vec<_>>();
+ addrs.shuffle(&mut rand::rng());
+ Ok(addrs)
+ });
+ tasks
+ .join_next()
+ .await
+ .expect("spawned one task")
+ .map_err(|err| Box::new(err) as DnsError)?
+ })
+ }
+ }
+}
- addrs.shuffle(&mut rand::rng());
+#[cfg(feature = "reqwest")]
+pub(crate) use reqwest_impl::{ReqwestResolver, ShuffleResolver};
- Ok(Box::new(addrs.into_iter()) as Addrs)
- });
+#[cfg(all(test, feature = "reqwest"))]
+mod tests {
+ use super::*;
+
+ #[tokio::test]
+ async fn shuffle_resolver_resolves_localhost() {
+ let ips = ShuffleResolver.resolve("localhost").await.unwrap();
+ assert!(!ips.is_empty());
+ assert!(ips.iter().all(|ip| ip.is_loopback()));
+ }
+
+ #[derive(Debug)]
+ struct FailingResolver;
+
+ impl DnsResolver for FailingResolver {
+ fn resolve(&self, _host: &str) -> DnsFuture {
+ Box::pin(async { Err("boom".into()) })
+ }
+ }
- tasks
- .join_next()
- .await
- .expect("spawned on task")
- .map_err(|err| Box::new(err) as DynErr)?
- })
+ #[tokio::test]
+ async fn adapter_propagates_errors() {
+ use reqwest::dns::Resolve;
+ use std::sync::Arc;
+ let adapter = ReqwestResolver(Arc::new(FailingResolver));
+ let err = match adapter.resolve("localhost".parse().unwrap()).await {
+ Ok(_) => panic!("expected resolution to fail"),
+ Err(e) => e,
+ };
+ assert!(err.to_string().contains("boom"));
}
}
diff --git a/src/client/mod.rs b/src/client/mod.rs
index a8d2ae5..dff11b9 100644
--- a/src/client/mod.rs
+++ b/src/client/mod.rs
@@ -21,8 +21,10 @@
pub(crate) mod backoff;
-#[cfg(all(feature = "reqwest", not(target_arch = "wasm32")))]
+#[cfg(not(target_arch = "wasm32"))]
mod dns;
+#[cfg(not(target_arch = "wasm32"))]
+pub use dns::{DnsError, DnsFuture, DnsResolver};
#[cfg(not(target_arch = "wasm32"))]
#[cfg(test)]
@@ -363,6 +365,8 @@ pub struct ClientOptions {
http1_only: ConfigValue<bool>,
http2_only: ConfigValue<bool>,
randomize_addresses: ConfigValue<bool>,
+ #[cfg(not(target_arch = "wasm32"))]
+ dns_resolver: Option<Arc<dyn DnsResolver>>,
}
impl Default for ClientOptions {
@@ -402,6 +406,8 @@ impl Default for ClientOptions {
http1_only: true.into(),
http2_only: Default::default(),
randomize_addresses: true.into(),
+ #[cfg(not(target_arch = "wasm32"))]
+ dns_resolver: Default::default(),
}
}
}
@@ -796,6 +802,32 @@ impl ClientOptions {
self
}
+ /// Override the default DNS resolution with a custom [`DnsResolver`]
+ ///
+ /// When set, [`ClientConfigKey::RandomizeAddresses`] is ignored: the
+ /// provided resolver is fully responsible for resolution, ordering,
+ /// shuffling, and caching.
+ ///
+ /// The built-in reqwest-based transport applies this automatically.
+ /// Custom [`HttpConnector`] implementations should read it via
+ /// [`Self::dns_resolver`] and honor it.
+ ///
+ /// Note: unlike other options, this cannot be configured via
+ /// [`ClientConfigKey`] / string configuration.
+ ///
+ /// [`HttpConnector`]: crate::client::HttpConnector
+ #[cfg(not(target_arch = "wasm32"))]
+ pub fn with_dns_resolver(mut self, resolver: Arc<dyn DnsResolver>) -> Self
{
+ self.dns_resolver = Some(resolver);
+ self
+ }
+
+ /// Return the custom [`DnsResolver`] if any, see
[`Self::with_dns_resolver`]
+ #[cfg(not(target_arch = "wasm32"))]
+ pub fn dns_resolver(&self) -> Option<&Arc<dyn DnsResolver>> {
+ self.dns_resolver.as_ref()
+ }
+
/// Get the default headers defined through
`ClientOptions::with_default_headers`
pub fn get_default_headers(&self) -> Option<&HeaderMap> {
self.default_headers.as_ref()
@@ -928,8 +960,13 @@ impl ClientOptions {
// size of objects.
builder = builder.no_gzip().no_brotli().no_zstd().no_deflate();
- if self.randomize_addresses.get()? {
- builder = builder.dns_resolver(Arc::new(dns::ShuffleResolver));
+ let resolver: Option<Arc<dyn DnsResolver>> = match &self.dns_resolver {
+ Some(resolver) => Some(Arc::clone(resolver)),
+ None if self.randomize_addresses.get()? =>
Some(Arc::new(dns::ShuffleResolver)),
+ None => None,
+ };
+ if let Some(resolver) = resolver {
+ builder =
builder.dns_resolver(Arc::new(dns::ReqwestResolver(resolver)));
}
builder
@@ -1263,4 +1300,40 @@ mod tests {
user_agent
);
}
+
+ #[tokio::test]
+ #[cfg(all(feature = "reqwest", not(target_arch = "wasm32")))]
+ async fn test_custom_dns_resolver() {
+ use crate::client::mock_server::MockServer;
+ use std::net::{IpAddr, Ipv4Addr};
+ use std::sync::atomic::{AtomicUsize, Ordering};
+
+ #[derive(Debug, Default)]
+ struct CountingResolver(AtomicUsize);
+
+ impl DnsResolver for CountingResolver {
+ fn resolve(&self, host: &str) -> DnsFuture {
+ assert_eq!(host, "localhost");
+ self.0.fetch_add(1, Ordering::SeqCst);
+ Box::pin(async { Ok(vec![IpAddr::V4(Ipv4Addr::LOCALHOST)]) })
+ }
+ }
+
+ let server = MockServer::new().await;
+ server.push(::http::Response::new("hello".to_string()));
+
+ let resolver = Arc::new(CountingResolver::default());
+ let client = ClientOptions::new()
+ .with_allow_http(true)
+ .with_dns_resolver(Arc::clone(&resolver) as Arc<dyn DnsResolver>)
+ .client()
+ .unwrap();
+
+ // Use a hostname (not an IP literal) so resolution actually runs
+ let parsed = url::Url::parse(server.url()).unwrap();
+ let url = format!("http://localhost:{}/";, parsed.port().unwrap());
+ let resp = client.get(url).send().await.unwrap();
+ assert!(resp.status().is_success());
+ assert_eq!(resolver.0.load(Ordering::SeqCst), 1);
+ }
}
