Retry PoolNeedsConnection Errors (#1010)

* big refator of connection

* just make the error retryable

* add new error that can be caught

* add retries to the places that kick out into signature requests

* Apply suggestions from code review

* fix parenthesis

---------

Co-authored-by: Andrew Plaza <[email protected]>
Co-authored-by: Andrew Plaza <[email protected]>

xmtp_mls/src/groups/mod.rs

@@ -188,6 +188,8 @@ pub enum GroupError {
     SqlKeyStore(#[from] sql_key_store::SqlKeyStoreError),
     #[error("No pending commit found")]
     MissingPendingCommit,
+    #[error("Sync failed to wait for intent")]
+    SyncFailedToWait,
 }
 
 impl RetryableError for GroupError {

xmtp_mls/src/groups/sync.rs

@@ -221,7 +221,7 @@ impl MlsGroup {
             num_attempts += 1;
         }
 
-        Err(last_err.unwrap_or(GroupError::Generic("failed to wait for intent".to_string())))
+        Err(last_err.unwrap_or(GroupError::SyncFailedToWait))
     }
 
     fn is_valid_epoch(

xmtp_mls/src/identity_updates.rs

@@ -1,6 +1,10 @@
 use std::collections::{HashMap, HashSet};
 
-use crate::{retry::RetryableError, retryable, storage::association_state::StoredAssociationState};
+use crate::{
+    retry::{Retry, RetryableError},
+    retry_async, retryable,
+    storage::association_state::StoredAssociationState,
+};
 use prost::Message;
 use thiserror::Error;
 use xmtp_id::associations::{
@@ -239,9 +243,13 @@ where
         wallets_to_revoke: Vec<String>,
     ) -> Result<SignatureRequest, ClientError> {
         let inbox_id = self.inbox_id();
-        let current_state = self
-            .get_association_state(&self.store().conn()?, &inbox_id, None)
-            .await?;
+        let current_state = retry_async!(
+            Retry::default(),
+            (async {
+                self.get_association_state(&self.store().conn()?, &inbox_id, None)
+                    .await
+            })
+        )?;
         let mut builder = SignatureRequestBuilder::new(inbox_id);
 
         for wallet in wallets_to_revoke {
@@ -259,9 +267,14 @@ where
         installation_ids: Vec<Vec<u8>>,
     ) -> Result<SignatureRequest, ClientError> {
         let inbox_id = self.inbox_id();
-        let current_state = self
-            .get_association_state(&self.store().conn()?, &inbox_id, None)
-            .await?;
+
+        let current_state = retry_async!(
+            Retry::default(),
+            (async {
+                self.get_association_state(&self.store().conn()?, &inbox_id, None)
+                    .await
+            })
+        )?;
 
         let mut builder = SignatureRequestBuilder::new(inbox_id);
 
@@ -291,7 +304,17 @@ where
             .await?;
 
         // Load the identity updates for the inbox so that we have a record in our DB
-        load_identity_updates(&self.api_client, &self.store().conn()?, vec![inbox_id]).await?;
+        retry_async!(
+            Retry::default(),
+            (async {
+                load_identity_updates(
+                    &self.api_client,
+                    &self.store().conn()?,
+                    vec![inbox_id.clone()],
+                )
+                .await
+            })
+        )?;
 
         Ok(())
     }

xmtp_mls/src/storage/errors.rs

@@ -68,6 +68,7 @@ impl RetryableError for StorageError {
             Self::Pool(_) => true,
             Self::Lock(_) => true,
             Self::SqlCipherNotLoaded => true,
+            Self::PoolNeedsConnection => true,
             _ => false,
         }
     }