From 289bf597196d6904a6556f1b0a23b7e221c09565 Mon Sep 17 00:00:00 2001
From: Ziyang Hu <hu.ziyang@cantab.net>
Date: Sun, 4 Dec 2022 21:59:36 +0800
Subject: [PATCH] retry if busy

---
 cozo-core/benches/pokec.rs | 53 +++++++++++++++++++++++++-------------
 1 file changed, 35 insertions(+), 18 deletions(-)

diff --git a/cozo-core/benches/pokec.rs b/cozo-core/benches/pokec.rs
index 5b9a2f3b..820d2999 100644
--- a/cozo-core/benches/pokec.rs
+++ b/cozo-core/benches/pokec.rs
@@ -227,12 +227,17 @@ fn single_vertex() {
 
 fn single_vertex_write() {
     let i = rand::thread_rng().gen_range(1..SIZES.0 * 10);
-    TEST_DB
-        .run_script(
-            "?[uid, cmpl_pct, gender, age] <- [[$id, 0, null, null]] :put user {uid => cmpl_pct, gender, age}",
-            BTreeMap::from([("id".to_string(), json!(i))]),
-        )
-        .unwrap();
+    for _ in 0..10 {
+        if TEST_DB
+            .run_script(
+                "?[uid, cmpl_pct, gender, age] <- [[$id, 0, null, null]] :put user {uid => cmpl_pct, gender, age}",
+                BTreeMap::from([("id".to_string(), json!(i))]),
+            )
+            .is_ok() {
+            return
+        }
+    }
+    panic!()
 }
 
 fn single_edge_write() {
@@ -241,28 +246,40 @@ fn single_edge_write() {
     while j == i {
         j = rand::thread_rng().gen_range(1..SIZES.0);
     }
-    TEST_DB
-        .run_script(
-            r#"
+    for _ in 0..10 {
+        if TEST_DB
+            .run_script(
+                r#"
             {?[fr, to] <- [[$i, $j]] :put friends {fr, to}}
             {?[fr, to] <- [[$i, $j]] :put friends.rev {fr, to}}
             "#,
-            BTreeMap::from([("i".to_string(), json!(i)), ("j".to_string(), json!(j))]),
-        )
-        .unwrap();
+                BTreeMap::from([("i".to_string(), json!(i)), ("j".to_string(), json!(j))]),
+            )
+            .is_ok()
+        {
+            return;
+        }
+    }
+    panic!()
 }
 
 fn single_vertex_update() {
     let i = rand::thread_rng().gen_range(1..SIZES.0);
-    TEST_DB
-        .run_script(
-            r#"
+    for _ in 0..10 {
+        if TEST_DB
+            .run_script(
+                r#"
             ?[uid, cmpl_pct, age, gender] := uid = $id, *user{uid, age, gender}, cmpl_pct = -1
             :put user {uid => cmpl_pct, age, gender}
             "#,
-            BTreeMap::from([("id".to_string(), json!(i))]),
-        )
-        .unwrap();
+                BTreeMap::from([("id".to_string(), json!(i))]),
+            )
+            .is_ok()
+        {
+            return;
+        }
+    }
+    panic!()
 }
 
 fn aggregation() {