From 00aafb216376440146636ef116d1c43c34e5b21a Mon Sep 17 00:00:00 2001
From: SHAcollision <127778313+SHAcollision@users.noreply.github.com>
Date: Thu, 8 May 2025 13:27:11 +0200
Subject: [PATCH] fix(client): signup-publish-retry-flexible (#120)

* feat(client): signup-publish-retry-flexible

* fix wasm

* Retry on known query errors fixed 3 attemps 1s

* simplify
---
 Cargo.lock                                | 13 ++++++
 pubky-client/Cargo.toml                   |  1 +
 pubky-client/src/native/internal/pkarr.rs | 57 ++++++++++++++++++-----
 3 files changed, 60 insertions(+), 11 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 4252aa6..40e6165 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1152,6 +1152,18 @@ version = "0.31.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "07e28edb80900c19c28f1072f2e8aeca7fa06b23cd4169cefe1af5aa3260783f"
 
+[[package]]
+name = "gloo-timers"
+version = "0.3.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "bbb143cf96099802033e0d4f4963b19fd2e0b728bcf076cd9cf7f6634f092994"
+dependencies = [
+ "futures-channel",
+ "futures-core",
+ "js-sys",
+ "wasm-bindgen",
+]
+
 [[package]]
 name = "governor"
 version = "0.8.0"
@@ -2280,6 +2292,7 @@ dependencies = [
  "futures-util",
  "getrandom 0.2.15",
  "getrandom 0.3.1",
+ "gloo-timers",
  "js-sys",
  "log",
  "mainline",
diff --git a/pubky-client/Cargo.toml b/pubky-client/Cargo.toml
index cf15624..1cfc604 100644
--- a/pubky-client/Cargo.toml
+++ b/pubky-client/Cargo.toml
@@ -49,6 +49,7 @@ wasm-bindgen = "0.2.100"
 wasm-bindgen-futures = "0.4.50"
 console_log = { version = "1.0.0", features = ["color"] }
 log = "0.4.25"
+gloo-timers = { version = "0.3", features = ["futures"] }
 
 js-sys = "0.3.77"
 web-sys = "0.3.77"
diff --git a/pubky-client/src/native/internal/pkarr.rs b/pubky-client/src/native/internal/pkarr.rs
index 336ad9e..1b37a12 100644
--- a/pubky-client/src/native/internal/pkarr.rs
+++ b/pubky-client/src/native/internal/pkarr.rs
@@ -1,6 +1,7 @@
 use anyhow::Result;
 use pkarr::{
     dns::rdata::{RData, SVCB},
+    errors::QueryError,
     Keypair, SignedPacket, Timestamp,
 };
 use std::convert::TryInto;
@@ -8,6 +9,28 @@ use std::time::Duration;
 
 use super::super::Client;
 
+// sleep for native
+#[cfg(not(wasm_browser))]
+use tokio::time::sleep;
+// sleep for wasm
+#[cfg(wasm_browser)]
+use gloo_timers::future::sleep;
+
+/// Helper returns true if this error (or any of its sources) is one of our
+/// three recoverable `QueryError`s with simple retrial.
+fn should_retry(err: &anyhow::Error) -> bool {
+    err.chain()
+        .filter_map(|cause| cause.downcast_ref::<QueryError>())
+        .any(|q| {
+            matches!(
+                q,
+                QueryError::Timeout
+                    | QueryError::NoClosestNodes
+                    | QueryError::DhtErrorResponse(_, _)
+            )
+        })
+}
+
 /// The strategy to decide whether to (re)publish a homeserver record.
 pub(crate) enum PublishStrategy {
     /// Always publish a new record (used on signup).
@@ -30,16 +53,16 @@ impl Client {
         host: Option<&str>,
         strategy: PublishStrategy,
     ) -> Result<()> {
-        // Resolve the most recent record.
+        // 1) Resolve the most recent record.
         let existing = self.pkarr.resolve_most_recent(&keypair.public_key()).await;
 
-        // Determine which host we should be using.
+        // 2) Determine which host we should be using.
         let host_str = match Self::determine_host(host, existing.as_ref()) {
             Some(host) => host,
             None => return Ok(()),
         };
 
-        // Calculate the age of the existing record.
+        // 3) Calculate age of the existing record.
         let packet_age = match existing {
             Some(ref record) => {
                 let elapsed = Timestamp::now() - record.timestamp();
@@ -48,15 +71,27 @@ impl Client {
             None => Duration::from_secs(u64::MAX), // Use max duration if no record exists.
         };
 
-        // Determine if we should publish based on the given strategy.
-        let should_publish = match strategy {
-            PublishStrategy::Force => true,
-            PublishStrategy::IfOlderThan => packet_age > self.max_record_age,
-        };
+        // 4) Should we publish?
+        let should_publish =
+            matches!(strategy, PublishStrategy::Force) || packet_age > self.max_record_age;
 
-        if should_publish {
-            self.publish_homeserver_inner(keypair, &host_str, existing)
-                .await?;
+        if !should_publish {
+            return Ok(());
+        }
+
+        // 5) Retry loop: up to 3 attempts, 1s back-off, only on specific QueryErrors.
+        for attempt in 1..=3 {
+            match self
+                .publish_homeserver_inner(keypair, &host_str, existing.clone())
+                .await
+            {
+                Ok(()) => break,
+                Err(e) if should_retry(&e) && attempt < 3 => {
+                    sleep(Duration::from_secs(1)).await;
+                    continue;
+                }
+                Err(e) => return Err(e),
+            }
         }
 
         Ok(())