Resume v2 polling on transient directory errors

payjoin-cli/src/app/v2/mod.rs

@@ -12,6 +12,7 @@ use payjoin::receive::v2::{
     ReceiverBuilder, SessionOutcome as ReceiverSessionOutcome, UncheckedOriginalPayload,
     WantsFeeRange, WantsInputs, WantsOutputs,
 };
+use payjoin::schedule::PollSchedule;
 use payjoin::send::v2::{
     replay_event_log as replay_sender_event_log, PendingFallback as SenderPendingFallback,
     PollingForProposal, SendSession, Sender, SenderBuilder, SessionOutcome as SenderSessionOutcome,
@@ -35,6 +36,7 @@ const W_ID: usize = 12;
 const W_ROLE: usize = 25;
 const W_DONE: usize = 15;
 const W_STATUS: usize = 15;
+const POLL_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(30);
 
 #[derive(Clone)]
 pub(crate) struct App {
@@ -704,27 +706,47 @@ impl App {
         sender: Sender<PollingForProposal>,
         persister: &SenderPersister,
     ) -> Result<()> {
-        let mut session = sender.clone();
-        // Long poll until we get a response
+        let session = sender;
+        let mut schedule = PollSchedule::new();
+        let mut polls = tokio::task::JoinSet::new();
+        let next = tokio::time::sleep(schedule.next_gap());
+        tokio::pin!(next);
         loop {
-            let (response, ctx) =
-                self.post_via_relay(|relay| session.create_poll_request(relay)).await?;
-            let res = session.process_response(&response.bytes().await?, ctx).save(persister);
-            match res {
-                Ok(OptionalTransitionOutcome::Progress(psbt)) => {
-                    println!("Proposal received. Processing...");
-                    self.process_pj_response(psbt)?;
-                    return Ok(());
-                }
-                Ok(OptionalTransitionOutcome::Stasis(current_state)) => {
-                    println!("No response yet.");
-                    session = current_state;
-                    continue;
+            tokio::select! {
+                Some(joined) = polls.join_next(), if !polls.is_empty() => {
+                    let (body, ctx): (Vec<u8>, _) = match joined {
+                        Ok(Ok(v)) => v,
+                        _ => continue,
+                    };
+                    match session.clone().process_response(&body, ctx).save(persister) {
+                        Ok(OptionalTransitionOutcome::Progress(psbt)) => {
+                            println!("Proposal received. Processing...");
+                            self.process_pj_response(psbt)?;
+                            return Ok(());
+                        }
+                        Ok(OptionalTransitionOutcome::Stasis(_)) => {
+                            println!("No response yet.");
+                        }
+                        Err(re) if re.is_transient() => {
+                            tracing::debug!("Transient directory error, retrying poll: {re:?}");
+                        }
+                        Err(re) => {
+                            println!("{re}");
+                            tracing::debug!("{re:?}");
+                            return Err(anyhow!("Response error").context(re));
+                        }
+                    }
                 }
-                Err(re) => {
-                    println!("{re}");
-                    tracing::debug!("{re:?}");
-                    return Err(anyhow!("Response error").context(re));
+                () = &mut next => {
+                    next.as_mut().reset(tokio::time::Instant::now() + schedule.next_gap());
+                    let relay = self.relay_manager.choose_relay()?;
+                    let (req, ctx) = session.create_poll_request(relay.as_str())?;
+                    let app = self.clone();
+                    polls.spawn(async move {
+                        let resp =
+                            tokio::time::timeout(POLL_TIMEOUT, app.post_request(req)).await??;
+                        Ok::<_, anyhow::Error>((resp.bytes().await?.to_vec(), ctx))
+                    });
                 }
             }
         }
@@ -735,24 +757,40 @@ impl App {
         session: Receiver<Initialized>,
         persister: &ReceiverPersister,
     ) -> Result<Receiver<UncheckedOriginalPayload>> {
-        let mut session = session;
+        let mut schedule = PollSchedule::new();
+        let mut polls = tokio::task::JoinSet::new();
+        let next = tokio::time::sleep(schedule.next_gap());
+        tokio::pin!(next);
         loop {
-            println!("Polling receive request...");
-            let (ohttp_response, context) =
-                self.post_via_relay(|relay| session.create_poll_request(relay)).await?;
-            let state_transition = session
-                .process_response(ohttp_response.bytes().await?.to_vec().as_slice(), context)
-                .save(persister);
-            match state_transition {
-                Ok(OptionalTransitionOutcome::Progress(next_state)) => {
-                    println!("Got a request from the sender. Responding with a Payjoin proposal.");
-                    return Ok(next_state);
+            tokio::select! {
+                Some(joined) = polls.join_next(), if !polls.is_empty() => {
+                    let (body, ctx): (Vec<u8>, _) = match joined {
+                        Ok(Ok(v)) => v,
+                        _ => continue,
+                    };
+                    match session.clone().process_response(&body, ctx).save(persister) {
+                        Ok(OptionalTransitionOutcome::Progress(next_state)) => {
+                            println!("Got a request from the sender. Responding with a Payjoin proposal.");
+                            return Ok(next_state);
+                        }
+                        Ok(OptionalTransitionOutcome::Stasis(_)) => {}
+                        Err(e) if e.is_transient() => {
+                            tracing::debug!("Transient directory error, retrying poll: {e:?}");
+                        }
+                        Err(e) => return Err(e.into()),
+                    }
                 }
-                Ok(OptionalTransitionOutcome::Stasis(current_state)) => {
-                    session = current_state;
-                    continue;
+                () = &mut next => {
+                    next.as_mut().reset(tokio::time::Instant::now() + schedule.next_gap());
+                    let relay = self.relay_manager.choose_relay()?;
+                    let (req, ctx) = session.create_poll_request(relay.as_str())?;
+                    let app = self.clone();
+                    polls.spawn(async move {
+                        let resp =
+                            tokio::time::timeout(POLL_TIMEOUT, app.post_request(req)).await??;
+                        Ok::<_, anyhow::Error>((resp.bytes().await?.to_vec(), ctx))
+                    });
                 }
-                Err(e) => return Err(e.into()),
             }
         }
     }

payjoin-cli/tests/e2e.rs

@@ -421,7 +421,7 @@ mod e2e {
         async fn respond_with_payjoin(mut cli_receive_resumer: Child) -> Result<()> {
             let mut stdout =
                 cli_receive_resumer.stdout.take().expect("Failed to take stdout of child process");
-            let timeout = tokio::time::Duration::from_secs(10);
+            let timeout = tokio::time::Duration::from_secs(45);
             let res = tokio::time::timeout(
                 timeout,
                 wait_for_stdout_match(&mut stdout, |line| line.contains("Response successful")),
@@ -436,7 +436,7 @@ mod e2e {
         async fn check_payjoin_sent(mut cli_send_resumer: Child) -> Result<()> {
             let mut stdout =
                 cli_send_resumer.stdout.take().expect("Failed to take stdout of child process");
-            let timeout = tokio::time::Duration::from_secs(10);
+            let timeout = tokio::time::Duration::from_secs(45);
             let res = tokio::time::timeout(
                 timeout,
                 wait_for_stdout_match(&mut stdout, |line| line.contains("Payjoin sent")),
@@ -466,7 +466,7 @@ mod e2e {
         async fn check_resume_completed(mut cli_resumer: Child) -> Result<()> {
             let mut stdout =
                 cli_resumer.stdout.take().expect("Failed to take stdout of child process");
-            let timeout = tokio::time::Duration::from_secs(10);
+            let timeout = tokio::time::Duration::from_secs(45);
             let res = tokio::time::timeout(
                 timeout,
                 wait_for_stdout_match(&mut stdout, |line| {

payjoin-mailroom/src/db/files.rs

@@ -238,6 +238,15 @@ impl DbTrait for FilesDb {
         Ok(guard.post_v2(id, payload).await?)
     }
 
+    async fn peek_v2_payload(
+        &self,
+        id: &ShortId,
+    ) -> Result<Option<Arc<Vec<u8>>>, DbError<Self::OperationalError>> {
+        let mut guard = self.mailboxes.lock().await;
+        Ok(guard.read(id).await?)
+    }
+
+    // Unused by GET after the non-blocking switch; v2 waitmap removal is a follow-up.
     async fn wait_for_v2_payload(
         &self,
         id: &ShortId,
@@ -1070,4 +1079,37 @@ mod tests {
 
         Ok(())
     }
+
+    #[tokio::test(start_paused = true)]
+    async fn peek_returns_immediately_on_empty_mailbox() {
+        let dir = tempfile::tempdir().unwrap();
+        let db = FilesDb::init(
+            Duration::from_secs(30),
+            dir.path().to_owned(),
+            Duration::from_secs(60 * 60 * 24 * 7),
+        )
+        .await
+        .unwrap();
+        let id = ShortId([0u8; 8]);
+        let start = tokio::time::Instant::now();
+        let got = db.peek_v2_payload(&id).await.expect("peek");
+        assert!(got.is_none());
+        assert_eq!(start.elapsed(), Duration::ZERO, "peek must not block");
+    }
+
+    #[tokio::test]
+    async fn peek_returns_present_payload() {
+        let dir = tempfile::tempdir().unwrap();
+        let db = FilesDb::init(
+            Duration::from_millis(10),
+            dir.path().to_owned(),
+            Duration::from_secs(60 * 60 * 24 * 7),
+        )
+        .await
+        .unwrap();
+        let id = ShortId([0u8; 8]);
+        db.post_v2_payload(&id, b"hi".to_vec()).await.unwrap().unwrap();
+        let got = db.peek_v2_payload(&id).await.expect("peek").expect("present");
+        assert_eq!(&got[..], b"hi");
+    }
 }

payjoin-mailroom/src/db/mod.rs

@@ -72,6 +72,12 @@ pub trait Db: Clone + Send + Sync + 'static {
         mailbox_id: &ShortId,
     ) -> impl Future<Output = Result<Arc<Vec<u8>>, Error<Self::OperationalError>>> + Send;
 
+    /// Read a stored v2 payload if present, without waiting.
+    fn peek_v2_payload(
+        &self,
+        mailbox_id: &ShortId,
+    ) -> impl Future<Output = Result<Option<Arc<Vec<u8>>>, Error<Self::OperationalError>>> + Send;
+
     /// Write a v1 response payload.
     fn post_v1_response(
         &self,
@@ -91,6 +97,7 @@ pub trait Db: Clone + Send + Sync + 'static {
 pub enum DbRequest {
     PostV2Payload { mailbox_id: ShortId, payload: Vec<u8> },
     WaitForV2Payload { mailbox_id: ShortId },
+    PeekV2Payload { mailbox_id: ShortId },
     PostV1Response { mailbox_id: ShortId, payload: Vec<u8> },
     PostV1RequestAndWaitForResponse { mailbox_id: ShortId, payload: Vec<u8> },
 }
@@ -99,6 +106,7 @@ pub enum DbRequest {
 pub enum DbResponse {
     PostV2Payload(Option<()>),
     WaitForV2Payload(Arc<Vec<u8>>),
+    PeekV2Payload(Option<Arc<Vec<u8>>>),
     PostV1Response(()),
     PostV1RequestAndWaitForResponse(Arc<Vec<u8>>),
 }
@@ -134,6 +142,8 @@ impl Service<DbRequest> for FilesDbService {
                     Ok(DbResponse::PostV2Payload(db.post_v2_payload(&mailbox_id, payload).await?)),
                 DbRequest::WaitForV2Payload { mailbox_id } =>
                     Ok(DbResponse::WaitForV2Payload(db.wait_for_v2_payload(&mailbox_id).await?)),
+                DbRequest::PeekV2Payload { mailbox_id } =>
+                    Ok(DbResponse::PeekV2Payload(db.peek_v2_payload(&mailbox_id).await?)),
                 DbRequest::PostV1Response { mailbox_id, payload } => {
                     db.post_v1_response(&mailbox_id, payload).await?;
                     Ok(DbResponse::PostV1Response(()))
@@ -199,6 +209,21 @@ impl Db for DbServiceAdapter {
         }
     }
 
+    async fn peek_v2_payload(
+        &self,
+        mailbox_id: &ShortId,
+    ) -> Result<Option<Arc<Vec<u8>>>, Error<Self::OperationalError>> {
+        let response = self
+            .inner
+            .clone()
+            .oneshot(DbRequest::PeekV2Payload { mailbox_id: *mailbox_id })
+            .await?;
+        match response {
+            DbResponse::PeekV2Payload(result) => Ok(result),
+            _ => Err(Self::invalid_response("peek_v2_payload")),
+        }
+    }
+
     async fn post_v1_response(
         &self,
         mailbox_id: &ShortId,
@@ -272,6 +297,14 @@ impl<D: Db> Db for MetricsDb<D> {
         self.inner.wait_for_v2_payload(mailbox_id).await
     }
 
+    async fn peek_v2_payload(
+        &self,
+        mailbox_id: &ShortId,
+    ) -> Result<Option<Arc<Vec<u8>>>, Error<Self::OperationalError>> {
+        self.metrics.record_short_id(mailbox_id);
+        self.inner.peek_v2_payload(mailbox_id).await
+    }
+
     async fn post_v1_response(
         &self,
         mailbox_id: &ShortId,

payjoin-mailroom/src/directory.rs

@@ -284,8 +284,16 @@ impl<D: Db> Service<D> {
 
     async fn get_mailbox(&self, id: &str) -> Result<Response<Body>, HandlerError> {
         let id = ShortId::from_str(id)?;
-        let timeout_response = Response::builder().status(StatusCode::ACCEPTED).body(empty())?;
-        handle_peek(self.db.wait_for_v2_payload(&id).await, timeout_response)
+        let empty_response = Response::builder().status(StatusCode::ACCEPTED).body(empty())?;
+        match self.db.peek_v2_payload(&id).await {
+            Ok(Some(payload)) => Ok(Response::new(full((*payload).clone()))),
+            Ok(None) => Ok(empty_response),
+            Err(DbError::Operational(err)) => {
+                error!("Storage error: {err}");
+                Err(HandlerError::InternalServerError(anyhow::Error::msg("Internal server error")))
+            }
+            Err(_) => Ok(empty_response),
+        }
     }
 
     /// Screen a V1 PSBT body against the address blocklist.
@@ -872,6 +880,25 @@ mod tests {
         }
     }
 
+    #[tokio::test(start_paused = true)]
+    async fn get_mailbox_returns_immediately_when_empty() {
+        let svc = test_service(None).await;
+        let id = valid_short_id_path();
+        let start = tokio::time::Instant::now();
+        let res = svc.get_mailbox(&id).await.expect("get_mailbox");
+        assert_eq!(res.status(), StatusCode::ACCEPTED);
+        assert_eq!(start.elapsed(), Duration::ZERO, "GET must not block");
+    }
+
+    #[tokio::test]
+    async fn get_mailbox_returns_payload_when_present() {
+        let svc = test_service(None).await;
+        let id = valid_short_id_path();
+        svc.post_mailbox(&id, Body::from(b"hi".to_vec())).await.expect("post");
+        let res = svc.get_mailbox(&id).await.expect("get_mailbox");
+        assert_eq!(res.status(), StatusCode::OK);
+    }
+
     #[tokio::test]
     async fn post_mailbox_records_short_id_cardinality() {
         use opentelemetry_sdk::metrics::{

payjoin/src/core/mod.rs

@@ -19,6 +19,8 @@ pub use into_url::{Error as IntoUrlError, IntoUrl};
 pub(crate) mod url;
 pub use url::{ParseError as UrlParseError, Url};
 #[cfg(feature = "v2")]
+pub mod schedule;
+#[cfg(feature = "v2")]
 pub mod time;
 pub mod uri;
 pub use uri::{PjParam, PjParseError, PjUri, Uri, UriExt};

payjoin/src/core/persist.rs

@@ -745,6 +745,10 @@ where
         }
     }
 
+    pub fn is_transient(&self) -> bool {
+        matches!(&self.0, InternalPersistedError::Api(ApiError::Transient(_)))
+    }
+
     pub fn error_state(self) -> Option<ErrorState> {
         match self.0 {
             InternalPersistedError::Api(ApiError::FatalWithState(_, state)) => Some(state),
@@ -1597,4 +1601,27 @@ mod tests {
         assert!(transient_error.storage_error_ref().is_none());
         assert!(transient_error.api_error_ref().is_some());
     }
+
+    #[test]
+    fn is_transient_only_for_transient_api_error() {
+        let transient = PersistedError::<InMemoryTestError, InMemoryTestError>(
+            InternalPersistedError::Api(ApiError::Transient(InMemoryTestError {})),
+        );
+        assert!(transient.is_transient());
+
+        let fatal = PersistedError::<InMemoryTestError, InMemoryTestError>(
+            InternalPersistedError::Api(ApiError::Fatal(InMemoryTestError {})),
+        );
+        assert!(!fatal.is_transient());
+
+        let storage = PersistedError::<InMemoryTestError, InMemoryTestError>(
+            InternalPersistedError::Storage(InMemoryTestError {}),
+        );
+        assert!(!storage.is_transient());
+
+        let fatal_with_state = PersistedError::<InMemoryTestError, InMemoryTestError>(
+            InternalPersistedError::Api(ApiError::FatalWithState(InMemoryTestError {}, ())),
+        );
+        assert!(!fatal_with_state.is_transient());
+    }
 }