From 9a283f9e1af7609811ceb2684989102c33982f77 Mon Sep 17 00:00:00 2001 From: Piotr Heilman Date: Wed, 10 Jul 2024 17:54:16 +0200 Subject: [PATCH 1/7] Add pre_root for identities. --- ...15_add_constraints_for_identities.down.sql | 7 + .../015_add_constraints_for_identities.up.sql | 52 ++++++++ src/database/mod.rs | 125 +++++++++++------- src/database/query.rs | 37 +++++- src/database/types.rs | 4 + src/task_monitor/tasks/create_batches.rs | 8 +- src/task_monitor/tasks/delete_identities.rs | 10 +- src/task_monitor/tasks/insert_identities.rs | 75 +++++------ 8 files changed, 220 insertions(+), 98 deletions(-) create mode 100644 schemas/database/015_add_constraints_for_identities.down.sql create mode 100644 schemas/database/015_add_constraints_for_identities.up.sql diff --git a/schemas/database/015_add_constraints_for_identities.down.sql b/schemas/database/015_add_constraints_for_identities.down.sql new file mode 100644 index 00000000..f6392d40 --- /dev/null +++ b/schemas/database/015_add_constraints_for_identities.down.sql @@ -0,0 +1,7 @@ +DROP UNIQUE INDEX idx_unique_insertion_leaf; +DROP UNIQUE INDEX idx_unique_deletion_leaf; + +DROP TRIGGER validate_pre_root_trigger; +DROP FUNCTION validate_pre_root(); + +ALTER TABLE identities DROP COLUMN pre_root; \ No newline at end of file diff --git a/schemas/database/015_add_constraints_for_identities.up.sql b/schemas/database/015_add_constraints_for_identities.up.sql new file mode 100644 index 00000000..028c4cca --- /dev/null +++ b/schemas/database/015_add_constraints_for_identities.up.sql @@ -0,0 +1,52 @@ +CREATE UNIQUE INDEX idx_unique_insertion_leaf on identities(leaf_index) WHERE commitment != E'\\x0000000000000000000000000000000000000000000000000000000000000000'; +CREATE UNIQUE INDEX idx_unique_deletion_leaf on identities(leaf_index) WHERE commitment = E'\\x0000000000000000000000000000000000000000000000000000000000000000'; + +-- Add the new 'prev_root' column +ALTER TABLE identities ADD COLUMN pre_root BYTEA; + +-- This constraint ensures that we have consistent database and changes to the tre are done in a valid sequence. +CREATE OR REPLACE FUNCTION validate_pre_root() returns trigger as $$ + DECLARE + last_id identities.id%type; + last_root identities.root%type; + BEGIN + SELECT id, root + INTO last_id, last_root + FROM identities + ORDER BY id DESC + LIMIT 1; + + -- When last_id is NULL that means there are no records in identities table. The first prev_root can + -- be a value not referencing previous root in database. + IF last_id IS NULL THEN RETURN NEW; + END IF; + + IF NEW.pre_root IS NULL THEN RAISE EXCEPTION 'Sent pre_root (%) can be null only for first record in table.', NEW.pre_root; + END IF; + + IF (last_root != NEW.pre_root) THEN RAISE EXCEPTION 'Sent pre_root (%) is different than last root (%) in database.', NEW.pre_root, last_root; + END IF; + + RETURN NEW; + END; +$$ language plpgsql; + +CREATE TRIGGER validate_pre_root_trigger BEFORE INSERT ON identities FOR EACH ROW EXECUTE PROCEDURE validate_pre_root(); + +-- Below function took around 10 minutes for 10 million records. +DO +$do$ +DECLARE + prev_root identities.pre_root%type := NULL; + identity identities%rowtype; +BEGIN + FOR identity IN SELECT * FROM identities ORDER BY id ASC + LOOP + IF identity.pre_root IS NULL THEN UPDATE identities SET pre_root = prev_root WHERE id = identity.id; + END IF; + prev_root = identity.root; + END LOOP; +END +$do$; + +CREATE UNIQUE INDEX idx_unique_pre_root on identities(pre_root) WHERE pre_root IS NOT NULL; diff --git a/src/database/mod.rs b/src/database/mod.rs index 61ee8cfa..a0cbf2bf 100644 --- a/src/database/mod.rs +++ b/src/database/mod.rs @@ -158,6 +158,7 @@ mod test { use ethers::types::U256; use postgres_docker_utils::DockerContainer; use ruint::Uint; + use semaphore::poseidon_tree::LazyPoseidonTree; use semaphore::Field; use testcontainers::clients::Cli; @@ -284,12 +285,15 @@ mod test { let (db, _db_container) = setup_db(&docker).await?; let zero: Hash = U256::zero().into(); + let initial_root = LazyPoseidonTree::new(4, zero).root(); let zero_root: Hash = U256::from_dec_str("6789")?.into(); let root: Hash = U256::from_dec_str("54321")?.into(); let commitment: Hash = U256::from_dec_str("12345")?.into(); - db.insert_pending_identity(0, &commitment, &root).await?; - db.insert_pending_identity(0, &zero, &zero_root).await?; + db.insert_pending_identity(0, &commitment, &root, &initial_root) + .await?; + db.insert_pending_identity(0, &zero, &zero_root, &root) + .await?; let leaf_index = db .get_identity_leaf_index(&commitment) @@ -557,23 +561,6 @@ mod test { Ok(()) } - #[tokio::test] - async fn test_update_insertion_timestamp() -> anyhow::Result<()> { - let docker = Cli::default(); - let (db, _db_container) = setup_db(&docker).await?; - - let insertion_timestamp = Utc::now(); - - db.update_latest_insertion_timestamp(insertion_timestamp) - .await?; - - let latest_insertion_timestamp = db.get_latest_insertion_timestamp().await?.unwrap(); - - assert!(latest_insertion_timestamp.timestamp() - insertion_timestamp.timestamp() <= 1); - - Ok(()) - } - #[tokio::test] async fn test_insert_deletion() -> anyhow::Result<()> { let docker = Cli::default(); @@ -637,6 +624,7 @@ mod test { let docker = Cli::default(); let (db, _db_container) = setup_db(&docker).await?; + let initial_root = LazyPoseidonTree::new(4, Hash::ZERO).root(); let identities = mock_identities(1); let roots = mock_roots(1); @@ -644,7 +632,7 @@ mod test { assert_eq!(next_leaf_index, 0, "Db should contain not leaf indexes"); - db.insert_pending_identity(0, &identities[0], &roots[0]) + db.insert_pending_identity(0, &identities[0], &roots[0], &initial_root) .await?; let next_leaf_index = db.get_next_leaf_index().await?; @@ -658,13 +646,16 @@ mod test { let docker = Cli::default(); let (db, _db_container) = setup_db(&docker).await?; + let initial_root = LazyPoseidonTree::new(4, Hash::ZERO).root(); let identities = mock_identities(5); let roots = mock_roots(5); + let mut pre_root = &initial_root; for i in 0..5 { - db.insert_pending_identity(i, &identities[i], &roots[i]) + db.insert_pending_identity(i, &identities[i], &roots[i], pre_root) .await .context("Inserting identity")?; + pre_root = &roots[i]; } db.mark_root_as_processed_tx(&roots[2]).await?; @@ -688,13 +679,16 @@ mod test { let docker = Cli::default(); let (db, _db_container) = setup_db(&docker).await?; + let initial_root = LazyPoseidonTree::new(4, Hash::ZERO).root(); let identities = mock_identities(5); let roots = mock_roots(5); + let mut pre_root = &initial_root; for i in 0..5 { - db.insert_pending_identity(i, &identities[i], &roots[i]) + db.insert_pending_identity(i, &identities[i], &roots[i], pre_root) .await .context("Inserting identity")?; + pre_root = &roots[i]; } db.mark_root_as_processed_tx(&roots[2]).await?; @@ -731,13 +725,16 @@ mod test { let docker = Cli::default(); let (db, _db_container) = setup_db(&docker).await?; + let initial_root = LazyPoseidonTree::new(4, Hash::ZERO).root(); let identities = mock_identities(5); let roots = mock_roots(5); + let mut pre_root = &initial_root; for i in 0..5 { - db.insert_pending_identity(i, &identities[i], &roots[i]) + db.insert_pending_identity(i, &identities[i], &roots[i], pre_root) .await .context("Inserting identity")?; + pre_root = &roots[i]; } db.mark_root_as_mined_tx(&roots[2]).await?; @@ -776,13 +773,16 @@ mod test { let num_identities = 6; + let initial_root = LazyPoseidonTree::new(4, Hash::ZERO).root(); let identities = mock_identities(num_identities); let roots = mock_roots(num_identities); + let mut pre_root = &initial_root; for i in 0..num_identities { - db.insert_pending_identity(i, &identities[i], &roots[i]) + db.insert_pending_identity(i, &identities[i], &roots[i], pre_root) .await .context("Inserting identity")?; + pre_root = &roots[i]; } println!("Marking roots up to 2nd as processed"); @@ -818,13 +818,16 @@ mod test { let docker = Cli::default(); let (db, _db_container) = setup_db(&docker).await?; + let initial_root = LazyPoseidonTree::new(4, Hash::ZERO).root(); let identities = mock_identities(5); let roots = mock_roots(5); + let mut pre_root = &initial_root; for i in 0..5 { - db.insert_pending_identity(i, &identities[i], &roots[i]) + db.insert_pending_identity(i, &identities[i], &roots[i], pre_root) .await .context("Inserting identity")?; + pre_root = &roots[i]; } // root[2] is somehow erroneously marked as mined @@ -862,6 +865,7 @@ mod test { let docker = Cli::default(); let (db, _db_container) = setup_db(&docker).await?; + let initial_root = LazyPoseidonTree::new(4, Hash::ZERO).root(); let identities = mock_identities(5); let roots = mock_roots(5); @@ -869,7 +873,7 @@ mod test { assert!(root.is_none(), "Root should not exist"); - db.insert_pending_identity(0, &identities[0], &roots[0]) + db.insert_pending_identity(0, &identities[0], &roots[0], &initial_root) .await?; let root = db @@ -920,14 +924,17 @@ mod test { let docker = Cli::default(); let (db, _db_container) = setup_db(&docker).await?; + let initial_root = LazyPoseidonTree::new(4, Hash::ZERO).root(); let identities = mock_identities(5); let roots = mock_roots(7); + let mut pre_root = &initial_root; for i in 0..5 { - db.insert_pending_identity(i, &identities[i], &roots[i]) + db.insert_pending_identity(i, &identities[i], &roots[i], pre_root) .await .context("Inserting identity")?; + pre_root = &roots[i]; } db.mark_root_as_processed_tx(&roots[2]).await?; @@ -959,20 +966,23 @@ mod test { let docker = Cli::default(); let (db, _db_container) = setup_db(&docker).await?; + let initial_root = LazyPoseidonTree::new(4, Hash::ZERO).root(); let identities = mock_identities(5); let roots = mock_roots(5); let zero_roots = mock_zero_roots(5); + let mut pre_root = &initial_root; for i in 0..5 { - db.insert_pending_identity(i, &identities[i], &roots[i]) + db.insert_pending_identity(i, &identities[i], &roots[i], pre_root) .await .context("Inserting identity")?; + pre_root = &roots[i]; } - db.insert_pending_identity(0, &Hash::ZERO, &zero_roots[0]) + db.insert_pending_identity(0, &Hash::ZERO, &zero_roots[0], &roots[4]) .await?; - db.insert_pending_identity(3, &Hash::ZERO, &zero_roots[3]) + db.insert_pending_identity(3, &Hash::ZERO, &zero_roots[3], &zero_roots[0]) .await?; let pending_tree_updates = db @@ -1014,10 +1024,11 @@ mod test { let docker = Cli::default(); let (db, _db_container) = setup_db(&docker).await?; + let initial_root = LazyPoseidonTree::new(4, Hash::ZERO).root(); let identities = mock_identities(5); let roots = mock_roots(5); - db.insert_pending_identity(0, &identities[0], &roots[0]) + db.insert_pending_identity(0, &identities[0], &roots[0], &initial_root) .await .context("Inserting identity 1")?; @@ -1034,9 +1045,9 @@ mod test { // Inserting a new pending root sets invalidation time for the // previous root - db.insert_pending_identity(1, &identities[1], &roots[1]) + db.insert_pending_identity(1, &identities[1], &roots[1], &roots[0]) .await?; - db.insert_pending_identity(2, &identities[2], &roots[2]) + db.insert_pending_identity(2, &identities[2], &roots[2], &roots[1]) .await?; let root_1_inserted_at = Utc::now(); @@ -1056,7 +1067,7 @@ mod test { assert_same_time!(root_item_1.pending_valid_as_of, root_1_inserted_at); // Test mined roots - db.insert_pending_identity(3, &identities[3], &roots[3]) + db.insert_pending_identity(3, &identities[3], &roots[3], &roots[2]) .await?; db.mark_root_as_processed_tx(&roots[0]) @@ -1091,6 +1102,7 @@ mod test { let docker = Cli::default(); let (db, _db_container) = setup_db(&docker).await?; + let initial_root = LazyPoseidonTree::new(4, Hash::ZERO).root(); let identities = mock_identities(2); let roots = mock_roots(1); @@ -1106,7 +1118,7 @@ mod test { assert!(db.identity_exists(identities[0]).await?); // When there's only processed identity - db.insert_pending_identity(0, &identities[1], &roots[0]) + db.insert_pending_identity(0, &identities[1], &roots[0], &initial_root) .await .context("Inserting identity")?; @@ -1148,7 +1160,35 @@ mod test { } #[tokio::test] - async fn test_latest_deletion_root() -> anyhow::Result<()> { + async fn test_latest_insertion() -> anyhow::Result<()> { + let docker = Cli::default(); + let (db, _db_container) = setup_db(&docker).await?; + + // Update with initial timestamp + let initial_timestamp = chrono::Utc::now(); + db.update_latest_insertion(initial_timestamp) + .await + .context("Inserting initial root")?; + + // Assert values + let initial_entry = db.get_latest_insertion().await?; + assert!(initial_entry.timestamp.timestamp() - initial_timestamp.timestamp() <= 1); + + // Update with a new timestamp + let new_timestamp = chrono::Utc::now(); + db.update_latest_insertion(new_timestamp) + .await + .context("Updating with new root")?; + + // Assert values + let new_entry = db.get_latest_insertion().await?; + assert!((new_entry.timestamp.timestamp() - new_timestamp.timestamp()) <= 1); + + Ok(()) + } + + #[tokio::test] + async fn test_latest_deletion() -> anyhow::Result<()> { let docker = Cli::default(); let (db, _db_container) = setup_db(&docker).await?; @@ -1179,25 +1219,20 @@ mod test { async fn can_not_insert_same_root_multiple_times() -> anyhow::Result<()> { let docker = Cli::default(); let (db, _db_container) = setup_db(&docker).await?; + + let initial_root = LazyPoseidonTree::new(4, Hash::ZERO).root(); let identities = mock_identities(2); let roots = mock_roots(2); - db.insert_pending_identity(0, &identities[0], &roots[0]) + db.insert_pending_identity(0, &identities[0], &roots[0], &initial_root) .await?; let res = db - .insert_pending_identity(1, &identities[1], &roots[0]) + .insert_pending_identity(1, &identities[1], &roots[0], &roots[0]) .await; assert!(res.is_err(), "Inserting duplicate root should fail"); - let root_state = db - .get_root_state(&roots[0]) - .await? - .context("Missing root")?; - - assert_eq!(root_state.status, ProcessedStatus::Pending); - Ok(()) } diff --git a/src/database/query.rs b/src/database/query.rs index 26121c88..01d2f1f9 100644 --- a/src/database/query.rs +++ b/src/database/query.rs @@ -6,7 +6,9 @@ use sqlx::{Executor, Postgres, Row}; use tracing::instrument; use types::{DeletionEntry, LatestDeletionEntry, RecoveryEntry}; -use crate::database::types::{BatchEntry, BatchEntryData, BatchType, TransactionEntry}; +use crate::database::types::{ + BatchEntry, BatchEntryData, BatchType, LatestInsertionEntry, TransactionEntry, +}; use crate::database::{types, Error}; use crate::identity_tree::{ Hash, ProcessedStatus, RootItem, TreeItem, TreeUpdate, UnprocessedStatus, @@ -25,17 +27,19 @@ pub trait DatabaseQuery<'a>: Executor<'a, Database = Postgres> { leaf_index: usize, identity: &Hash, root: &Hash, + pre_root: &Hash, ) -> Result<(), Error> { let insert_pending_identity_query = sqlx::query( r#" - INSERT INTO identities (leaf_index, commitment, root, status, pending_as_of) - VALUES ($1, $2, $3, $4, CURRENT_TIMESTAMP) + INSERT INTO identities (leaf_index, commitment, root, status, pending_as_of, pre_root) + VALUES ($1, $2, $3, $4, CURRENT_TIMESTAMP, $5) "#, ) .bind(leaf_index as i64) .bind(identity) .bind(root) - .bind(<&str>::from(ProcessedStatus::Pending)); + .bind(<&str>::from(ProcessedStatus::Pending)) + .bind(pre_root); self.execute(insert_pending_identity_query).await?; @@ -182,6 +186,17 @@ pub trait DatabaseQuery<'a>: Executor<'a, Database = Postgres> { .collect()) } + async fn get_latest_root(self) -> Result, Error> { + Ok(sqlx::query( + r#" + SELECT root FROM identities ORDER BY id DESC LIMIT 1 + "#, + ) + .fetch_optional(self) + .await? + .map(|r| r.get::(0))) + } + async fn get_latest_root_by_status( self, status: ProcessedStatus, @@ -218,7 +233,7 @@ pub trait DatabaseQuery<'a>: Executor<'a, Database = Postgres> { .await?) } - async fn get_latest_insertion_timestamp(self) -> Result>, Error> { + async fn get_latest_insertion(self) -> Result { let query = sqlx::query( r#" SELECT insertion_timestamp @@ -228,7 +243,15 @@ pub trait DatabaseQuery<'a>: Executor<'a, Database = Postgres> { let row = self.fetch_optional(query).await?; - Ok(row.map(|r| r.get::, _>(0))) + if let Some(row) = row { + Ok(LatestInsertionEntry { + timestamp: row.get(0), + }) + } else { + Ok(LatestInsertionEntry { + timestamp: Utc::now(), + }) + } } async fn count_unprocessed_identities(self) -> Result { @@ -383,7 +406,7 @@ pub trait DatabaseQuery<'a>: Executor<'a, Database = Postgres> { } } - async fn update_latest_insertion_timestamp( + async fn update_latest_insertion( self, insertion_timestamp: DateTime, ) -> Result<(), Error> { diff --git a/src/database/types.rs b/src/database/types.rs index 768f6be6..a0fbe585 100644 --- a/src/database/types.rs +++ b/src/database/types.rs @@ -24,6 +24,10 @@ pub struct RecoveryEntry { pub new_commitment: Hash, } +pub struct LatestInsertionEntry { + pub timestamp: DateTime, +} + pub struct LatestDeletionEntry { pub timestamp: DateTime, } diff --git a/src/task_monitor/tasks/create_batches.rs b/src/task_monitor/tasks/create_batches.rs index 1df6eebc..4f16e36f 100644 --- a/src/task_monitor/tasks/create_batches.rs +++ b/src/task_monitor/tasks/create_batches.rs @@ -50,11 +50,7 @@ pub async fn create_batches( // inserted. If we have an incomplete batch but are within a small delta of the // tick happening anyway in the wake branch, we insert the current // (possibly-incomplete) batch anyway. - let mut last_batch_time: DateTime = app - .database - .get_latest_insertion_timestamp() - .await? - .unwrap_or(Utc::now()); + let mut last_batch_time: DateTime = app.database.get_latest_insertion().await?.timestamp; loop { // We wait either for a timer tick or a full batch @@ -127,7 +123,7 @@ pub async fn create_batches( timer.reset(); last_batch_time = Utc::now(); app.database - .update_latest_insertion_timestamp(last_batch_time) + .update_latest_insertion(last_batch_time) .await?; } else { // Check if the next batch after the current insertion batch is diff --git a/src/task_monitor/tasks/delete_identities.rs b/src/task_monitor/tasks/delete_identities.rs index a499553c..35aca551 100644 --- a/src/task_monitor/tasks/delete_identities.rs +++ b/src/task_monitor/tasks/delete_identities.rs @@ -13,6 +13,12 @@ use crate::database::query::DatabaseQuery; use crate::database::types::DeletionEntry; use crate::identity_tree::{Hash, TreeVersionReadOps}; +// Deletion here differs from insert_identites task. This is because two +// different flows are created for both tasks. Due to how our prover works +// (can handle only a batch of same operations types - insertion or deletion) +// we want to group together insertions and deletions. We are doing it by +// grouping deletions (as the not need to be put into tree immediately as +// insertions) and putting them into the tree pub async fn delete_identities( app: Arc, pending_insertions_mutex: Arc>, @@ -70,6 +76,7 @@ pub async fn delete_identities( } } + let mut pre_root = app.tree_state()?.latest_tree().get_root(); // Delete the commitments at the target leaf indices in the latest tree, // generating the proof for each update let data = app.tree_state()?.latest_tree().delete_many(&leaf_indices); @@ -84,8 +91,9 @@ pub async fn delete_identities( let items = data.into_iter().zip(leaf_indices); for ((root, _proof), leaf_index) in items { app.database - .insert_pending_identity(leaf_index, &Hash::ZERO, &root) + .insert_pending_identity(leaf_index, &Hash::ZERO, &root, &pre_root) .await?; + pre_root = root; } // Remove the previous commitments from the deletions table diff --git a/src/task_monitor/tasks/insert_identities.rs b/src/task_monitor/tasks/insert_identities.rs index e7afd568..bbd680a7 100644 --- a/src/task_monitor/tasks/insert_identities.rs +++ b/src/task_monitor/tasks/insert_identities.rs @@ -1,6 +1,7 @@ use std::sync::Arc; use std::time::Duration; +use sqlx::{Postgres, Transaction}; use tokio::sync::{Mutex, Notify}; use tokio::time; use tracing::info; @@ -8,10 +9,14 @@ use tracing::info; use crate::app::App; use crate::database::query::DatabaseQuery as _; use crate::database::types::UnprocessedCommitment; -use crate::database::Database; use crate::identity_tree::{Latest, TreeVersion, TreeVersionReadOps, UnprocessedStatus}; use crate::retry_tx; +// Insertion here differs from delete_identities task. This is because two +// different flows are created for both tasks. We need to insert identities as +// fast as possible to the tree to be able to return inclusion proof as our +// customers depend on it. Flow here is to rewrite from unprocessed_identities +// into identities every 5 seconds. pub async fn insert_identities( app: Arc, pending_insertions_mutex: Arc>, @@ -34,12 +39,12 @@ pub async fn insert_identities( continue; } - insert_identities_batch( - &app.database, - app.tree_state()?.latest_tree(), - &unprocessed, - &pending_insertions_mutex, - ) + let _guard = pending_insertions_mutex.lock().await; + let latest_tree = app.tree_state()?.latest_tree(); + + retry_tx!(&app.database, tx, { + insert_identities_batch(&mut tx, latest_tree, &unprocessed).await + }) .await?; // Notify the identity processing task, that there are new identities @@ -48,35 +53,28 @@ pub async fn insert_identities( } pub async fn insert_identities_batch( - database: &Database, + tx: &mut Transaction<'_, Postgres>, latest_tree: &TreeVersion, identities: &[UnprocessedCommitment], - pending_insertions_mutex: &Mutex<()>, -) -> anyhow::Result<()> { - let filtered_identities = retry_tx!(database, tx, { - // Filter out any identities that are already in the `identities` table - let mut filtered_identities = vec![]; - for identity in identities { - if tx - .get_identity_leaf_index(&identity.commitment) - .await? - .is_some() - { - tracing::warn!(?identity.commitment, "Duplicate identity"); - tx.remove_unprocessed_identity(&identity.commitment).await?; - } else { - filtered_identities.push(identity.commitment); - } +) -> Result<(), anyhow::Error> { + // Filter out any identities that are already in the `identities` table + let mut filtered_identities = vec![]; + for identity in identities { + if tx + .get_identity_leaf_index(&identity.commitment) + .await? + .is_some() + { + tracing::warn!(?identity.commitment, "Duplicate identity"); + tx.remove_unprocessed_identity(&identity.commitment).await?; + } else { + filtered_identities.push(identity.commitment); } - Result::<_, anyhow::Error>::Ok(filtered_identities) - }) - .await?; - - let _guard = pending_insertions_mutex.lock().await; + } let next_leaf = latest_tree.next_leaf(); - let next_db_index = retry_tx!(database, tx, tx.get_next_leaf_index().await).await?; + let next_db_index = tx.get_next_leaf_index().await?; assert_eq!( next_leaf, next_db_index, @@ -84,6 +82,7 @@ pub async fn insert_identities_batch( {next_db_index}" ); + let mut pre_root = &latest_tree.get_root(); let data = latest_tree.append_many(&filtered_identities); assert_eq!( @@ -92,15 +91,13 @@ pub async fn insert_identities_batch( "Length mismatch when appending identities to tree" ); - retry_tx!(database, tx, { - for ((root, _proof, leaf_index), identity) in data.iter().zip(&filtered_identities) { - tx.insert_pending_identity(*leaf_index, identity, root) - .await?; + for ((root, _proof, leaf_index), identity) in data.iter().zip(&filtered_identities) { + tx.insert_pending_identity(*leaf_index, identity, root, pre_root) + .await?; + pre_root = root; - tx.remove_unprocessed_identity(identity).await?; - } + tx.remove_unprocessed_identity(identity).await?; + } - Result::<_, anyhow::Error>::Ok(()) - }) - .await + Ok(()) } From e458afb2fb0361b205cd5103a105153588848187 Mon Sep 17 00:00:00 2001 From: Piotr Heilman Date: Wed, 17 Jul 2024 09:45:23 +0200 Subject: [PATCH 2/7] Join insert/delete tasks. Fix flaky tests. --- src/database/query.rs | 11 + src/task_monitor/mod.rs | 44 +--- src/task_monitor/tasks/create_batches.rs | 3 + src/task_monitor/tasks/delete_identities.rs | 103 -------- src/task_monitor/tasks/insert_identities.rs | 103 -------- src/task_monitor/tasks/mod.rs | 3 +- src/task_monitor/tasks/modify_tree.rs | 235 ++++++++++++++++++ tests/common/mod.rs | 27 +- tests/tree_restore_with_root_back_to_init.rs | 13 +- .../tree_restore_with_root_back_to_middle.rs | 15 +- tests/validate_proof_with_age.rs | 2 +- 11 files changed, 301 insertions(+), 258 deletions(-) delete mode 100644 src/task_monitor/tasks/delete_identities.rs delete mode 100644 src/task_monitor/tasks/insert_identities.rs create mode 100644 src/task_monitor/tasks/modify_tree.rs diff --git a/src/database/query.rs b/src/database/query.rs index 01d2f1f9..54cb7a51 100644 --- a/src/database/query.rs +++ b/src/database/query.rs @@ -486,6 +486,17 @@ pub trait DatabaseQuery<'a>: Executor<'a, Database = Postgres> { Ok(()) } + async fn count_deletions(self) -> Result { + let query = sqlx::query( + r#" + SELECT COUNT(*) + FROM deletions + "#, + ); + let result = self.fetch_one(query).await?; + Ok(result.get::(0) as i32) + } + // TODO: consider using a larger value than i64 for leaf index, ruint should // have postgres compatibility for u256 async fn get_deletions(self) -> Result, Error> { diff --git a/src/task_monitor/mod.rs b/src/task_monitor/mod.rs index 6318c21e..e9ffceea 100644 --- a/src/task_monitor/mod.rs +++ b/src/task_monitor/mod.rs @@ -18,8 +18,7 @@ const TREE_INIT_BACKOFF: Duration = Duration::from_secs(5); const PROCESS_IDENTITIES_BACKOFF: Duration = Duration::from_secs(5); const FINALIZE_IDENTITIES_BACKOFF: Duration = Duration::from_secs(5); const QUEUE_MONITOR_BACKOFF: Duration = Duration::from_secs(5); -const INSERT_IDENTITIES_BACKOFF: Duration = Duration::from_secs(5); -const DELETE_IDENTITIES_BACKOFF: Duration = Duration::from_secs(5); +const MODIFY_TREE_BACKOFF: Duration = Duration::from_secs(5); struct RunningInstance { handles: Vec>, @@ -204,46 +203,19 @@ impl TaskMonitor { ); handles.push(monitor_txs_handle); - let pending_insertion_mutex = Arc::new(Mutex::new(())); - - // Insert identities - let app = self.app.clone(); - let wake_up_notify = base_wake_up_notify.clone(); - let insertion_mutex = pending_insertion_mutex.clone(); - let insert_identities = move || { - self::tasks::insert_identities::insert_identities( - app.clone(), - insertion_mutex.clone(), - wake_up_notify.clone(), - ) - }; - - let insert_identities_handle = crate::utils::spawn_monitored_with_backoff( - insert_identities, - shutdown_sender.clone(), - INSERT_IDENTITIES_BACKOFF, - self.shutdown.clone(), - ); - handles.push(insert_identities_handle); - - // Delete identities + // Modify tree let app = self.app.clone(); let wake_up_notify = base_wake_up_notify.clone(); - let delete_identities = move || { - self::tasks::delete_identities::delete_identities( - app.clone(), - pending_insertion_mutex.clone(), - wake_up_notify.clone(), - ) - }; + let modify_tree = + move || tasks::modify_tree::modify_tree(app.clone(), wake_up_notify.clone()); - let delete_identities_handle = crate::utils::spawn_monitored_with_backoff( - delete_identities, + let modify_tree_handle = crate::utils::spawn_monitored_with_backoff( + modify_tree, shutdown_sender.clone(), - DELETE_IDENTITIES_BACKOFF, + MODIFY_TREE_BACKOFF, self.shutdown.clone(), ); - handles.push(delete_identities_handle); + handles.push(modify_tree_handle); // Create the instance *instance = Some(RunningInstance { diff --git a/src/task_monitor/tasks/create_batches.rs b/src/task_monitor/tasks/create_batches.rs index 4f16e36f..27e479ac 100644 --- a/src/task_monitor/tasks/create_batches.rs +++ b/src/task_monitor/tasks/create_batches.rs @@ -34,6 +34,9 @@ pub async fn create_batches( tracing::info!("Awaiting for a clean slate"); app.identity_processor.await_clean_slate().await?; + tracing::info!("Awaiting for initialized tree"); + app.tree_state()?; + tracing::info!("Starting batch creator."); ensure_batch_chain_initialized(&app).await?; diff --git a/src/task_monitor/tasks/delete_identities.rs b/src/task_monitor/tasks/delete_identities.rs deleted file mode 100644 index 35aca551..00000000 --- a/src/task_monitor/tasks/delete_identities.rs +++ /dev/null @@ -1,103 +0,0 @@ -use std::collections::HashSet; -use std::sync::Arc; -use std::time::Duration; - -use anyhow::Context; -use chrono::Utc; -use tokio::sync::{Mutex, Notify}; -use tokio::time; -use tracing::info; - -use crate::app::App; -use crate::database::query::DatabaseQuery; -use crate::database::types::DeletionEntry; -use crate::identity_tree::{Hash, TreeVersionReadOps}; - -// Deletion here differs from insert_identites task. This is because two -// different flows are created for both tasks. Due to how our prover works -// (can handle only a batch of same operations types - insertion or deletion) -// we want to group together insertions and deletions. We are doing it by -// grouping deletions (as the not need to be put into tree immediately as -// insertions) and putting them into the tree -pub async fn delete_identities( - app: Arc, - pending_insertions_mutex: Arc>, - wake_up_notify: Arc, -) -> anyhow::Result<()> { - info!("Starting deletion processor."); - - let batch_deletion_timeout = chrono::Duration::from_std(app.config.app.batch_deletion_timeout) - .context("Invalid batch deletion timeout duration")?; - - let mut timer = time::interval(Duration::from_secs(5)); - - loop { - _ = timer.tick().await; - info!("Deletion processor woken due to timeout"); - - let deletions = app.database.get_deletions().await?; - if deletions.is_empty() { - continue; - } - - let last_deletion_timestamp = app.database.get_latest_deletion().await?.timestamp; - - // If the minimum deletions batch size is not reached and the deletion time - // interval has not elapsed then we can skip - if deletions.len() < app.config.app.min_batch_deletion_size - && Utc::now() - last_deletion_timestamp <= batch_deletion_timeout - { - continue; - } - - // Dedup deletion entries - let deletions = deletions.into_iter().collect::>(); - - let (leaf_indices, previous_commitments): (Vec, Vec) = deletions - .iter() - .map(|d| (d.leaf_index, d.commitment)) - .unzip(); - - let _guard = pending_insertions_mutex.lock().await; - - // Check if the deletion batch could potentially create a duplicate root batch - if let Some(last_leaf_index) = app.tree_state()?.latest_tree().next_leaf().checked_sub(1) { - let mut sorted_indices = leaf_indices.clone(); - sorted_indices.sort(); - - let indices_are_continuous = sorted_indices.windows(2).all(|w| w[1] == w[0] + 1); - - if indices_are_continuous && sorted_indices.last().unwrap() == &last_leaf_index { - tracing::warn!( - "Deletion batch could potentially create a duplicate root batch. Deletion \ - batch will be postponed" - ); - continue; - } - } - - let mut pre_root = app.tree_state()?.latest_tree().get_root(); - // Delete the commitments at the target leaf indices in the latest tree, - // generating the proof for each update - let data = app.tree_state()?.latest_tree().delete_many(&leaf_indices); - - assert_eq!( - data.len(), - leaf_indices.len(), - "Length mismatch when appending identities to tree" - ); - - // Insert the new items into pending identities - let items = data.into_iter().zip(leaf_indices); - for ((root, _proof), leaf_index) in items { - app.database - .insert_pending_identity(leaf_index, &Hash::ZERO, &root, &pre_root) - .await?; - pre_root = root; - } - - // Remove the previous commitments from the deletions table - app.database.remove_deletions(&previous_commitments).await?; - wake_up_notify.notify_one(); - } -} diff --git a/src/task_monitor/tasks/insert_identities.rs b/src/task_monitor/tasks/insert_identities.rs deleted file mode 100644 index bbd680a7..00000000 --- a/src/task_monitor/tasks/insert_identities.rs +++ /dev/null @@ -1,103 +0,0 @@ -use std::sync::Arc; -use std::time::Duration; - -use sqlx::{Postgres, Transaction}; -use tokio::sync::{Mutex, Notify}; -use tokio::time; -use tracing::info; - -use crate::app::App; -use crate::database::query::DatabaseQuery as _; -use crate::database::types::UnprocessedCommitment; -use crate::identity_tree::{Latest, TreeVersion, TreeVersionReadOps, UnprocessedStatus}; -use crate::retry_tx; - -// Insertion here differs from delete_identities task. This is because two -// different flows are created for both tasks. We need to insert identities as -// fast as possible to the tree to be able to return inclusion proof as our -// customers depend on it. Flow here is to rewrite from unprocessed_identities -// into identities every 5 seconds. -pub async fn insert_identities( - app: Arc, - pending_insertions_mutex: Arc>, - wake_up_notify: Arc, -) -> anyhow::Result<()> { - info!("Starting insertion processor task."); - - let mut timer = time::interval(Duration::from_secs(5)); - - loop { - _ = timer.tick().await; - info!("Insertion processor woken due to timeout."); - - // get commits from database - let unprocessed = app - .database - .get_eligible_unprocessed_commitments(UnprocessedStatus::New) - .await?; - if unprocessed.is_empty() { - continue; - } - - let _guard = pending_insertions_mutex.lock().await; - let latest_tree = app.tree_state()?.latest_tree(); - - retry_tx!(&app.database, tx, { - insert_identities_batch(&mut tx, latest_tree, &unprocessed).await - }) - .await?; - - // Notify the identity processing task, that there are new identities - wake_up_notify.notify_one(); - } -} - -pub async fn insert_identities_batch( - tx: &mut Transaction<'_, Postgres>, - latest_tree: &TreeVersion, - identities: &[UnprocessedCommitment], -) -> Result<(), anyhow::Error> { - // Filter out any identities that are already in the `identities` table - let mut filtered_identities = vec![]; - for identity in identities { - if tx - .get_identity_leaf_index(&identity.commitment) - .await? - .is_some() - { - tracing::warn!(?identity.commitment, "Duplicate identity"); - tx.remove_unprocessed_identity(&identity.commitment).await?; - } else { - filtered_identities.push(identity.commitment); - } - } - - let next_leaf = latest_tree.next_leaf(); - - let next_db_index = tx.get_next_leaf_index().await?; - - assert_eq!( - next_leaf, next_db_index, - "Database and tree are out of sync. Next leaf index in tree is: {next_leaf}, in database: \ - {next_db_index}" - ); - - let mut pre_root = &latest_tree.get_root(); - let data = latest_tree.append_many(&filtered_identities); - - assert_eq!( - data.len(), - filtered_identities.len(), - "Length mismatch when appending identities to tree" - ); - - for ((root, _proof, leaf_index), identity) in data.iter().zip(&filtered_identities) { - tx.insert_pending_identity(*leaf_index, identity, root, pre_root) - .await?; - pre_root = root; - - tx.remove_unprocessed_identity(identity).await?; - } - - Ok(()) -} diff --git a/src/task_monitor/tasks/mod.rs b/src/task_monitor/tasks/mod.rs index fbbee77e..d6c44db6 100644 --- a/src/task_monitor/tasks/mod.rs +++ b/src/task_monitor/tasks/mod.rs @@ -1,7 +1,6 @@ pub mod create_batches; -pub mod delete_identities; pub mod finalize_identities; -pub mod insert_identities; +pub mod modify_tree; pub mod monitor_queue; pub mod monitor_txs; pub mod process_batches; diff --git a/src/task_monitor/tasks/modify_tree.rs b/src/task_monitor/tasks/modify_tree.rs new file mode 100644 index 00000000..454d1e80 --- /dev/null +++ b/src/task_monitor/tasks/modify_tree.rs @@ -0,0 +1,235 @@ +use std::sync::Arc; +use std::time::Duration; + +use anyhow::Context; +use chrono::Utc; +use sqlx::{Postgres, Transaction}; +use tokio::sync::Notify; +use tokio::{select, time}; +use tracing::info; + +use crate::app::App; +use crate::database::query::DatabaseQuery as _; +use crate::database::types::DeletionEntry; +use crate::identity_tree::{Hash, TreeState, TreeVersionReadOps, UnprocessedStatus}; +use crate::retry_tx; + +// Because tree operations are single threaded (done one by one) we are running +// them from single task that determines which type of operations to run. It is +// done that way to reduce number of used mutexes and eliminate the risk of some +// tasks not being run at all as mutex is not preserving unlock order. +pub async fn modify_tree(app: Arc, wake_up_notify: Arc) -> anyhow::Result<()> { + info!("Starting modify tree task."); + + let batch_deletion_timeout = chrono::Duration::from_std(app.config.app.batch_deletion_timeout) + .context("Invalid batch deletion timeout duration")?; + let min_batch_deletion_size = app.config.app.min_batch_deletion_size; + + let mut timer = time::interval(Duration::from_secs(5)); + + loop { + // We wait either for a timer tick or a full batch + select! { + _ = timer.tick() => { + info!("Modify tree task woken due to timeout"); + } + + () = wake_up_notify.notified() => { + info!("Modify tree task woken due to request"); + }, + } + + let tree_state = app.tree_state()?; + + retry_tx!(&app.database, tx, { + do_modify_tree( + &mut tx, + batch_deletion_timeout, + min_batch_deletion_size, + tree_state, + &wake_up_notify, + ) + .await + }) + .await?; + + // wake_up_notify.notify_one(); + } +} + +async fn do_modify_tree( + tx: &mut Transaction<'_, Postgres>, + batch_deletion_timeout: chrono::Duration, + min_batch_deletion_size: usize, + tree_state: &TreeState, + wake_up_notify: &Arc, +) -> anyhow::Result<()> { + let deletions = tx.get_deletions().await?; + + // Deleting identities has precedence over inserting them. + if should_run_deletion( + tx, + batch_deletion_timeout, + min_batch_deletion_size, + tree_state, + &deletions, + ) + .await? + { + run_deletions(tx, tree_state, deletions, wake_up_notify).await?; + } else { + run_insertions(tx, tree_state, wake_up_notify).await?; + } + + Ok(()) +} + +pub async fn should_run_deletion( + tx: &mut Transaction<'_, Postgres>, + batch_deletion_timeout: chrono::Duration, + min_batch_deletion_size: usize, + tree_state: &TreeState, + deletions: &[DeletionEntry], +) -> anyhow::Result { + let last_deletion_timestamp = tx.get_latest_deletion().await?.timestamp; + + if deletions.is_empty() { + return Ok(false); + } + + // If min batch size is not reached and batch deletion timeout not elapsed + if deletions.len() < min_batch_deletion_size + && Utc::now() - last_deletion_timestamp <= batch_deletion_timeout + { + return Ok(false); + } + + // Now also check if the deletion batch could potentially create a duplicate + // root batch + if let Some(last_leaf_index) = tree_state.latest_tree().next_leaf().checked_sub(1) { + let mut sorted_indices: Vec = deletions.iter().map(|v| v.leaf_index).collect(); + sorted_indices.sort(); + + let indices_are_continuous = sorted_indices.windows(2).all(|w| w[1] == w[0] + 1); + + if indices_are_continuous && sorted_indices.last().unwrap() == &last_leaf_index { + tracing::warn!( + "Deletion batch could potentially create a duplicate root batch. Deletion batch \ + will be postponed." + ); + return Ok(false); + } + } + + Ok(true) +} + +pub async fn run_insertions( + tx: &mut Transaction<'_, Postgres>, + tree_state: &TreeState, + wake_up_notify: &Arc, +) -> anyhow::Result<()> { + let unprocessed = tx + .get_eligible_unprocessed_commitments(UnprocessedStatus::New) + .await?; + if unprocessed.is_empty() { + return Ok(()); + } + + let latest_tree = tree_state.latest_tree(); + + // Filter out any identities that are already in the `identities` table + let mut filtered_identities = vec![]; + for identity in unprocessed { + if tx + .get_identity_leaf_index(&identity.commitment) + .await? + .is_some() + { + tracing::warn!(?identity.commitment, "Duplicate identity"); + tx.remove_unprocessed_identity(&identity.commitment).await?; + } else { + filtered_identities.push(identity.commitment); + } + } + + let next_leaf = latest_tree.next_leaf(); + + let next_db_index = tx.get_next_leaf_index().await?; + + assert_eq!( + next_leaf, next_db_index, + "Database and tree are out of sync. Next leaf index in tree is: {next_leaf}, in database: \ + {next_db_index}" + ); + + let mut pre_root = &latest_tree.get_root(); + let data = latest_tree.append_many(&filtered_identities); + + assert_eq!( + data.len(), + filtered_identities.len(), + "Length mismatch when appending identities to tree" + ); + + for ((root, _proof, leaf_index), identity) in data.iter().zip(&filtered_identities) { + tx.insert_pending_identity(*leaf_index, identity, root, pre_root) + .await?; + pre_root = root; + + tx.remove_unprocessed_identity(identity).await?; + } + + // Immediately look for next operations + wake_up_notify.notify_one(); + + Ok(()) +} + +pub async fn run_deletions( + tx: &mut Transaction<'_, Postgres>, + tree_state: &TreeState, + mut deletions: Vec, + wake_up_notify: &Arc, +) -> anyhow::Result<()> { + if deletions.is_empty() { + return Ok(()); + } + + // This sorting is very important. It ensures that we will create a unique root + // after deletion. It mostly ensures that we won't delete things in reverse + // order as they were added. + deletions.sort_by(|v1, v2| v1.leaf_index.cmp(&v2.leaf_index)); + + let (leaf_indices, previous_commitments): (Vec, Vec) = deletions + .iter() + .map(|d| (d.leaf_index, d.commitment)) + .unzip(); + + let mut pre_root = tree_state.latest_tree().get_root(); + // Delete the commitments at the target leaf indices in the latest tree, + // generating the proof for each update + let data = tree_state.latest_tree().delete_many(&leaf_indices); + + assert_eq!( + data.len(), + leaf_indices.len(), + "Length mismatch when appending identities to tree" + ); + + // Insert the new items into pending identities + let items = data.into_iter().zip(leaf_indices); + for ((root, _proof), leaf_index) in items { + tx.insert_pending_identity(leaf_index, &Hash::ZERO, &root, &pre_root) + .await?; + pre_root = root; + } + + // Remove the previous commitments from the deletions table + tx.remove_deletions(&previous_commitments).await?; + + // Immediately look for next operations + wake_up_notify.notify_one(); + + Ok(()) +} diff --git a/tests/common/mod.rs b/tests/common/mod.rs index 32d43492..544c3fb5 100644 --- a/tests/common/mod.rs +++ b/tests/common/mod.rs @@ -714,6 +714,22 @@ fn construct_verify_proof_body( pub async fn spawn_app( config: Config, ) -> anyhow::Result<(Arc, JoinHandle<()>, SocketAddr, Arc)> { + spawn_app_returning_initialized_tree(config).await.map( + |(app, join_handle, socket_addr, shutdown, _)| (app, join_handle, socket_addr, shutdown), + ) +} + +#[instrument(skip_all)] +#[allow(clippy::type_complexity)] +pub async fn spawn_app_returning_initialized_tree( + config: Config, +) -> anyhow::Result<( + Arc, + JoinHandle<()>, + SocketAddr, + Arc, + TreeState, +)> { let server_config = config.server.clone(); let app = App::new(config).await.expect("Failed to create App"); let shutdown = Arc::new(Shutdown::new()); @@ -728,8 +744,9 @@ pub async fn spawn_app( // For our tests to work we need the tree to be initialized. while app.tree_state().is_err() { trace!("Waiting for the tree to be initialized"); - tokio::time::sleep(Duration::from_millis(250)).await; + tokio::time::sleep(Duration::from_millis(100)).await; } + let initialized_tree_state = app.tree_state()?.clone(); let app_clone = app.clone(); let shutdown_clone = shutdown.clone(); @@ -756,7 +773,13 @@ pub async fn spawn_app( info!("App ready"); - Ok((app, app_handle, local_addr, shutdown)) + Ok(( + app, + app_handle, + local_addr, + shutdown, + initialized_tree_state, + )) } pub async fn check_metrics(socket_addr: &SocketAddr) -> anyhow::Result<()> { diff --git a/tests/tree_restore_with_root_back_to_init.rs b/tests/tree_restore_with_root_back_to_init.rs index f81a7dd7..884c78ac 100644 --- a/tests/tree_restore_with_root_back_to_init.rs +++ b/tests/tree_restore_with_root_back_to_init.rs @@ -2,6 +2,8 @@ mod common; use common::prelude::*; +use crate::common::spawn_app_returning_initialized_tree; + const IDLE_TIME: u64 = 7; #[tokio::test] @@ -109,7 +111,7 @@ async fn tree_restore_with_root_back_to_init(offchain_mode_enabled: bool) -> any ) .await; - tokio::time::sleep(Duration::from_secs(2)).await; + tokio::time::sleep(Duration::from_secs(5)).await; let tree_state = app.tree_state()?.clone(); @@ -144,13 +146,14 @@ async fn tree_restore_with_root_back_to_init(offchain_mode_enabled: bool) -> any info!("Starting the app again for testing purposes"); - let (app, app_handle, local_addr, shutdown) = spawn_app(config.clone()) - .await - .expect("Failed to spawn app."); + let (_, app_handle, local_addr, shutdown, initialized_tree_state) = + spawn_app_returning_initialized_tree(config.clone()) + .await + .expect("Failed to spawn app."); let uri = "http://".to_owned() + &local_addr.to_string(); - let restored_tree_state = app.tree_state()?.clone(); + let restored_tree_state = initialized_tree_state; assert_eq!( restored_tree_state.latest_tree().get_root(), diff --git a/tests/tree_restore_with_root_back_to_middle.rs b/tests/tree_restore_with_root_back_to_middle.rs index 812017ef..b35f510f 100644 --- a/tests/tree_restore_with_root_back_to_middle.rs +++ b/tests/tree_restore_with_root_back_to_middle.rs @@ -2,6 +2,8 @@ mod common; use common::prelude::*; +use crate::common::spawn_app_returning_initialized_tree; + const IDLE_TIME: u64 = 7; #[tokio::test] @@ -109,7 +111,7 @@ async fn tree_restore_with_root_back_to_middle(offchain_mode_enabled: bool) -> a ) .await; - tokio::time::sleep(Duration::from_secs(2)).await; + tokio::time::sleep(Duration::from_secs(5)).await; let mid_root: U256 = ref_tree.root().into(); @@ -157,7 +159,7 @@ async fn tree_restore_with_root_back_to_middle(offchain_mode_enabled: bool) -> a ) .await; - tokio::time::sleep(Duration::from_secs(2)).await; + tokio::time::sleep(Duration::from_secs(5)).await; let tree_state = app.tree_state()?.clone(); @@ -190,13 +192,14 @@ async fn tree_restore_with_root_back_to_middle(offchain_mode_enabled: bool) -> a .offchain_mode(offchain_mode_enabled) .build()?; - let (app, app_handle, local_addr, shutdown) = spawn_app(config.clone()) - .await - .expect("Failed to spawn app."); + let (_, app_handle, local_addr, shutdown, initialized_tree) = + spawn_app_returning_initialized_tree(config.clone()) + .await + .expect("Failed to spawn app."); let uri = "http://".to_owned() + &local_addr.to_string(); - let restored_tree_state = app.tree_state()?.clone(); + let restored_tree_state = initialized_tree; assert_eq!( restored_tree_state.latest_tree().get_root(), diff --git a/tests/validate_proof_with_age.rs b/tests/validate_proof_with_age.rs index 58417df5..ff5e23c6 100644 --- a/tests/validate_proof_with_age.rs +++ b/tests/validate_proof_with_age.rs @@ -140,7 +140,7 @@ async fn validate_proof_with_age(offchain_mode_enabled: bool) -> anyhow::Result< ) .await; - let max_age_of_proof = (Instant::now() - time_of_identity_insertion).as_secs(); + let max_age_of_proof = (Instant::now() - time_of_identity_insertion).as_secs() + 2; assert!( max_age_of_proof > sleep_duration_seconds * 2, "Proof age should be at least 2 batch times" From 2858723728af7ea4d4ce819bbc07e9f7741e8798 Mon Sep 17 00:00:00 2001 From: Piotr Heilman Date: Wed, 24 Jul 2024 12:01:12 +0200 Subject: [PATCH 3/7] Sync in-memory trees using DB (DB is source of truth now). --- Dockerfile | 2 +- e2e_tests/docker-compose/compose.yml | 14 +- .../scenarios/tests/common/docker_compose.rs | 18 + e2e_tests/scenarios/tests/insert_100.rs | 2 +- src/database/query.rs | 51 +- src/identity/processor.rs | 29 +- src/identity_tree/builder.rs | 198 ++++++ src/identity_tree/initializer.rs | 86 +-- src/identity_tree/mod.rs | 654 +++++++++--------- src/identity_tree/state.rs | 67 ++ src/task_monitor/mod.rs | 68 +- src/task_monitor/tasks/create_batches.rs | 44 +- src/task_monitor/tasks/finalize_identities.rs | 6 +- src/task_monitor/tasks/mod.rs | 1 + src/task_monitor/tasks/modify_tree.rs | 62 +- src/task_monitor/tasks/process_batches.rs | 13 +- .../tasks/sync_tree_state_with_db.rs | 123 ++++ src/utils/tree_updates.rs | 58 +- 18 files changed, 1032 insertions(+), 464 deletions(-) create mode 100644 src/identity_tree/builder.rs create mode 100644 src/identity_tree/state.rs create mode 100644 src/task_monitor/tasks/sync_tree_state_with_db.rs diff --git a/Dockerfile b/Dockerfile index 206a97e8..6cb2a197 100644 --- a/Dockerfile +++ b/Dockerfile @@ -1,4 +1,4 @@ -FROM debian:12 as build-env +FROM debian:12 AS build-env WORKDIR /src diff --git a/e2e_tests/docker-compose/compose.yml b/e2e_tests/docker-compose/compose.yml index ef0e1cb8..05af04c7 100644 --- a/e2e_tests/docker-compose/compose.yml +++ b/e2e_tests/docker-compose/compose.yml @@ -140,21 +140,21 @@ services: command: [ "/config.toml" ] environment: - RUST_LOG=debug -# signup-sequencer-1: -# <<: *signup-sequencer-def -# hostname: signup-sequencer-1 -# ports: -# - ${SIGNUP_SEQUENCER_0_PORT:-9081}:8080 + signup-sequencer-1: + <<: *signup-sequencer-def + hostname: signup-sequencer-1 + ports: + - ${SIGNUP_SEQUENCER_1_PORT:-9081}:8080 # signup-sequencer-2: # <<: *signup-sequencer-def # hostname: signup-sequencer-2 # ports: -# - ${SIGNUP_SEQUENCER_0_PORT:-9082}:8080 +# - ${SIGNUP_SEQUENCER_2_PORT:-9082}:8080 # signup-sequencer-3: # <<: *signup-sequencer-def # hostname: signup-sequencer-3 # ports: -# - ${SIGNUP_SEQUENCER_0_PORT:-9083}:8080 +# - ${SIGNUP_SEQUENCER_3_PORT:-9083}:8080 volumes: tx_sitter_db_data: driver: local diff --git a/e2e_tests/scenarios/tests/common/docker_compose.rs b/e2e_tests/scenarios/tests/common/docker_compose.rs index 0fa4fa60..ca534ef8 100644 --- a/e2e_tests/scenarios/tests/common/docker_compose.rs +++ b/e2e_tests/scenarios/tests/common/docker_compose.rs @@ -26,6 +26,9 @@ pub struct DockerComposeGuard<'a> { semaphore_insertion_port: u32, semaphore_deletion_port: u32, signup_sequencer_0_port: u32, + signup_sequencer_1_port: u32, + signup_sequencer_2_port: u32, + signup_sequencer_3_port: u32, signup_sequencer_balancer_port: u32, } @@ -82,6 +85,18 @@ impl<'a> DockerComposeGuard<'a> { String::from("SIGNUP_SEQUENCER_0_PORT"), self.signup_sequencer_0_port.to_string(), ); + res.insert( + String::from("SIGNUP_SEQUENCER_1_PORT"), + self.signup_sequencer_1_port.to_string(), + ); + res.insert( + String::from("SIGNUP_SEQUENCER_2_PORT"), + self.signup_sequencer_2_port.to_string(), + ); + res.insert( + String::from("SIGNUP_SEQUENCER_3_PORT"), + self.signup_sequencer_3_port.to_string(), + ); res.insert( String::from("SIGNUP_SEQUENCER_BALANCER_PORT"), self.signup_sequencer_balancer_port.to_string(), @@ -152,6 +167,9 @@ pub async fn setup(cwd: &str) -> anyhow::Result { semaphore_insertion_port: 0, semaphore_deletion_port: 0, signup_sequencer_0_port: 0, + signup_sequencer_1_port: 0, + signup_sequencer_2_port: 0, + signup_sequencer_3_port: 0, signup_sequencer_balancer_port: 0, }; diff --git a/e2e_tests/scenarios/tests/insert_100.rs b/e2e_tests/scenarios/tests/insert_100.rs index 10e72eb2..a96588f0 100644 --- a/e2e_tests/scenarios/tests/insert_100.rs +++ b/e2e_tests/scenarios/tests/insert_100.rs @@ -16,7 +16,7 @@ async fn insert_100() -> anyhow::Result<()> { let uri = format!("http://{}", docker_compose.get_local_addr()); let client = Client::new(); - let identities = generate_test_commitments(10); + let identities = generate_test_commitments(100); for commitment in identities.iter() { insert_identity_with_retries(&client, &uri, commitment, 10, 3.0).await?; diff --git a/src/database/query.rs b/src/database/query.rs index 54cb7a51..b8b33e30 100644 --- a/src/database/query.rs +++ b/src/database/query.rs @@ -134,7 +134,7 @@ pub trait DatabaseQuery<'a>: Executor<'a, Database = Postgres> { ) -> Result, Error> { Ok(sqlx::query_as::<_, TreeUpdate>( r#" - SELECT leaf_index, commitment as element + SELECT id as sequence_id, leaf_index, commitment as element, root as post_root FROM identities WHERE status = $1 ORDER BY id ASC; @@ -152,7 +152,7 @@ pub trait DatabaseQuery<'a>: Executor<'a, Database = Postgres> { let statuses: Vec<&str> = statuses.into_iter().map(<&str>::from).collect(); Ok(sqlx::query_as::<_, TreeUpdate>( r#" - SELECT leaf_index, commitment as element + SELECT id as sequence_id, leaf_index, commitment as element, root as post_root FROM identities WHERE status = ANY($1) ORDER BY id ASC; @@ -163,6 +163,20 @@ pub trait DatabaseQuery<'a>: Executor<'a, Database = Postgres> { .await?) } + async fn get_commitments_after_id(self, id: usize) -> Result, Error> { + Ok(sqlx::query_as::<_, TreeUpdate>( + r#" + SELECT id as sequence_id, leaf_index, commitment as element, root as post_root + FROM identities + WHERE id > $1 + ORDER BY id ASC; + "#, + ) + .bind(id as i64) + .fetch_all(self) + .await?) + } + async fn get_non_zero_commitments_by_leaf_indexes>( self, leaf_indexes: I, @@ -233,6 +247,39 @@ pub trait DatabaseQuery<'a>: Executor<'a, Database = Postgres> { .await?) } + async fn get_latest_tree_update_by_statuses( + self, + statuses: Vec, + ) -> Result, Error> { + let statuses: Vec<&str> = statuses.into_iter().map(<&str>::from).collect(); + Ok(sqlx::query_as::<_, TreeUpdate>( + r#" + SELECT id as sequence_id, leaf_index, commitment as element, root as post_root + FROM identities + WHERE status = ANY($1) + ORDER BY id DESC + LIMIT 1; + "#, + ) + .bind(&statuses[..]) // Official workaround https://github.com/launchbadge/sqlx/blob/main/FAQ.md#how-can-i-do-a-select--where-foo-in--query + .fetch_optional(self) + .await?) + } + + async fn get_tree_update_by_root(self, root: &Hash) -> Result, Error> { + Ok(sqlx::query_as::<_, TreeUpdate>( + r#" + SELECT id as sequence_id, leaf_index, commitment as element, root as post_root + FROM identities + WHERE root = $1 + LIMIT 1; + "#, + ) + .bind(root) // Official workaround https://github.com/launchbadge/sqlx/blob/main/FAQ.md#how-can-i-do-a-select--where-foo-in--query + .fetch_optional(self) + .await?) + } + async fn get_latest_insertion(self) -> Result { let query = sqlx::query( r#" diff --git a/src/identity/processor.rs b/src/identity/processor.rs index 41c3e1e5..63c32517 100644 --- a/src/identity/processor.rs +++ b/src/identity/processor.rs @@ -10,6 +10,7 @@ use ethers::addressbook::Address; use ethers::contract::EthEvent; use ethers::middleware::Middleware; use ethers::prelude::{Log, Topic, ValueOrArray, U256}; +use tokio::sync::Notify; use tracing::{error, info, instrument}; use crate::config::Config; @@ -20,7 +21,7 @@ use crate::database::query::DatabaseQuery; use crate::database::types::{BatchEntry, BatchType}; use crate::database::{Database, Error}; use crate::ethereum::{Ethereum, ReadProvider}; -use crate::identity_tree::{Canonical, Hash, TreeVersion, TreeWithNextVersion}; +use crate::identity_tree::Hash; use crate::prover::identity::Identity; use crate::prover::repository::ProverRepository; use crate::prover::Prover; @@ -33,10 +34,7 @@ pub type TransactionId = String; pub trait IdentityProcessor: Send + Sync + 'static { async fn commit_identities(&self, batch: &BatchEntry) -> anyhow::Result; - async fn finalize_identities( - &self, - processed_tree: &TreeVersion, - ) -> anyhow::Result<()>; + async fn finalize_identities(&self, sync_tree_notify: &Arc) -> anyhow::Result<()>; async fn await_clean_slate(&self) -> anyhow::Result<()>; @@ -89,14 +87,11 @@ impl IdentityProcessor for OnChainIdentityProcessor { } } - async fn finalize_identities( - &self, - processed_tree: &TreeVersion, - ) -> anyhow::Result<()> { + async fn finalize_identities(&self, sync_tree_notify: &Arc) -> anyhow::Result<()> { let mainnet_logs = self.fetch_mainnet_logs().await?; self.finalize_mainnet_roots( - processed_tree, + sync_tree_notify, &mainnet_logs, self.config.app.max_epoch_duration, ) @@ -400,7 +395,7 @@ impl OnChainIdentityProcessor { #[instrument(level = "info", skip_all)] async fn finalize_mainnet_roots( &self, - processed_tree: &TreeVersion, + sync_tree_notify: &Arc, logs: &[Log], max_epoch_duration: Duration, ) -> Result<(), anyhow::Error> { @@ -434,9 +429,7 @@ impl OnChainIdentityProcessor { .await?; } - let updates_count = processed_tree.apply_updates_up_to(post_root.into()); - - info!(updates_count, ?pre_root, ?post_root, "Mined tree updated"); + sync_tree_notify.notify_one(); } Ok(()) @@ -559,10 +552,7 @@ impl IdentityProcessor for OffChainIdentityProcessor { Ok(batch.id.to_string()) } - async fn finalize_identities( - &self, - processed_tree: &TreeVersion, - ) -> anyhow::Result<()> { + async fn finalize_identities(&self, sync_tree_notify: &Arc) -> anyhow::Result<()> { let batches = { let mut committed_batches = self.committed_batches.lock().unwrap(); let copied = committed_batches.clone(); @@ -583,7 +573,8 @@ impl IdentityProcessor for OffChainIdentityProcessor { self.database .mark_root_as_mined_tx(&batch.next_root) .await?; - processed_tree.apply_updates_up_to(batch.next_root); + + sync_tree_notify.notify_one(); } Ok(()) diff --git a/src/identity_tree/builder.rs b/src/identity_tree/builder.rs new file mode 100644 index 00000000..cf2ac613 --- /dev/null +++ b/src/identity_tree/builder.rs @@ -0,0 +1,198 @@ +use std::cmp::min; +use std::sync::{Arc, Mutex}; + +use semaphore::lazy_merkle_tree::LazyMerkleTree; +use semaphore::poseidon_tree::PoseidonHash; +use semaphore::{lazy_merkle_tree, Field}; +use tracing::warn; + +use crate::identity_tree::{ + BasicTreeOps, Canonical, CanonicalTreeMetadata, DerivedTreeMetadata, Intermediate, Latest, + PoseidonTree, TreeUpdate, TreeVersion, TreeVersionData, TreeVersionState, Version, +}; + +/// A helper for building the first tree version. Exposes a type-safe API over +/// building a sequence of tree versions efficiently. +pub struct CanonicalTreeBuilder(TreeVersionData); + +impl CanonicalTreeBuilder { + /// Creates a new builder with the given parameters. + /// * `tree_depth`: The depth of the tree. + /// * `dense_prefix_depth`: The depth of the dense prefix – i.e. the prefix + /// of the tree that will be stored in a vector rather than in a + /// pointer-based structure. + /// * `flattening_threshold`: The number of updates that can be applied to + /// this tree before garbage collection is triggered. GC is quite + /// expensive time-wise, so this value should be chosen carefully to make + /// sure it pays off in memory savings. A rule of thumb is that GC will + /// free up roughly `Depth * Number of Versions * Flattening Threshold` + /// nodes in the tree. + /// * `initial_leaf`: The default value of the tree leaves. + /// * `initial_leaves`: The initial values of the tree leaves. Index in + /// array is a leaf index in the tree. + /// * `mmap_file_path`: Path to file where data are stored on disk. + #[must_use] + pub fn new( + tree_depth: usize, + dense_prefix_depth: usize, + flattening_threshold: usize, + default_leaf: Field, + initial_leaves: &[Option], + mmap_file_path: &str, + ) -> Self { + let initial_leaves_in_dense_count = min(initial_leaves.len(), 1 << dense_prefix_depth); + let (initial_leaves_in_dense, leftover_initial_leaves) = + initial_leaves.split_at(initial_leaves_in_dense_count); + + let tree = + PoseidonTree::::new_mmapped_with_dense_prefix_with_init_values( + tree_depth, + dense_prefix_depth, + &default_leaf, + &initial_leaves_in_dense + .iter() + .map(|tree_update| tree_update.as_ref().map(|v| v.element).unwrap_or(default_leaf)) + .collect::>(), + mmap_file_path + ).unwrap(); + let metadata = CanonicalTreeMetadata { + flatten_threshold: flattening_threshold, + count_since_last_flatten: 0, + }; + + let last_dense_leaf = initial_leaves_in_dense + .iter() + .last() + .unwrap_or(&None) + .as_ref(); + let mut builder = Self(TreeVersionData { + state: TreeVersionState { + tree, + next_leaf: last_dense_leaf.map(|v| v.leaf_index + 1).unwrap_or(0), + last_sequence_id: last_dense_leaf.map(|v| v.sequence_id).unwrap_or(0), + }, + next: None, + metadata, + }); + for tree_update in leftover_initial_leaves.iter().flatten() { + builder.update(tree_update); + } + builder + } + + pub fn restore( + tree_depth: usize, + dense_prefix_depth: usize, + default_leaf: &Field, + last_dense_leaf: Option, + leftover_items: &[TreeUpdate], + flattening_threshold: usize, + mmap_file_path: &str, + ) -> Option { + let tree: LazyMerkleTree = + match PoseidonTree::::attempt_dense_mmap_restore( + tree_depth, + dense_prefix_depth, + default_leaf, + mmap_file_path, + ) { + Ok(tree) => tree, + Err(error) => { + warn!("Tree wasn't restored. Reason: {}", error.to_string()); + return None; + } + }; + + let metadata = CanonicalTreeMetadata { + flatten_threshold: flattening_threshold, + count_since_last_flatten: 0, + }; + let mut builder = Self(TreeVersionData { + state: TreeVersionState { + tree, + next_leaf: last_dense_leaf + .as_ref() + .map(|v| v.leaf_index + 1) + .unwrap_or(0), + last_sequence_id: last_dense_leaf.as_ref().map(|v| v.sequence_id).unwrap_or(0), + }, + next: None, + metadata, + }); + + for tree_update in leftover_items.iter() { + builder.update(tree_update); + } + + Some(builder) + } + + /// Updates a leaf in the resulting tree. + pub fn update(&mut self, update: &TreeUpdate) { + self.0 + .update(update.sequence_id, update.leaf_index, update.element); + } + + /// Seals this version and returns a builder for the next version. + #[must_use] + pub fn seal(self) -> (TreeVersion, DerivedTreeBuilder) { + let state = self.0.state.derived(); + let sealed = TreeVersion(Arc::new(Mutex::new(self.0))); + let next = DerivedTreeBuilder::::new(state, sealed.clone()); + (sealed, next) + } +} + +/// A helper for building successive tree versions. Exposes a type-safe API over +/// building a sequence of tree versions efficiently. +pub struct DerivedTreeBuilder { + prev: TreeVersion

, + current: TreeVersionData, +} + +impl DerivedTreeBuilder

{ + #[must_use] + fn new( + state: TreeVersionState, + prev: TreeVersion, + ) -> DerivedTreeBuilder { + let metadata = DerivedTreeMetadata { + diff: vec![], + ref_state: state.clone(), + }; + DerivedTreeBuilder { + prev, + current: TreeVersionData { + state, + next: None, + metadata, + }, + } + } + + /// Updates a leaf in the resulting tree. + pub fn update(&mut self, update: &TreeUpdate) { + self.current + .update(update.sequence_id, update.leaf_index, update.element); + } + + /// Seals this version and returns a builder for the next version. + #[must_use] + pub fn seal_and_continue( + self, + ) -> (TreeVersion, DerivedTreeBuilder) { + let state = self.current.state.clone(); + let sealed = TreeVersion(Arc::new(Mutex::new(self.current))); + let next = Self::new(state, sealed.clone()); + self.prev.get_data().next = Some(sealed.as_derived()); + (sealed, next) + } + + /// Seals this version and finishes the building process. + #[must_use] + pub fn seal(self) -> TreeVersion { + let sealed = TreeVersion(Arc::new(Mutex::new(self.current))); + self.prev.get_data().next = Some(sealed.as_derived()); + sealed + } +} diff --git a/src/identity_tree/initializer.rs b/src/identity_tree/initializer.rs index 772dc268..65596e92 100644 --- a/src/identity_tree/initializer.rs +++ b/src/identity_tree/initializer.rs @@ -8,9 +8,9 @@ use crate::config::TreeConfig; use crate::database::query::DatabaseQuery; use crate::database::Database; use crate::identity::processor::IdentityProcessor; +use crate::identity_tree::builder::CanonicalTreeBuilder; use crate::identity_tree::{ - CanonicalTreeBuilder, Hash, ProcessedStatus, TreeState, TreeUpdate, TreeVersionReadOps, - TreeWithNextVersion, + Hash, ProcessedStatus, TreeState, TreeUpdate, TreeVersionReadOps, TreeWithNextVersion, }; use crate::utils::tree_updates::dedup_tree_updates; @@ -98,32 +98,32 @@ impl TreeInitializer { } pub fn get_leftover_leaves_and_update_index( - index: &mut Option, + last_dense_leaf: &mut Option, dense_prefix_depth: usize, - mined_items: &[TreeUpdate], - ) -> Vec> { - let leftover_items = if mined_items.is_empty() { + mined_or_processed_items: &[TreeUpdate], + ) -> Vec { + let leftover_items = if mined_or_processed_items.is_empty() { vec![] } else { - let max_leaf = mined_items.last().map(|item| item.leaf_index).unwrap(); - // if the last index is greater then dense_prefix_depth, 1 << dense_prefix_depth + let max_leaf = mined_or_processed_items + .last() + .map(|item| item.leaf_index) + .unwrap(); + // if the last index is greater than dense_prefix_depth, 1 << dense_prefix_depth // should be the last index in restored tree - let last_index = std::cmp::min(max_leaf, (1 << dense_prefix_depth) - 1); - *index = Some(last_index); - - if max_leaf - last_index == 0 { - return vec![]; - } - - let mut leaves = Vec::with_capacity(max_leaf - last_index); + let max_dense_leaf_index = std::cmp::min(max_leaf, (1 << dense_prefix_depth) - 1); + let last_dense_leaf_index = mined_or_processed_items + .iter() + .rposition(|v| v.leaf_index <= max_dense_leaf_index); - let leftover = &mined_items[(last_index + 1)..]; + *last_dense_leaf = last_dense_leaf_index + .and_then(|v| mined_or_processed_items.get(v)) + .cloned(); - for item in leftover { - leaves.push(item.element); - } + println!("{:?}", mined_or_processed_items); + println!("{:?}", last_dense_leaf_index); - leaves + mined_or_processed_items[last_dense_leaf_index.map(|v| v + 1).unwrap_or(0)..].to_vec() }; leftover_items @@ -134,9 +134,9 @@ impl TreeInitializer { mined_or_processed_items: &[TreeUpdate], initial_root_hash: Hash, ) -> anyhow::Result> { - let mut last_mined_or_processed_index_in_dense: Option = None; + let mut last_mined_or_processed_leaf_in_dense: Option = None; let leftover_items = Self::get_leftover_leaves_and_update_index( - &mut last_mined_or_processed_index_in_dense, + &mut last_mined_or_processed_leaf_in_dense, self.config.dense_tree_prefix_depth, mined_or_processed_items, ); @@ -145,7 +145,7 @@ impl TreeInitializer { self.config.tree_depth, self.config.dense_tree_prefix_depth, &self.config.initial_leaf_value, - last_mined_or_processed_index_in_dense, + last_mined_or_processed_leaf_in_dense, &leftover_items, self.config.tree_gc_threshold, &self.config.cache_file, @@ -208,10 +208,11 @@ impl TreeInitializer { .last() .map(|item| item.leaf_index) .unwrap(); - let mut leaves = vec![initial_leaf_value; max_leaf + 1]; + let mut leaves = vec![None; max_leaf + 1]; for item in mined_or_processed_items { - leaves[item.leaf_index] = item.element; + let i = item.leaf_index; + leaves[i] = Some(item); } leaves @@ -279,13 +280,15 @@ mod test { pub fn generate_test_identities_with_index(identity_count: usize) -> Vec { let mut identities = vec![]; - for i in 1..=identity_count { + for i in 0..identity_count { let bytes: [u8; 32] = U256::from(rand::random::()).into(); let identity = Uint::<256, 4>::from_le_bytes(bytes); identities.push(TreeUpdate { - leaf_index: i, - element: identity, + sequence_id: i + 1, + leaf_index: i, + element: identity, + post_root: identity, }); } @@ -304,7 +307,7 @@ mod test { // indecies at all) let identities: Vec = vec![]; - let mut last_mined_index_in_dense: Option = None; + let mut last_mined_index_in_dense: Option = None; let leaves = TreeInitializer::get_leftover_leaves_and_update_index( &mut last_mined_index_in_dense, dense_prefix_depth, @@ -317,7 +320,7 @@ mod test { // since there are no identities at all the leaves should be 0 assert_eq!(leaves.len(), 0); - // first test with less then dense prefix + // first test with less than dense prefix let identities = generate_test_identities_with_index(less_identities_count); last_mined_index_in_dense = None; @@ -329,12 +332,15 @@ mod test { ); // check if the index is correct - assert_eq!(last_mined_index_in_dense, Some(identities.len())); - // since there are less identities then dense prefix, the leavs should be empty - // vector + assert_eq!( + last_mined_index_in_dense.unwrap().leaf_index, + identities.len() - 1 + ); + // since there are fewer identities than dense prefix, the leaves should be + // empty vector assert!(leaves.is_empty()); - // lets try now with more identities then dense prefix supports + // let's try now with more identities than dense prefix supports // this should generate 2^dense_prefix + 2 let identities = generate_test_identities_with_index(more_identities_count); @@ -348,16 +354,16 @@ mod test { // check if the index is correct assert_eq!( - last_mined_index_in_dense, - Some((1 << dense_prefix_depth) - 1) + last_mined_index_in_dense.unwrap().leaf_index, + (1 << dense_prefix_depth) - 1 ); - // since there are more identities then dense prefix, the leavs should be 2 + // since there are more identities than dense prefix, the leaves should be 2 assert_eq!(leaves.len(), 2); // additional check for correctness - assert_eq!(leaves[0], identities[8].element); - assert_eq!(leaves[1], identities[9].element); + assert_eq!(leaves[0].element, identities[8].element); + assert_eq!(leaves[1].element, identities[9].element); Ok(()) } diff --git a/src/identity_tree/mod.rs b/src/identity_tree/mod.rs index d0b4ce16..04ee364f 100644 --- a/src/identity_tree/mod.rs +++ b/src/identity_tree/mod.rs @@ -1,8 +1,7 @@ -use std::cmp::min; use std::sync::{Arc, Mutex, MutexGuard}; use chrono::Utc; -use semaphore::lazy_merkle_tree::{Derived, LazyMerkleTree}; +use semaphore::lazy_merkle_tree::LazyMerkleTree; use semaphore::merkle_tree::Hasher; use semaphore::poseidon_tree::{PoseidonHash, Proof}; use semaphore::{lazy_merkle_tree, Field}; @@ -10,27 +9,40 @@ use serde::{Deserialize, Serialize}; use sqlx::prelude::FromRow; use tracing::{info, warn}; +pub mod builder; pub mod initializer; +pub mod state; mod status; pub type PoseidonTree = LazyMerkleTree; pub type Hash = ::Hash; +pub use self::state::TreeState; pub use self::status::{ProcessedStatus, Status, UnknownStatus, UnprocessedStatus}; #[derive(Clone, Eq, PartialEq, Hash, Debug, FromRow)] pub struct TreeUpdate { #[sqlx(try_from = "i64")] - pub leaf_index: usize, - pub element: Hash, + pub sequence_id: usize, + #[sqlx(try_from = "i64")] + pub leaf_index: usize, + pub element: Hash, + pub post_root: Hash, } impl TreeUpdate { #[must_use] - pub const fn new(leaf_index: usize, element: Hash) -> Self { + pub const fn new( + sequence_id: usize, + leaf_index: usize, + element: Hash, + post_root: Hash, + ) -> Self { Self { + sequence_id, leaf_index, element, + post_root, } } } @@ -69,13 +81,14 @@ pub struct CanonicalTreeMetadata { /// Additional data held by any derived tree version. Includes the list of /// updates performed since previous version. pub struct DerivedTreeMetadata { - diff: Vec, + diff: Vec, + ref_state: TreeVersionState, } #[derive(Clone)] pub struct AppliedTreeUpdate { - pub update: TreeUpdate, - pub result: PoseidonTree, + pub update: TreeUpdate, + pub post_state: TreeVersionState, } /// Trait used to associate a version marker with its metadata type. @@ -94,20 +107,46 @@ impl AllowedTreeVersionMarker for lazy_merkle_tree::Derived { type Metadata = DerivedTreeMetadata; } +pub struct TreeVersionState { + pub tree: PoseidonTree, + pub next_leaf: usize, + pub last_sequence_id: usize, +} + +impl TreeVersionState { + fn derived(&self) -> TreeVersionState { + TreeVersionState { + tree: self.tree.derived(), + next_leaf: self.next_leaf, + last_sequence_id: self.last_sequence_id, + } + } +} + +impl Clone for TreeVersionState { + fn clone(&self) -> Self { + TreeVersionState { + tree: self.tree.clone(), + next_leaf: self.next_leaf, + last_sequence_id: self.last_sequence_id, + } + } +} + /// Underlying data structure for a tree version. It holds the tree itself, the -/// next leaf (only used in the latest tree), a pointer to the next version (if -/// exists) and the metadata specified by the version marker. +/// next leaf (only used in the latest tree), a last sequence id from database +/// indicating order of operations, a pointer to the next version (if exists) +/// and the metadata specified by the version marker. struct TreeVersionData { - tree: PoseidonTree, - next_leaf: usize, - next: Option>, - metadata: V::Metadata, + state: TreeVersionState, + next: Option>, + metadata: V::Metadata, } /// Basic operations that should be available for all tree versions. trait BasicTreeOps { /// Updates the tree with the given element at the given leaf index. - fn update(&mut self, leaf_index: usize, element: Hash); + fn update(&mut self, sequence_id: usize, leaf_index: usize, element: Hash); fn apply_diffs(&mut self, diffs: Vec); @@ -124,18 +163,18 @@ where { /// Gets the current tree root. fn get_root(&self) -> Hash { - self.tree.root() + self.state.tree.root() } /// Gets the leaf value at a given index. fn get_leaf(&self, leaf: usize) -> Hash { - self.tree.get_leaf(leaf) + self.state.tree.get_leaf(leaf) } /// Gets the proof of the given leaf index element fn get_proof(&self, leaf: usize) -> (Hash, Proof) { - let proof = self.tree.proof(leaf); - (self.tree.root(), proof) + let proof = self.state.tree.proof(leaf); + (self.state.tree.root(), proof) } /// Returns _up to_ `maximum_update_count` contiguous deletion or insertion @@ -172,6 +211,7 @@ where .collect() } + /// Applies updates _up to_ `root`. Returns zero when root was not found. fn apply_updates_up_to(&mut self, root: Hash) -> usize { let Some(next) = self.next.clone() else { return 0; @@ -179,26 +219,83 @@ where let num_updates; { - // Acquire the exclusive write lock on the next version. - let mut next = next.get_data(); + let applied_updates = { + // Acquire the exclusive write lock on the next version. + let mut next = next.get_data(); + + let index_of_root = next + .metadata + .diff + .iter() + .position(|update| update.post_state.tree.root() == root); + + let Some(index_of_root) = index_of_root else { + warn!(?root, "Root not found in the diff"); + return 0; + }; + + next.metadata + .diff + .drain(..=index_of_root) + .collect::>() + }; + + num_updates = applied_updates.len(); + + self.apply_diffs(applied_updates); + } + + self.garbage_collect(); - let index_of_root = next + num_updates + } +} + +impl TreeVersionData +where + Self: BasicTreeOps, +{ + /// Rewinds updates _up to_ `root`. Returns zero when root was not found. + pub fn rewind_updates_up_to(&mut self, root: Hash) -> usize { + let mut rest = if root == self.metadata.ref_state.tree.root() { + self.state = self.metadata.ref_state.clone(); + + self.metadata.diff.drain(..).collect::>() + } else { + let index_of_root = self .metadata .diff .iter() - .position(|update| update.result.root() == root); + .position(|update| update.post_state.tree.root() == root); let Some(index_of_root) = index_of_root else { warn!(?root, "Root not found in the diff"); return 0; }; - let applied_updates: Vec<_> = next.metadata.diff.drain(..=index_of_root).collect(); + let Some(root_update) = self.metadata.diff.get(index_of_root) else { + warn!(?root, "Root position not found in the diff"); + return 0; + }; - num_updates = applied_updates.len(); + self.state = root_update.post_state.clone(); - self.apply_diffs(applied_updates); - } + self.metadata + .diff + .drain((index_of_root + 1)..) + .collect::>() + }; + + let num_updates = rest.len(); + + if let Some(next) = self.next.clone() { + let mut next = next.get_data(); + rest.append(&mut next.metadata.diff); + next.metadata.diff = rest; + next.metadata.ref_state = self.state.clone(); + + next.garbage_collect(); + }; self.garbage_collect(); @@ -207,20 +304,21 @@ where } impl BasicTreeOps for TreeVersionData { - fn update(&mut self, leaf_index: usize, element: Hash) { - take_mut::take(&mut self.tree, |tree| { + fn update(&mut self, sequence_id: usize, leaf_index: usize, element: Hash) { + take_mut::take(&mut self.state.tree, |tree| { tree.update_with_mutation(leaf_index, &element) }); if element != Hash::ZERO { - self.next_leaf = leaf_index + 1; + self.state.next_leaf = leaf_index + 1; } self.metadata.count_since_last_flatten += 1; + self.state.last_sequence_id = sequence_id; } fn apply_diffs(&mut self, diffs: Vec) { for applied_update in &diffs { let update = &applied_update.update; - self.update(update.leaf_index, update.element); + self.update(update.sequence_id, update.leaf_index, update.element); } } @@ -238,7 +336,7 @@ impl BasicTreeOps for TreeVersionData { self.metadata.count_since_last_flatten = 0; let next = &self.next; if let Some(next) = next { - next.get_data().rebuild_on(self.tree.derived()); + next.get_data().rebuild_on(self.state.tree.derived()); } info!("Tree versions rebuilt"); } @@ -246,35 +344,51 @@ impl BasicTreeOps for TreeVersionData { } impl TreeVersionData { + /// This method recalculate tree to use different tree version as a base. + /// The tree itself is same in terms of root hash but differs how + /// internally is stored in memory and on disk. fn rebuild_on(&mut self, mut tree: PoseidonTree) { + self.metadata.ref_state.tree = tree.clone(); for update in &mut self.metadata.diff { tree = tree.update(update.update.leaf_index, &update.update.element); - update.result = tree.clone(); + update.post_state.tree = tree.clone(); } - self.tree = tree; + self.state.tree = tree.clone(); let next = &self.next; if let Some(next) = next { - next.get_data().rebuild_on(self.tree.clone()); + next.get_data().rebuild_on(self.state.tree.clone()); } } } impl BasicTreeOps for TreeVersionData { - fn update(&mut self, leaf_index: usize, element: Hash) { - let updated_tree = self.tree.update(leaf_index, &element); - - self.tree = updated_tree.clone(); + fn update(&mut self, sequence_id: usize, leaf_index: usize, element: Hash) { + let updated_tree = self.state.tree.update(leaf_index, &element); + let updated_next_leaf = if element != Hash::ZERO { + leaf_index + 1 + } else { + self.state.next_leaf + }; - if element != Hash::ZERO { - self.next_leaf = leaf_index + 1; - } + self.state = TreeVersionState { + tree: updated_tree.clone(), + next_leaf: updated_next_leaf, + last_sequence_id: sequence_id, + }; self.metadata.diff.push(AppliedTreeUpdate { - update: TreeUpdate { + update: TreeUpdate { + sequence_id, leaf_index, element, + post_root: updated_tree.root(), }, - result: updated_tree, + post_state: self.state.clone(), }); + + if let Some(next) = &self.next { + let mut next = next.get_data(); + next.metadata.ref_state = self.state.clone(); + } } fn apply_diffs(&mut self, mut diffs: Vec) { @@ -283,11 +397,12 @@ impl BasicTreeOps for TreeVersionData { self.metadata.diff.append(&mut diffs); if let Some(last) = last { - self.tree = last.result.clone(); + self.state = last.post_state.clone(); + } - if last.update.element != Hash::ZERO { - self.next_leaf = last.update.leaf_index + 1; - } + if let Some(next) = &self.next { + let mut next = next.get_data(); + next.metadata.ref_state = self.state.clone(); } } @@ -379,6 +494,10 @@ pub trait TreeVersionReadOps { fn get_proof(&self, leaf: usize) -> (Hash, Proof); /// Gets the leaf value at a given index. fn get_leaf(&self, leaf: usize) -> Hash; + /// Gets commitments at given leaf values + fn commitments_by_leaves(&self, leaves: impl IntoIterator) -> Vec; + /// Gets last sequence id + fn get_last_sequence_id(&self) -> usize; } impl TreeVersionReadOps for TreeVersion @@ -390,7 +509,7 @@ where } fn next_leaf(&self) -> usize { - self.get_data().next_leaf + self.get_data().state.next_leaf } fn get_leaf_and_proof(&self, leaf: usize) -> (Hash, Hash, Proof) { @@ -411,6 +530,22 @@ where let tree = self.get_data(); tree.get_leaf(leaf) } + + fn commitments_by_leaves(&self, leaves: impl IntoIterator) -> Vec { + let tree = self.get_data(); + + let mut commitments = vec![]; + + for leaf in leaves { + commitments.push(tree.state.tree.get_leaf(leaf)); + } + + commitments + } + + fn get_last_sequence_id(&self) -> usize { + self.get_data().state.last_sequence_id + } } impl TreeVersion { @@ -420,20 +555,23 @@ impl TreeVersion { } impl TreeVersion { - /// Appends many identities to the tree, returns a list with the root, proof - /// of inclusion and leaf index + /// Simulate appending many identities to the tree by copying it underneath, + /// returns a list with the root, proof of inclusion and leaf index. No + /// changes are made to the tree. #[must_use] - pub fn append_many(&self, identities: &[Hash]) -> Vec<(Hash, Proof, usize)> { - let mut data = self.get_data(); - let next_leaf = data.next_leaf; + pub fn simulate_append_many(&self, identities: &[Hash]) -> Vec<(Hash, Proof, usize)> { + let data = self.get_data(); + let mut tree = data.state.tree.clone(); + let next_leaf = data.state.next_leaf; let mut output = Vec::with_capacity(identities.len()); for (idx, identity) in identities.iter().enumerate() { let leaf_index = next_leaf + idx; - data.update(leaf_index, *identity); - let (root, proof) = data.get_proof(leaf_index); + tree = tree.update(leaf_index, identity); + let root = tree.root(); + let proof = tree.proof(leaf_index); output.push((root, proof, leaf_index)); } @@ -441,39 +579,43 @@ impl TreeVersion { output } - /// Deletes many identities from the tree, returns a list with the root - /// and proof of inclusion + /// Simulates deleting many identities from the tree by copying it + /// underneath, returns a list with the root and proof of inclusion. No + /// changes are made to the tree. #[must_use] - pub fn delete_many(&self, leaf_indices: &[usize]) -> Vec<(Hash, Proof)> { - let mut data = self.get_data(); + pub fn simulate_delete_many(&self, leaf_indices: &[usize]) -> Vec<(Hash, Proof)> { + let mut tree = self.get_data().state.tree.clone(); let mut output = Vec::with_capacity(leaf_indices.len()); for leaf_index in leaf_indices { - data.update(*leaf_index, Hash::ZERO); - let (root, proof) = data.get_proof(*leaf_index); + tree = tree.update(*leaf_index, &Hash::ZERO); + let root = tree.root(); + let proof = tree.proof(*leaf_index); output.push((root, proof)); } output } -} -impl TreeVersion -where - T: Version, -{ - pub fn commitments_by_indices(&self, indices: impl IntoIterator) -> Vec { - let tree = self.get_data(); + /// Latest tree is the only way to apply new updates. Other versions may + /// only move on the chain of changes by passing desired root. + pub fn apply_updates(&self, tree_updates: &[TreeUpdate]) -> Vec { + let mut data = self.get_data(); - let mut commitments = vec![]; + let mut output = Vec::with_capacity(tree_updates.len()); - for idx in indices { - commitments.push(tree.tree.get_leaf(idx)); + for tree_update in tree_updates { + data.update( + tree_update.sequence_id, + tree_update.leaf_index, + tree_update.element, + ); + output.push(data.get_root()); } - commitments + output } } @@ -498,242 +640,22 @@ where } } -#[derive(Clone)] -pub struct TreeState { - processed: TreeVersion, - batching: TreeVersion, - latest: TreeVersion, +/// Public API for working with versions that can rollback updates. Rollback is +/// only possible up to previous tree root. +pub trait ReversibleVersion { + fn rewind_updates_up_to(&self, root: Hash) -> usize; } -impl TreeState { - #[must_use] - pub const fn new( - processed: TreeVersion, - batching: TreeVersion, - latest: TreeVersion, - ) -> Self { - Self { - processed, - batching, - latest, - } - } - - pub fn latest_tree(&self) -> &TreeVersion { - &self.latest - } - - #[must_use] - pub fn get_latest_tree(&self) -> TreeVersion { - self.latest.clone() - } - - #[must_use] - pub fn get_processed_tree(&self) -> TreeVersion { - self.processed.clone() - } - - pub fn processed_tree(&self) -> &TreeVersion { - &self.processed - } - - #[must_use] - pub fn get_batching_tree(&self) -> TreeVersion { - self.batching.clone() - } - - pub fn batching_tree(&self) -> &TreeVersion { - &self.batching - } - - #[must_use] - pub fn get_proof_for(&self, item: &TreeItem) -> (Field, InclusionProof) { - let (leaf, root, proof) = self.latest.get_leaf_and_proof(item.leaf_index); - - let proof = InclusionProof { - root: Some(root), - proof: Some(proof), - message: None, - }; - - (leaf, proof) - } -} - -/// A helper for building the first tree version. Exposes a type-safe API over -/// building a sequence of tree versions efficiently. -pub struct CanonicalTreeBuilder(TreeVersionData); -impl CanonicalTreeBuilder { - /// Creates a new builder with the given parameters. - /// * `tree_depth`: The depth of the tree. - /// * `dense_prefix_depth`: The depth of the dense prefix – i.e. the prefix - /// of the tree that will be stored in a vector rather than in a - /// pointer-based structure. - /// * `flattening_threshold`: The number of updates that can be applied to - /// this tree before garbage collection is triggered. GC is quite - /// expensive time-wise, so this value should be chosen carefully to make - /// sure it pays off in memory savings. A rule of thumb is that GC will - /// free up roughly `Depth * Number of Versions * Flattening Threshold` - /// nodes in the tree. - /// * `initial_leaf`: The initial value of the tree leaves. - #[must_use] - pub fn new( - tree_depth: usize, - dense_prefix_depth: usize, - flattening_threshold: usize, - initial_leaf: Field, - initial_leaves: &[Field], - mmap_file_path: &str, - ) -> Self { - let initial_leaves_in_dense_count = min(initial_leaves.len(), 1 << dense_prefix_depth); - let (initial_leaves_in_dense, leftover_initial_leaves) = - initial_leaves.split_at(initial_leaves_in_dense_count); - - let tree = - PoseidonTree::::new_mmapped_with_dense_prefix_with_init_values( - tree_depth, - dense_prefix_depth, - &initial_leaf, - initial_leaves_in_dense, - mmap_file_path - ).unwrap(); - let metadata = CanonicalTreeMetadata { - flatten_threshold: flattening_threshold, - count_since_last_flatten: 0, - }; - let mut builder = Self(TreeVersionData { - tree, - next_leaf: initial_leaves_in_dense_count, - metadata, - next: None, - }); - for (index, leaf) in leftover_initial_leaves.iter().enumerate() { - builder.update(&TreeUpdate { - leaf_index: index + initial_leaves_in_dense_count, - element: *leaf, - }); - } - builder - } - - pub fn restore( - tree_depth: usize, - dense_prefix_depth: usize, - initial_leaf: &Field, - last_index: Option, - leftover_items: &[ruint::Uint<256, 4>], - flattening_threshold: usize, - mmap_file_path: &str, - ) -> Option { - let tree: LazyMerkleTree = - match PoseidonTree::::attempt_dense_mmap_restore( - tree_depth, - dense_prefix_depth, - initial_leaf, - mmap_file_path, - ) { - Ok(tree) => tree, - Err(error) => { - warn!("Tree wasn't restored. Reason: {}", error.to_string()); - return None; - } - }; - - let metadata = CanonicalTreeMetadata { - flatten_threshold: flattening_threshold, - count_since_last_flatten: 0, - }; - let next_leaf = last_index.map(|v| v + 1).unwrap_or(0); - let mut builder = Self(TreeVersionData { - tree, - next_leaf, - metadata, - next: None, - }); - - for (index, leaf) in leftover_items.iter().enumerate() { - builder.update(&TreeUpdate { - leaf_index: next_leaf + index, - element: *leaf, - }); - } - - Some(builder) - } - - /// Updates a leaf in the resulting tree. - pub fn update(&mut self, update: &TreeUpdate) { - self.0.update(update.leaf_index, update.element); - } - - /// Seals this version and returns a builder for the next version. - #[must_use] - pub fn seal(self) -> (TreeVersion, DerivedTreeBuilder) { - let next_tree = self.0.tree.derived(); - let next_leaf = self.0.next_leaf; - let sealed = TreeVersion(Arc::new(Mutex::new(self.0))); - let next = DerivedTreeBuilder::::new(next_tree, next_leaf, sealed.clone()); - (sealed, next) - } -} - -/// A helper for building successive tree versions. Exposes a type-safe API over -/// building a sequence of tree versions efficiently. -pub struct DerivedTreeBuilder { - prev: TreeVersion

, - current: TreeVersionData, -} - -impl DerivedTreeBuilder

{ - #[must_use] - const fn new( - tree: PoseidonTree, - next_leaf: usize, - prev: TreeVersion, - ) -> DerivedTreeBuilder { - let metadata = DerivedTreeMetadata { diff: vec![] }; - DerivedTreeBuilder { - prev, - current: TreeVersionData { - tree, - next_leaf, - metadata, - next: None, - }, - } - } - - /// Updates a leaf in the resulting tree. - pub fn update(&mut self, update: &TreeUpdate) { - self.current.update(update.leaf_index, update.element); - } - - /// Seals this version and returns a builder for the next version. - #[must_use] - pub fn seal_and_continue( - self, - ) -> (TreeVersion, DerivedTreeBuilder) { - let next_tree = self.current.tree.clone(); - let next_leaf = self.current.next_leaf; - let sealed = TreeVersion(Arc::new(Mutex::new(self.current))); - let next = Self::new(next_tree, next_leaf, sealed.clone()); - self.prev.get_data().next = Some(sealed.as_derived()); - (sealed, next) - } - - /// Seals this version and finishes the building process. - #[must_use] - pub fn seal(self) -> TreeVersion { - let sealed = TreeVersion(Arc::new(Mutex::new(self.current))); - self.prev.get_data().next = Some(sealed.as_derived()); - sealed +impl> ReversibleVersion for TreeVersion { + fn rewind_updates_up_to(&self, root: Hash) -> usize { + self.get_data().rewind_updates_up_to(root) } } #[cfg(test)] mod tests { - - use super::{CanonicalTreeBuilder, Hash, TreeWithNextVersion}; + use super::builder::CanonicalTreeBuilder; + use super::{Hash, ReversibleVersion, TreeUpdate, TreeVersionReadOps, TreeWithNextVersion}; #[test] fn test_peek_next_updates() { @@ -749,7 +671,8 @@ mod tests { ) .seal(); let processed_tree = processed_builder.seal(); - let insertion_updates = processed_tree.append_many(&vec![ + + let insertions = [ Hash::from(1), Hash::from(2), Hash::from(3), @@ -757,29 +680,132 @@ mod tests { Hash::from(5), Hash::from(6), Hash::from(7), - ]); - - let _deletion_updates = processed_tree.delete_many(&[0, 1, 2]); + ]; + let updates = processed_tree.simulate_append_many(&insertions); + let insertion_updates = (0..7) + .zip(updates) + .map(|(i, (root, _, leaf_index))| { + TreeUpdate::new(i, leaf_index, *insertions.get(i).unwrap(), root) + }) + .collect::>(); + _ = processed_tree.apply_updates(&insertion_updates); + + let deletions = [0, 1, 2]; + let updates = processed_tree.simulate_delete_many(&deletions); + let deletion_updates = (7..10) + .zip(updates) + .map(|(i, (root, _))| { + TreeUpdate::new(i, *deletions.get(i - 7).unwrap(), Hash::ZERO, root) + }) + .collect::>(); + _ = processed_tree.apply_updates(&deletion_updates); let next_updates = canonical_tree.peek_next_updates(10); assert_eq!(next_updates.len(), 7); canonical_tree.apply_updates_up_to( - insertion_updates + next_updates .last() .expect("Could not get insertion updates") - .0, + .update + .post_root, ); - let _ = processed_tree.append_many(&[ - Hash::from(5), - Hash::from(6), - Hash::from(7), - Hash::from(8), - ]); + let insertions = [Hash::from(8), Hash::from(9), Hash::from(10), Hash::from(11)]; + let updates = processed_tree.simulate_append_many(&insertions); + let insertion_updates = (10..14) + .zip(updates) + .map(|(i, (root, _, leaf_index))| { + TreeUpdate::new(i, leaf_index, *insertions.get(i - 10).unwrap(), root) + }) + .collect::>(); + let _ = processed_tree.apply_updates(&insertion_updates); let next_updates = canonical_tree.peek_next_updates(10); assert_eq!(next_updates.len(), 3); } + + #[test] + fn test_rewind_up_to_root() { + let temp_dir = tempfile::tempdir().unwrap(); + + let (processed_tree, batching_tree_builder) = CanonicalTreeBuilder::new( + 10, + 10, + 0, + Hash::ZERO, + &[], + temp_dir.path().join("testfile").to_str().unwrap(), + ) + .seal(); + let (batching_tree, latest_tree_builder) = batching_tree_builder.seal_and_continue(); + let latest_tree = latest_tree_builder.seal(); + + let insertions = (1..=30).map(Hash::from).collect::>(); + let updates = latest_tree.simulate_append_many(&insertions); + let insertion_updates = (0..30) + .zip(updates) + .map(|(i, (root, _, leaf_index))| { + TreeUpdate::new(i, leaf_index, *insertions.get(i).unwrap(), root) + }) + .collect::>(); + _ = latest_tree.apply_updates(&insertion_updates); + + batching_tree.apply_updates_up_to(insertion_updates.get(19).unwrap().post_root); + processed_tree.apply_updates_up_to(insertion_updates.get(9).unwrap().post_root); + + assert_eq!(processed_tree.next_leaf(), 10); + assert_eq!( + processed_tree.get_root(), + insertion_updates.get(9).unwrap().post_root + ); + assert_eq!(batching_tree.next_leaf(), 20); + assert_eq!( + batching_tree.get_root(), + insertion_updates.get(19).unwrap().post_root + ); + assert_eq!(latest_tree.next_leaf(), 30); + assert_eq!( + latest_tree.get_root(), + insertion_updates.get(29).unwrap().post_root + ); + + batching_tree.rewind_updates_up_to(insertion_updates.get(15).unwrap().post_root); + latest_tree.rewind_updates_up_to(insertion_updates.get(25).unwrap().post_root); + + assert_eq!(processed_tree.next_leaf(), 10); + assert_eq!( + processed_tree.get_root(), + insertion_updates.get(9).unwrap().post_root + ); + assert_eq!(batching_tree.next_leaf(), 16); + assert_eq!( + batching_tree.get_root(), + insertion_updates.get(15).unwrap().post_root + ); + assert_eq!(latest_tree.next_leaf(), 26); + assert_eq!( + latest_tree.get_root(), + insertion_updates.get(25).unwrap().post_root + ); + + latest_tree.rewind_updates_up_to(insertion_updates.get(15).unwrap().post_root); + + assert_eq!(processed_tree.next_leaf(), 10); + assert_eq!( + processed_tree.get_root(), + insertion_updates.get(9).unwrap().post_root + ); + assert_eq!(batching_tree.next_leaf(), 16); + assert_eq!( + batching_tree.get_root(), + insertion_updates.get(15).unwrap().post_root + ); + assert_eq!(latest_tree.next_leaf(), 16); + assert_eq!( + latest_tree.get_root(), + insertion_updates.get(15).unwrap().post_root + ); + } } diff --git a/src/identity_tree/state.rs b/src/identity_tree/state.rs new file mode 100644 index 00000000..9f62d514 --- /dev/null +++ b/src/identity_tree/state.rs @@ -0,0 +1,67 @@ +use semaphore::Field; + +use crate::identity_tree::{ + Canonical, InclusionProof, Intermediate, Latest, TreeItem, TreeVersion, TreeVersionReadOps, +}; + +#[derive(Clone)] +pub struct TreeState { + processed: TreeVersion, + batching: TreeVersion, + latest: TreeVersion, +} + +impl TreeState { + #[must_use] + pub const fn new( + processed: TreeVersion, + batching: TreeVersion, + latest: TreeVersion, + ) -> Self { + Self { + processed, + batching, + latest, + } + } + + pub fn latest_tree(&self) -> &TreeVersion { + &self.latest + } + + #[must_use] + pub fn get_latest_tree(&self) -> TreeVersion { + self.latest.clone() + } + + #[must_use] + pub fn get_processed_tree(&self) -> TreeVersion { + self.processed.clone() + } + + pub fn processed_tree(&self) -> &TreeVersion { + &self.processed + } + + #[must_use] + pub fn get_batching_tree(&self) -> TreeVersion { + self.batching.clone() + } + + pub fn batching_tree(&self) -> &TreeVersion { + &self.batching + } + + #[must_use] + pub fn get_proof_for(&self, item: &TreeItem) -> (Field, InclusionProof) { + let (leaf, root, proof) = self.latest.get_leaf_and_proof(item.leaf_index); + + let proof = InclusionProof { + root: Some(root), + proof: Some(proof), + message: None, + }; + + (leaf, proof) + } +} diff --git a/src/task_monitor/mod.rs b/src/task_monitor/mod.rs index e9ffceea..2b30ca38 100644 --- a/src/task_monitor/mod.rs +++ b/src/task_monitor/mod.rs @@ -19,6 +19,7 @@ const PROCESS_IDENTITIES_BACKOFF: Duration = Duration::from_secs(5); const FINALIZE_IDENTITIES_BACKOFF: Duration = Duration::from_secs(5); const QUEUE_MONITOR_BACKOFF: Duration = Duration::from_secs(5); const MODIFY_TREE_BACKOFF: Duration = Duration::from_secs(5); +const SYNC_TREE_STATE_WITH_DB_BACKOFF: Duration = Duration::from_secs(5); struct RunningInstance { handles: Vec>, @@ -106,15 +107,23 @@ impl TaskMonitor { let monitored_txs_sender = Arc::new(monitored_txs_sender); let monitored_txs_receiver = Arc::new(Mutex::new(monitored_txs_receiver)); - let base_wake_up_notify = Arc::new(Notify::new()); - // Immediately notify so we can start processing if we have pending identities - // in the database - base_wake_up_notify.notify_one(); + let base_next_batch_notify = Arc::new(Notify::new()); + // Immediately notify, so we can start processing if we have pending operations + base_next_batch_notify.notify_one(); + + let base_sync_tree_notify = Arc::new(Notify::new()); + // Immediately notify, so we can start processing if we have pending operations + base_sync_tree_notify.notify_one(); + + let base_tree_synced_notify = Arc::new(Notify::new()); + // Immediately notify, so we can start processing if we have pending operations + base_tree_synced_notify.notify_one(); let mut handles = Vec::new(); // Initialize the Tree let app = self.app.clone(); + let tree_init = move || app.clone().init_tree(); let tree_init_handle = crate::utils::spawn_monitored_with_backoff( tree_init, @@ -127,7 +136,11 @@ impl TaskMonitor { // Finalize identities let app = self.app.clone(); - let finalize_identities = move || tasks::finalize_identities::finalize_roots(app.clone()); + let sync_tree_notify = base_sync_tree_notify.clone(); + + let finalize_identities = move || { + tasks::finalize_identities::finalize_roots(app.clone(), sync_tree_notify.clone()) + }; let finalize_identities_handle = crate::utils::spawn_monitored_with_backoff( finalize_identities, shutdown_sender.clone(), @@ -138,6 +151,7 @@ impl TaskMonitor { // Report length of the queue of identities let app = self.app.clone(); + let queue_monitor = move || tasks::monitor_queue::monitor_queue(app.clone()); let queue_monitor_handle = crate::utils::spawn_monitored_with_backoff( queue_monitor, @@ -147,19 +161,18 @@ impl TaskMonitor { ); handles.push(queue_monitor_handle); - // Process identities - let base_next_batch_notify = Arc::new(Notify::new()); - // Create batches let app = self.app.clone(); let next_batch_notify = base_next_batch_notify.clone(); - let wake_up_notify = base_wake_up_notify.clone(); + let sync_tree_notify = base_sync_tree_notify.clone(); + let tree_synced_notify = base_tree_synced_notify.clone(); let create_batches = move || { tasks::create_batches::create_batches( app.clone(), next_batch_notify.clone(), - wake_up_notify.clone(), + sync_tree_notify.clone(), + tree_synced_notify.clone(), ) }; let create_batches_handle = crate::utils::spawn_monitored_with_backoff( @@ -173,14 +186,12 @@ impl TaskMonitor { // Process batches let app = self.app.clone(); let next_batch_notify = base_next_batch_notify.clone(); - let wake_up_notify = base_wake_up_notify.clone(); let process_identities = move || { tasks::process_batches::process_batches( app.clone(), monitored_txs_sender.clone(), next_batch_notify.clone(), - wake_up_notify.clone(), ) }; let process_identities_handle = crate::utils::spawn_monitored_with_backoff( @@ -193,6 +204,7 @@ impl TaskMonitor { // Monitor transactions let app = self.app.clone(); + let monitor_txs = move || tasks::monitor_txs::monitor_txs(app.clone(), monitored_txs_receiver.clone()); let monitor_txs_handle = crate::utils::spawn_monitored_with_backoff( @@ -205,10 +217,16 @@ impl TaskMonitor { // Modify tree let app = self.app.clone(); - let wake_up_notify = base_wake_up_notify.clone(); - let modify_tree = - move || tasks::modify_tree::modify_tree(app.clone(), wake_up_notify.clone()); + let sync_tree_notify = base_sync_tree_notify.clone(); + let tree_synced_notify = base_tree_synced_notify.clone(); + let modify_tree = move || { + tasks::modify_tree::modify_tree( + app.clone(), + sync_tree_notify.clone(), + tree_synced_notify.clone(), + ) + }; let modify_tree_handle = crate::utils::spawn_monitored_with_backoff( modify_tree, shutdown_sender.clone(), @@ -217,6 +235,26 @@ impl TaskMonitor { ); handles.push(modify_tree_handle); + // Sync tree state with DB + let app = self.app.clone(); + let sync_tree_notify = base_sync_tree_notify.clone(); + let tree_synced_notify = base_tree_synced_notify.clone(); + + let sync_tree_state_with_db = move || { + tasks::sync_tree_state_with_db::sync_tree_state_with_db( + app.clone(), + sync_tree_notify.clone(), + tree_synced_notify.clone(), + ) + }; + let sync_tree_state_with_db_handle = crate::utils::spawn_monitored_with_backoff( + sync_tree_state_with_db, + shutdown_sender.clone(), + SYNC_TREE_STATE_WITH_DB_BACKOFF, + self.shutdown.clone(), + ); + handles.push(sync_tree_state_with_db_handle); + // Create the instance *instance = Some(RunningInstance { handles, diff --git a/src/task_monitor/tasks/create_batches.rs b/src/task_monitor/tasks/create_batches.rs index 27e479ac..6a0febda 100644 --- a/src/task_monitor/tasks/create_batches.rs +++ b/src/task_monitor/tasks/create_batches.rs @@ -29,7 +29,8 @@ const DEBOUNCE_THRESHOLD_SECS: i64 = 1; pub async fn create_batches( app: Arc, next_batch_notify: Arc, - wake_up_notify: Arc, + sync_tree_notify: Arc, + tree_synced_notify: Arc, ) -> anyhow::Result<()> { tracing::info!("Awaiting for a clean slate"); app.identity_processor.await_clean_slate().await?; @@ -55,15 +56,21 @@ pub async fn create_batches( // (possibly-incomplete) batch anyway. let mut last_batch_time: DateTime = app.database.get_latest_insertion().await?.timestamp; + let check_next_batch_notify = Notify::new(); + loop { // We wait either for a timer tick or a full batch select! { _ = timer.tick() => { - tracing::info!("Identity batch insertion woken due to timeout"); + tracing::info!("Create batches woken due to timeout"); } - () = wake_up_notify.notified() => { - tracing::trace!("Identity batch insertion woken due to request"); + () = tree_synced_notify.notified() => { + tracing::trace!("Create batches woken due tree synced event"); + }, + + () = check_next_batch_notify.notified() => { + tracing::trace!("Create batches woken due instant check for next batch"); }, } @@ -94,6 +101,7 @@ pub async fn create_batches( &app.prover_repository, app.tree_state()?.batching_tree(), &next_batch_notify, + &sync_tree_notify, &updates, ) .await?; @@ -116,6 +124,7 @@ pub async fn create_batches( &app.prover_repository, app.tree_state()?.batching_tree(), &next_batch_notify, + &sync_tree_notify, &updates, ) .await?; @@ -152,6 +161,7 @@ pub async fn create_batches( &app.prover_repository, app.tree_state()?.batching_tree(), &next_batch_notify, + &sync_tree_notify, &updates, ) .await?; @@ -169,7 +179,7 @@ pub async fn create_batches( } // We want to check if there's a full batch available immediately - wake_up_notify.notify_one(); + check_next_batch_notify.notify_one(); } } @@ -188,6 +198,7 @@ async fn commit_identities( prover_repository: &Arc, batching_tree: &TreeVersion, next_batch_notify: &Arc, + sync_tree_notify: &Arc, updates: &[AppliedTreeUpdate], ) -> anyhow::Result<()> { // If the update is an insertion @@ -208,6 +219,7 @@ async fn commit_identities( database, batching_tree, next_batch_notify, + sync_tree_notify, updates, batch_size, ) @@ -223,6 +235,7 @@ async fn commit_identities( database, batching_tree, next_batch_notify, + sync_tree_notify, updates, batch_size, ) @@ -235,6 +248,7 @@ pub async fn insert_identities( database: &Database, batching_tree: &TreeVersion, next_batch_notify: &Arc, + sync_tree_notify: &Arc, updates: &[AppliedTreeUpdate], batch_size: usize, ) -> anyhow::Result<()> { @@ -250,14 +264,15 @@ pub async fn insert_identities( let latest_tree_from_updates = updates .last() .expect("Updates is non empty.") - .result + .post_state + .tree .clone(); // Next get merkle proofs for each update - note the proofs are acquired from // intermediate versions of the tree let mut merkle_proofs: Vec<_> = updates .iter() - .map(|update| update.result.proof(update.update.leaf_index)) + .map(|update| update.post_state.tree.proof(update.update.leaf_index)) .collect(); // Grab some variables for sizes to make querying easier. @@ -338,7 +353,8 @@ pub async fn insert_identities( TaskMonitor::log_batch_size(updates.len()); - batching_tree.apply_updates_up_to(post_root); + sync_tree_notify.notify_one(); + // batching_tree.apply_updates_up_to(post_root); Ok(()) } @@ -371,6 +387,7 @@ pub async fn delete_identities( database: &Database, batching_tree: &TreeVersion, next_batch_notify: &Arc, + sync_tree_notify: &Arc, updates: &[AppliedTreeUpdate], batch_size: usize, ) -> anyhow::Result<()> { @@ -379,13 +396,14 @@ pub async fn delete_identities( let mut deletion_indices: Vec<_> = updates.iter().map(|f| f.update.leaf_index).collect(); - let commitments = batching_tree.commitments_by_indices(deletion_indices.iter().copied()); + let commitments = batching_tree.commitments_by_leaves(deletion_indices.iter().copied()); let mut commitments: Vec = commitments.into_iter().map(U256::from).collect(); let latest_tree_from_updates = updates .last() .expect("Updates is non empty.") - .result + .post_state + .tree .clone(); // Next get merkle proofs for each update - note the proofs are acquired from @@ -394,7 +412,8 @@ pub async fn delete_identities( .iter() .map(|update_with_tree| { update_with_tree - .result + .post_state + .tree .proof(update_with_tree.update.leaf_index) }) .collect(); @@ -459,7 +478,8 @@ pub async fn delete_identities( TaskMonitor::log_batch_size(updates.len()); - batching_tree.apply_updates_up_to(post_root); + sync_tree_notify.notify_one(); + // batching_tree.apply_updates_up_to(post_root); Ok(()) } diff --git a/src/task_monitor/tasks/finalize_identities.rs b/src/task_monitor/tasks/finalize_identities.rs index 14237f4f..3a9d94cc 100644 --- a/src/task_monitor/tasks/finalize_identities.rs +++ b/src/task_monitor/tasks/finalize_identities.rs @@ -1,11 +1,13 @@ use std::sync::Arc; +use tokio::sync::Notify; + use crate::app::App; -pub async fn finalize_roots(app: Arc) -> anyhow::Result<()> { +pub async fn finalize_roots(app: Arc, sync_tree_notify: Arc) -> anyhow::Result<()> { loop { app.identity_processor - .finalize_identities(app.tree_state()?.processed_tree()) + .finalize_identities(&sync_tree_notify) .await?; tokio::time::sleep(app.config.app.time_between_scans).await; diff --git a/src/task_monitor/tasks/mod.rs b/src/task_monitor/tasks/mod.rs index d6c44db6..fc93bc46 100644 --- a/src/task_monitor/tasks/mod.rs +++ b/src/task_monitor/tasks/mod.rs @@ -4,3 +4,4 @@ pub mod modify_tree; pub mod monitor_queue; pub mod monitor_txs; pub mod process_batches; +pub mod sync_tree_state_with_db; diff --git a/src/task_monitor/tasks/modify_tree.rs b/src/task_monitor/tasks/modify_tree.rs index 454d1e80..72aef2c5 100644 --- a/src/task_monitor/tasks/modify_tree.rs +++ b/src/task_monitor/tasks/modify_tree.rs @@ -18,7 +18,11 @@ use crate::retry_tx; // them from single task that determines which type of operations to run. It is // done that way to reduce number of used mutexes and eliminate the risk of some // tasks not being run at all as mutex is not preserving unlock order. -pub async fn modify_tree(app: Arc, wake_up_notify: Arc) -> anyhow::Result<()> { +pub async fn modify_tree( + app: Arc, + sync_tree_notify: Arc, + tree_synced_notify: Arc, +) -> anyhow::Result<()> { info!("Starting modify tree task."); let batch_deletion_timeout = chrono::Duration::from_std(app.config.app.batch_deletion_timeout) @@ -34,36 +38,41 @@ pub async fn modify_tree(app: Arc, wake_up_notify: Arc) -> anyhow:: info!("Modify tree task woken due to timeout"); } - () = wake_up_notify.notified() => { - info!("Modify tree task woken due to request"); + () = tree_synced_notify.notified() => { + info!("Modify tree task woken due to tree synced event"); }, } let tree_state = app.tree_state()?; - retry_tx!(&app.database, tx, { + let tree_modified = retry_tx!(&app.database, tx, { do_modify_tree( &mut tx, batch_deletion_timeout, min_batch_deletion_size, tree_state, - &wake_up_notify, ) .await }) .await?; - // wake_up_notify.notify_one(); + // It is very important to generate that event AFTER transaction is committed to + // database. Otherwise, notified task may not see changes as transaction was not + // committed yet. + if tree_modified { + sync_tree_notify.notify_one(); + } } } +/// Looks for any pending changes to the tree. Returns true if there were any +/// changes applied to the tree. async fn do_modify_tree( tx: &mut Transaction<'_, Postgres>, batch_deletion_timeout: chrono::Duration, min_batch_deletion_size: usize, tree_state: &TreeState, - wake_up_notify: &Arc, -) -> anyhow::Result<()> { +) -> anyhow::Result { let deletions = tx.get_deletions().await?; // Deleting identities has precedence over inserting them. @@ -76,12 +85,10 @@ async fn do_modify_tree( ) .await? { - run_deletions(tx, tree_state, deletions, wake_up_notify).await?; + run_deletions(tx, tree_state, deletions).await } else { - run_insertions(tx, tree_state, wake_up_notify).await?; + run_insertions(tx, tree_state).await } - - Ok(()) } pub async fn should_run_deletion( @@ -124,16 +131,16 @@ pub async fn should_run_deletion( Ok(true) } +/// Run insertions and returns true if there were any changes to the tree. pub async fn run_insertions( tx: &mut Transaction<'_, Postgres>, tree_state: &TreeState, - wake_up_notify: &Arc, -) -> anyhow::Result<()> { +) -> anyhow::Result { let unprocessed = tx .get_eligible_unprocessed_commitments(UnprocessedStatus::New) .await?; if unprocessed.is_empty() { - return Ok(()); + return Ok(false); } let latest_tree = tree_state.latest_tree(); @@ -164,7 +171,9 @@ pub async fn run_insertions( ); let mut pre_root = &latest_tree.get_root(); - let data = latest_tree.append_many(&filtered_identities); + let data = latest_tree + .clone() + .simulate_append_many(&filtered_identities); assert_eq!( data.len(), @@ -180,20 +189,17 @@ pub async fn run_insertions( tx.remove_unprocessed_identity(identity).await?; } - // Immediately look for next operations - wake_up_notify.notify_one(); - - Ok(()) + Ok(true) } +/// Run deletions and returns true if there were any changes to the tree. pub async fn run_deletions( tx: &mut Transaction<'_, Postgres>, tree_state: &TreeState, mut deletions: Vec, - wake_up_notify: &Arc, -) -> anyhow::Result<()> { +) -> anyhow::Result { if deletions.is_empty() { - return Ok(()); + return Ok(false); } // This sorting is very important. It ensures that we will create a unique root @@ -209,7 +215,10 @@ pub async fn run_deletions( let mut pre_root = tree_state.latest_tree().get_root(); // Delete the commitments at the target leaf indices in the latest tree, // generating the proof for each update - let data = tree_state.latest_tree().delete_many(&leaf_indices); + let data = tree_state + .latest_tree() + .clone() + .simulate_delete_many(&leaf_indices); assert_eq!( data.len(), @@ -228,8 +237,5 @@ pub async fn run_deletions( // Remove the previous commitments from the deletions table tx.remove_deletions(&previous_commitments).await?; - // Immediately look for next operations - wake_up_notify.notify_one(); - - Ok(()) + Ok(true) } diff --git a/src/task_monitor/tasks/process_batches.rs b/src/task_monitor/tasks/process_batches.rs index 95fb6055..67fac774 100644 --- a/src/task_monitor/tasks/process_batches.rs +++ b/src/task_monitor/tasks/process_batches.rs @@ -12,7 +12,6 @@ pub async fn process_batches( app: Arc, monitored_txs_sender: Arc>, next_batch_notify: Arc, - wake_up_notify: Arc, ) -> anyhow::Result<()> { tracing::info!("Awaiting for a clean slate"); app.identity_processor.await_clean_slate().await?; @@ -24,19 +23,21 @@ pub async fn process_batches( let mut timer = time::interval(Duration::from_secs(5)); + let check_next_batch_notify = Notify::new(); + loop { // We wait either for a timer tick or a full batch select! { _ = timer.tick() => { - tracing::info!("Identity processor woken due to timeout"); + tracing::info!("Process batches woken due to timeout"); } () = next_batch_notify.notified() => { - tracing::trace!("Identity processor woken due to next batch creation"); + tracing::trace!("Process batches woken due to next batch request"); }, - () = wake_up_notify.notified() => { - tracing::trace!("Identity processor woken due to request"); + () = check_next_batch_notify.notified() => { + tracing::trace!("Process batches woken due instant check for next batch"); }, } @@ -57,6 +58,6 @@ pub async fn process_batches( .await?; // We want to check if there's a full batch available immediately - wake_up_notify.notify_one(); + check_next_batch_notify.notify_one(); } } diff --git a/src/task_monitor/tasks/sync_tree_state_with_db.rs b/src/task_monitor/tasks/sync_tree_state_with_db.rs new file mode 100644 index 00000000..022fbc8b --- /dev/null +++ b/src/task_monitor/tasks/sync_tree_state_with_db.rs @@ -0,0 +1,123 @@ +use std::sync::Arc; + +use sqlx::{Postgres, Transaction}; +use tokio::sync::Notify; +use tokio::time::Duration; +use tokio::{select, time}; + +use crate::database::query::DatabaseQuery; +use crate::identity_tree::{ + ProcessedStatus, ReversibleVersion, TreeState, TreeVersionReadOps, TreeWithNextVersion, +}; +use crate::retry_tx; +use crate::task_monitor::App; + +pub async fn sync_tree_state_with_db( + app: Arc, + sync_tree_notify: Arc, + tree_synced_notify: Arc, +) -> anyhow::Result<()> { + tracing::info!("Awaiting for a clean slate"); + app.identity_processor.await_clean_slate().await?; + + tracing::info!("Awaiting for initialized tree"); + app.tree_state()?; + + let mut timer = time::interval(Duration::from_secs(5)); + + loop { + // We wait either for a timer tick or a full batch + select! { + _ = timer.tick() => { + tracing::info!("Sync TreeState with DB task woken due to timeout"); + } + + () = sync_tree_notify.notified() => { + tracing::info!("Sync TreeState with DB task woken due to sync request"); + }, + } + + let tree_state = app.tree_state()?; + + retry_tx!(&app.database, tx, sync_tree(&mut tx, tree_state).await).await?; + + tree_synced_notify.notify_one(); + } +} + +/// Order of operations in sync tree is very important as it ensures we can +/// apply new updates or rewind them properly. +async fn sync_tree( + tx: &mut Transaction<'_, Postgres>, + tree_state: &TreeState, +) -> anyhow::Result<()> { + let latest_processed_tree_update = tx + .get_latest_tree_update_by_statuses(vec![ + ProcessedStatus::Processed, + ProcessedStatus::Mined, + ]) + .await?; + + let processed_tree = tree_state.processed_tree(); + let batching_tree = tree_state.batching_tree(); + let latest_tree = tree_state.latest_tree(); + + // First check if processed tree needs to be rolled back. If so then we must + // panic to quit to rebuild the tree on startup. This is a time-consuming + // operation. + if let Some(ref tree_update) = latest_processed_tree_update { + let last_sequence_id = processed_tree.get_last_sequence_id(); + assert!( + tree_update.sequence_id >= last_sequence_id, + "Processed tree needs to be rolled back." + ); + }; + + let latest_pending_tree_update = tx + .get_latest_tree_update_by_statuses(vec![ + ProcessedStatus::Pending, + ProcessedStatus::Processed, + ProcessedStatus::Mined, + ]) + .await?; + + let latest_batch = tx.get_latest_batch().await?; + let latest_batching_tree_update = if let Some(latest_batch) = latest_batch { + tx.get_tree_update_by_root(&latest_batch.next_root).await? + } else { + None + }; + + // Then check if latest tree can be updated forward. + if let Some(latest_tree_update) = latest_pending_tree_update { + if latest_tree_update.sequence_id >= latest_tree.get_last_sequence_id() { + let tree_updates = tx + .get_commitments_after_id(latest_tree.get_last_sequence_id()) + .await?; + latest_tree.apply_updates(&tree_updates); + + if let Some(batching_tree_update) = latest_batching_tree_update { + if batching_tree_update.sequence_id >= batching_tree.get_last_sequence_id() { + batching_tree.apply_updates_up_to(batching_tree_update.post_root); + } else { + batching_tree.rewind_updates_up_to(batching_tree_update.post_root); + } + } + } else { + if let Some(batching_tree_update) = latest_batching_tree_update { + if batching_tree_update.sequence_id >= batching_tree.get_last_sequence_id() { + batching_tree.apply_updates_up_to(batching_tree_update.post_root); + } else { + batching_tree.rewind_updates_up_to(batching_tree_update.post_root); + } + } + latest_tree.rewind_updates_up_to(latest_tree_update.post_root); + } + } + + if let Some(ref processed_tree_update) = latest_processed_tree_update { + processed_tree.apply_updates_up_to(processed_tree_update.post_root); + } + + Ok(()) +} diff --git a/src/utils/tree_updates.rs b/src/utils/tree_updates.rs index 32a30fc3..276beb22 100644 --- a/src/utils/tree_updates.rs +++ b/src/utils/tree_updates.rs @@ -1,5 +1,7 @@ use crate::identity_tree::TreeUpdate; +/// Deduplicates changes to same leaf. Requires as input updates sorted by leaf +/// index and also for same leaf index sorted in chronological order. #[must_use] pub fn dedup_tree_updates(updates: Vec) -> Vec { let mut deduped = Vec::new(); @@ -7,15 +9,11 @@ pub fn dedup_tree_updates(updates: Vec) -> Vec { for update in updates { if let Some(prev) = temp { - if prev.leaf_index == update.leaf_index { - temp = Some(update); - } else { + if prev.leaf_index != update.leaf_index { deduped.push(prev); - temp = Some(update); } - } else { - temp = Some(update); } + temp = Some(update); } if let Some(item) = temp { @@ -37,19 +35,45 @@ mod tests { let hashes: Vec = (0..10).map(Field::from).collect(); let updates = vec![ - TreeUpdate::new(0, hashes[0]), - TreeUpdate::new(1, hashes[1]), - TreeUpdate::new(1, hashes[2]), - TreeUpdate::new(1, hashes[3]), - TreeUpdate::new(2, hashes[4]), - TreeUpdate::new(2, hashes[5]), - TreeUpdate::new(3, hashes[6]), + TreeUpdate::new(1, 0, hashes[0], hashes[0]), + TreeUpdate::new(2, 1, hashes[1], hashes[1]), + TreeUpdate::new(3, 1, hashes[2], hashes[2]), + TreeUpdate::new(4, 1, hashes[3], hashes[3]), + TreeUpdate::new(5, 2, hashes[4], hashes[4]), + TreeUpdate::new(6, 2, hashes[5], hashes[5]), + TreeUpdate::new(7, 3, hashes[6], hashes[6]), ]; let expected = vec![ - TreeUpdate::new(0, hashes[0]), - TreeUpdate::new(1, hashes[3]), - TreeUpdate::new(2, hashes[5]), - TreeUpdate::new(3, hashes[6]), + TreeUpdate::new(1, 0, hashes[0], hashes[0]), + TreeUpdate::new(4, 1, hashes[3], hashes[3]), + TreeUpdate::new(6, 2, hashes[5], hashes[5]), + TreeUpdate::new(7, 3, hashes[6], hashes[6]), + ]; + + let deduped = dedup_tree_updates(updates); + + assert_eq!(expected, deduped); + } + + #[test] + fn deduplicates_tree_updates_with_same_last() { + let hashes: Vec = (0..10).map(Field::from).collect(); + + let updates = vec![ + TreeUpdate::new(1, 0, hashes[0], hashes[0]), + TreeUpdate::new(2, 1, hashes[1], hashes[1]), + TreeUpdate::new(3, 1, hashes[2], hashes[2]), + TreeUpdate::new(4, 1, hashes[3], hashes[3]), + TreeUpdate::new(5, 2, hashes[4], hashes[4]), + TreeUpdate::new(6, 2, hashes[5], hashes[5]), + TreeUpdate::new(7, 3, hashes[6], hashes[6]), + TreeUpdate::new(8, 3, hashes[7], hashes[7]), + ]; + let expected = vec![ + TreeUpdate::new(1, 0, hashes[0], hashes[0]), + TreeUpdate::new(4, 1, hashes[3], hashes[3]), + TreeUpdate::new(6, 2, hashes[5], hashes[5]), + TreeUpdate::new(8, 3, hashes[7], hashes[7]), ]; let deduped = dedup_tree_updates(updates); From c7b3ae34d019bcc3f56e26e722ed5eb3d67e58ec Mon Sep 17 00:00:00 2001 From: Piotr Heilman Date: Tue, 6 Aug 2024 12:16:23 +0200 Subject: [PATCH 4/7] working HA --- crates/tx-sitter-client/src/lib.rs | 35 +++-- e2e_tests/docker-compose/compose.yml | 23 ++-- e2e_tests/scenarios/tests/common/api.rs | 10 +- .../scenarios/tests/common/docker_compose.rs | 44 +++--- e2e_tests/scenarios/tests/common/mod.rs | 24 ++-- e2e_tests/scenarios/tests/insert_100.rs | 4 +- src/contracts/mod.rs | 17 ++- src/database/query.rs | 17 +++ src/database/transaction.rs | 2 +- src/ethereum/mod.rs | 5 +- src/ethereum/write_provider/inner.rs | 1 + src/ethereum/write_provider/mod.rs | 3 +- src/ethereum/write_provider/openzeppelin.rs | 4 +- src/ethereum/write_provider/tx_sitter.rs | 27 +++- src/identity/processor.rs | 127 ++++++++++++------ src/task_monitor/tasks/create_batches.rs | 62 ++++++--- src/task_monitor/tasks/process_batches.rs | 42 ++++-- .../tasks/sync_tree_state_with_db.rs | 12 +- 18 files changed, 324 insertions(+), 135 deletions(-) diff --git a/crates/tx-sitter-client/src/lib.rs b/crates/tx-sitter-client/src/lib.rs index 82680cd5..8241833e 100644 --- a/crates/tx-sitter-client/src/lib.rs +++ b/crates/tx-sitter-client/src/lib.rs @@ -1,5 +1,8 @@ +use std::fmt; + +use anyhow::bail; use data::{GetTxResponse, SendTxRequest, SendTxResponse, TxStatus}; -use reqwest::Response; +use reqwest::{Response, StatusCode}; use tracing::instrument; pub mod data; @@ -9,6 +12,22 @@ pub struct TxSitterClient { url: String, } +#[derive(Debug)] +pub struct HttpError { + pub status: StatusCode, + pub body: String, +} + +impl fmt::Display for HttpError { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + write!( + f, + "Response failed with status {} - {}", + self.status, self.body + ) + } +} + impl TxSitterClient { pub fn new(url: impl ToString) -> Self { Self { @@ -46,9 +65,7 @@ impl TxSitterClient { let body = response.text().await?; tracing::error!("Response failed with status {} - {}", status, body); - return Err(anyhow::anyhow!( - "Response failed with status {status} - {body}" - )); + bail!(HttpError { body, status }); } Ok(response) @@ -56,19 +73,19 @@ impl TxSitterClient { #[instrument(skip(self))] pub async fn send_tx(&self, req: &SendTxRequest) -> anyhow::Result { - self.json_post(&format!("{}/tx", self.url), req).await + Ok(self.json_post(&format!("{}/tx", self.url), req).await?) } #[instrument(skip(self))] pub async fn get_tx(&self, tx_id: &str) -> anyhow::Result { - self.json_get(&format!("{}/tx/{}", self.url, tx_id)).await + Ok(self.json_get(&format!("{}/tx/{}", self.url, tx_id)).await?) } #[instrument(skip(self))] pub async fn get_txs(&self) -> anyhow::Result> { let url = format!("{}/txs", self.url); - self.json_get(&url).await + Ok(self.json_get(&url).await?) } #[instrument(skip(self))] @@ -78,14 +95,14 @@ impl TxSitterClient { ) -> anyhow::Result> { let url = format!("{}/txs?status={}", self.url, tx_status); - self.json_get(&url).await + Ok(self.json_get(&url).await?) } #[instrument(skip(self))] pub async fn get_unsent_txs(&self) -> anyhow::Result> { let url = format!("{}/txs?unsent=true", self.url); - self.json_get(&url).await + Ok(self.json_get(&url).await?) } pub fn rpc_url(&self) -> String { diff --git a/e2e_tests/docker-compose/compose.yml b/e2e_tests/docker-compose/compose.yml index 05af04c7..8a610d98 100644 --- a/e2e_tests/docker-compose/compose.yml +++ b/e2e_tests/docker-compose/compose.yml @@ -28,7 +28,8 @@ services: volumes: - sequencer_db_data:/var/lib/postgresql/data tx-sitter: - image: ghcr.io/worldcoin/tx-sitter-monolith:latest + #image: ghcr.io/worldcoin/tx-sitter-monolith:latest + image: tx-sitter-monolith hostname: tx-sitter depends_on: - tx-sitter-db @@ -145,16 +146,16 @@ services: hostname: signup-sequencer-1 ports: - ${SIGNUP_SEQUENCER_1_PORT:-9081}:8080 -# signup-sequencer-2: -# <<: *signup-sequencer-def -# hostname: signup-sequencer-2 -# ports: -# - ${SIGNUP_SEQUENCER_2_PORT:-9082}:8080 -# signup-sequencer-3: -# <<: *signup-sequencer-def -# hostname: signup-sequencer-3 -# ports: -# - ${SIGNUP_SEQUENCER_3_PORT:-9083}:8080 + signup-sequencer-2: + <<: *signup-sequencer-def + hostname: signup-sequencer-2 + ports: + - ${SIGNUP_SEQUENCER_2_PORT:-9082}:8080 + signup-sequencer-3: + <<: *signup-sequencer-def + hostname: signup-sequencer-3 + ports: + - ${SIGNUP_SEQUENCER_3_PORT:-9083}:8080 volumes: tx_sitter_db_data: driver: local diff --git a/e2e_tests/scenarios/tests/common/api.rs b/e2e_tests/scenarios/tests/common/api.rs index eee76f2d..866a5c87 100644 --- a/e2e_tests/scenarios/tests/common/api.rs +++ b/e2e_tests/scenarios/tests/common/api.rs @@ -8,7 +8,7 @@ use signup_sequencer::identity_tree::Hash; use signup_sequencer::server::data::{ DeletionRequest, InclusionProofRequest, InclusionProofResponse, InsertCommitmentRequest, }; -use tracing::debug; +use tracing::{debug, info}; use crate::common::prelude::StatusCode; @@ -126,11 +126,15 @@ pub async fn inclusion_proof( client: &Client, uri: &String, commitment: &Hash, -) -> anyhow::Result { +) -> anyhow::Result<(StatusCode, Option)> { let result = inclusion_proof_raw(client, uri, commitment).await?; + if !result.status_code.is_success() { + return Ok((result.status_code, None)); + } + let result_json = serde_json::from_str::(&result.body) .expect("Failed to parse response as json"); - Ok(result_json) + Ok((result.status_code, Some(result_json))) } diff --git a/e2e_tests/scenarios/tests/common/docker_compose.rs b/e2e_tests/scenarios/tests/common/docker_compose.rs index ca534ef8..a21762b9 100644 --- a/e2e_tests/scenarios/tests/common/docker_compose.rs +++ b/e2e_tests/scenarios/tests/common/docker_compose.rs @@ -3,7 +3,7 @@ use std::process::{Command, Stdio}; use std::sync::atomic::{AtomicU32, Ordering}; use std::time::{Duration, Instant}; -use anyhow::{Context, Error}; +use anyhow::{anyhow, Context, Error}; use hyper::{Body, Client}; use rand::distributions::Alphanumeric; use rand::{thread_rng, Rng}; @@ -133,7 +133,7 @@ impl<'a> DockerComposeGuard<'a> { stdout, stderr ); - Ok(parse_exposed_port(stdout)) + parse_exposed_port(stdout) } } @@ -189,11 +189,27 @@ pub async fn setup(cwd: &str) -> anyhow::Result { tokio::time::sleep(Duration::from_secs(1)).await; - let balancer_port = res.get_mapped_port("signup-sequencer-balancer", 8080)?; - res.update_balancer_port(balancer_port); + let mut balancer_port = Err(anyhow!("Balancer port not queried.")); + for _ in 0..3 { + balancer_port = res.get_mapped_port("signup-sequencer-balancer", 8080); + if balancer_port.is_ok() { + break; + } + + tokio::time::sleep(Duration::from_secs(1)).await; + } + res.update_balancer_port(balancer_port?); + + let mut chain_port = Err(anyhow!("Chain port not queried.")); + for _ in 0..3 { + chain_port = res.get_mapped_port("chain", 8545); + if chain_port.is_ok() { + break; + } - let chain_port = res.get_mapped_port("chain", 8545)?; - res.update_chain_port(chain_port); + tokio::time::sleep(Duration::from_secs(1)).await; + } + res.update_chain_port(chain_port?); await_running(&res).await?; @@ -284,19 +300,17 @@ fn run_cmd(cwd: &str, envs: HashMap, cmd_str: String) -> anyhow: Ok(()) } -fn parse_exposed_port(s: String) -> u32 { +fn parse_exposed_port(s: String) -> anyhow::Result { let parts: Vec<_> = s .split_whitespace() .map(|s| s.trim()) .filter(|s| !s.is_empty()) .collect(); - parts - .last() - .unwrap() - .split(':') - .last() - .unwrap() - .parse::() - .unwrap() + let port = parts.last().map(|v| v.split(':').last()).flatten(); + + match port { + Some(port) => port.parse::().map_err(|err| anyhow!(err)), + None => Err(anyhow!("Port not found in string.")), + } } diff --git a/e2e_tests/scenarios/tests/common/mod.rs b/e2e_tests/scenarios/tests/common/mod.rs index 512c43be..b144eefd 100644 --- a/e2e_tests/scenarios/tests/common/mod.rs +++ b/e2e_tests/scenarios/tests/common/mod.rs @@ -146,16 +146,20 @@ pub async fn mined_inclusion_proof_with_retries( for _i in 0..retries_count { last_res = Some(inclusion_proof(client, uri, commitment).await?); - if let Some(ref inclusion_proof_json) = last_res { - if let Some(root) = inclusion_proof_json.0.root { - let (root, ..) = chain - .identity_manager - .query_root(root.into()) - .call() - .await?; - - if root != U256::zero() { - return Ok(()); + if let Some((status_code, ref inclusion_proof_json)) = last_res { + if status_code.is_success() { + if let Some(inclusion_proof_json) = inclusion_proof_json { + if let Some(root) = inclusion_proof_json.0.root { + let (root, ..) = chain + .identity_manager + .query_root(root.into()) + .call() + .await?; + + if root != U256::zero() { + return Ok(()); + } + } } } }; diff --git a/e2e_tests/scenarios/tests/insert_100.rs b/e2e_tests/scenarios/tests/insert_100.rs index a96588f0..db04d152 100644 --- a/e2e_tests/scenarios/tests/insert_100.rs +++ b/e2e_tests/scenarios/tests/insert_100.rs @@ -16,14 +16,14 @@ async fn insert_100() -> anyhow::Result<()> { let uri = format!("http://{}", docker_compose.get_local_addr()); let client = Client::new(); - let identities = generate_test_commitments(100); + let identities = generate_test_commitments(1000); for commitment in identities.iter() { insert_identity_with_retries(&client, &uri, commitment, 10, 3.0).await?; } for commitment in identities.iter() { - mined_inclusion_proof_with_retries(&client, &uri, &chain, commitment, 60, 10.0).await?; + mined_inclusion_proof_with_retries(&client, &uri, &chain, commitment, 60, 45.0).await?; } Ok(()) diff --git a/src/contracts/mod.rs b/src/contracts/mod.rs index 5e021733..13ee3681 100644 --- a/src/contracts/mod.rs +++ b/src/contracts/mod.rs @@ -3,6 +3,7 @@ pub mod abi; pub mod scanner; use anyhow::{anyhow, bail, Context}; +use ethers::abi::AbiEncode; use ethers::providers::Middleware; use ethers::types::{H256, U256}; use tracing::{error, info, instrument}; @@ -131,8 +132,14 @@ impl IdentityManager { ) .tx; + let tx_id = format!( + "tx-{}-{}", + hex::encode(pre_root.encode()), + hex::encode(post_root.encode()) + ); + self.ethereum - .send_transaction(register_identities_transaction, true) + .send_transaction(register_identities_transaction, true, Some(tx_id)) .await .map_err(|tx_err| anyhow!("{}", tx_err.to_string())) } @@ -158,8 +165,14 @@ impl IdentityManager { ) .tx; + let tx_id = format!( + "tx-{}-{}", + hex::encode(pre_root.encode()), + hex::encode(post_root.encode()) + ); + self.ethereum - .send_transaction(delete_identities_transaction, true) + .send_transaction(delete_identities_transaction, true, Some(tx_id)) .await .map_err(|tx_err| anyhow!("{}", tx_err.to_string())) } diff --git a/src/database/query.rs b/src/database/query.rs index b8b33e30..d087a8e9 100644 --- a/src/database/query.rs +++ b/src/database/query.rs @@ -782,6 +782,23 @@ pub trait DatabaseQuery<'a>: Executor<'a, Database = Postgres> { Ok(res) } + async fn count_not_finalized_batches(self) -> Result { + let res = sqlx::query( + r#" + SELECT COUNT(*) + FROM batches + LEFT JOIN transactions ON batches.next_root = transactions.batch_next_root + LEFT JOIN identities ON batches.next_root = identities.root + WHERE transactions.batch_next_root IS NOT NULL AND batches.prev_root IS NOT NULL AND identities.status = $1 + "#, + ) + .bind(<&str>::from(ProcessedStatus::Pending)) + .fetch_one(self) + .await?; + + Ok(res.get::(0) as i32) + } + async fn get_next_batch_without_transaction(self) -> Result, Error> { let res = sqlx::query_as::<_, BatchEntry>( r#" diff --git a/src/database/transaction.rs b/src/database/transaction.rs index c2e98e11..fbb9078f 100644 --- a/src/database/transaction.rs +++ b/src/database/transaction.rs @@ -6,7 +6,7 @@ use crate::database::{Database, Error}; use crate::identity_tree::{Hash, ProcessedStatus}; use crate::retry_tx; -async fn mark_root_as_processed( +pub async fn mark_root_as_processed( tx: &mut Transaction<'_, Postgres>, root: &Hash, ) -> Result<(), Error> { diff --git a/src/ethereum/mod.rs b/src/ethereum/mod.rs index e4447ce6..1fbd6c62 100644 --- a/src/ethereum/mod.rs +++ b/src/ethereum/mod.rs @@ -78,9 +78,12 @@ impl Ethereum { &self, tx: TypedTransaction, only_once: bool, + tx_id: Option, ) -> Result { tracing::info!(?tx, "Sending transaction"); - self.write_provider.send_transaction(tx, only_once).await + self.write_provider + .send_transaction(tx, only_once, tx_id) + .await } pub async fn fetch_pending_transactions(&self) -> Result, TxError> { diff --git a/src/ethereum/write_provider/inner.rs b/src/ethereum/write_provider/inner.rs index 7ac6c64a..f21be988 100644 --- a/src/ethereum/write_provider/inner.rs +++ b/src/ethereum/write_provider/inner.rs @@ -10,6 +10,7 @@ pub trait Inner: Send + Sync + 'static { &self, tx: TypedTransaction, only_once: bool, + tx_id: Option, ) -> Result; async fn fetch_pending_transactions(&self) -> Result, TxError>; diff --git a/src/ethereum/write_provider/mod.rs b/src/ethereum/write_provider/mod.rs index b80194ca..e0d61a36 100644 --- a/src/ethereum/write_provider/mod.rs +++ b/src/ethereum/write_provider/mod.rs @@ -60,8 +60,9 @@ impl WriteProvider { &self, tx: TypedTransaction, only_once: bool, + tx_id: Option, ) -> Result { - self.inner.send_transaction(tx, only_once).await + self.inner.send_transaction(tx, only_once, tx_id).await } pub async fn fetch_pending_transactions(&self) -> Result, TxError> { diff --git a/src/ethereum/write_provider/openzeppelin.rs b/src/ethereum/write_provider/openzeppelin.rs index 7222b39b..2e5945e2 100644 --- a/src/ethereum/write_provider/openzeppelin.rs +++ b/src/ethereum/write_provider/openzeppelin.rs @@ -133,6 +133,7 @@ impl OzRelay { &self, mut tx: TypedTransaction, only_once: bool, + _tx_id: Option, ) -> Result { if let Some(gas_limit) = self.gas_limit { tx.set_gas(gas_limit); @@ -220,8 +221,9 @@ impl Inner for OzRelay { &self, tx: TypedTransaction, only_once: bool, + tx_id: Option, ) -> Result { - self.send_transaction(tx, only_once).await + self.send_transaction(tx, only_once, tx_id).await } async fn fetch_pending_transactions(&self) -> Result, TxError> { diff --git a/src/ethereum/write_provider/tx_sitter.rs b/src/ethereum/write_provider/tx_sitter.rs index c986e3b7..3108c264 100644 --- a/src/ethereum/write_provider/tx_sitter.rs +++ b/src/ethereum/write_provider/tx_sitter.rs @@ -4,8 +4,9 @@ use anyhow::Context; use async_trait::async_trait; use ethers::types::transaction::eip2718::TypedTransaction; use ethers::types::U256; +use reqwest::StatusCode; use tx_sitter_client::data::{SendTxRequest, TransactionPriority, TxStatus}; -use tx_sitter_client::TxSitterClient; +use tx_sitter_client::{HttpError, TxSitterClient}; use super::inner::{Inner, TransactionResult}; use crate::config::TxSitterConfig; @@ -50,7 +51,7 @@ impl TxSitter { }); } - tokio::time::sleep(std::time::Duration::from_secs(1)).await; + tokio::time::sleep(Duration::from_secs(1)).await; } } } @@ -61,13 +62,14 @@ impl Inner for TxSitter { &self, mut tx: TypedTransaction, _only_once: bool, + tx_id: Option, ) -> Result { if let Some(gas_limit) = self.gas_limit { tx.set_gas(gas_limit); } // TODO: Handle only_once - let tx = self + let res = self .client .send_tx(&SendTxRequest { to: *tx @@ -81,9 +83,24 @@ impl Inner for TxSitter { .context("Missing tx gas limit") .map_err(TxError::Send)?, priority: TransactionPriority::Regular, - tx_id: None, + tx_id: tx_id.clone(), }) - .await + .await; + + let res = match res { + Err(err) => match err.downcast_ref::() { + Some(http_err) if http_err.status == StatusCode::CONFLICT => { + if let Some(tx_id) = tx_id { + return Ok(tx_id); + } + Err(err) + } + _ => Err(err), + }, + res => res, + }; + + let tx = res .context("Error sending transaction") .map_err(TxError::Send)?; diff --git a/src/identity/processor.rs b/src/identity/processor.rs index 63c32517..0d854ca3 100644 --- a/src/identity/processor.rs +++ b/src/identity/processor.rs @@ -9,7 +9,7 @@ use ethers::abi::RawLog; use ethers::addressbook::Address; use ethers::contract::EthEvent; use ethers::middleware::Middleware; -use ethers::prelude::{Log, Topic, ValueOrArray, U256}; +use ethers::prelude::{Log, Topic, ValueOrArray}; use tokio::sync::Notify; use tracing::{error, info, instrument}; @@ -18,10 +18,11 @@ use crate::contracts::abi::{BridgedWorldId, RootAddedFilter, TreeChangeKind, Tre use crate::contracts::scanner::BlockScanner; use crate::contracts::IdentityManager; use crate::database::query::DatabaseQuery; +use crate::database::transaction::{mark_root_as_mined, mark_root_as_processed}; use crate::database::types::{BatchEntry, BatchType}; use crate::database::{Database, Error}; use crate::ethereum::{Ethereum, ReadProvider}; -use crate::identity_tree::Hash; +use crate::identity_tree::{Hash, ProcessedStatus}; use crate::prover::identity::Identity; use crate::prover::repository::ProverRepository; use crate::prover::Prover; @@ -222,16 +223,16 @@ impl OnChainIdentityProcessor { ) -> anyhow::Result { self.validate_merkle_proofs(&batch.data.0.identities)?; let start_index = *batch.data.0.indexes.first().expect("Should exist."); - let pre_root: U256 = batch.prev_root.expect("Should exist.").into(); - let post_root: U256 = batch.next_root.into(); + let pre_root: Hash = batch.prev_root.expect("Should exist."); + let post_root: Hash = batch.next_root; // We prepare the proof before reserving a slot in the pending identities let proof = crate::prover::proof::prepare_insertion_proof( prover, start_index, - pre_root, + pre_root.into(), &batch.data.0.identities, - post_root, + post_root.into(), ) .await?; @@ -248,8 +249,8 @@ impl OnChainIdentityProcessor { .identity_manager .register_identities( start_index, - pre_root, - post_root, + pre_root.into(), + post_root.into(), batch.data.0.identities.clone(), proof, ) @@ -277,17 +278,17 @@ impl OnChainIdentityProcessor { batch: &BatchEntry, ) -> anyhow::Result { self.validate_merkle_proofs(&batch.data.0.identities)?; - let pre_root: U256 = batch.prev_root.expect("Should exist.").into(); - let post_root: U256 = batch.next_root.into(); + let pre_root: Hash = batch.prev_root.expect("Should exist."); + let post_root: Hash = batch.next_root; let deletion_indices: Vec<_> = batch.data.0.indexes.iter().map(|&v| v as u32).collect(); // We prepare the proof before reserving a slot in the pending identities let proof = crate::prover::proof::prepare_deletion_proof( prover, - pre_root, + pre_root.into(), deletion_indices.clone(), batch.data.0.identities.clone(), - post_root, + post_root.into(), ) .await?; @@ -299,7 +300,12 @@ impl OnChainIdentityProcessor { // identity manager and wait for that transaction to be mined. let transaction_id = self .identity_manager - .delete_identities(proof, packed_deletion_indices, pre_root, post_root) + .delete_identities( + proof, + packed_deletion_indices, + pre_root.into(), + post_root.into(), + ) .await .map_err(|e| { error!(?e, "Failed to insert identity to contract."); @@ -362,7 +368,7 @@ impl OnChainIdentityProcessor { Ok(mainnet_logs) } - async fn fetch_secondary_logs(&self) -> anyhow::Result> + async fn fetch_secondary_logs(&self) -> anyhow::Result> where ::Error: 'static, { @@ -404,20 +410,37 @@ impl OnChainIdentityProcessor { continue; }; - let pre_root = event.pre_root; - let post_root = event.post_root; + let pre_root: Hash = event.pre_root.into(); + let post_root: Hash = event.post_root.into(); let kind = TreeChangeKind::from(event.kind); info!(?pre_root, ?post_root, ?kind, "Mining batch"); // Double check - if !self.identity_manager.is_root_mined(post_root).await? { + if !self + .identity_manager + .is_root_mined(post_root.into()) + .await? + { continue; } - self.database - .mark_root_as_processed_tx(&post_root.into()) - .await?; + retry_tx!(self.database.pool, tx, { + // With current flow it is required to mark root as processed first as this is + // how required mined_at field is set, We set proper state only if not set + // previously. + let root_state = tx.get_root_state(&post_root.into()).await?; + match root_state { + Some(root_state) if root_state.status == ProcessedStatus::Processed => {} + Some(root_state) if root_state.status == ProcessedStatus::Mined => {} + _ => { + mark_root_as_processed(&mut tx, &post_root.into()).await?; + } + } + + Ok::<(), anyhow::Error>(()) + }) + .await?; info!(?pre_root, ?post_root, ?kind, "Batch mined"); @@ -436,20 +459,34 @@ impl OnChainIdentityProcessor { } #[instrument(level = "info", skip_all)] - async fn finalize_secondary_roots(&self, roots: Vec) -> Result<(), anyhow::Error> { + async fn finalize_secondary_roots(&self, roots: Vec) -> Result<(), anyhow::Error> { for root in roots { info!(?root, "Finalizing root"); // Check if mined on all L2s if !self .identity_manager - .is_root_mined_multi_chain(root) + .is_root_mined_multi_chain(root.into()) .await? { continue; } - self.database.mark_root_as_mined_tx(&root.into()).await?; + retry_tx!(self.database.pool, tx, { + // With current flow it is required to mark root as processed first as this is + // how required mined_at field is set, We set proper state only if not set + // previously. + let root_state = tx.get_root_state(&root.into()).await?; + match root_state { + Some(root_state) if root_state.status == ProcessedStatus::Mined => {} + _ => { + mark_root_as_mined(&mut tx, &root.into()).await?; + } + } + + Ok::<(), anyhow::Error>(()) + }) + .await?; info!(?root, "Root finalized"); } @@ -457,7 +494,7 @@ impl OnChainIdentityProcessor { Ok(()) } - fn extract_roots_from_mainnet_logs(mainnet_logs: Vec) -> Vec { + fn extract_roots_from_mainnet_logs(mainnet_logs: Vec) -> Vec { let mut roots = vec![]; for log in mainnet_logs { let Some(event) = Self::raw_log_to_tree_changed(&log) else { @@ -466,7 +503,7 @@ impl OnChainIdentityProcessor { let post_root = event.post_root; - roots.push(post_root); + roots.push(post_root.into()); } roots } @@ -477,13 +514,13 @@ impl OnChainIdentityProcessor { TreeChangedFilter::decode_log(&raw_log).ok() } - fn extract_roots_from_secondary_logs(logs: &[Log]) -> Vec { + fn extract_roots_from_secondary_logs(logs: &[Log]) -> Vec { let mut roots = vec![]; for log in logs { let raw_log = RawLog::from((log.topics.clone(), log.data.to_vec())); if let Ok(event) = RootAddedFilter::decode_log(&raw_log) { - roots.push(event.root); + roots.push(event.root.into()); } } @@ -507,7 +544,6 @@ impl OnChainIdentityProcessor { .database .get_non_zero_commitments_by_leaf_indexes(commitments.iter().copied()) .await?; - let commitments: Vec = commitments.into_iter().map(Into::into).collect(); // Fetch the root history expiry time on chain let root_history_expiry = self.identity_manager.root_history_expiry().await?; @@ -565,14 +601,25 @@ impl IdentityProcessor for OffChainIdentityProcessor { self.update_eligible_recoveries(batch).await?; } - // With current flow it is required to mark root as processed first as this is - // how required mined_at field is set - self.database - .mark_root_as_processed_tx(&batch.next_root) - .await?; - self.database - .mark_root_as_mined_tx(&batch.next_root) - .await?; + retry_tx!(self.database.pool, tx, { + // With current flow it is required to mark root as processed first as this is + // how required mined_at field is set, We set proper state only if not set + // previously. + let root_state = tx.get_root_state(&batch.next_root).await?; + match root_state { + Some(root_state) if root_state.status == ProcessedStatus::Processed => { + mark_root_as_mined(&mut tx, &batch.next_root).await?; + } + Some(root_state) if root_state.status == ProcessedStatus::Mined => {} + _ => { + mark_root_as_processed(&mut tx, &batch.next_root).await?; + mark_root_as_mined(&mut tx, &batch.next_root).await?; + } + } + + Ok::<(), anyhow::Error>(()) + }) + .await?; sync_tree_notify.notify_one(); } @@ -613,8 +660,12 @@ impl OffChainIdentityProcessor { async fn update_eligible_recoveries(&self, batch: &BatchEntry) -> anyhow::Result<()> { retry_tx!(self.database.pool, tx, { - let commitments: Vec = - batch.data.identities.iter().map(|v| v.commitment).collect(); + let commitments: Vec = batch + .data + .identities + .iter() + .map(|v| v.commitment.into()) + .collect(); let eligibility_timestamp = Utc::now(); // Check if any deleted commitments correspond with entries in the diff --git a/src/task_monitor/tasks/create_batches.rs b/src/task_monitor/tasks/create_batches.rs index 6a0febda..d11cd1f7 100644 --- a/src/task_monitor/tasks/create_batches.rs +++ b/src/task_monitor/tasks/create_batches.rs @@ -255,6 +255,17 @@ pub async fn insert_identities( assert_updates_are_consecutive(updates); let pre_root = batching_tree.get_root(); + + let mut tx = database.pool.begin().await?; + let latest_batch = tx.get_latest_batch().await?; + if let Some(latest_batch) = latest_batch { + if pre_root != latest_batch.next_root { + // Tree not synced + sync_tree_notify.notify_one(); + return Ok(()); + } + } + let mut insertion_indices: Vec<_> = updates.iter().map(|f| f.update.leaf_index).collect(); let mut commitments: Vec = updates .iter() @@ -332,15 +343,17 @@ pub async fn insert_identities( ); // With all the data prepared we can submit the batch to database. - database - .insert_new_batch( - &post_root, - &pre_root, - database::types::BatchType::Insertion, - &identity_commitments, - &insertion_indices, - ) - .await?; + tx.insert_new_batch( + &post_root, + &pre_root, + database::types::BatchType::Insertion, + &identity_commitments, + &insertion_indices, + ) + .await?; + + // It is important to commit transaction as soon as possible. + tx.commit().await?; tracing::info!( start_index, @@ -394,6 +407,16 @@ pub async fn delete_identities( // Grab the initial conditions before the updates are applied to the tree. let pre_root = batching_tree.get_root(); + let mut tx = database.pool.begin().await?; + let latest_batch = tx.get_latest_batch().await?; + if let Some(latest_batch) = latest_batch { + if pre_root != latest_batch.next_root { + // Tree not synced + sync_tree_notify.notify_one(); + return Ok(()); + } + } + let mut deletion_indices: Vec<_> = updates.iter().map(|f| f.update.leaf_index).collect(); let commitments = batching_tree.commitments_by_leaves(deletion_indices.iter().copied()); @@ -462,15 +485,17 @@ pub async fn delete_identities( tracing::info!(?pre_root, ?post_root, "Submitting deletion batch to DB"); // With all the data prepared we can submit the batch to database. - database - .insert_new_batch( - &post_root, - &pre_root, - database::types::BatchType::Deletion, - &identity_commitments, - &deletion_indices, - ) - .await?; + tx.insert_new_batch( + &post_root, + &pre_root, + database::types::BatchType::Deletion, + &identity_commitments, + &deletion_indices, + ) + .await?; + + // It is important to commit transaction as soon as possible. + tx.commit().await?; tracing::info!(?pre_root, ?post_root, "Deletion batch submitted to DB"); @@ -479,7 +504,6 @@ pub async fn delete_identities( TaskMonitor::log_batch_size(updates.len()); sync_tree_notify.notify_one(); - // batching_tree.apply_updates_up_to(post_root); Ok(()) } diff --git a/src/task_monitor/tasks/process_batches.rs b/src/task_monitor/tasks/process_batches.rs index 67fac774..7eb2e6ce 100644 --- a/src/task_monitor/tasks/process_batches.rs +++ b/src/task_monitor/tasks/process_batches.rs @@ -8,6 +8,8 @@ use crate::app::App; use crate::database::query::DatabaseQuery as _; use crate::identity::processor::TransactionId; +const MAX_BUFFERED_TRANSACTIONS: i32 = 5; + pub async fn process_batches( app: Arc, monitored_txs_sender: Arc>, @@ -41,21 +43,37 @@ pub async fn process_batches( }, } - let next_batch = app.database.get_next_batch_without_transaction().await?; - let Some(next_batch) = next_batch else { - continue; - }; + { + let mut tx = app.database.pool.begin().await?; + + sqlx::query("LOCK TABLE transactions IN ACCESS EXCLUSIVE MODE;") + .execute(&mut *tx) + .await?; + + let buffered_transactions = tx.count_not_finalized_batches().await?; + if buffered_transactions >= MAX_BUFFERED_TRANSACTIONS { + tx.commit().await?; + continue; + } - let tx_id = app - .identity_processor - .commit_identities(&next_batch) - .await?; + let next_batch = tx.get_next_batch_without_transaction().await?; + let Some(next_batch) = next_batch else { + tx.commit().await?; + continue; + }; - monitored_txs_sender.send(tx_id.clone()).await?; + let tx_id = app + .identity_processor + .commit_identities(&next_batch) + .await?; - app.database - .insert_new_transaction(&tx_id, &next_batch.next_root) - .await?; + monitored_txs_sender.send(tx_id.clone()).await?; + + tx.insert_new_transaction(&tx_id, &next_batch.next_root) + .await?; + + tx.commit().await?; + } // We want to check if there's a full batch available immediately check_next_batch_notify.notify_one(); diff --git a/src/task_monitor/tasks/sync_tree_state_with_db.rs b/src/task_monitor/tasks/sync_tree_state_with_db.rs index 022fbc8b..c70d282e 100644 --- a/src/task_monitor/tasks/sync_tree_state_with_db.rs +++ b/src/task_monitor/tasks/sync_tree_state_with_db.rs @@ -97,17 +97,17 @@ async fn sync_tree( latest_tree.apply_updates(&tree_updates); if let Some(batching_tree_update) = latest_batching_tree_update { - if batching_tree_update.sequence_id >= batching_tree.get_last_sequence_id() { + if batching_tree_update.sequence_id > batching_tree.get_last_sequence_id() { batching_tree.apply_updates_up_to(batching_tree_update.post_root); - } else { + } else if batching_tree_update.sequence_id < batching_tree.get_last_sequence_id() { batching_tree.rewind_updates_up_to(batching_tree_update.post_root); } } } else { if let Some(batching_tree_update) = latest_batching_tree_update { - if batching_tree_update.sequence_id >= batching_tree.get_last_sequence_id() { + if batching_tree_update.sequence_id > batching_tree.get_last_sequence_id() { batching_tree.apply_updates_up_to(batching_tree_update.post_root); - } else { + } else if batching_tree_update.sequence_id < batching_tree.get_last_sequence_id() { batching_tree.rewind_updates_up_to(batching_tree_update.post_root); } } @@ -116,7 +116,9 @@ async fn sync_tree( } if let Some(ref processed_tree_update) = latest_processed_tree_update { - processed_tree.apply_updates_up_to(processed_tree_update.post_root); + if processed_tree_update.sequence_id > processed_tree.get_last_sequence_id() { + processed_tree.apply_updates_up_to(processed_tree_update.post_root); + } } Ok(()) From 34938c1e0bf8981e9d765b8fce40c2c5a0adb2ee Mon Sep 17 00:00:00 2001 From: Piotr Heilman Date: Tue, 13 Aug 2024 09:45:35 +0200 Subject: [PATCH 5/7] Refactor. --- crates/tx-sitter-client/src/lib.rs | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/crates/tx-sitter-client/src/lib.rs b/crates/tx-sitter-client/src/lib.rs index 8241833e..ea462c82 100644 --- a/crates/tx-sitter-client/src/lib.rs +++ b/crates/tx-sitter-client/src/lib.rs @@ -73,19 +73,19 @@ impl TxSitterClient { #[instrument(skip(self))] pub async fn send_tx(&self, req: &SendTxRequest) -> anyhow::Result { - Ok(self.json_post(&format!("{}/tx", self.url), req).await?) + self.json_post(&format!("{}/tx", self.url), req).await } #[instrument(skip(self))] pub async fn get_tx(&self, tx_id: &str) -> anyhow::Result { - Ok(self.json_get(&format!("{}/tx/{}", self.url, tx_id)).await?) + self.json_get(&format!("{}/tx/{}", self.url, tx_id)).await } #[instrument(skip(self))] pub async fn get_txs(&self) -> anyhow::Result> { let url = format!("{}/txs", self.url); - Ok(self.json_get(&url).await?) + self.json_get(&url).await } #[instrument(skip(self))] @@ -95,14 +95,14 @@ impl TxSitterClient { ) -> anyhow::Result> { let url = format!("{}/txs?status={}", self.url, tx_status); - Ok(self.json_get(&url).await?) + self.json_get(&url).await } #[instrument(skip(self))] pub async fn get_unsent_txs(&self) -> anyhow::Result> { let url = format!("{}/txs?unsent=true", self.url); - Ok(self.json_get(&url).await?) + self.json_get(&url).await } pub fn rpc_url(&self) -> String { From 6e1f300fdef75e9a02d8e4cf3669753bd70fe70a Mon Sep 17 00:00:00 2001 From: Piotr Heilman Date: Tue, 13 Aug 2024 09:48:55 +0200 Subject: [PATCH 6/7] Refactor. --- src/task_monitor/tasks/create_batches.rs | 1 - 1 file changed, 1 deletion(-) diff --git a/src/task_monitor/tasks/create_batches.rs b/src/task_monitor/tasks/create_batches.rs index d11cd1f7..ba5161fd 100644 --- a/src/task_monitor/tasks/create_batches.rs +++ b/src/task_monitor/tasks/create_batches.rs @@ -367,7 +367,6 @@ pub async fn insert_identities( TaskMonitor::log_batch_size(updates.len()); sync_tree_notify.notify_one(); - // batching_tree.apply_updates_up_to(post_root); Ok(()) } From 5551a33fd69077be89d225c7441e4913d2510ded Mon Sep 17 00:00:00 2001 From: Piotr Heilman Date: Wed, 21 Aug 2024 15:04:50 +0200 Subject: [PATCH 7/7] Revert test to original value. --- e2e_tests/scenarios/tests/insert_100.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/e2e_tests/scenarios/tests/insert_100.rs b/e2e_tests/scenarios/tests/insert_100.rs index db04d152..c0236f47 100644 --- a/e2e_tests/scenarios/tests/insert_100.rs +++ b/e2e_tests/scenarios/tests/insert_100.rs @@ -16,7 +16,7 @@ async fn insert_100() -> anyhow::Result<()> { let uri = format!("http://{}", docker_compose.get_local_addr()); let client = Client::new(); - let identities = generate_test_commitments(1000); + let identities = generate_test_commitments(100); for commitment in identities.iter() { insert_identity_with_retries(&client, &uri, commitment, 10, 3.0).await?;