fix: concurrent client registration errors (#341)

WalletConnect · Jun 5, 2024 · 85b9997 · 85b9997
1 parent 3a70f80
commit 85b9997
Show file tree

Hide file tree

Showing 2 changed files with 54 additions and 1 deletion.
diff --git a/src/stores/client.rs b/src/stores/client.rs
@@ -59,6 +59,21 @@ impl ClientStore for sqlx::PgPool {
         if let Some(metrics) = metrics {
             metrics.postgres_query("create_client_begin", start);
         }
+        // Lock the records in-case of fast concurrent requests
+        let start = Instant::now();
+        let query = "
+            SELECT
+            pg_advisory_xact_lock(abs(hashtext($1::text))),
+            pg_advisory_xact_lock(abs(hashtext($2::text)))
+        ";
+        sqlx::query(query)
+            .bind(id)
+            .bind(client.token.clone())
+            .execute(&mut transaction)
+            .await?;
+        if let Some(metrics) = metrics {
+            metrics.postgres_query("create_client_pg_advisory_xact_lock", start);
+        }
 
         let query = "
             SELECT *
@@ -79,9 +94,12 @@ impl ClientStore for sqlx::PgPool {
                 e => Err(e),
             })?;
         if let Some(metrics) = metrics {
-            metrics.postgres_query("create_client_delete", start);
+            metrics.postgres_query("create_client_select", start);
         }
 
+        #[cfg(feature = "functional_tests")]
+        tokio::time::sleep(std::time::Duration::from_millis(100)).await;
+
         if let Some(existing_client) = existing_client {
             if existing_client.id == id && existing_client.device_token != client.token {
                 let query = "

diff --git a/tests/functional/stores/client.rs b/tests/functional/stores/client.rs
@@ -32,6 +32,41 @@ async fn client_creation(ctx: &mut StoreContext) {
     ctx.clients.delete_client(TENANT_ID, &id).await.unwrap();
 }
 
+#[test_context(StoreContext)]
+#[tokio::test]
+async fn client_creation_concurrent(ctx: &mut StoreContext) {
+    let id = format!("id-{}", gen_id());
+    let token = format!("token-{}", gen_id());
+    let futures = (0..2).map(|_| {
+        // tokio::task::spawn({
+        let id = id.clone();
+        let token = token.clone();
+        let clients = ctx.clients.clone();
+        async move {
+            clients
+                .create_client(
+                    TENANT_ID,
+                    &id,
+                    Client {
+                        tenant_id: TENANT_ID.to_string(),
+                        push_type: ProviderKind::Noop,
+                        token,
+                        always_raw: false,
+                    },
+                    None,
+                )
+                .await
+        }
+        // })
+    });
+    let results = futures_util::future::join_all(futures).await;
+    for result in results {
+        result.unwrap(); //.unwrap();
+    }
+    // Cleaning up records
+    ctx.clients.delete_client(TENANT_ID, &id).await.unwrap();
+}
+
 #[test_context(StoreContext)]
 #[tokio::test]
 async fn client_creation_fcm(ctx: &mut StoreContext) {