Set up DNS cache

Author: magec

src/client.rs

@@ -59,14 +59,14 @@ pub struct Client<S, T> {
     /// to connect and cancel a query.
     client_server_map: ClientServerMap,
 
-    /// Client parameters, e.g. user, client_encoding, etc.
+    /// Client parameters, e.g. user, client_encoding, etc1
     #[allow(dead_code)]
     parameters: HashMap<String, String>,
 
     /// Statistics
     stats: Reporter,
 
-    /// Clients want to talk to admin database.
+    /// Clients want to talk to admin database0
     admin: bool,
 
     /// Last address the client talked to.

src/config.rs

@@ -208,7 +208,7 @@ impl General {
     }
 
     pub fn default_dns_max_ttl() -> u64 {
-	30
+        30
     }
 
     pub fn default_healthcheck_timeout() -> u64 {
@@ -238,8 +238,8 @@ impl Default for General {
             ban_time: Self::default_ban_time(),
             log_client_connections: false,
             log_client_disconnections: false,
-	    dns_cache_enabled: false,
-	    dns_max_ttl: Self::default_dns_max_ttl(),
+            dns_cache_enabled: false,
+            dns_max_ttl: Self::default_dns_max_ttl(),
             autoreload: false,
             tls_certificate: None,
             tls_private_key: None,

src/dns_cache.rs

@@ -1,9 +1,9 @@
 use log::{debug, error};
 use std::collections::{HashMap, HashSet};
+use std::io;
 use std::net::IpAddr;
 use std::sync::Arc;
 use std::sync::RwLock;
-use std::io;
 use tokio::time::{sleep, Duration};
 use trust_dns_resolver::lookup_ip::LookupIp;
 use trust_dns_resolver::TokioAsyncResolver;
@@ -55,7 +55,7 @@ impl From<LookupIp> for AddrSet {
 // // calls will be returned from cache.
 // resolver.has_changed("www.example.com.", addrset)
 // ```
-#[derive(Clone)]
+#[derive(Clone, Debug)]
 pub struct CachedResolver {
     // The configuration of the cached_resolver.
     pub config: CachedResolverConfig,
@@ -91,7 +91,7 @@ impl CachedResolver {
     //
     pub fn new(config: CachedResolverConfig) -> io::Result<Self> {
         // Construct a new Resolver with default configuration options
-	let resolver = Arc::new(TokioAsyncResolver::tokio_from_system_conf()?);
+        let resolver = Arc::new(TokioAsyncResolver::tokio_from_system_conf()?);
         let data = Arc::new(RwLock::new(HashMap::new()));
 
         Ok(Self {
@@ -103,46 +103,53 @@ impl CachedResolver {
 
     // Schedules the refresher
     pub async fn refresh_dns_entries_loop(&mut self) {
-	let data = self.data.clone();
-	let resolver = TokioAsyncResolver::tokio_from_system_conf().unwrap();
+        let data = self.data.clone();
+        let resolver = TokioAsyncResolver::tokio_from_system_conf().unwrap();
         let interval = Duration::from_secs(self.config.dns_max_ttl);
         loop {
-	    debug!("Begin refreshing cached DNS addresses.");
-	    // To minimize the time we hold the lock, we first create
-	    // an array with keys.
-	    let mut hostnames: Vec<String> = Vec::new();
-	    {
+            debug!("Begin refreshing cached DNS addresses.");
+            // To minimize the time we hold the lock, we first create
+            // an array with keys.
+            let mut hostnames: Vec<String> = Vec::new();
+            {
                 for hostname in data.read().unwrap().keys() {
-		    hostnames.push(hostname.clone());
+                    hostnames.push(hostname.clone());
                 }
-	    }
+            }
 
-	    for hostname in hostnames.iter() {
+            for hostname in hostnames.iter() {
                 match CachedResolver::fetch_from_data_cache(data.clone(), hostname.as_str()) {
-		    Some(addrset) => {
-                        match resolver.lookup_ip(hostname).await {
-			    Ok(lookup_ip) => {
-				let new_addrset = AddrSet::from(lookup_ip);
-				debug!("Obtained address for host ({}) -> ({:?})", hostname, new_addrset);
+                    Some(addrset) => match resolver.lookup_ip(hostname).await {
+                        Ok(lookup_ip) => {
+                            let new_addrset = AddrSet::from(lookup_ip);
+                            debug!(
+                                "Obtained address for host ({}) -> ({:?})",
+                                hostname, new_addrset
+                            );
 
-				if addrset != new_addrset {
-				    debug!("Addr changed from {:?} to {:?} updating cache.", addrset, new_addrset);
-				    CachedResolver::store_in_cache(data.clone(), hostname, new_addrset);
-				}
-			    },
-			    Err(err) => {
-				error!("There was an error trying to resolv {}: ({}).", hostname, err);
-			    }
+                            if addrset != new_addrset {
+                                debug!(
+                                    "Addr changed from {:?} to {:?} updating cache.",
+                                    addrset, new_addrset
+                                );
+                                CachedResolver::store_in_cache(data.clone(), hostname, new_addrset);
+                            }
+                        }
+                        Err(err) => {
+                            error!(
+                                "There was an error trying to resolv {}: ({}).",
+                                hostname, err
+                            );
                         }
-		    }
-		    None => {
+                    },
+                    None => {
                         error!("Could not obtain expected address from cache, this should not happen, \
 				as this cache does not allow deleting addresses.");
-		    }
+                    }
                 }
-	    }
-	    debug!("Finished refreshing cached DNS addresses.");
-	    sleep(interval).await;
+            }
+            debug!("Finished refreshing cached DNS addresses.");
+            sleep(interval).await;
         }
     }
 
@@ -166,16 +173,16 @@ impl CachedResolver {
     // ```
     //
     pub async fn lookup_ip(&mut self, host: &str) -> ResolveResult<AddrSet> {
-	debug!("Lookup up {} in cache", host);
+        debug!("Lookup up {} in cache", host);
         match self.fetch_from_cache(host) {
             Some(addr_set) => {
-		debug!("Cache hit!");
-		Ok(addr_set)
-	    },
+                debug!("Cache hit!");
+                Ok(addr_set)
+            }
             None => {
-		debug!("Not found, executing a dns query!");
+                debug!("Not found, executing a dns query!");
                 let addr_set = AddrSet::from(self.resolver.clone().lookup_ip(host).await?);
-		debug!("Obtained: {:?}", addr_set);
+                debug!("Obtained: {:?}", addr_set);
                 CachedResolver::store_in_cache(self.data.clone(), host, addr_set.clone());
                 Ok(addr_set)
             }
@@ -247,7 +254,7 @@ mod tests {
         let mut resolver = CachedResolver::new(config).unwrap();
         let hostname = "www.idontexists.";
         let response = resolver.lookup_ip(hostname).await;
-	assert!(matches!(response, Err(ResolveError{ .. })));
+        assert!(matches!(response, Err(ResolveError { .. })));
     }
 
     #[tokio::test]
@@ -256,26 +263,26 @@ mod tests {
         let mut resolver = CachedResolver::new(config).unwrap();
         let hostname = "w  ww.idontexists.";
         let response = resolver.lookup_ip(hostname).await;
-	assert!(matches!(response, Err(ResolveError{ .. })));
-	assert!(!resolver.has_changed(hostname, &AddrSet::new()));
+        assert!(matches!(response, Err(ResolveError { .. })));
+        assert!(!resolver.has_changed(hostname, &AddrSet::new()));
     }
 
     #[tokio::test]
     // Ok, this test is based on the fact that google does DNS RR
     // and does not responds with every available ip everytime, so
     // if I cache here, it will miss after one cache iteration or two.
     async fn thread() {
-	env_logger::init();
+        env_logger::init();
         let config = CachedResolverConfig { dns_max_ttl: 10 };
         let mut resolver = CachedResolver::new(config).unwrap();
         let hostname = "www.google.com.";
         let response = resolver.lookup_ip(hostname).await;
         let addr_set = response.unwrap();
         assert!(!resolver.has_changed(hostname, &addr_set));
-	let mut resolver_for_refresher = resolver.clone();
-	let _thread_handle = tokio::task::spawn(async move {
-	    resolver_for_refresher.refresh_dns_entries_loop().await;
-	});
-	assert!(!resolver.has_changed(hostname, &addr_set));
+        let mut resolver_for_refresher = resolver.clone();
+        let _thread_handle = tokio::task::spawn(async move {
+            resolver_for_refresher.refresh_dns_entries_loop().await;
+        });
+        assert!(!resolver.has_changed(hostname, &addr_set));
     }
 }

src/lib.rs

@@ -1,5 +1,6 @@
 pub mod config;
 pub mod constants;
+pub mod dns_cache;
 pub mod errors;
 pub mod messages;
 pub mod pool;

src/main.rs

@@ -64,6 +64,7 @@ mod admin;
 mod client;
 mod config;
 mod constants;
+mod dns_cache;
 mod errors;
 mod messages;
 mod pool;
@@ -74,7 +75,6 @@ mod server;
 mod sharding;
 mod stats;
 mod tls;
-mod dns_cache;
 
 use crate::config::{get_config, reload_config, VERSION};
 use crate::pool::{ClientServerMap, ConnectionPool};

src/pool.rs

@@ -15,6 +15,7 @@ use std::time::Instant;
 use crate::config::{get_config, Address, General, PoolMode, Role, User};
 use crate::errors::Error;
 
+use crate::dns_cache::{CachedResolver, CachedResolverConfig};
 use crate::server::Server;
 use crate::sharding::ShardingFunction;
 use crate::stats::{get_reporter, Reporter};
@@ -132,13 +133,40 @@ pub struct ConnectionPool {
 
     /// Pool configuration.
     pub settings: PoolSettings,
+
+    /// CachedResolver to use when dns_cache is enabled
+    pub cached_resolver: Arc<Option<CachedResolver>>,
 }
 
 impl ConnectionPool {
     /// Construct the connection pool from the configuration.
     pub async fn from_config(client_server_map: ClientServerMap) -> Result<(), Error> {
         let config = get_config();
 
+        // Configure dns_cache if enabled
+        let mut cached_resolver: Arc<Option<CachedResolver>> = Arc::new(None);
+        if config.general.dns_cache_enabled {
+            info!("Starting Dns cache");
+            let cached_resolver_config = CachedResolverConfig {
+                dns_max_ttl: config.general.dns_max_ttl,
+            };
+            cached_resolver = match CachedResolver::new(cached_resolver_config) {
+                Ok(ok) => {
+                    let mut refresher_cached_resolver = ok.clone();
+                    info!("Scheduling DNS refresh loop");
+                    tokio::task::spawn(async move {
+                        refresher_cached_resolver.refresh_dns_entries_loop().await;
+                    });
+                    Arc::new(Some(ok))
+                }
+                Err(err) => {
+                    error!("Error Starting cached_resolver error: {:?}", err);
+                    error!("Will continue without this feature");
+                    Arc::new(None)
+                }
+            };
+        };
+
         let mut new_pools = HashMap::new();
         let mut address_id = 0;
 
@@ -246,6 +274,7 @@ impl ConnectionPool {
                 let mut pool = ConnectionPool {
                     databases: shards,
                     addresses,
+                    cached_resolver: cached_resolver.clone(),
                     banlist: Arc::new(RwLock::new(banlist)),
                     stats: get_reporter(),
                     server_info: BytesMut::new(),
@@ -596,6 +625,12 @@ impl ManageConnection for ServerPool {
     /// Attempts to create a new connection.
     async fn connect(&self) -> Result<Self::Connection, Self::Error> {
         info!("Creating a new server connection {:?}", self.address);
+
+        let cached_resolver = match get_pool(&self.database, &self.user.username) {
+            Some(pool) => pool.cached_resolver,
+            None => Arc::new(None),
+        };
+
         let server_id = rand::random::<i32>();
 
         self.stats.server_register(
@@ -614,6 +649,7 @@ impl ManageConnection for ServerPool {
             &self.user,
             &self.database,
             self.client_server_map.clone(),
+            cached_resolver,
             self.stats.clone(),
         )
         .await

src/server.rs

@@ -1,8 +1,9 @@
-/// Implementation of the PostgreSQL server (database) protocol.
-/// Here we are pretending to the a Postgres client.
 use bytes::{Buf, BufMut, BytesMut};
 use log::{debug, error, info, trace, warn};
 use std::io::Read;
+/// Implementation of the PostgreSQL server (database) protocol.
+/// Here we are pretending to the a Postgres client.
+use std::sync::Arc;
 use std::time::SystemTime;
 use tokio::io::{AsyncReadExt, BufReader};
 use tokio::net::{
@@ -12,6 +13,7 @@ use tokio::net::{
 
 use crate::config::{Address, User};
 use crate::constants::*;
+use crate::dns_cache::{AddrSet, CachedResolver};
 use crate::errors::Error;
 use crate::messages::*;
 use crate::pool::ClientServerMap;
@@ -68,6 +70,12 @@ pub struct Server {
 
     // Last time that a successful server send or response happened
     last_activity: SystemTime,
+
+    // Cached Resolver to be used to resolve addresses with hostnames
+    cached_resolver: Arc<Option<CachedResolver>>,
+
+    // Associated addresses used
+    addr_set: Option<AddrSet>,
 }
 
 impl Server {
@@ -79,8 +87,22 @@ impl Server {
         user: &User,
         database: &str,
         client_server_map: ClientServerMap,
+        cached_resolver: Arc<Option<CachedResolver>>,
         stats: Reporter,
     ) -> Result<Server, Error> {
+        let addr_set = match *cached_resolver.clone() {
+            Some(ref cached_resolver) => {
+                match cached_resolver.clone().lookup_ip(&address.host).await {
+                    Ok(ok) => Some(ok),
+                    Err(err) => {
+                        warn!("Error trying to resolve {}, ({:?})", &address.host, err);
+                        None
+                    }
+                }
+            }
+            None => None,
+        };
+
         let mut stream =
             match TcpStream::connect(&format!("{}:{}", &address.host, address.port)).await {
                 Ok(stream) => stream,
@@ -329,6 +351,8 @@ impl Server {
                         bad: false,
                         needs_cleanup: false,
                         client_server_map,
+                        cached_resolver,
+                        addr_set,
                         connected_at: chrono::offset::Utc::now().naive_utc(),
                         stats,
                         application_name: String::new(),
@@ -554,6 +578,12 @@ impl Server {
     /// Server & client are out of sync, we must discard this connection.
     /// This happens with clients that misbehave.
     pub fn is_bad(&self) -> bool {
+        if let Some(cached_resolver) = &(*(self.cached_resolver.clone())) {
+            if let Some(addr_set) = &self.addr_set {
+                cached_resolver.has_changed(self.address.host.as_str(), addr_set);
+                return true;
+            }
+        }
         self.bad
     }