feat: retry all networked operations with exponential backoff

src/main.rs

@@ -6,15 +6,16 @@ use time::Duration;
 
 use hickory_client::{
     client::{Client, SyncClient},
+    error::ClientErrorKind,
     op::{DnsResponse, ResponseCode},
-    rr::{rdata::TXT, Name, RData, Record, RecordType, DNSClass},
+    rr::{rdata::TXT, DNSClass, Name, RData, Record, RecordType},
     udp::UdpClientConnection,
 };
+use hickory_proto::error::ProtoErrorKind;
 
-use tracing::{debug, info, trace, warn, error, instrument};
+use tracing::{debug, error, info, instrument, trace, warn};
 use tracing_subscriber::{EnvFilter, FmtSubscriber};
 
-use backoff::{retry, ExponentialBackoff};
 
 use once_cell::sync::OnceCell;
 
@@ -23,6 +24,8 @@ mod config;
 
 static RESOLVER: OnceCell<std::net::SocketAddr> = OnceCell::new();
 
+
+
 fn main() -> anyhow::Result<()> {
     let subscriber = FmtSubscriber::builder()
         .with_env_filter(EnvFilter::from_default_env())
@@ -33,7 +36,9 @@ fn main() -> anyhow::Result<()> {
     let cli = cli::parse();
 
     let config = read_config(&cli.config)?;
-    RESOLVER.set(config.resolver).map_err(| _| anyhow::format_err!("failed to set global recursor"))?;
+    RESOLVER
+        .set(config.resolver)
+        .map_err(|_| anyhow::format_err!("failed to set global recursor"))?;
 
     debug!("initializing with zones {:?}", config.zones.keys());
 
@@ -47,9 +52,7 @@ fn main() -> anyhow::Result<()> {
     let primary_ns_client = &config
         .zones
         .get(&zone)
-        .context(
-            "Couldn't find challenge zone in config, maybe you forgot the trailing dot?",
-        )?
+        .context("Couldn't find challenge zone in config, maybe you forgot the trailing dot?")?
         .create_client()?;
 
     let challenge_record = Record::from_rdata(
@@ -60,8 +63,7 @@ fn main() -> anyhow::Result<()> {
 
     match &cli.command {
         cli::Commands::Set => {
-            let result = primary_ns_client
-                .create(challenge_record.clone(), zone)
+            let result = retry_op(primary_ns_client.create(challenge_record.clone(), zone))
                 .context("Couldn't set record")?;
             match result.response_code() {
                 ResponseCode::NoError => {
@@ -75,8 +77,7 @@ fn main() -> anyhow::Result<()> {
             Ok(())
         }
         cli::Commands::Unset => {
-            let result = primary_ns_client
-                .delete_rrset(challenge_record.clone(), zone)
+            let result = retry_op(primary_ns_client.delete_rrset(challenge_record.clone(), zone))
                 .context("Couldn't remove record")?;
             match result.response_code() {
                 ResponseCode::NoError => {
@@ -111,7 +112,7 @@ fn find_record(identifier: &Name) -> anyhow::Result<Name> {
                 last_record.name_labels
             );
             Ok(last_record.name_labels)
-        },
+        }
         _ => Result::Err(anyhow::format_err!("unexpected record type")),
     }
 }
@@ -132,7 +133,6 @@ fn find_zone(identifier: &Name) -> anyhow::Result<Name> {
     }
 }
 
-#[instrument]
 pub fn search(identifier: Name, rtype: RecordType) -> anyhow::Result<DnsResponse> {
     let client = SyncClient::new(
         UdpClientConnection::new(
@@ -143,7 +143,42 @@ pub fn search(identifier: Name, rtype: RecordType) -> anyhow::Result<DnsResponse
         .context("failed to connect to recursor")?,
     );
 
-    let query = || Ok(client.query(&identifier, DNSClass::IN, rtype)?);
+    Ok(retry_op(client.query(
+        &identifier,
+        DNSClass::IN,
+        rtype
+    ))?)
+}
 
-    Ok(retry(ExponentialBackoff::default(), query)?)
+#[instrument]
+pub fn retry_op(
+    payload: Result<DnsResponse, hickory_client::error::ClientError>,
+) -> Result<DnsResponse, backoff::Error<hickory_client::error::ClientError>> {
+    backoff::retry(
+        backoff::ExponentialBackoff::default(),
+        || {
+        match payload.clone() {
+            Ok(response) => Ok(response),
+            Err(e) => Result::Err({
+                let transient = backoff::Error::Transient {
+                    err: e.clone(),
+                    retry_after: None,
+                };
+                let permanent = backoff::Error::Permanent(e.clone());
+
+                match &e.kind() {
+                    ClientErrorKind::Proto(proto_error) => match proto_error.kind() {
+                        ProtoErrorKind::Busy => transient,
+                        ProtoErrorKind::Io(_) => transient,
+                        ProtoErrorKind::Timeout => transient,
+                        _ => permanent,
+                    },
+                    ClientErrorKind::SendError(_) => transient,
+                    ClientErrorKind::Io(_) => transient,
+                    ClientErrorKind::Timeout => transient,
+                    _ => permanent,
+                }
+            }),
+        }
+    })
 }