Migrate Redis from r2d2 to bb8 (fully async now!)

valeriansaliou · Jul 20, 2024 · 4d326ae · 4d326ae
1 parent caf7475
commit 4d326ae
Show file tree

Hide file tree

Showing 9 changed files with 133 additions and 119 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -27,7 +27,7 @@ serde_json = "1.0"
 url_serde = { version = "0.2", default-features = false }
 regex = "1.10"
 async-trait = "0.1"
-r2d2_redis = "0.14"
+bb8-redis = "0.15"
 actix-web = { version = "3.3", default-features = false }
 actix-web-httpauth = "0.5"
 hickory-server = { version = "0.24", default-features = false }

diff --git a/src/dns/handler.rs b/src/dns/handler.rs
@@ -155,8 +155,9 @@ impl DNSHandler {
 
         // #4. Resolve from remote store
         // Notice: this is used to serve all records set with the HTTP API.
-        // TODO: this is a blocking code path (records_from_store() must be made async)
-        return match Self::records_from_store(authority, &zone_name, request.src().ip(), query) {
+        return match Self::records_from_store(authority, &zone_name, request.src().ip(), query)
+            .await
+        {
             Ok(records_remote) => {
                 // Serve response data?
                 if let Some(records_remote_inner) = records_remote {
@@ -336,7 +337,7 @@ impl DNSHandler {
         None
     }
 
-    fn records_from_store(
+    async fn records_from_store(
         authority: &DNSAuthority,
         zone_name: &Option<ZoneName>,
         source: IpAddr,
@@ -351,7 +352,6 @@ impl DNSHandler {
         }
 
         // Attempt with requested domain
-        // TODO: this is a blocking code path (method must be made async)
         let mut records = Self::records_from_store_attempt(
             authority,
             source,
@@ -360,7 +360,8 @@ impl DNSHandler {
             &query_name,
             &query_type,
             &record_type,
-        )?;
+        )
+        .await?;
 
         // Check if 'records' is empty
         let is_records_empty = if let Some(ref records_inner) = records {
@@ -383,7 +384,6 @@ impl DNSHandler {
                     let wildcard_name_lower = LowerName::new(&wildcard_name);
 
                     if &wildcard_name_lower != query_name {
-                        // TODO: this is a blocking code path (method must be made async)
                         let records_wildcard = Self::records_from_store_attempt(
                             authority,
                             source,
@@ -392,7 +392,8 @@ impl DNSHandler {
                             &wildcard_name_lower,
                             &query_type,
                             &record_type,
-                        )?;
+                        )
+                        .await?;
 
                         // Assign non-none wildcard records? (retain any NOERROR from 'records')
                         if records_wildcard.is_none() == false {
@@ -406,7 +407,7 @@ impl DNSHandler {
         Ok(records)
     }
 
-    fn records_from_store_attempt(
+    async fn records_from_store_attempt(
         authority: &DNSAuthority,
         source: IpAddr,
         zone_name: &Option<ZoneName>,
@@ -427,13 +428,15 @@ impl DNSHandler {
                 let mut records = Vec::new();
 
                 if let &Some(ref record_type_inner) = record_type {
-                    // TODO: this is a blocking code path (method must be made async)
-                    match APP_STORE.get(
-                        &zone_name,
-                        &record_name,
-                        record_type_inner,
-                        StoreAccessOrigin::External,
-                    ) {
+                    match APP_STORE
+                        .get(
+                            &zone_name,
+                            &record_name,
+                            record_type_inner,
+                            StoreAccessOrigin::External,
+                        )
+                        .await
+                    {
                         Ok(record) => {
                             debug!(
                                 "found record in store for query: {} {}; got: {:?}",
@@ -461,13 +464,15 @@ impl DNSHandler {
 
                     // Look for a CNAME result? (if no records were acquired)
                     if record_type_inner != &RecordType::CNAME && records.is_empty() {
-                        // TODO: this is a blocking code path (method must be made async)
-                        match APP_STORE.get(
-                            &zone_name,
-                            &record_name,
-                            &RecordType::CNAME,
-                            StoreAccessOrigin::External,
-                        ) {
+                        match APP_STORE
+                            .get(
+                                &zone_name,
+                                &record_name,
+                                &RecordType::CNAME,
+                                StoreAccessOrigin::External,
+                            )
+                            .await
+                        {
                             Ok(record_cname) => {
                                 debug!(
                                     "found cname hint record in store for query: {} {}; got: {:?}",
@@ -502,8 +507,8 @@ impl DNSHandler {
 
                 // No record found, exhaust all record types to check if name exists
                 // Notice: a DNS server must return NOERROR if name exists, else NXDOMAIN
-                // TODO: this is a blocking code path (method must be made async)
-                if Self::check_name_exists(&zone_name, &record_name, StoreAccessOrigin::External)?
+                if Self::check_name_exists(&zone_name, &record_name, StoreAccessOrigin::External)
+                    .await?
                     == true
                 {
                     // Name exists, return empty records (ie. NOERROR)
@@ -778,7 +783,7 @@ impl DNSHandler {
         }
     }
 
-    fn check_name_exists(
+    async fn check_name_exists(
         zone_name: &ZoneName,
         record_name: &RecordName,
         origin: StoreAccessOrigin,
@@ -789,8 +794,10 @@ impl DNSHandler {
             // Notice: instead of performing a simple exist check, we acquire full record data, \
             //   as this lets us use the local store and therefore prevent non-existing domain \
             //   attacks on the remote store.
-            // TODO: this is a blocking code path (method must be made async)
-            match APP_STORE.get(zone_name, record_name, &record_type, origin) {
+            match APP_STORE
+                .get(zone_name, record_name, &record_type, origin)
+                .await
+            {
                 Ok(_) => {
                     // Record exists for name and type; abort there.
                     return Ok(true);

diff --git a/src/dns/health.rs b/src/dns/health.rs
@@ -156,17 +156,18 @@ impl DNSHealth {
 }
 
 impl DNSHealthHTTP {
-    fn run(notifier: &mut DNSHealthNotify) {
+    #[tokio::main]
+    async fn run(notifier: &mut DNSHealthNotify) {
         debug!("running dns health checks for the http protocol...");
 
         for domain in &APP_CONF.dns.health.http {
-            Self::check_domain(domain, notifier);
+            Self::check_domain(domain, notifier).await;
         }
 
         debug!("ran dns health checks for the http protocol");
     }
 
-    fn check_domain(domain: &ConfigDNSHealthHTTP, notifier: &mut DNSHealthNotify) {
+    async fn check_domain(domain: &ConfigDNSHealthHTTP, notifier: &mut DNSHealthNotify) {
         for record_type in HEALTH_CHECK_RECORD_TYPES.iter() {
             debug!(
                 "checking dns health for target: {} on zone: {} with type: {:?}",
@@ -175,13 +176,15 @@ impl DNSHealthHTTP {
                 record_type
             );
 
-            // TODO: this is a blocking code path (method must be made async)
-            if let Ok(record) = APP_STORE.get(
-                &domain.zone,
-                &domain.name,
-                record_type,
-                StoreAccessOrigin::Internal,
-            ) {
+            if let Ok(record) = APP_STORE
+                .get(
+                    &domain.zone,
+                    &domain.name,
+                    record_type,
+                    StoreAccessOrigin::Internal,
+                )
+                .await
+            {
                 let unique_values = record.list_record_values();
 
                 for record_value in unique_values {