Improved consistency check + docs

Cargo.toml

@@ -59,7 +59,7 @@ hex-literal = "0.4.1"
 indicatif = "0.17.8"
 reqwest = { version = "0.11.22", features = ["json"] }
 statrs = "0.17.1"
-
+tokio-retry = "0.3.0"
 
 [dev-dependencies]
 indoc = "2.0.5"

docs/consistency_check.md

@@ -0,0 +1,28 @@
+To verify that the world tree service and signup sequencer are in sync we can run a consistency check.
+
+On a very basic level this is as simple as running
+```
+./utils consistency
+```
+
+This will fetch inclusion proofs for a set of randomly selected hardcoded identities. Correct output should look like this:
+```
+================================================= 39/39 [00:00:02]Matches: 39
+Mismatches: 0
+Missing: 0
+Failures: 0
+```
+
+# With DB access
+If you have access to the sequencer database you can run a more thorough consistency check by fetching a new set of identities.
+
+To do that connect to the database and fetch new identities, assuming you're connecting to a localhost db the command would look like this
+```
+psql postgresql://postgres:postgres@localhost:5432/db  -c "WITH latest_leaf_id as (SELECT MAX(id) as id, leaf_index FROM identities GROUP BY leaf_index) SELECT encode(commitment, 'hex') FROM identities INNER JOIN latest_leaf_id ON latest_leaf_id.id = identities.id WHERE commitment <> '\x0000000000000000000000000000000000000000000000000000000000000000' AND status = 'mined' ORDER BY RANDOM() LIMIT 2400" -t -A > identities
+```
+then run the consistency check using this newly created file.
+
+## Note
+Make sure to include `commitment <> '\x0000000000000000000000000000000000000000000000000000000000000000'` in the query, otherwise the output will try to fetch inclusion proofs for zero identities which is invalid.
+
+If you omit `status = 'mined'` from the query the output might include fresh identities which haven't been mined yet - these will identities will be present in the sequncer but not known yet to the world-tree service. But it's a useful test to run anyway, as running the same check in ~1-2 hours should match all identities.

src/bin/utils.rs

@@ -1,9 +1,10 @@
 use std::path::{Path, PathBuf};
-use std::sync::atomic::{AtomicUsize, Ordering};
+use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
 use std::sync::Arc;
 use std::time::{Duration, Instant};
 
 use clap::{Args, Parser, Subcommand};
+use futures::stream::{self, StreamExt};
 use indicatif::{ProgressBar, ProgressStyle};
 use rand::seq::SliceRandom;
 use rand::{thread_rng, Rng};
@@ -12,6 +13,9 @@ use ruint::aliases::U256;
 use serde::{Deserialize, Serialize};
 use serde_json::json;
 use statrs::statistics::Statistics;
+use tokio::sync::Mutex;
+use tokio_retry::strategy::{jitter, ExponentialBackoff};
+use tokio_retry::RetryIf;
 use world_tree::tree::Hash;
 
 macro_rules! hash {
@@ -88,6 +92,10 @@ struct Common {
         default_value = "https://world-tree.crypto.worldcoin.org/inclusionProof"
     )]
     world_tree_endpoint: String,
+
+    /// Number of concurrent jobs to run
+    #[clap(short, long, default_value = "10")]
+    jobs: usize,
 }
 
 #[derive(Debug, Clone, Subcommand)]
@@ -150,6 +158,7 @@ async fn main() -> eyre::Result<()> {
                 args.common.world_tree_endpoint,
                 sequencer_endpoint,
                 &identities,
+                args.common.jobs,
             )
             .await
         }
@@ -280,52 +289,186 @@ async fn run_consistency_check(
     world_tree_endpoint: String,
     sequencer_endpoint: String,
     identities: &[Hash],
+    jobs: usize,
 ) -> eyre::Result<()> {
     let client = Client::new();
 
-    let progress_bar = ProgressBar::new(identities.len() as u64);
-    progress_bar.set_style(
+    let total_progress_bar = ProgressBar::new(identities.len() as u64);
+    total_progress_bar.set_style(
         ProgressStyle::default_bar()
             .template("{wide_bar} {pos}/{len} [{elapsed_precise}]")?
             .progress_chars("=> "),
     );
 
-    for identity in identities {
-        let world_tree_response = client
-            .post(&world_tree_endpoint)
-            .json(&json!({
-                "identityCommitment": identity.to_string(),
-            }))
-            .send()
-            .await?
-            .error_for_status()?;
-
-        let world_tree_response: InclusionProof =
-            world_tree_response.json().await?;
-
-        let sequencer_response = client
-            .post(&sequencer_endpoint)
-            .json(&json!({
-                "identityCommitment": identity.to_string(),
-            }))
-            .send()
-            .await?
-            .error_for_status()?;
-
-        let sequencer_response: InclusionProof =
-            sequencer_response.json().await?;
-
-        assert_eq!(world_tree_response.root, sequencer_response.root);
-        assert_eq!(world_tree_response.proof, sequencer_response.proof);
-
-        progress_bar.inc(1);
+    let matches_count = Arc::new(AtomicU64::new(0));
+    let mismatches_count = Arc::new(AtomicU64::new(0));
+    let missing_count = Arc::new(AtomicU64::new(0));
+    let failures_count = Arc::new(AtomicU64::new(0));
+
+    let missing_identities: Arc<Mutex<Vec<Hash>>> =
+        Arc::new(Mutex::new(Vec::new()));
+    let failures: Arc<Mutex<Vec<Hash>>> = Arc::new(Mutex::new(Vec::new()));
+
+    let concurrency_limit = jobs; // Adjust as needed
+
+    stream::iter(identities.iter().cloned())
+        .map(|identity| {
+            let client = client.clone();
+            let world_tree_endpoint = world_tree_endpoint.clone();
+            let sequencer_endpoint = sequencer_endpoint.clone();
+            let matches_count = matches_count.clone();
+            let mismatches_count = mismatches_count.clone();
+            let missing_count = missing_count.clone();
+            let failures_count = failures_count.clone();
+            let total_progress_bar = total_progress_bar.clone();
+            let missing_identities = missing_identities.clone();
+            let failures = failures.clone();
+
+            async move {
+                match get_world_tree_inclusion_proof(
+                    &client,
+                    &world_tree_endpoint,
+                    &identity,
+                )
+                .await
+                {
+                    Ok(Some(world_tree_response)) => {
+                        match get_sequencer_inclusion_proof(
+                            &client,
+                            &sequencer_endpoint,
+                            &identity,
+                        )
+                        .await
+                        {
+                            Ok(sequencer_response) => {
+                                if world_tree_response.root
+                                    == sequencer_response.root
+                                    && world_tree_response.proof
+                                        == sequencer_response.proof
+                                {
+                                    matches_count
+                                        .fetch_add(1, Ordering::SeqCst);
+                                } else {
+                                    mismatches_count
+                                        .fetch_add(1, Ordering::SeqCst);
+                                }
+                            }
+                            Err(_) => {
+                                failures_count.fetch_add(1, Ordering::SeqCst);
+
+                                failures.lock().await.push(identity);
+                            }
+                        }
+                    }
+                    Ok(None) => {
+                        // Identity missing on world-tree
+                        missing_count.fetch_add(1, Ordering::SeqCst);
+
+                        missing_identities.lock().await.push(identity);
+                    }
+                    Err(_) => {
+                        failures_count.fetch_add(1, Ordering::SeqCst);
+
+                        failures.lock().await.push(identity);
+                    }
+                }
+                total_progress_bar.inc(1);
+            }
+        })
+        .buffer_unordered(concurrency_limit)
+        .collect::<()>()
+        .await;
+
+    total_progress_bar.finish_with_message("Done!");
+
+    println!("Matches: {}", matches_count.load(Ordering::SeqCst));
+    println!("Mismatches: {}", mismatches_count.load(Ordering::SeqCst));
+    println!("Missing: {}", missing_count.load(Ordering::SeqCst));
+    println!("Failures: {}", failures_count.load(Ordering::SeqCst));
+
+    let missing_identities = missing_identities.lock().await;
+    if !missing_identities.is_empty() {
+        println!("Missing identities:");
+        for identity in missing_identities.iter() {
+            println!("{}", identity);
+        }
     }
 
-    progress_bar.finish_with_message("Done!");
+    let failures = failures.lock().await;
+    if !failures.is_empty() {
+        println!("Failures:");
+        for identity in failures.iter() {
+            println!("{}", identity);
+        }
+    }
 
     Ok(())
 }
 
+async fn get_world_tree_inclusion_proof(
+    client: &Client,
+    endpoint: &str,
+    identity: &Hash,
+) -> Result<Option<InclusionProof>, reqwest::Error> {
+    let retry_strategy =
+        ExponentialBackoff::from_millis(10).map(jitter).take(5);
+
+    RetryIf::spawn(
+        retry_strategy,
+        || async {
+            let response = client
+                .post(endpoint)
+                .json(&json!({
+                    "identityCommitment": identity.to_string(),
+                }))
+                .send()
+                .await?;
+
+            if response.status() == 404 {
+                Ok(None)
+            } else {
+                let response = response.error_for_status()?;
+                let inclusion_proof = response.json().await?;
+                Ok(Some(inclusion_proof))
+            }
+        },
+        is_retryable_error,
+    )
+    .await
+}
+
+async fn get_sequencer_inclusion_proof(
+    client: &Client,
+    endpoint: &str,
+    identity: &Hash,
+) -> Result<InclusionProof, reqwest::Error> {
+    let retry_strategy =
+        ExponentialBackoff::from_millis(10).map(jitter).take(5);
+
+    RetryIf::spawn(
+        retry_strategy,
+        || async {
+            let response = client
+                .post(endpoint)
+                .json(&json!({
+                    "identityCommitment": identity.to_string(),
+                }))
+                .send()
+                .await?;
+
+            let response = response.error_for_status()?;
+            let inclusion_proof = response.json().await?;
+            Ok(inclusion_proof)
+        },
+        is_retryable_error,
+    )
+    .await
+}
+
+fn is_retryable_error(err: &reqwest::Error) -> bool {
+    err.is_timeout() || err.is_connect() || err.is_request()
+}
+
 impl Common {
     fn identities(&self) -> Vec<Hash> {
         if let Some(identities_file) = &self.identities_file {