Skip to content

Commit

Permalink
Optimize TrustDnsResolver (#1967)
Browse files Browse the repository at this point in the history
* Optimize initialization of `TrustDnsResolver`

Replace use of `tokio::sync::Mutex` with `once_cell::sync::OnceCell`
since there's no need for it to be `async` and thus it can just use
`OnceCell` instead of `Mutex`.

Signed-off-by: Jiahao XU <Jiahao_XU@outlook.com>

* Optimize `TrustDnsResolver`: Rm use of `Lazy`

Signed-off-by: Jiahao XU <Jiahao_XU@outlook.com>

* Optimize `TrustDnsResolver`: Rm unnecessary `Arc`

Signed-off-by: Jiahao XU <Jiahao_XU@outlook.com>

* Derive `Default` on `TrustDnsResolver`

instead of implementing `TrustDnsResolver::new()`

Signed-off-by: Jiahao XU <Jiahao_XU@outlook.com>

---------

Signed-off-by: Jiahao XU <Jiahao_XU@outlook.com>
  • Loading branch information
NobodyXu authored Sep 4, 2023
1 parent 0292486 commit d3d95a5
Show file tree
Hide file tree
Showing 2 changed files with 14 additions and 57 deletions.
2 changes: 1 addition & 1 deletion src/async_impl/client.rs
Original file line number Diff line number Diff line change
Expand Up @@ -266,7 +266,7 @@ impl ClientBuilder {
let mut resolver: Arc<dyn Resolve> = match config.trust_dns {
false => Arc::new(GaiResolver::new()),
#[cfg(feature = "trust-dns")]
true => Arc::new(TrustDnsResolver::new().map_err(crate::error::builder)?),
true => Arc::new(TrustDnsResolver::default()),
#[cfg(not(feature = "trust-dns"))]
true => unreachable!("trust-dns shouldn't be enabled unless the feature is"),
};
Expand Down
69 changes: 13 additions & 56 deletions src/dns/trust_dns.rs
Original file line number Diff line number Diff line change
@@ -1,8 +1,7 @@
//! DNS resolution via the [trust_dns_resolver](https://github.com/bluejekyll/trust-dns) crate

use hyper::client::connect::dns::Name;
use once_cell::sync::Lazy;
use tokio::sync::Mutex;
use once_cell::sync::OnceCell;
pub use trust_dns_resolver::config::{ResolverConfig, ResolverOpts};
use trust_dns_resolver::{lookup_ip::LookupIpIntoIter, system_conf, TokioAsyncResolver};

Expand All @@ -12,62 +11,24 @@ use std::sync::Arc;

use super::{Addrs, Resolve, Resolving};

type SharedResolver = Arc<TokioAsyncResolver>;

static SYSTEM_CONF: Lazy<io::Result<(ResolverConfig, ResolverOpts)>> =
Lazy::new(|| system_conf::read_system_conf().map_err(io::Error::from));

/// Wrapper around an `AsyncResolver`, which implements the `Resolve` trait.
#[derive(Debug, Clone)]
#[derive(Debug, Default, Clone)]
pub(crate) struct TrustDnsResolver {
state: Arc<Mutex<State>>,
/// Since we might not have been called in the context of a
/// Tokio Runtime in initialization, so we must delay the actual
/// construction of the resolver.
state: Arc<OnceCell<TokioAsyncResolver>>,
}

struct SocketAddrs {
iter: LookupIpIntoIter,
}

#[derive(Debug)]
enum State {
Init,
Ready(SharedResolver),
}

impl TrustDnsResolver {
/// Create a new resolver with the default configuration,
/// which reads from `/etc/resolve.conf`.
pub fn new() -> io::Result<Self> {
SYSTEM_CONF.as_ref().map_err(|e| {
io::Error::new(e.kind(), format!("error reading DNS system conf: {}", e))
})?;

// At this stage, we might not have been called in the context of a
// Tokio Runtime, so we must delay the actual construction of the
// resolver.
Ok(TrustDnsResolver {
state: Arc::new(Mutex::new(State::Init)),
})
}
}

impl Resolve for TrustDnsResolver {
fn resolve(&self, name: Name) -> Resolving {
let resolver = self.clone();
Box::pin(async move {
let mut lock = resolver.state.lock().await;

let resolver = match &*lock {
State::Init => {
let resolver = new_resolver().await;
*lock = State::Ready(resolver.clone());
resolver
}
State::Ready(resolver) => resolver.clone(),
};

// Don't keep lock once the resolver is constructed, otherwise
// only one lookup could be done at a time.
drop(lock);
let resolver = resolver.state.get_or_try_init(new_resolver)?;

let lookup = resolver.lookup_ip(name.as_str()).await?;
let addrs: Addrs = Box::new(SocketAddrs {
Expand All @@ -86,14 +47,10 @@ impl Iterator for SocketAddrs {
}
}

async fn new_resolver() -> SharedResolver {
let (config, opts) = SYSTEM_CONF
.as_ref()
.expect("can't construct TrustDnsResolver if SYSTEM_CONF is error")
.clone();
new_resolver_with_config(config, opts)
}

fn new_resolver_with_config(config: ResolverConfig, opts: ResolverOpts) -> SharedResolver {
Arc::new(TokioAsyncResolver::tokio(config, opts))
/// Create a new resolver with the default configuration,
/// which reads from `/etc/resolve.conf`.
fn new_resolver() -> io::Result<TokioAsyncResolver> {
let (config, opts) = system_conf::read_system_conf()
.map_err(|e| io::Error::new(e.kind(), format!("error reading DNS system conf: {}", e)))?;
Ok(TokioAsyncResolver::tokio(config, opts))
}

0 comments on commit d3d95a5

Please sign in to comment.