From f9e2ac0929e2f340f6ef9104712a57a2fd1b1029 Mon Sep 17 00:00:00 2001 From: Mariano Nicolini Date: Wed, 25 Sep 2024 17:12:55 -0300 Subject: [PATCH 01/47] Save work in progress --- batcher/aligned-batcher/src/lib.rs | 23 +++++++++++++---------- batcher/aligned-batcher/src/types/mod.rs | 2 +- 2 files changed, 14 insertions(+), 11 deletions(-) diff --git a/batcher/aligned-batcher/src/lib.rs b/batcher/aligned-batcher/src/lib.rs index 3e2ceed89..0f87453fd 100644 --- a/batcher/aligned-batcher/src/lib.rs +++ b/batcher/aligned-batcher/src/lib.rs @@ -7,6 +7,7 @@ use dotenvy::dotenv; use ethers::contract::ContractError; use ethers::signers::Signer; use serde::Serialize; +use types::user_state::UserState; use std::collections::hash_map::Entry; use std::collections::HashMap; @@ -65,22 +66,12 @@ const DEFAULT_AGGREGATOR_FEE_DIVIDER: u128 = 2; struct BatchState { batch_queue: BatchQueue, - user_nonces: HashMap, - /// The minimum fee of a pending proof for a user. - /// This should always be the fee of the biggest pending nonce by the user. - /// This is used to check if a user is submitting a proof with a higher nonce and higher fee, - /// which is invalid and should be rejected. - user_min_fee: HashMap, - user_proof_count_in_batch: HashMap, } impl BatchState { fn new() -> Self { Self { batch_queue: BatchQueue::new(), - user_nonces: HashMap::new(), - user_min_fee: HashMap::new(), - user_proof_count_in_batch: HashMap::new(), } } @@ -212,6 +203,7 @@ pub struct Batcher { pre_verification_is_enabled: bool, non_paying_config: Option, posting_batch: Mutex, + user_states: HashMap>, } impl Batcher { @@ -322,6 +314,7 @@ impl Batcher { pre_verification_is_enabled: config.batcher.pre_verification_is_enabled, non_paying_config, posting_batch: Mutex::new(false), + user_states: HashMap::new(), } } @@ -419,6 +412,8 @@ impl Batcher { Ok(()) } + async fn get_or_insert(&self) -> Mutex<> + /// Handle an individual message from the client. async fn handle_message( self: Arc, @@ -461,6 +456,14 @@ impl Batcher { self.handle_nonpaying_msg(ws_conn_sink.clone(), client_msg) .await } else { + // Check that we had a user state entry for this user and insert it if not. + if !self.user_states.contains_key(&addr) { + let new_user_state = Mutex::new(UserState::new()); + self.user_states.insert(addr.clone(), new_user_state); + } + + + if !self .check_user_balance_and_increment_proof_count(&addr) .await diff --git a/batcher/aligned-batcher/src/types/mod.rs b/batcher/aligned-batcher/src/types/mod.rs index 438bf92f8..cb2327411 100644 --- a/batcher/aligned-batcher/src/types/mod.rs +++ b/batcher/aligned-batcher/src/types/mod.rs @@ -1,3 +1,3 @@ pub(crate) mod batch_queue; - pub mod errors; +pub(crate) mod user_state; From 0faa5b420ba9b3713be0fca14e0912b908825828 Mon Sep 17 00:00:00 2001 From: Mariano Nicolini Date: Wed, 25 Sep 2024 18:22:40 -0300 Subject: [PATCH 02/47] Save work in progress - Adding user state validations --- batcher/aligned-batcher/src/lib.rs | 97 +++++++++++++++++------------ operator/merkle_tree/lib/Cargo.lock | 8 +-- 2 files changed, 62 insertions(+), 43 deletions(-) diff --git a/batcher/aligned-batcher/src/lib.rs b/batcher/aligned-batcher/src/lib.rs index 0f87453fd..b68a6ff55 100644 --- a/batcher/aligned-batcher/src/lib.rs +++ b/batcher/aligned-batcher/src/lib.rs @@ -412,8 +412,6 @@ impl Batcher { Ok(()) } - async fn get_or_insert(&self) -> Mutex<> - /// Handle an individual message from the client. async fn handle_message( self: Arc, @@ -462,18 +460,28 @@ impl Batcher { self.user_states.insert(addr.clone(), new_user_state); } + // At this point, we will have a user state for sure, since we have inserted it + // if not already present + let user_state = self.user_states.get(&addr).unwrap().lock().await; - if !self - .check_user_balance_and_increment_proof_count(&addr) - .await - { + // Perform validations on user state + if self.user_balance_is_unlocked(&addr).await { send_message( ws_conn_sink.clone(), ValidityResponseMessage::InsufficientBalance(addr), ) .await; + return Ok(()); + } + let updated_user_proofs = user_state.proofs_in_batch + 1; + if !self.check_user_balance(&addr, updated_user_proofs).await { + send_message( + ws_conn_sink.clone(), + ValidityResponseMessage::InsufficientBalance(addr), + ) + .await; return Ok(()); } @@ -497,7 +505,6 @@ impl Batcher { // Nonce and max fee verification let nonce = nonced_verification_data.nonce; let max_fee = nonced_verification_data.max_fee; - if max_fee < U256::from(MIN_FEE_PER_PROOF) { error!("The max fee signed in the message is less than the accepted minimum fee to be included in the batch."); send_message(ws_conn_sink.clone(), ValidityResponseMessage::InvalidMaxFee) @@ -505,30 +512,52 @@ impl Batcher { return Ok(()); } - let mut batch_state = self.batch_state.lock().await; - - let expected_user_nonce = match batch_state.user_nonces.get(&addr) { - Some(nonce) => *nonce, - None => { - let user_nonce = match self.get_user_nonce(addr).await { - Ok(nonce) => nonce, - Err(e) => { - error!("Failed to get user nonce for address {:?}: {:?}", addr, e); - send_message( - ws_conn_sink.clone(), - ValidityResponseMessage::InvalidNonce, - ) - .await; + let expected_nonce = if let Some(cached_nonce) = user_state.nonce { + cached_nonce + } else { + let ethereum_user_nonce = match self.get_user_nonce_from_ethereum(addr).await { + Ok(ethereum_user_nonce) => ethereum_user_nonce, + Err(e) => { + error!( + "Failed to get user nonce from Ethereum for address {:?}. Error: {:?}", + addr, e + ); + send_message( + ws_conn_sink.clone(), + ValidityResponseMessage::InvalidNonce, + ) + .await; - return Ok(()); - } - }; + return Ok(()); + } + }; - batch_state.user_nonces.insert(addr, user_nonce); - user_nonce - } + user_state.nonce = Some(ethereum_user_nonce); + ethereum_user_nonce }; + // let expected_user_nonce = match user_state.nonce { + // Some(nonce) => nonce, + // None => { + // let user_nonce = match self.get_user_nonce_from_ethereum(addr).await { + // Ok(nonce) => nonce, + // Err(e) => { + // error!("Failed to get user nonce for address {:?}: {:?}", addr, e); + // send_message( + // ws_conn_sink.clone(), + // ValidityResponseMessage::InvalidNonce, + // ) + // .await; + + // return Ok(()); + // } + // }; + + // batch_state.user_nonces.insert(addr, user_nonce); + // user_nonce + // } + // }; + let min_fee = match batch_state.user_min_fee.get(&addr) { Some(fee) => *fee, None => U256::max_value(), @@ -602,22 +631,12 @@ impl Batcher { // Checks user has sufficient balance // If user has sufficient balance, increments the user's proof count in the batch - async fn check_user_balance_and_increment_proof_count(&self, addr: &Address) -> bool { - if self.user_balance_is_unlocked(addr).await { - return false; - } - let mut batch_state = self.batch_state.lock().await; - - let user_proofs_in_batch = batch_state.get_user_proof_count(addr) + 1; - + async fn check_user_balance(&self, addr: &Address, user_proofs_in_batch: u64) -> bool { let user_balance = self.get_user_balance(addr).await; - let min_balance = U256::from(user_proofs_in_batch) * U256::from(MIN_FEE_PER_PROOF); if user_balance < min_balance { return false; } - - batch_state.increment_user_proof_count(addr); true } @@ -743,7 +762,7 @@ impl Batcher { true } - async fn get_user_nonce( + async fn get_user_nonce_from_ethereum( &self, addr: Address, ) -> Result> { diff --git a/operator/merkle_tree/lib/Cargo.lock b/operator/merkle_tree/lib/Cargo.lock index 079878cdb..47549f3ba 100644 --- a/operator/merkle_tree/lib/Cargo.lock +++ b/operator/merkle_tree/lib/Cargo.lock @@ -1677,9 +1677,9 @@ dependencies = [ [[package]] name = "lambdaworks-crypto" -version = "0.7.0" +version = "0.10.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7fb5d4f22241504f7c7b8d2c3a7d7835d7c07117f10bff2a7d96a9ef6ef217c3" +checksum = "bbc2a4da0d9e52ccfe6306801a112e81a8fc0c76aa3e4449fefeda7fef72bb34" dependencies = [ "lambdaworks-math", "serde", @@ -1689,9 +1689,9 @@ dependencies = [ [[package]] name = "lambdaworks-math" -version = "0.7.0" +version = "0.10.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "358e172628e713b80a530a59654154bfc45783a6ed70ea284839800cebdf8f97" +checksum = "d1bd2632acbd9957afc5aeec07ad39f078ae38656654043bf16e046fa2730e23" dependencies = [ "serde", "serde_json", From 304e8624db7e9822259e38fe64304a29716f2899 Mon Sep 17 00:00:00 2001 From: Mariano Nicolini Date: Wed, 25 Sep 2024 19:07:44 -0300 Subject: [PATCH 03/47] Save work in progress - Refactoring message handlers --- batcher/aligned-batcher/src/lib.rs | 148 +++++++++++++++++++---------- 1 file changed, 100 insertions(+), 48 deletions(-) diff --git a/batcher/aligned-batcher/src/lib.rs b/batcher/aligned-batcher/src/lib.rs index b68a6ff55..cd60d3a60 100644 --- a/batcher/aligned-batcher/src/lib.rs +++ b/batcher/aligned-batcher/src/lib.rs @@ -503,7 +503,7 @@ impl Batcher { } // Nonce and max fee verification - let nonce = nonced_verification_data.nonce; + let msg_nonce = nonced_verification_data.nonce; let max_fee = nonced_verification_data.max_fee; if max_fee < U256::from(MIN_FEE_PER_PROOF) { error!("The max fee signed in the message is less than the accepted minimum fee to be included in the batch."); @@ -558,61 +558,112 @@ impl Batcher { // } // }; - let min_fee = match batch_state.user_min_fee.get(&addr) { - Some(fee) => *fee, - None => U256::max_value(), - }; + // let min_fee = match batch_state.user_min_fee.get(&addr) { + // Some(fee) => *fee, + // None => U256::max_value(), + // }; - match expected_user_nonce.cmp(&nonce) { - std::cmp::Ordering::Less => { - // invalid, expected user nonce < nonce + let min_fee = user_state.min_fee; + + if expected_nonce < msg_nonce { + warn!( + "Invalid nonce for address {addr}, had nonce {:?} < {:?}", + expected_nonce, msg_nonce + ); + send_message(ws_conn_sink.clone(), ValidityResponseMessage::InvalidNonce).await; + return Ok(()); + } else if expected_nonce == msg_nonce { + let max_fee = nonced_verification_data.max_fee; + if max_fee > min_fee { warn!( - "Invalid nonce for address {addr}, had nonce {:?} < {:?}", - expected_user_nonce, nonce + "Invalid max fee for address {addr}, had fee {:?} < {:?}", + min_fee, max_fee ); - send_message(ws_conn_sink.clone(), ValidityResponseMessage::InvalidNonce) + send_message(ws_conn_sink.clone(), ValidityResponseMessage::InvalidMaxFee) .await; return Ok(()); } - std::cmp::Ordering::Equal => { - // if we are here nonce == expected_user_nonce - if !self - .handle_expected_nonce_message( - batch_state, - min_fee, - nonced_verification_data, - ws_conn_sink.clone(), - client_msg.signature, - addr, - ) - .await - { - // message should not be added to batch - return Ok(()); - }; - } - std::cmp::Ordering::Greater => { - // might be replacement message - // if the message is already in the batch - // we can check if we need to increment the fee - // get the entry with the same sender and nonce - if !self - .handle_replacement_message( - batch_state, - nonced_verification_data, - ws_conn_sink.clone(), - client_msg.signature, - addr, - expected_user_nonce, - ) - .await - { - // message should not be added to batch - return Ok(()); - } + + user_state.nonce = Some(msg_nonce + U256::one()); + user_state.min_fee = max_fee; + self.add_to_batch( + nonced_verification_data, + ws_conn_sink.clone(), + client_msg.signature, + addr, + ) + .await; + } else { + // might be replacement message + // if the message is already in the batch + // we can check if we need to increment the fee + // get the entry with the same sender and nonce + if !self + .handle_replacement_message( + batch_state, + nonced_verification_data, + ws_conn_sink.clone(), + client_msg.signature, + addr, + expected_user_nonce, + ) + .await + { + // message should not be added to batch + return Ok(()); } } + // match expected_user_nonce.cmp(&nonce) { + // std::cmp::Ordering::Less => { + // // invalid, expected user nonce < nonce + // warn!( + // "Invalid nonce for address {addr}, had nonce {:?} < {:?}", + // expected_user_nonce, nonce + // ); + // send_message(ws_conn_sink.clone(), ValidityResponseMessage::InvalidNonce) + // .await; + // return Ok(()); + // } + // std::cmp::Ordering::Equal => { + // // if we are here nonce == expected_user_nonce + // if !self + // .handle_expected_nonce_message( + // batch_state, + // min_fee, + // nonced_verification_data, + // ws_conn_sink.clone(), + // client_msg.signature, + // addr, + // ) + // .await + // { + // // message should not be added to batch + // return Ok(()); + // }; + // } + // std::cmp::Ordering::Greater => { + // // might be replacement message + // // if the message is already in the batch + // // we can check if we need to increment the fee + // // get the entry with the same sender and nonce + // if !self + // .handle_replacement_message( + // batch_state, + // nonced_verification_data, + // ws_conn_sink.clone(), + // client_msg.signature, + // addr, + // expected_user_nonce, + // ) + // .await + // { + // // message should not be added to batch + // return Ok(()); + // } + // } + // } + info!("Verification data message handled"); send_message(ws_conn_sink, ValidityResponseMessage::Valid).await; @@ -775,7 +826,7 @@ impl Batcher { /// Adds verification data to the current batch queue. async fn add_to_batch( &self, - mut batch_state: tokio::sync::MutexGuard<'_, BatchState>, + // mut batch_state: tokio::sync::MutexGuard<'_, BatchState>, verification_data: NoncedVerificationData, ws_conn_sink: Arc, Message>>>, proof_submitter_sig: Signature, @@ -788,6 +839,7 @@ impl Batcher { let max_fee = verification_data.max_fee; let nonce = verification_data.nonce; + let batch_state = self.batch_state.lock().await; batch_state.batch_queue.push( BatchQueueEntry::new( verification_data, From 4010b807dcd8eef51b80839e0c0541445ae3d261 Mon Sep 17 00:00:00 2001 From: Mariano Nicolini Date: Thu, 26 Sep 2024 12:20:54 -0300 Subject: [PATCH 04/47] Save work in progress - More refactoring message handlers --- batcher/aligned-batcher/src/lib.rs | 11 +++++------ 1 file changed, 5 insertions(+), 6 deletions(-) diff --git a/batcher/aligned-batcher/src/lib.rs b/batcher/aligned-batcher/src/lib.rs index cd60d3a60..28c808360 100644 --- a/batcher/aligned-batcher/src/lib.rs +++ b/batcher/aligned-batcher/src/lib.rs @@ -463,7 +463,7 @@ impl Batcher { // At this point, we will have a user state for sure, since we have inserted it // if not already present - let user_state = self.user_states.get(&addr).unwrap().lock().await; + let mut user_state = self.user_states.get(&addr).unwrap().lock().await; // Perform validations on user state if self.user_balance_is_unlocked(&addr).await { @@ -564,7 +564,6 @@ impl Batcher { // }; let min_fee = user_state.min_fee; - if expected_nonce < msg_nonce { warn!( "Invalid nonce for address {addr}, had nonce {:?} < {:?}", @@ -600,12 +599,12 @@ impl Batcher { // get the entry with the same sender and nonce if !self .handle_replacement_message( - batch_state, + // batch_state, nonced_verification_data, ws_conn_sink.clone(), client_msg.signature, addr, - expected_user_nonce, + expected_nonce, ) .await { @@ -665,7 +664,6 @@ impl Batcher { // } info!("Verification data message handled"); - send_message(ws_conn_sink, ValidityResponseMessage::Valid).await; Ok(()) } @@ -740,7 +738,7 @@ impl Batcher { /// Returns true if the message was replaced in the batch, false otherwise async fn handle_replacement_message( &self, - mut batch_state: tokio::sync::MutexGuard<'_, BatchState>, + // mut batch_state: tokio::sync::MutexGuard<'_, BatchState>, nonced_verification_data: NoncedVerificationData, ws_conn_sink: Arc, Message>>>, signature: Signature, @@ -749,6 +747,7 @@ impl Batcher { ) -> bool { let replacement_max_fee = nonced_verification_data.max_fee; let nonce = nonced_verification_data.nonce; + let mut batch_state = self.batch_state.lock().await; let mut replacement_entry = match batch_state.get_entry(addr, nonce) { Some(entry) => { From 02582e8f68917610cb91d24c20254ebecf09add3 Mon Sep 17 00:00:00 2001 From: Mariano Nicolini Date: Thu, 26 Sep 2024 17:29:48 -0300 Subject: [PATCH 05/47] Save work in progress - More refactoring replacement message function --- batcher/aligned-batcher/src/lib.rs | 480 +++++++----------- .../aligned-batcher/src/types/batch_queue.rs | 1 + batcher/aligned-batcher/src/types/mod.rs | 1 + 3 files changed, 194 insertions(+), 288 deletions(-) diff --git a/batcher/aligned-batcher/src/lib.rs b/batcher/aligned-batcher/src/lib.rs index 28c808360..a43376096 100644 --- a/batcher/aligned-batcher/src/lib.rs +++ b/batcher/aligned-batcher/src/lib.rs @@ -7,9 +7,9 @@ use dotenvy::dotenv; use ethers::contract::ContractError; use ethers::signers::Signer; use serde::Serialize; +use types::batch_state::BatchState; use types::user_state::UserState; -use std::collections::hash_map::Entry; use std::collections::HashMap; use std::env; use std::iter::repeat; @@ -35,7 +35,7 @@ use tokio::net::{TcpListener, TcpStream}; use tokio::sync::{Mutex, RwLock}; use tokio_tungstenite::tungstenite::{Error, Message}; use tokio_tungstenite::WebSocketStream; -use types::batch_queue::{self, BatchQueue, BatchQueueEntry, BatchQueueEntryPriority}; +use types::batch_queue::{self, BatchQueueEntry, BatchQueueEntryPriority}; use types::errors::{BatcherError, BatcherSendError}; use crate::config::{ConfigFromYaml, ContractDeploymentOutput}; @@ -64,127 +64,6 @@ const RESPOND_TO_TASK_FEE_LIMIT_DIVIDER: u128 = 2; const DEFAULT_AGGREGATOR_FEE_MULTIPLIER: u128 = 3; // to set the feeForAggregator variable higher than what was calculated const DEFAULT_AGGREGATOR_FEE_DIVIDER: u128 = 2; -struct BatchState { - batch_queue: BatchQueue, -} - -impl BatchState { - fn new() -> Self { - Self { - batch_queue: BatchQueue::new(), - } - } - - fn get_user_proof_count(&self, addr: &Address) -> u64 { - *self.user_proof_count_in_batch.get(addr).unwrap_or(&0) - } - - /* - Increments the user proof count in the batch, if the user is already in the hashmap. - If the user is not in the hashmap, it adds the user to the hashmap with a count of 1 to represent the first proof. - */ - fn increment_user_proof_count(&mut self, addr: &Address) { - self.user_proof_count_in_batch - .entry(*addr) - .and_modify(|count| *count += 1) - .or_insert(1); - } - - fn get_entry(&self, sender: Address, nonce: U256) -> Option<&BatchQueueEntry> { - self.batch_queue - .iter() - .map(|(entry, _)| entry) - .find(|entry| entry.sender == sender && entry.nonced_verification_data.nonce == nonce) - } - - /// Checks if the entry is valid - /// An entry is valid if there is no entry with the same sender, - /// lower nonce and a lower fee - /// If the entry is valid, it replaces the entry in the queue - /// to increment the max fee, then it updates the user min fee if necessary - /// If the entry is invalid, it returns a validity response message. - /// If the entry is valid, it returns None. - fn validate_and_increment_max_fee( - &mut self, - replacement_entry: BatchQueueEntry, - ) -> Option { - let replacement_max_fee = replacement_entry.nonced_verification_data.max_fee; - let nonce = replacement_entry.nonced_verification_data.nonce; - let sender = replacement_entry.sender; - - debug!( - "Checking validity of entry with sender: {:?}, nonce: {:?}, max_fee: {:?}", - sender, nonce, replacement_max_fee - ); - - // it is a valid entry only if there is no entry with the same sender, lower nonce and a lower fee - let is_valid = !self.batch_queue.iter().any(|(entry, _)| { - entry.sender == sender - && entry.nonced_verification_data.nonce < nonce - && entry.nonced_verification_data.max_fee < replacement_max_fee - }); - - if !is_valid { - return Some(ValidityResponseMessage::InvalidReplacementMessage); - } - - info!( - "Entry is valid, incrementing fee for sender: {:?}, nonce: {:?}, max_fee: {:?}", - sender, nonce, replacement_max_fee - ); - - // remove the old entry and insert the new one - // note that the entries are considered equal for the priority queue - // if they have the same nonce and sender, so we can remove the old entry - // by calling remove with the new entry - self.batch_queue.remove(&replacement_entry); - self.batch_queue.push( - replacement_entry.clone(), - BatchQueueEntryPriority::new(replacement_max_fee, nonce), - ); - - let user_min_fee = self - .batch_queue - .iter() - .filter(|(e, _)| e.sender == sender) - .map(|(e, _)| e.nonced_verification_data.max_fee) - .min() - .unwrap_or(U256::max_value()); - - self.user_min_fee.insert(sender, user_min_fee); - - None - } - - /// Updates: - /// * The user proof count in batch - /// * The user min fee pending in batch (which is the one with the highest nonce) - /// based on whats currenlty in the batch queue. - /// This is necessary because the whole batch may not be included in the finalized batch, - /// This caches are needed to validate user messages. - fn update_user_proofs_in_batch_and_min_fee(&mut self) { - let mut updated_user_min_fee = HashMap::new(); - let mut updated_user_proof_count_in_batch = HashMap::new(); - - for (entry, _) in self.batch_queue.iter() { - *updated_user_proof_count_in_batch - .entry(entry.sender) - .or_insert(0) += 1; - - let min_fee = updated_user_min_fee - .entry(entry.sender) - .or_insert(entry.nonced_verification_data.max_fee); - - if entry.nonced_verification_data.max_fee < *min_fee { - *min_fee = entry.nonced_verification_data.max_fee; - } - } - - self.user_proof_count_in_batch = updated_user_proof_count_in_batch; - self.user_min_fee = updated_user_min_fee; - } -} - pub struct Batcher { s3_client: S3Client, s3_bucket_name: String, @@ -203,7 +82,7 @@ pub struct Batcher { pre_verification_is_enabled: bool, non_paying_config: Option, posting_batch: Mutex, - user_states: HashMap>, + user_states: RwLock>>, } impl Batcher { @@ -288,10 +167,25 @@ impl Batcher { .await .expect("Failed to get fallback Batcher Payment Service contract"); + let user_states = RwLock::new(HashMap::new()); let non_paying_config = if let Some(non_paying_config) = config.batcher.non_paying { warn!("Non-paying address configuration detected. Will replace non-paying address {} with configured address.", non_paying_config.address); - Some(NonPayingConfig::from_yaml_config(non_paying_config).await) + + let non_paying_config = NonPayingConfig::from_yaml_config(non_paying_config).await; + let nonpaying_nonce = payment_service + .user_nonces(non_paying_config.replacement.address()) + .call() + .await + .expect("Could not get non-paying nonce from Ethereum"); + + let non_paying_user_state = Mutex::new(UserState::new_non_paying(nonpaying_nonce)); + user_states.write().await.insert( + non_paying_config.replacement.address(), + non_paying_user_state, + ); + + Some(non_paying_config) } else { None }; @@ -314,7 +208,7 @@ impl Batcher { pre_verification_is_enabled: config.batcher.pre_verification_is_enabled, non_paying_config, posting_batch: Mutex::new(false), - user_states: HashMap::new(), + user_states, } } @@ -455,15 +349,18 @@ impl Batcher { .await } else { // Check that we had a user state entry for this user and insert it if not. - if !self.user_states.contains_key(&addr) { + if !self.user_states.read().await.contains_key(&addr) { let new_user_state = Mutex::new(UserState::new()); - self.user_states.insert(addr.clone(), new_user_state); + self.user_states + .write() + .await + .insert(addr.clone(), new_user_state); } // At this point, we will have a user state for sure, since we have inserted it - // if not already present - - let mut user_state = self.user_states.get(&addr).unwrap().lock().await; + // if not already present. + let user_state_read_lock = self.user_states.read().await; + let mut user_state = user_state_read_lock.get(&addr).unwrap().lock().await; // Perform validations on user state if self.user_balance_is_unlocked(&addr).await { @@ -689,45 +586,45 @@ impl Batcher { true } - /// Handles a message with an expected nonce. - /// If the max_fee is valid, it is added to the batch. - /// If the max_fee is invalid, a message is sent to the client. - /// Returns true if the message was added to the batch, false otherwise. - async fn handle_expected_nonce_message( - &self, - mut batch_state: tokio::sync::MutexGuard<'_, BatchState>, - min_fee: U256, - nonced_verification_data: NoncedVerificationData, - ws_conn_sink: Arc, Message>>>, - signature: Signature, - addr: Address, - ) -> bool { - let max_fee = nonced_verification_data.max_fee; - if max_fee > min_fee { - warn!( - "Invalid max fee for address {addr}, had fee {:?} < {:?}", - min_fee, max_fee - ); - send_message(ws_conn_sink.clone(), ValidityResponseMessage::InvalidMaxFee).await; - return false; - } - - let nonce = nonced_verification_data.nonce; - - batch_state.user_nonces.insert(addr, nonce + U256::one()); - batch_state.user_min_fee.insert(addr, max_fee); - - self.add_to_batch( - batch_state, - nonced_verification_data, - ws_conn_sink.clone(), - signature, - addr, - ) - .await; - - true - } + // /// Handles a message with an expected nonce. + // /// If the max_fee is valid, it is added to the batch. + // /// If the max_fee is invalid, a message is sent to the client. + // /// Returns true if the message was added to the batch, false otherwise. + // async fn handle_expected_nonce_message( + // &self, + // mut batch_state: tokio::sync::MutexGuard<'_, BatchState>, + // min_fee: U256, + // nonced_verification_data: NoncedVerificationData, + // ws_conn_sink: Arc, Message>>>, + // signature: Signature, + // addr: Address, + // ) -> bool { + // let max_fee = nonced_verification_data.max_fee; + // if max_fee > min_fee { + // warn!( + // "Invalid max fee for address {addr}, had fee {:?} < {:?}", + // min_fee, max_fee + // ); + // send_message(ws_conn_sink.clone(), ValidityResponseMessage::InvalidMaxFee).await; + // return false; + // } + + // let nonce = nonced_verification_data.nonce; + + // batch_state.user_nonces.insert(addr, nonce + U256::one()); + // batch_state.user_min_fee.insert(addr, max_fee); + + // self.add_to_batch( + // batch_state, + // nonced_verification_data, + // ws_conn_sink.clone(), + // signature, + // addr, + // ) + // .await; + + // true + // } /// Handles a replacement message /// First checks if the message is already in the batch @@ -738,7 +635,6 @@ impl Batcher { /// Returns true if the message was replaced in the batch, false otherwise async fn handle_replacement_message( &self, - // mut batch_state: tokio::sync::MutexGuard<'_, BatchState>, nonced_verification_data: NoncedVerificationData, ws_conn_sink: Arc, Message>>>, signature: Signature, @@ -746,40 +642,63 @@ impl Batcher { expected_user_nonce: U256, ) -> bool { let replacement_max_fee = nonced_verification_data.max_fee; - let nonce = nonced_verification_data.nonce; + let msg_nonce = nonced_verification_data.nonce; let mut batch_state = self.batch_state.lock().await; - let mut replacement_entry = match batch_state.get_entry(addr, nonce) { - Some(entry) => { - if entry.nonced_verification_data.max_fee < replacement_max_fee { - entry.clone() - } else { - warn!( - "Invalid replacement message for address {addr}, had fee {:?} < {:?}", - entry.nonced_verification_data.max_fee, replacement_max_fee - ); - send_message( - ws_conn_sink.clone(), - ValidityResponseMessage::InvalidReplacementMessage, - ) - .await; - - return false; - } - } - None => { - warn!( - "Invalid nonce for address {addr} Expected: {:?}, got: {:?}", - expected_user_nonce, nonce - ); - send_message(ws_conn_sink.clone(), ValidityResponseMessage::InvalidNonce).await; - return false; - } + let Some(entry) = batch_state.get_entry(addr, msg_nonce) else { + warn!( + "Invalid nonce for address {addr} Expected: {:?}, got: {:?}", + expected_user_nonce, msg_nonce + ); + send_message(ws_conn_sink.clone(), ValidityResponseMessage::InvalidNonce).await; + return false; }; + if entry.nonced_verification_data.max_fee > replacement_max_fee { + warn!( + "Invalid replacement message for address {addr}, had fee {:?} < {:?}", + entry.nonced_verification_data.max_fee, replacement_max_fee + ); + send_message( + ws_conn_sink.clone(), + ValidityResponseMessage::InvalidReplacementMessage, + ) + .await; + + return false; + } + + // let mut replacement_entry = match batch_state.get_entry(addr, msg_nonce) { + // Some(entry) => { + // if entry.nonced_verification_data.max_fee < replacement_max_fee { + // entry.clone() + // } else { + // warn!( + // "Invalid replacement message for address {addr}, had fee {:?} < {:?}", + // entry.nonced_verification_data.max_fee, replacement_max_fee + // ); + // send_message( + // ws_conn_sink.clone(), + // ValidityResponseMessage::InvalidReplacementMessage, + // ) + // .await; + + // return false; + // } + // } + // None => { + // warn!( + // "Invalid nonce for address {addr} Expected: {:?}, got: {:?}", + // expected_user_nonce, msg_nonce + // ); + // send_message(ws_conn_sink.clone(), ValidityResponseMessage::InvalidNonce).await; + // return false; + // } + // }; + info!( "Replacing message for address {} with nonce {} and max fee {}", - addr, nonce, replacement_max_fee + addr, msg_nonce, replacement_max_fee ); replacement_entry.signature = signature; @@ -838,7 +757,7 @@ impl Batcher { let max_fee = verification_data.max_fee; let nonce = verification_data.nonce; - let batch_state = self.batch_state.lock().await; + let mut batch_state = self.batch_state.lock().await; batch_state.batch_queue.push( BatchQueueEntry::new( verification_data, @@ -873,8 +792,9 @@ impl Batcher { gas_price: U256, ) -> Option> { let mut batch_state = self.batch_state.lock().await; - let current_batch_len = batch_state.batch_queue.len(); + let user_states = self.user_states.write().await; + let current_batch_len = batch_state.batch_queue.len(); let last_uploaded_batch_block_lock = self.last_uploaded_batch_block.lock().await; // FIXME(marian): This condition should be changed to current_batch_size == 0 @@ -905,17 +825,23 @@ impl Batcher { // Set the batch posting flag to true *batch_posting = true; - let batch_queue_copy = batch_state.batch_queue.clone(); match batch_queue::try_build_batch(batch_queue_copy, gas_price, self.max_batch_size) { Ok((resulting_batch_queue, finalized_batch)) => { - // Set the batch queue to batch queue copy batch_state.batch_queue = resulting_batch_queue; - batch_state.update_user_proofs_in_batch_and_min_fee(); + let updated_user_proof_count_and_min_fee = + batch_state.get_user_proofs_in_batch_and_min_fee(); + + for (addr, (proof_count, min_fee)) in updated_user_proof_count_and_min_fee.iter() { + let mut user_state = user_states.get(addr).unwrap().lock().await; + user_state.proofs_in_batch = *proof_count; + user_state.min_fee = *min_fee; + } + Some(finalized_batch) } Err(BatcherError::BatchCostTooHigh) => { - // We cant post a batch since users are not willing to pay the needed fee, wait for more proofs + // We can't post a batch since users are not willing to pay the needed fee, wait for more proofs info!("No working batch found. Waiting for more proofs..."); *batch_posting = false; None @@ -1022,9 +948,10 @@ impl Batcher { } batch_state.batch_queue.clear(); - batch_state.user_nonces.clear(); - batch_state.user_proof_count_in_batch.clear(); - batch_state.user_min_fee.clear(); + self.user_states.write().await.clear(); + // batch_state.user_nonces.clear(); + // batch_state.user_proof_count_in_batch.clear(); + // batch_state.user_min_fee.clear(); } /// Receives new block numbers, checks if conditions are met for submission and @@ -1175,8 +1102,7 @@ impl Batcher { { Ok(receipt) => Ok(receipt), Err(BatcherSendError::TransactionReverted(err)) => { - // dont retry with fallback - // just return the error + // Since transaction was reverted, we don't want to retry with fallback. warn!("Transaction reverted {:?}", err); Err(BatcherError::TransactionSendError) @@ -1213,101 +1139,79 @@ impl Batcher { .is_some_and(|non_paying_config| non_paying_config.address == *addr) } - /// Only relevant for testing and for users to easily use Aligned + /// Only relevant for testing and for users to easily use Aligned in testnet. async fn handle_nonpaying_msg( - self: Arc, - ws_conn_sink: Arc, Message>>>, + &self, + ws_sink: Arc, Message>>>, client_msg: ClientMessage, ) -> Result<(), Error> { - let non_paying_config = self.non_paying_config.as_ref().unwrap(); - let addr = non_paying_config.replacement.address(); + let Some(non_paying_config) = self.non_paying_config.as_ref() else { + warn!("There isn't a non-paying configuration loaded. This message will be ignored"); + send_message(ws_sink.clone(), ValidityResponseMessage::InvalidNonce).await; + return Ok(()); + }; - let user_balance = self.get_user_balance(&addr).await; + if client_msg.verification_data.verification_data.proof.len() > self.max_proof_size { + error!("Proof is too large"); + send_message(ws_sink.clone(), ValidityResponseMessage::ProofTooLarge).await; + return Ok(()); // Send error message to the client and return + } - if user_balance == U256::from(0) { - error!("Insufficient funds for address {:?}", addr); + let replacement_addr = non_paying_config.replacement.address(); + let replacement_user_balance = self.get_user_balance(&replacement_addr).await; + if replacement_user_balance == U256::from(0) { + error!("Insufficient funds for address {:?}", replacement_addr); send_message( - ws_conn_sink.clone(), - ValidityResponseMessage::InsufficientBalance(addr), + ws_sink.clone(), + ValidityResponseMessage::InsufficientBalance(replacement_addr), ) .await; return Ok(()); // Send error message to the client and return } - if client_msg.verification_data.verification_data.proof.len() <= self.max_proof_size { - // When pre-verification is enabled, batcher will verify proofs for faster feedback with clients - if self.pre_verification_is_enabled - && !zk_utils::verify(&client_msg.verification_data.verification_data).await - { - error!("Invalid proof detected. Verification failed."); - send_message(ws_conn_sink.clone(), ValidityResponseMessage::InvalidProof).await; - return Ok(()); // Send error message to the client and return - } - - let nonced_verification_data = { - let mut batch_state = self.batch_state.lock().await; - - let nonpaying_nonce = match batch_state.user_nonces.entry(addr) { - Entry::Occupied(o) => o.into_mut(), - Entry::Vacant(vacant) => { - let nonce = match self.payment_service.user_nonces(addr).call().await { - Ok(nonce) => nonce, - Err(e) => { - error!("Failed to get nonce for address {:?}: {:?}", addr, e); - send_message( - ws_conn_sink.clone(), - ValidityResponseMessage::InvalidNonce, - ) - .await; - - return Ok(()); - } - }; - - vacant.insert(nonce) - } - }; - - info!("non paying nonce: {:?}", nonpaying_nonce); - - let nonce_value = *nonpaying_nonce; - - *nonpaying_nonce += U256::one(); + // When pre-verification is enabled, batcher will verify proofs for faster feedback with clients + if self.pre_verification_is_enabled + && !zk_utils::verify(&client_msg.verification_data.verification_data).await + { + error!("Invalid proof detected. Verification failed."); + send_message(ws_sink.clone(), ValidityResponseMessage::InvalidProof).await; + return Ok(()); // Send error message to the client and return + } - NoncedVerificationData::new( - client_msg.verification_data.verification_data.clone(), - nonce_value, - DEFAULT_MAX_FEE_PER_PROOF.into(), // 13_000 gas per proof * 100 gwei gas price (upper bound) - self.chain_id, - self.payment_service.address(), - ) - }; + let user_states = self.user_states.read().await; + + // Safe to unwrap here since at this point we have a non-paying configuration loaded for sure + // and the non-paying address nonce cached in the user states. + let non_paying_user_state = user_states.get(&replacement_addr).unwrap(); + let non_paying_nonce = non_paying_user_state.lock().await.nonce.unwrap(); + debug!("Non-paying nonce: {:?}", non_paying_nonce); + let updated_non_paying_nonce = non_paying_nonce + U256::one(); + + let nonced_verification_data = NoncedVerificationData::new( + client_msg.verification_data.verification_data.clone(), + updated_non_paying_nonce, + DEFAULT_MAX_FEE_PER_PROOF.into(), // 13_000 gas per proof * 100 gwei gas price (upper bound) + self.chain_id, + self.payment_service.address(), + ); - let client_msg = ClientMessage::new( - nonced_verification_data.clone(), - non_paying_config.replacement.clone(), - ) - .await; + let client_msg = ClientMessage::new( + nonced_verification_data.clone(), + non_paying_config.replacement.clone(), + ) + .await; - let batch_state = self.batch_state.lock().await; - self.clone() - .add_to_batch( - batch_state, - nonced_verification_data, - ws_conn_sink.clone(), - client_msg.signature, - non_paying_config.address, - ) - .await; - } else { - error!("Proof is too large"); - send_message(ws_conn_sink.clone(), ValidityResponseMessage::ProofTooLarge).await; - return Ok(()); // Send error message to the client and return - }; + self.add_to_batch( + nonced_verification_data, + ws_sink.clone(), + client_msg.signature, + non_paying_config.address, + ) + .await; info!("Verification data message handled"); + send_message(ws_sink, ValidityResponseMessage::Valid).await; - send_message(ws_conn_sink, ValidityResponseMessage::Valid).await; Ok(()) } diff --git a/batcher/aligned-batcher/src/types/batch_queue.rs b/batcher/aligned-batcher/src/types/batch_queue.rs index 9427cc9db..59eede378 100644 --- a/batcher/aligned-batcher/src/types/batch_queue.rs +++ b/batcher/aligned-batcher/src/types/batch_queue.rs @@ -2,6 +2,7 @@ use ethers::types::{Address, Signature, U256}; use futures_util::stream::SplitSink; use priority_queue::PriorityQueue; use std::{ + collections::HashMap, hash::{Hash, Hasher}, ops::ControlFlow, sync::Arc, diff --git a/batcher/aligned-batcher/src/types/mod.rs b/batcher/aligned-batcher/src/types/mod.rs index cb2327411..893882786 100644 --- a/batcher/aligned-batcher/src/types/mod.rs +++ b/batcher/aligned-batcher/src/types/mod.rs @@ -1,3 +1,4 @@ pub(crate) mod batch_queue; +pub(crate) mod batch_state; pub mod errors; pub(crate) mod user_state; From 02316a44671a91a238dbf6eb35b0bb9b5893686b Mon Sep 17 00:00:00 2001 From: Mariano Nicolini Date: Thu, 26 Sep 2024 17:52:36 -0300 Subject: [PATCH 06/47] Save work in progress - Batcher compiling --- batcher/aligned-batcher/src/lib.rs | 57 ++++++++----------- .../aligned-batcher/src/types/batch_queue.rs | 1 - 2 files changed, 25 insertions(+), 33 deletions(-) diff --git a/batcher/aligned-batcher/src/lib.rs b/batcher/aligned-batcher/src/lib.rs index a43376096..752f3ff5a 100644 --- a/batcher/aligned-batcher/src/lib.rs +++ b/batcher/aligned-batcher/src/lib.rs @@ -577,7 +577,7 @@ impl Batcher { // Checks user has sufficient balance // If user has sufficient balance, increments the user's proof count in the batch - async fn check_user_balance(&self, addr: &Address, user_proofs_in_batch: u64) -> bool { + async fn check_user_balance(&self, addr: &Address, user_proofs_in_batch: usize) -> bool { let user_balance = self.get_user_balance(addr).await; let min_balance = U256::from(user_proofs_in_batch) * U256::from(MIN_FEE_PER_PROOF); if user_balance < min_balance { @@ -668,45 +668,19 @@ impl Batcher { return false; } - // let mut replacement_entry = match batch_state.get_entry(addr, msg_nonce) { - // Some(entry) => { - // if entry.nonced_verification_data.max_fee < replacement_max_fee { - // entry.clone() - // } else { - // warn!( - // "Invalid replacement message for address {addr}, had fee {:?} < {:?}", - // entry.nonced_verification_data.max_fee, replacement_max_fee - // ); - // send_message( - // ws_conn_sink.clone(), - // ValidityResponseMessage::InvalidReplacementMessage, - // ) - // .await; - - // return false; - // } - // } - // None => { - // warn!( - // "Invalid nonce for address {addr} Expected: {:?}, got: {:?}", - // expected_user_nonce, msg_nonce - // ); - // send_message(ws_conn_sink.clone(), ValidityResponseMessage::InvalidNonce).await; - // return false; - // } - // }; - info!( "Replacing message for address {} with nonce {} and max fee {}", addr, msg_nonce, replacement_max_fee ); + // The replacement entry is built from the old entry and validated for then to be replaced + let mut replacement_entry = entry.clone(); replacement_entry.signature = signature; replacement_entry.verification_data_commitment = nonced_verification_data.verification_data.clone().into(); replacement_entry.nonced_verification_data = nonced_verification_data; - // close old sink and replace with new one + // Close old sink in old entry replace it with new one { if let Some(messaging_sink) = replacement_entry.messaging_sink { let mut old_sink = messaging_sink.write().await; @@ -722,12 +696,31 @@ impl Batcher { } replacement_entry.messaging_sink = Some(ws_conn_sink.clone()); - if let Some(msg) = batch_state.validate_and_increment_max_fee(replacement_entry) { + if !batch_state.replacement_entry_is_valid(&replacement_entry) { warn!("Invalid max fee"); - send_message(ws_conn_sink.clone(), msg).await; + send_message( + ws_conn_sink.clone(), + ValidityResponseMessage::InvalidReplacementMessage, + ) + .await; return false; } + info!( + "Replacement entry is valid, incrementing fee for sender: {:?}, nonce: {:?}, max_fee: {:?}", + replacement_entry.sender, replacement_entry.nonced_verification_data.nonce, replacement_max_fee + ); + + // remove the old entry and insert the new one + // note that the entries are considered equal for the priority queue + // if they have the same nonce and sender, so we can remove the old entry + // by calling remove with the new entry + batch_state.batch_queue.remove(&replacement_entry); + batch_state.batch_queue.push( + replacement_entry.clone(), + BatchQueueEntryPriority::new(replacement_max_fee, msg_nonce), + ); + true } diff --git a/batcher/aligned-batcher/src/types/batch_queue.rs b/batcher/aligned-batcher/src/types/batch_queue.rs index 59eede378..9427cc9db 100644 --- a/batcher/aligned-batcher/src/types/batch_queue.rs +++ b/batcher/aligned-batcher/src/types/batch_queue.rs @@ -2,7 +2,6 @@ use ethers::types::{Address, Signature, U256}; use futures_util::stream::SplitSink; use priority_queue::PriorityQueue; use std::{ - collections::HashMap, hash::{Hash, Hasher}, ops::ControlFlow, sync::Arc, From da389d11b1fd27e6a04ab4c0023e126f75ce1a35 Mon Sep 17 00:00:00 2001 From: Mariano Nicolini Date: Thu, 26 Sep 2024 18:49:41 -0300 Subject: [PATCH 07/47] Full flow working, fix non-paying message handling bug --- batcher/aligned-batcher/src/lib.rs | 24 ++++++++++++++---------- operator/merkle_tree/lib/Cargo.lock | 8 ++++---- 2 files changed, 18 insertions(+), 14 deletions(-) diff --git a/batcher/aligned-batcher/src/lib.rs b/batcher/aligned-batcher/src/lib.rs index 752f3ff5a..fa2b9914f 100644 --- a/batcher/aligned-batcher/src/lib.rs +++ b/batcher/aligned-batcher/src/lib.rs @@ -761,6 +761,7 @@ impl Batcher { ), BatchQueueEntryPriority::new(max_fee, nonce), ); + info!( "Current batch queue length: {}", batch_state.batch_queue.len() @@ -1147,7 +1148,7 @@ impl Batcher { if client_msg.verification_data.verification_data.proof.len() > self.max_proof_size { error!("Proof is too large"); send_message(ws_sink.clone(), ValidityResponseMessage::ProofTooLarge).await; - return Ok(()); // Send error message to the client and return + return Ok(()); } let replacement_addr = non_paying_config.replacement.address(); @@ -1159,7 +1160,7 @@ impl Batcher { ValidityResponseMessage::InsufficientBalance(replacement_addr), ) .await; - return Ok(()); // Send error message to the client and return + return Ok(()); } // When pre-verification is enabled, batcher will verify proofs for faster feedback with clients @@ -1173,16 +1174,17 @@ impl Batcher { let user_states = self.user_states.read().await; - // Safe to unwrap here since at this point we have a non-paying configuration loaded for sure - // and the non-paying address nonce cached in the user states. - let non_paying_user_state = user_states.get(&replacement_addr).unwrap(); - let non_paying_nonce = non_paying_user_state.lock().await.nonce.unwrap(); - debug!("Non-paying nonce: {:?}", non_paying_nonce); - let updated_non_paying_nonce = non_paying_nonce + U256::one(); + // Safe to call `unwrap()` here since at this point we have a non-paying configuration loaded + // for sure and the non-paying address nonce cached in the user states. + let non_paying_user_state_lock = user_states.get(&replacement_addr).unwrap(); + let mut non_paying_user_state = non_paying_user_state_lock.lock().await; + let non_paying_nonce = non_paying_user_state.nonce.unwrap(); + + info!("Non-paying nonce: {:?}", non_paying_nonce); let nonced_verification_data = NoncedVerificationData::new( client_msg.verification_data.verification_data.clone(), - updated_non_paying_nonce, + non_paying_nonce, DEFAULT_MAX_FEE_PER_PROOF.into(), // 13_000 gas per proof * 100 gwei gas price (upper bound) self.chain_id, self.payment_service.address(), @@ -1202,7 +1204,9 @@ impl Batcher { ) .await; - info!("Verification data message handled"); + // Non-paying user nonce is updated + (*non_paying_user_state).nonce = Some(non_paying_nonce + U256::one()); + info!("Non-paying verification data message handled"); send_message(ws_sink, ValidityResponseMessage::Valid).await; Ok(()) diff --git a/operator/merkle_tree/lib/Cargo.lock b/operator/merkle_tree/lib/Cargo.lock index 47549f3ba..079878cdb 100644 --- a/operator/merkle_tree/lib/Cargo.lock +++ b/operator/merkle_tree/lib/Cargo.lock @@ -1677,9 +1677,9 @@ dependencies = [ [[package]] name = "lambdaworks-crypto" -version = "0.10.0" +version = "0.7.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bbc2a4da0d9e52ccfe6306801a112e81a8fc0c76aa3e4449fefeda7fef72bb34" +checksum = "7fb5d4f22241504f7c7b8d2c3a7d7835d7c07117f10bff2a7d96a9ef6ef217c3" dependencies = [ "lambdaworks-math", "serde", @@ -1689,9 +1689,9 @@ dependencies = [ [[package]] name = "lambdaworks-math" -version = "0.10.0" +version = "0.7.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d1bd2632acbd9957afc5aeec07ad39f078ae38656654043bf16e046fa2730e23" +checksum = "358e172628e713b80a530a59654154bfc45783a6ed70ea284839800cebdf8f97" dependencies = [ "serde", "serde_json", From 2fac2b0c8996b7efda2e54511a4a351f57f663d0 Mon Sep 17 00:00:00 2001 From: Mariano Nicolini Date: Thu, 26 Sep 2024 21:07:57 -0300 Subject: [PATCH 08/47] Polish code, remove commented legacy code --- batcher/aligned-batcher/src/lib.rs | 426 +++++++----------- .../aligned-batcher/src/types/batch_queue.rs | 2 +- 2 files changed, 152 insertions(+), 276 deletions(-) diff --git a/batcher/aligned-batcher/src/lib.rs b/batcher/aligned-batcher/src/lib.rs index fa2b9914f..c1da10eb1 100644 --- a/batcher/aligned-batcher/src/lib.rs +++ b/batcher/aligned-batcher/src/lib.rs @@ -26,7 +26,6 @@ use eth::{try_create_new_task, BatcherPaymentService, CreateNewTaskFeeParams, Si use ethers::prelude::{Middleware, Provider}; use ethers::providers::Ws; use ethers::types::{Address, Signature, TransactionReceipt, U256}; -use futures_util::stream::SplitSink; use futures_util::{future, SinkExt, StreamExt, TryStreamExt}; use lambdaworks_crypto::merkle_tree::merkle::MerkleTree; use lambdaworks_crypto::merkle_tree::traits::IsMerkleTreeBackend; @@ -34,8 +33,7 @@ use log::{debug, error, info, warn}; use tokio::net::{TcpListener, TcpStream}; use tokio::sync::{Mutex, RwLock}; use tokio_tungstenite::tungstenite::{Error, Message}; -use tokio_tungstenite::WebSocketStream; -use types::batch_queue::{self, BatchQueueEntry, BatchQueueEntryPriority}; +use types::batch_queue::{self, BatchQueueEntry, BatchQueueEntryPriority, WsMessageSink}; use types::errors::{BatcherError, BatcherSendError}; use crate::config::{ConfigFromYaml, ContractDeploymentOutput}; @@ -310,7 +308,7 @@ impl Batcher { async fn handle_message( self: Arc, message: Message, - ws_conn_sink: Arc, Message>>>, + ws_conn_sink: WsMessageSink, ) -> Result<(), Error> { // Deserialize verification data from message let client_msg: ClientMessage = match cbor_deserialize(message.into_data().as_slice()) { @@ -342,237 +340,157 @@ impl Batcher { } info!("Verifying message signature..."); - if let Ok(addr) = client_msg.verify_signature() { - info!("Message signature verified"); - if self.is_nonpaying(&addr) { - self.handle_nonpaying_msg(ws_conn_sink.clone(), client_msg) - .await - } else { - // Check that we had a user state entry for this user and insert it if not. - if !self.user_states.read().await.contains_key(&addr) { - let new_user_state = Mutex::new(UserState::new()); - self.user_states - .write() - .await - .insert(addr.clone(), new_user_state); - } + let Ok(addr) = client_msg.verify_signature() else { + error!("Signature verification error"); + send_message( + ws_conn_sink.clone(), + ValidityResponseMessage::InvalidSignature, + ) + .await; + return Ok(()); + }; + info!("Message signature verified"); - // At this point, we will have a user state for sure, since we have inserted it - // if not already present. - let user_state_read_lock = self.user_states.read().await; - let mut user_state = user_state_read_lock.get(&addr).unwrap().lock().await; + if self.is_nonpaying(&addr) { + return self + .handle_nonpaying_msg(ws_conn_sink.clone(), client_msg) + .await; + } + // Check that we had a user state entry for this user and insert it if not. + if !self.user_states.read().await.contains_key(&addr) { + let new_user_state = Mutex::new(UserState::new()); + self.user_states + .write() + .await + .insert(addr.clone(), new_user_state); + } - // Perform validations on user state - if self.user_balance_is_unlocked(&addr).await { - send_message( - ws_conn_sink.clone(), - ValidityResponseMessage::InsufficientBalance(addr), - ) - .await; - return Ok(()); - } + // At this point, we will have a user state for sure, since we have inserted it + // if not already present. + let user_state_read_lock = self.user_states.read().await; + let mut user_state = user_state_read_lock.get(&addr).unwrap().lock().await; - let updated_user_proofs = user_state.proofs_in_batch + 1; - if !self.check_user_balance(&addr, updated_user_proofs).await { - send_message( - ws_conn_sink.clone(), - ValidityResponseMessage::InsufficientBalance(addr), - ) - .await; - return Ok(()); - } + // Perform validations on user state + if self.user_balance_is_unlocked(&addr).await { + send_message( + ws_conn_sink.clone(), + ValidityResponseMessage::InsufficientBalance(addr), + ) + .await; + return Ok(()); + } - let nonced_verification_data = client_msg.verification_data; - if nonced_verification_data.verification_data.proof.len() > self.max_proof_size { - error!("Proof size exceeds the maximum allowed size."); - send_message(ws_conn_sink.clone(), ValidityResponseMessage::ProofTooLarge) - .await; - return Ok(()); - } + let updated_user_proofs = user_state.proofs_in_batch + 1; + if !self.check_user_balance(&addr, updated_user_proofs).await { + send_message( + ws_conn_sink.clone(), + ValidityResponseMessage::InsufficientBalance(addr), + ) + .await; + return Ok(()); + } - // When pre-verification is enabled, batcher will verify proofs for faster feedback with clients - if self.pre_verification_is_enabled - && !zk_utils::verify(&nonced_verification_data.verification_data).await - { - error!("Invalid proof detected. Verification failed."); - send_message(ws_conn_sink.clone(), ValidityResponseMessage::InvalidProof).await; - return Ok(()); // Send error message to the client and return - } + let nonced_verification_data = client_msg.verification_data; + if nonced_verification_data.verification_data.proof.len() > self.max_proof_size { + error!("Proof size exceeds the maximum allowed size."); + send_message(ws_conn_sink.clone(), ValidityResponseMessage::ProofTooLarge).await; + return Ok(()); + } - // Nonce and max fee verification - let msg_nonce = nonced_verification_data.nonce; - let max_fee = nonced_verification_data.max_fee; - if max_fee < U256::from(MIN_FEE_PER_PROOF) { - error!("The max fee signed in the message is less than the accepted minimum fee to be included in the batch."); - send_message(ws_conn_sink.clone(), ValidityResponseMessage::InvalidMaxFee) - .await; - return Ok(()); - } + // When pre-verification is enabled, batcher will verify proofs for faster feedback with clients + if self.pre_verification_is_enabled + && !zk_utils::verify(&nonced_verification_data.verification_data).await + { + error!("Invalid proof detected. Verification failed."); + send_message(ws_conn_sink.clone(), ValidityResponseMessage::InvalidProof).await; + return Ok(()); // Send error message to the client and return + } - let expected_nonce = if let Some(cached_nonce) = user_state.nonce { - cached_nonce - } else { - let ethereum_user_nonce = match self.get_user_nonce_from_ethereum(addr).await { - Ok(ethereum_user_nonce) => ethereum_user_nonce, - Err(e) => { - error!( - "Failed to get user nonce from Ethereum for address {:?}. Error: {:?}", - addr, e - ); - send_message( - ws_conn_sink.clone(), - ValidityResponseMessage::InvalidNonce, - ) - .await; - - return Ok(()); - } - }; - - user_state.nonce = Some(ethereum_user_nonce); - ethereum_user_nonce - }; + // Nonce and max fee verification + let msg_nonce = nonced_verification_data.nonce; + let max_fee = nonced_verification_data.max_fee; + if max_fee < U256::from(MIN_FEE_PER_PROOF) { + error!("The max fee signed in the message is less than the accepted minimum fee to be included in the batch."); + send_message(ws_conn_sink.clone(), ValidityResponseMessage::InvalidMaxFee).await; + return Ok(()); + } - // let expected_user_nonce = match user_state.nonce { - // Some(nonce) => nonce, - // None => { - // let user_nonce = match self.get_user_nonce_from_ethereum(addr).await { - // Ok(nonce) => nonce, - // Err(e) => { - // error!("Failed to get user nonce for address {:?}: {:?}", addr, e); - // send_message( - // ws_conn_sink.clone(), - // ValidityResponseMessage::InvalidNonce, - // ) - // .await; - - // return Ok(()); - // } - // }; - - // batch_state.user_nonces.insert(addr, user_nonce); - // user_nonce - // } - // }; - - // let min_fee = match batch_state.user_min_fee.get(&addr) { - // Some(fee) => *fee, - // None => U256::max_value(), - // }; - - let min_fee = user_state.min_fee; - if expected_nonce < msg_nonce { - warn!( - "Invalid nonce for address {addr}, had nonce {:?} < {:?}", - expected_nonce, msg_nonce + let expected_nonce = if let Some(cached_nonce) = user_state.nonce { + cached_nonce + } else { + let ethereum_user_nonce = match self.get_user_nonce_from_ethereum(addr).await { + Ok(ethereum_user_nonce) => ethereum_user_nonce, + Err(e) => { + error!( + "Failed to get user nonce from Ethereum for address {:?}. Error: {:?}", + addr, e ); send_message(ws_conn_sink.clone(), ValidityResponseMessage::InvalidNonce).await; + return Ok(()); - } else if expected_nonce == msg_nonce { - let max_fee = nonced_verification_data.max_fee; - if max_fee > min_fee { - warn!( - "Invalid max fee for address {addr}, had fee {:?} < {:?}", - min_fee, max_fee - ); - send_message(ws_conn_sink.clone(), ValidityResponseMessage::InvalidMaxFee) - .await; - return Ok(()); - } - - user_state.nonce = Some(msg_nonce + U256::one()); - user_state.min_fee = max_fee; - self.add_to_batch( - nonced_verification_data, - ws_conn_sink.clone(), - client_msg.signature, - addr, - ) - .await; - } else { - // might be replacement message - // if the message is already in the batch - // we can check if we need to increment the fee - // get the entry with the same sender and nonce - if !self - .handle_replacement_message( - // batch_state, - nonced_verification_data, - ws_conn_sink.clone(), - client_msg.signature, - addr, - expected_nonce, - ) - .await - { - // message should not be added to batch - return Ok(()); - } } + }; - // match expected_user_nonce.cmp(&nonce) { - // std::cmp::Ordering::Less => { - // // invalid, expected user nonce < nonce - // warn!( - // "Invalid nonce for address {addr}, had nonce {:?} < {:?}", - // expected_user_nonce, nonce - // ); - // send_message(ws_conn_sink.clone(), ValidityResponseMessage::InvalidNonce) - // .await; - // return Ok(()); - // } - // std::cmp::Ordering::Equal => { - // // if we are here nonce == expected_user_nonce - // if !self - // .handle_expected_nonce_message( - // batch_state, - // min_fee, - // nonced_verification_data, - // ws_conn_sink.clone(), - // client_msg.signature, - // addr, - // ) - // .await - // { - // // message should not be added to batch - // return Ok(()); - // }; - // } - // std::cmp::Ordering::Greater => { - // // might be replacement message - // // if the message is already in the batch - // // we can check if we need to increment the fee - // // get the entry with the same sender and nonce - // if !self - // .handle_replacement_message( - // batch_state, - // nonced_verification_data, - // ws_conn_sink.clone(), - // client_msg.signature, - // addr, - // expected_user_nonce, - // ) - // .await - // { - // // message should not be added to batch - // return Ok(()); - // } - // } - // } - - info!("Verification data message handled"); - send_message(ws_conn_sink, ValidityResponseMessage::Valid).await; - Ok(()) + user_state.nonce = Some(ethereum_user_nonce); + ethereum_user_nonce + }; + + let min_fee = user_state.min_fee; + if expected_nonce < msg_nonce { + warn!( + "Invalid nonce for address {addr}, had nonce {:?} < {:?}", + expected_nonce, msg_nonce + ); + send_message(ws_conn_sink.clone(), ValidityResponseMessage::InvalidNonce).await; + return Ok(()); + } else if expected_nonce == msg_nonce { + let max_fee = nonced_verification_data.max_fee; + if max_fee > min_fee { + warn!( + "Invalid max fee for address {addr}, had fee {:?} < {:?}", + min_fee, max_fee + ); + send_message(ws_conn_sink.clone(), ValidityResponseMessage::InvalidMaxFee).await; + return Ok(()); } - } else { - error!("Signature verification error"); - send_message( + + user_state.nonce = Some(msg_nonce + U256::one()); + user_state.min_fee = max_fee; + self.add_to_batch( + nonced_verification_data, ws_conn_sink.clone(), - ValidityResponseMessage::InvalidSignature, + client_msg.signature, + addr, ) .await; - Ok(()) // Send error message to the client and return + } else { + // In this case, the message might be a replacement one. If it is valid, + // we replace the old entry with the new from the replacement message. + if !self + .handle_replacement_message( + nonced_verification_data, + ws_conn_sink.clone(), + client_msg.signature, + addr, + expected_nonce, + ) + .await + { + // message should not be added to batch + return Ok(()); + } + let user_min_fee = self + .batch_state + .lock() + .await + .get_user_min_fee_in_batch(&addr); + + user_state.min_fee = user_min_fee; } + + info!("Verification data message handled"); + send_message(ws_conn_sink, ValidityResponseMessage::Valid).await; + Ok(()) } // Checks user has sufficient balance @@ -586,46 +504,6 @@ impl Batcher { true } - // /// Handles a message with an expected nonce. - // /// If the max_fee is valid, it is added to the batch. - // /// If the max_fee is invalid, a message is sent to the client. - // /// Returns true if the message was added to the batch, false otherwise. - // async fn handle_expected_nonce_message( - // &self, - // mut batch_state: tokio::sync::MutexGuard<'_, BatchState>, - // min_fee: U256, - // nonced_verification_data: NoncedVerificationData, - // ws_conn_sink: Arc, Message>>>, - // signature: Signature, - // addr: Address, - // ) -> bool { - // let max_fee = nonced_verification_data.max_fee; - // if max_fee > min_fee { - // warn!( - // "Invalid max fee for address {addr}, had fee {:?} < {:?}", - // min_fee, max_fee - // ); - // send_message(ws_conn_sink.clone(), ValidityResponseMessage::InvalidMaxFee).await; - // return false; - // } - - // let nonce = nonced_verification_data.nonce; - - // batch_state.user_nonces.insert(addr, nonce + U256::one()); - // batch_state.user_min_fee.insert(addr, max_fee); - - // self.add_to_batch( - // batch_state, - // nonced_verification_data, - // ws_conn_sink.clone(), - // signature, - // addr, - // ) - // .await; - - // true - // } - /// Handles a replacement message /// First checks if the message is already in the batch /// If the message is in the batch, checks if the max fee is higher @@ -636,19 +514,20 @@ impl Batcher { async fn handle_replacement_message( &self, nonced_verification_data: NoncedVerificationData, - ws_conn_sink: Arc, Message>>>, + ws_conn_sink: WsMessageSink, signature: Signature, addr: Address, expected_user_nonce: U256, ) -> bool { - let replacement_max_fee = nonced_verification_data.max_fee; - let msg_nonce = nonced_verification_data.nonce; let mut batch_state = self.batch_state.lock().await; - let Some(entry) = batch_state.get_entry(addr, msg_nonce) else { + let replacement_max_fee = nonced_verification_data.max_fee; + let nonce = nonced_verification_data.nonce; + + let Some(entry) = batch_state.get_entry(addr, nonce) else { warn!( "Invalid nonce for address {addr} Expected: {:?}, got: {:?}", - expected_user_nonce, msg_nonce + expected_user_nonce, nonce ); send_message(ws_conn_sink.clone(), ValidityResponseMessage::InvalidNonce).await; return false; @@ -670,7 +549,7 @@ impl Batcher { info!( "Replacing message for address {} with nonce {} and max fee {}", - addr, msg_nonce, replacement_max_fee + addr, nonce, replacement_max_fee ); // The replacement entry is built from the old entry and validated for then to be replaced @@ -716,11 +595,13 @@ impl Batcher { // if they have the same nonce and sender, so we can remove the old entry // by calling remove with the new entry batch_state.batch_queue.remove(&replacement_entry); - batch_state.batch_queue.push( + let result = batch_state.batch_queue.push( replacement_entry.clone(), - BatchQueueEntryPriority::new(replacement_max_fee, msg_nonce), + BatchQueueEntryPriority::new(replacement_max_fee, nonce), ); + assert!(result.is_none()); + true } @@ -737,9 +618,8 @@ impl Batcher { /// Adds verification data to the current batch queue. async fn add_to_batch( &self, - // mut batch_state: tokio::sync::MutexGuard<'_, BatchState>, verification_data: NoncedVerificationData, - ws_conn_sink: Arc, Message>>>, + ws_conn_sink: WsMessageSink, proof_submitter_sig: Signature, proof_submiter_addr: Address, ) { @@ -898,6 +778,8 @@ impl Batcher { .map(VerificationCommitmentBatch::hash_data) .collect(); + println!("ENTERING SUBMIT BATCH..."); + if let Err(e) = self .submit_batch( &batch_bytes, @@ -943,9 +825,6 @@ impl Batcher { batch_state.batch_queue.clear(); self.user_states.write().await.clear(); - // batch_state.user_nonces.clear(); - // batch_state.user_proof_count_in_batch.clear(); - // batch_state.user_min_fee.clear(); } /// Receives new block numbers, checks if conditions are met for submission and @@ -1136,7 +1015,7 @@ impl Batcher { /// Only relevant for testing and for users to easily use Aligned in testnet. async fn handle_nonpaying_msg( &self, - ws_sink: Arc, Message>>>, + ws_sink: WsMessageSink, client_msg: ClientMessage, ) -> Result<(), Error> { let Some(non_paying_config) = self.non_paying_config.as_ref() else { @@ -1297,10 +1176,7 @@ async fn send_batch_inclusion_data_responses( Ok(()) } -async fn send_message( - ws_conn_sink: Arc, Message>>>, - message: T, -) { +async fn send_message(ws_conn_sink: WsMessageSink, message: T) { match cbor_serialize(&message) { Ok(serialized_response) => { if let Err(err) = ws_conn_sink diff --git a/batcher/aligned-batcher/src/types/batch_queue.rs b/batcher/aligned-batcher/src/types/batch_queue.rs index 9427cc9db..2e78dff78 100644 --- a/batcher/aligned-batcher/src/types/batch_queue.rs +++ b/batcher/aligned-batcher/src/types/batch_queue.rs @@ -16,7 +16,7 @@ use aligned_sdk::{ use super::errors::BatcherError; -type WsMessageSink = Arc, Message>>>; +pub(crate) type WsMessageSink = Arc, Message>>>; #[derive(Clone)] pub(crate) struct BatchQueueEntry { From e82fc66ee58a811b7ca6e11a41e87ba92817cfec Mon Sep 17 00:00:00 2001 From: Mariano Nicolini Date: Thu, 26 Sep 2024 21:32:16 -0300 Subject: [PATCH 09/47] Some more code polishing --- batcher/aligned-batcher/src/lib.rs | 77 +++---------------- .../aligned-batcher/src/types/batch_queue.rs | 16 ++-- 2 files changed, 16 insertions(+), 77 deletions(-) diff --git a/batcher/aligned-batcher/src/lib.rs b/batcher/aligned-batcher/src/lib.rs index c1da10eb1..92eb3b170 100644 --- a/batcher/aligned-batcher/src/lib.rs +++ b/batcher/aligned-batcher/src/lib.rs @@ -3,10 +3,10 @@ extern crate core; use aligned_sdk::communication::serialization::{cbor_deserialize, cbor_serialize}; use aligned_sdk::eth::batcher_payment_service::SignatureData; use config::NonPayingConfig; +use connection::{send_message, WsMessageSink}; use dotenvy::dotenv; use ethers::contract::ContractError; use ethers::signers::Signer; -use serde::Serialize; use types::batch_state::BatchState; use types::user_state::UserState; @@ -17,9 +17,8 @@ use std::net::SocketAddr; use std::sync::Arc; use aligned_sdk::core::types::{ - BatchInclusionData, ClientMessage, NoncedVerificationData, ResponseMessage, - ValidityResponseMessage, VerificationCommitmentBatch, VerificationData, - VerificationDataCommitment, + ClientMessage, NoncedVerificationData, ResponseMessage, ValidityResponseMessage, + VerificationCommitmentBatch, VerificationData, VerificationDataCommitment, }; use aws_sdk_s3::client::Client as S3Client; use eth::{try_create_new_task, BatcherPaymentService, CreateNewTaskFeeParams, SignerMiddlewareT}; @@ -33,12 +32,13 @@ use log::{debug, error, info, warn}; use tokio::net::{TcpListener, TcpStream}; use tokio::sync::{Mutex, RwLock}; use tokio_tungstenite::tungstenite::{Error, Message}; -use types::batch_queue::{self, BatchQueueEntry, BatchQueueEntryPriority, WsMessageSink}; +use types::batch_queue::{self, BatchQueueEntry, BatchQueueEntryPriority}; use types::errors::{BatcherError, BatcherSendError}; use crate::config::{ConfigFromYaml, ContractDeploymentOutput}; mod config; +mod connection; mod eth; pub mod gnark; pub mod halo2; @@ -435,7 +435,6 @@ impl Batcher { ethereum_user_nonce }; - let min_fee = user_state.min_fee; if expected_nonce < msg_nonce { warn!( "Invalid nonce for address {addr}, had nonce {:?} < {:?}", @@ -445,10 +444,10 @@ impl Batcher { return Ok(()); } else if expected_nonce == msg_nonce { let max_fee = nonced_verification_data.max_fee; - if max_fee > min_fee { + if max_fee > user_state.min_fee { warn!( "Invalid max fee for address {addr}, had fee {:?} < {:?}", - min_fee, max_fee + user_state.min_fee, max_fee ); send_message(ws_conn_sink.clone(), ValidityResponseMessage::InvalidMaxFee).await; return Ok(()); @@ -479,13 +478,10 @@ impl Batcher { // message should not be added to batch return Ok(()); } - let user_min_fee = self - .batch_state - .lock() - .await - .get_user_min_fee_in_batch(&addr); - user_state.min_fee = user_min_fee; + let batch_state = self.batch_state.lock().await; + let updated_min_fee_in_batch = batch_state.get_user_min_fee_in_batch(&addr); + user_state.min_fee = updated_min_fee_in_batch; } info!("Verification data message handled"); @@ -778,8 +774,6 @@ impl Batcher { .map(VerificationCommitmentBatch::hash_data) .collect(); - println!("ENTERING SUBMIT BATCH..."); - if let Err(e) = self .submit_batch( &batch_bytes, @@ -808,7 +802,7 @@ impl Batcher { return Err(e); }; - send_batch_inclusion_data_responses(finalized_batch, &batch_merkle_tree).await + connection::send_batch_inclusion_data_responses(finalized_batch, &batch_merkle_tree).await } async fn flush_queue_and_clear_nonce_cache(&self) { @@ -1142,52 +1136,3 @@ impl Batcher { } } } - -async fn send_batch_inclusion_data_responses( - finalized_batch: Vec, - batch_merkle_tree: &MerkleTree, -) -> Result<(), BatcherError> { - for (vd_batch_idx, entry) in finalized_batch.iter().enumerate() { - let batch_inclusion_data = BatchInclusionData::new(vd_batch_idx, batch_merkle_tree); - let response = ResponseMessage::BatchInclusionData(batch_inclusion_data); - - let serialized_response = cbor_serialize(&response) - .map_err(|e| BatcherError::SerializationError(e.to_string()))?; - - let Some(ws_sink) = entry.messaging_sink.as_ref() else { - return Err(BatcherError::WsSinkEmpty); - }; - - let sending_result = ws_sink - .write() - .await - .send(Message::binary(serialized_response)) - .await; - - match sending_result { - Err(Error::AlreadyClosed) => (), - Err(e) => error!("Error while sending batch inclusion data response: {}", e), - Ok(_) => (), - } - - info!("Response sent"); - } - - Ok(()) -} - -async fn send_message(ws_conn_sink: WsMessageSink, message: T) { - match cbor_serialize(&message) { - Ok(serialized_response) => { - if let Err(err) = ws_conn_sink - .write() - .await - .send(Message::binary(serialized_response)) - .await - { - error!("Error while sending message: {}", err) - } - } - Err(e) => error!("Error while serializing message: {}", e), - } -} diff --git a/batcher/aligned-batcher/src/types/batch_queue.rs b/batcher/aligned-batcher/src/types/batch_queue.rs index 2e78dff78..c90219d23 100644 --- a/batcher/aligned-batcher/src/types/batch_queue.rs +++ b/batcher/aligned-batcher/src/types/batch_queue.rs @@ -1,22 +1,16 @@ +use aligned_sdk::{ + communication::serialization::cbor_serialize, + core::types::{NoncedVerificationData, VerificationDataCommitment}, +}; use ethers::types::{Address, Signature, U256}; -use futures_util::stream::SplitSink; use priority_queue::PriorityQueue; use std::{ hash::{Hash, Hasher}, ops::ControlFlow, - sync::Arc, -}; -use tokio::{net::TcpStream, sync::RwLock}; -use tokio_tungstenite::{tungstenite::Message, WebSocketStream}; - -use aligned_sdk::{ - communication::serialization::cbor_serialize, - core::types::{NoncedVerificationData, VerificationDataCommitment}, }; use super::errors::BatcherError; - -pub(crate) type WsMessageSink = Arc, Message>>>; +use crate::connection::WsMessageSink; #[derive(Clone)] pub(crate) struct BatchQueueEntry { From f7c060638b5bb81531efaf0fd66253f113708469 Mon Sep 17 00:00:00 2001 From: Mariano Nicolini Date: Thu, 26 Sep 2024 21:34:38 -0300 Subject: [PATCH 10/47] Add new entities files --- batcher/aligned-batcher/src/connection.rs | 68 ++++++++++ .../aligned-batcher/src/types/batch_state.rs | 121 ++++++++++++++++++ .../aligned-batcher/src/types/user_state.rs | 29 +++++ 3 files changed, 218 insertions(+) create mode 100644 batcher/aligned-batcher/src/connection.rs create mode 100644 batcher/aligned-batcher/src/types/batch_state.rs create mode 100644 batcher/aligned-batcher/src/types/user_state.rs diff --git a/batcher/aligned-batcher/src/connection.rs b/batcher/aligned-batcher/src/connection.rs new file mode 100644 index 000000000..65bbcf384 --- /dev/null +++ b/batcher/aligned-batcher/src/connection.rs @@ -0,0 +1,68 @@ +use std::sync::Arc; + +use aligned_sdk::{ + communication::serialization::cbor_serialize, + core::types::{BatchInclusionData, ResponseMessage, VerificationCommitmentBatch}, +}; +use futures_util::{stream::SplitSink, SinkExt}; +use lambdaworks_crypto::merkle_tree::merkle::MerkleTree; +use log::{error, info}; +use serde::Serialize; +use tokio::{net::TcpStream, sync::RwLock}; +use tokio_tungstenite::{ + tungstenite::{Error, Message}, + WebSocketStream, +}; + +use crate::types::{batch_queue::BatchQueueEntry, errors::BatcherError}; + +pub(crate) type WsMessageSink = Arc, Message>>>; + +pub(crate) async fn send_batch_inclusion_data_responses( + finalized_batch: Vec, + batch_merkle_tree: &MerkleTree, +) -> Result<(), BatcherError> { + for (vd_batch_idx, entry) in finalized_batch.iter().enumerate() { + let batch_inclusion_data = BatchInclusionData::new(vd_batch_idx, batch_merkle_tree); + let response = ResponseMessage::BatchInclusionData(batch_inclusion_data); + + let serialized_response = cbor_serialize(&response) + .map_err(|e| BatcherError::SerializationError(e.to_string()))?; + + let Some(ws_sink) = entry.messaging_sink.as_ref() else { + return Err(BatcherError::WsSinkEmpty); + }; + + let sending_result = ws_sink + .write() + .await + .send(Message::binary(serialized_response)) + .await; + + match sending_result { + Err(Error::AlreadyClosed) => (), + Err(e) => error!("Error while sending batch inclusion data response: {}", e), + Ok(_) => (), + } + + info!("Response sent"); + } + + Ok(()) +} + +pub(crate) async fn send_message(ws_conn_sink: WsMessageSink, message: T) { + match cbor_serialize(&message) { + Ok(serialized_response) => { + if let Err(err) = ws_conn_sink + .write() + .await + .send(Message::binary(serialized_response)) + .await + { + error!("Error while sending message: {}", err) + } + } + Err(e) => error!("Error while serializing message: {}", e), + } +} diff --git a/batcher/aligned-batcher/src/types/batch_state.rs b/batcher/aligned-batcher/src/types/batch_state.rs new file mode 100644 index 000000000..b5139f858 --- /dev/null +++ b/batcher/aligned-batcher/src/types/batch_state.rs @@ -0,0 +1,121 @@ +use std::collections::HashMap; + +use super::batch_queue::{BatchQueue, BatchQueueEntry}; +use ethers::types::{Address, U256}; +use log::debug; + +pub(crate) struct BatchState { + pub(crate) batch_queue: BatchQueue, +} + +impl BatchState { + pub(crate) fn new() -> Self { + Self { + batch_queue: BatchQueue::new(), + } + } + + // fn get_user_proof_count(&self, addr: &Address) -> u64 { + // *self.user_proof_count_in_batch.get(addr).unwrap_or(&0) + // } + + // /* + // Increments the user proof count in the batch, if the user is already in the hashmap. + // If the user is not in the hashmap, it adds the user to the hashmap with a count of 1 to represent the first proof. + // */ + // fn increment_user_proof_count(&mut self, addr: &Address) { + // self.user_proof_count_in_batch + // .entry(*addr) + // .and_modify(|count| *count += 1) + // .or_insert(1); + // } + + pub(crate) fn get_entry(&self, sender: Address, nonce: U256) -> Option<&BatchQueueEntry> { + self.batch_queue + .iter() + .map(|(entry, _)| entry) + .find(|entry| entry.sender == sender && entry.nonced_verification_data.nonce == nonce) + } + + /// Checks if the entry is valid + /// An entry is valid if there is no entry with the same sender, + /// lower nonce and a lower fee + /// If the entry is valid, it replaces the entry in the queue + /// to increment the max fee, then it updates the user min fee if necessary + /// If the entry is invalid, it returns a validity response message. + /// If the entry is valid, it returns None. + pub(crate) fn replacement_entry_is_valid( + &mut self, + replacement_entry: &BatchQueueEntry, + ) -> bool { + let replacement_max_fee = replacement_entry.nonced_verification_data.max_fee; + let nonce = replacement_entry.nonced_verification_data.nonce; + let sender = replacement_entry.sender; + + debug!( + "Checking validity of entry with sender: {:?}, nonce: {:?}, max_fee: {:?}", + sender, nonce, replacement_max_fee + ); + + // it is a valid entry only if there is no entry with the same sender, lower nonce and a lower fee + !self.batch_queue.iter().any(|(entry, _)| { + entry.sender == sender + && entry.nonced_verification_data.nonce < nonce + && entry.nonced_verification_data.max_fee < replacement_max_fee + }) + } + + // if !is_valid { + // return Some(ValidityResponseMessage::InvalidReplacementMessage); + // } + + // // remove the old entry and insert the new one + // // note that the entries are considered equal for the priority queue + // // if they have the same nonce and sender, so we can remove the old entry + // // by calling remove with the new entry + // self.batch_queue.remove(&replacement_entry); + // self.batch_queue.push( + // replacement_entry.clone(), + // BatchQueueEntryPriority::new(replacement_max_fee, nonce), + // ); + + // let user_min_fee = self + // .batch_queue + // .iter() + // .filter(|(e, _)| e.sender == sender) + // .map(|(e, _)| e.nonced_verification_data.max_fee) + // .min() + // .unwrap_or(U256::max_value()); + + // self.user_min_fee.insert(sender, user_min_fee); + + // None + // } + + pub(crate) fn get_user_min_fee_in_batch(&self, addr: &Address) -> U256 { + self.batch_queue + .iter() + .filter(|(e, _)| &e.sender == addr) + .map(|(e, _)| e.nonced_verification_data.max_fee) + .min() + .unwrap_or(U256::max_value()) + } + + pub(crate) fn get_user_proofs_in_batch_and_min_fee(&self) -> HashMap { + let mut updated_user_states = HashMap::new(); + for (entry, _) in self.batch_queue.iter() { + let addr = entry.sender; + let user_min_fee = entry.nonced_verification_data.max_fee; + + let (proof_count, min_fee) = + updated_user_states.entry(addr).or_insert((0, user_min_fee)); + + *proof_count += 1; + if entry.nonced_verification_data.max_fee < *min_fee { + *min_fee = entry.nonced_verification_data.max_fee; + } + } + + updated_user_states + } +} diff --git a/batcher/aligned-batcher/src/types/user_state.rs b/batcher/aligned-batcher/src/types/user_state.rs new file mode 100644 index 000000000..c5c650ade --- /dev/null +++ b/batcher/aligned-batcher/src/types/user_state.rs @@ -0,0 +1,29 @@ +use ethers::types::U256; + +pub(crate) struct UserState { + pub nonce: Option, + /// The minimum fee of a pending proof for a user. + /// This should always be the fee of the biggest pending nonce by the user. + /// This is used to check if a user is submitting a proof with a higher nonce and higher fee, + /// which is invalid and should be rejected. + pub min_fee: U256, + pub proofs_in_batch: usize, +} + +impl UserState { + pub(crate) fn new() -> Self { + UserState { + nonce: None, + min_fee: U256::max_value(), + proofs_in_batch: 0, + } + } + + pub(crate) fn new_non_paying(nonce: U256) -> Self { + UserState { + nonce: Some(nonce), + min_fee: U256::max_value(), + proofs_in_batch: 0, + } + } +} From 26d66ccc89691319de540035c916c1108b6f27b9 Mon Sep 17 00:00:00 2001 From: Mariano Nicolini Date: Thu, 26 Sep 2024 21:37:43 -0300 Subject: [PATCH 11/47] Remoe legacy code --- .../aligned-batcher/src/types/batch_state.rs | 49 +------------------ 1 file changed, 1 insertion(+), 48 deletions(-) diff --git a/batcher/aligned-batcher/src/types/batch_state.rs b/batcher/aligned-batcher/src/types/batch_state.rs index b5139f858..497aaebe3 100644 --- a/batcher/aligned-batcher/src/types/batch_state.rs +++ b/batcher/aligned-batcher/src/types/batch_state.rs @@ -15,21 +15,6 @@ impl BatchState { } } - // fn get_user_proof_count(&self, addr: &Address) -> u64 { - // *self.user_proof_count_in_batch.get(addr).unwrap_or(&0) - // } - - // /* - // Increments the user proof count in the batch, if the user is already in the hashmap. - // If the user is not in the hashmap, it adds the user to the hashmap with a count of 1 to represent the first proof. - // */ - // fn increment_user_proof_count(&mut self, addr: &Address) { - // self.user_proof_count_in_batch - // .entry(*addr) - // .and_modify(|count| *count += 1) - // .or_insert(1); - // } - pub(crate) fn get_entry(&self, sender: Address, nonce: U256) -> Option<&BatchQueueEntry> { self.batch_queue .iter() @@ -38,12 +23,7 @@ impl BatchState { } /// Checks if the entry is valid - /// An entry is valid if there is no entry with the same sender, - /// lower nonce and a lower fee - /// If the entry is valid, it replaces the entry in the queue - /// to increment the max fee, then it updates the user min fee if necessary - /// If the entry is invalid, it returns a validity response message. - /// If the entry is valid, it returns None. + /// An entry is valid if there is no entry with the same sender, lower nonce and a lower fee pub(crate) fn replacement_entry_is_valid( &mut self, replacement_entry: &BatchQueueEntry, @@ -65,33 +45,6 @@ impl BatchState { }) } - // if !is_valid { - // return Some(ValidityResponseMessage::InvalidReplacementMessage); - // } - - // // remove the old entry and insert the new one - // // note that the entries are considered equal for the priority queue - // // if they have the same nonce and sender, so we can remove the old entry - // // by calling remove with the new entry - // self.batch_queue.remove(&replacement_entry); - // self.batch_queue.push( - // replacement_entry.clone(), - // BatchQueueEntryPriority::new(replacement_max_fee, nonce), - // ); - - // let user_min_fee = self - // .batch_queue - // .iter() - // .filter(|(e, _)| e.sender == sender) - // .map(|(e, _)| e.nonced_verification_data.max_fee) - // .min() - // .unwrap_or(U256::max_value()); - - // self.user_min_fee.insert(sender, user_min_fee); - - // None - // } - pub(crate) fn get_user_min_fee_in_batch(&self, addr: &Address) -> U256 { self.batch_queue .iter() From 9996df5b59182eb979897c938d5a4327f5de312e Mon Sep 17 00:00:00 2001 From: Mariano Nicolini Date: Fri, 27 Sep 2024 13:31:02 -0300 Subject: [PATCH 12/47] Save work in progress - Refactor user states into batch state --- batcher/aligned-batcher/src/lib.rs | 157 +++++++++++++++++++++-------- 1 file changed, 117 insertions(+), 40 deletions(-) diff --git a/batcher/aligned-batcher/src/lib.rs b/batcher/aligned-batcher/src/lib.rs index 92eb3b170..e98a28c65 100644 --- a/batcher/aligned-batcher/src/lib.rs +++ b/batcher/aligned-batcher/src/lib.rs @@ -30,7 +30,7 @@ use lambdaworks_crypto::merkle_tree::merkle::MerkleTree; use lambdaworks_crypto::merkle_tree::traits::IsMerkleTreeBackend; use log::{debug, error, info, warn}; use tokio::net::{TcpListener, TcpStream}; -use tokio::sync::{Mutex, RwLock}; +use tokio::sync::{Mutex, MutexGuard, RwLock}; use tokio_tungstenite::tungstenite::{Error, Message}; use types::batch_queue::{self, BatchQueueEntry, BatchQueueEntryPriority}; use types::errors::{BatcherError, BatcherSendError}; @@ -352,41 +352,19 @@ impl Batcher { info!("Message signature verified"); if self.is_nonpaying(&addr) { + println!("Handling nonpaying MESSAGEEE!"); return self .handle_nonpaying_msg(ws_conn_sink.clone(), client_msg) .await; } - // Check that we had a user state entry for this user and insert it if not. - if !self.user_states.read().await.contains_key(&addr) { - let new_user_state = Mutex::new(UserState::new()); - self.user_states - .write() - .await - .insert(addr.clone(), new_user_state); - } - // At this point, we will have a user state for sure, since we have inserted it - // if not already present. - let user_state_read_lock = self.user_states.read().await; - let mut user_state = user_state_read_lock.get(&addr).unwrap().lock().await; - - // Perform validations on user state if self.user_balance_is_unlocked(&addr).await { send_message( ws_conn_sink.clone(), ValidityResponseMessage::InsufficientBalance(addr), ) .await; - return Ok(()); - } - - let updated_user_proofs = user_state.proofs_in_batch + 1; - if !self.check_user_balance(&addr, updated_user_proofs).await { - send_message( - ws_conn_sink.clone(), - ValidityResponseMessage::InsufficientBalance(addr), - ) - .await; + println!("UNLOCKING THE USER {addr} STATEEE!"); return Ok(()); } @@ -415,6 +393,39 @@ impl Batcher { return Ok(()); } + // Check that we had a user state entry for this user and insert it if not. + if !self.user_states.read().await.contains_key(&addr) { + let new_user_state = Mutex::new(UserState::new()); + self.user_states + .write() + .await + .insert(addr.clone(), new_user_state); + } + + // At this point, we will have a user state for sure, since we have inserted it + // if not already present. + println!("LOCKING USER STATE FOR ADDR: {addr}"); + let user_state_read_lock = self.user_states.read().await; + let mut user_state = user_state_read_lock.get(&addr).unwrap().lock().await; + println!("USER STATE FOR {addr} LOCKED BABY!"); + + println!("LOCKING BATCH STATE IN HANDLE MESSAGE"); + let batch_state_lock = self.batch_state.lock().await; + println!("BATCH STATE LOCKEDIN HANDLE MESSAGE"); + + if !self + .check_user_balance(&addr, user_state.proofs_in_batch + 1) + .await + { + send_message( + ws_conn_sink.clone(), + ValidityResponseMessage::InsufficientBalance(addr), + ) + .await; + println!("UNLOCKING THE USER {addr} STATEEE!"); + return Ok(()); + } + let expected_nonce = if let Some(cached_nonce) = user_state.nonce { cached_nonce } else { @@ -440,6 +451,7 @@ impl Batcher { "Invalid nonce for address {addr}, had nonce {:?} < {:?}", expected_nonce, msg_nonce ); + println!("UNLOCKING THE USER {addr} STATEEE!"); send_message(ws_conn_sink.clone(), ValidityResponseMessage::InvalidNonce).await; return Ok(()); } else if expected_nonce == msg_nonce { @@ -455,18 +467,30 @@ impl Batcher { user_state.nonce = Some(msg_nonce + U256::one()); user_state.min_fee = max_fee; + user_state.proofs_in_batch += 1; + + println!( + "USER {addr} PROOFS IN BATCH: {}", + user_state.proofs_in_batch + ); + + println!("ADDING NORMAL ENTRY TO BATCH"); self.add_to_batch( + batch_state_lock, nonced_verification_data, ws_conn_sink.clone(), client_msg.signature, addr, ) .await; + println!("NORMAL ENTRY ADDED SUCCESSFULLY TO BATCH"); } else { // In this case, the message might be a replacement one. If it is valid, // we replace the old entry with the new from the replacement message. if !self .handle_replacement_message( + batch_state_lock, + user_state, nonced_verification_data, ws_conn_sink.clone(), client_msg.signature, @@ -475,15 +499,20 @@ impl Batcher { ) .await { + println!("UNLOCKING THE USER {addr} STATEEE!"); + // message should not be added to batch return Ok(()); } - let batch_state = self.batch_state.lock().await; - let updated_min_fee_in_batch = batch_state.get_user_min_fee_in_batch(&addr); - user_state.min_fee = updated_min_fee_in_batch; + // println!("LOCKING BATCH STATE - HANDLE MESSAGE - REPLACEMENT MESSAGE"); + // let batch_state = self.batch_state.lock().await; + // println!("BATCH STATE LOCKED - HANDLE MESSAGE - REPLACEMENT MESSAGE"); + // let updated_min_fee_in_batch = batch_state.get_user_min_fee_in_batch(&addr); + // user_state.min_fee = updated_min_fee_in_batch; } - + println!("BATCH STATE UNLOCKED LOCKED - HANDLE MESSAGE"); + println!("UNLOCKING THE USER {addr} STATEEE!"); info!("Verification data message handled"); send_message(ws_conn_sink, ValidityResponseMessage::Valid).await; Ok(()) @@ -509,18 +538,22 @@ impl Batcher { /// Returns true if the message was replaced in the batch, false otherwise async fn handle_replacement_message( &self, + mut batch_state_lock: MutexGuard<'_, BatchState>, + mut user_state_lock: MutexGuard<'_, UserState>, nonced_verification_data: NoncedVerificationData, ws_conn_sink: WsMessageSink, signature: Signature, addr: Address, expected_user_nonce: U256, ) -> bool { - let mut batch_state = self.batch_state.lock().await; + // println!("LOCKING BATCH STATE - HANDLE MESSAGE - REPLACEMENT MESSAGE FUNCTION"); + // let mut batch_state = self.batch_state.lock().await; + // println!("BATCH STATE LOCKED - HANDLE MESSAGE - REPLACEMENT MESSAGE FUNCTION"); let replacement_max_fee = nonced_verification_data.max_fee; let nonce = nonced_verification_data.nonce; - let Some(entry) = batch_state.get_entry(addr, nonce) else { + let Some(entry) = batch_state_lock.get_entry(addr, nonce) else { warn!( "Invalid nonce for address {addr} Expected: {:?}, got: {:?}", expected_user_nonce, nonce @@ -571,13 +604,14 @@ impl Batcher { } replacement_entry.messaging_sink = Some(ws_conn_sink.clone()); - if !batch_state.replacement_entry_is_valid(&replacement_entry) { + if !batch_state_lock.replacement_entry_is_valid(&replacement_entry) { warn!("Invalid max fee"); send_message( ws_conn_sink.clone(), ValidityResponseMessage::InvalidReplacementMessage, ) .await; + println!("BATCH STATE UNLOCKED - HANDLE MESSAGE - REPLACEMENT MESSAGE FUNCTION"); return false; } @@ -590,14 +624,17 @@ impl Batcher { // note that the entries are considered equal for the priority queue // if they have the same nonce and sender, so we can remove the old entry // by calling remove with the new entry - batch_state.batch_queue.remove(&replacement_entry); - let result = batch_state.batch_queue.push( + batch_state_lock.batch_queue.remove(&replacement_entry); + let result = batch_state_lock.batch_queue.push( replacement_entry.clone(), BatchQueueEntryPriority::new(replacement_max_fee, nonce), ); + let updated_min_fee_in_batch = batch_state_lock.get_user_min_fee_in_batch(&addr); + user_state_lock.min_fee = updated_min_fee_in_batch; assert!(result.is_none()); + println!("BATCH STATE UNLOCKED - HANDLE MESSAGE - REPLACEMENT MESSAGE FUNCTION"); true } @@ -614,6 +651,7 @@ impl Batcher { /// Adds verification data to the current batch queue. async fn add_to_batch( &self, + mut batch_state_lock: MutexGuard<'_, BatchState>, verification_data: NoncedVerificationData, ws_conn_sink: WsMessageSink, proof_submitter_sig: Signature, @@ -626,8 +664,10 @@ impl Batcher { let max_fee = verification_data.max_fee; let nonce = verification_data.nonce; - let mut batch_state = self.batch_state.lock().await; - batch_state.batch_queue.push( + // println!("LOCKING BATCH STATE - ADD TO BATCH FUNCTION"); + // let mut batch_state = self.batch_state.lock().await; + // println!("BATCH STATE LOCKED - ADD TO BATCH FUNCTION"); + batch_state_lock.batch_queue.push( BatchQueueEntry::new( verification_data, verification_data_comm, @@ -640,8 +680,9 @@ impl Batcher { info!( "Current batch queue length: {}", - batch_state.batch_queue.len() + batch_state_lock.batch_queue.len() ); + println!("UNLOCKING BATCH STATE - ADD TO BATCH FUNCTION"); } /// Given a new block number listened from the blockchain, checks if the current batch is ready to be posted. @@ -661,8 +702,9 @@ impl Batcher { block_number: u64, gas_price: U256, ) -> Option> { + println!("LOCKING BATCH STATE - IS BATCH READY FUNCTION"); let mut batch_state = self.batch_state.lock().await; - let user_states = self.user_states.write().await; + println!("BATCH STATE LOCKED - IS BATCH READY FUNCTION"); let current_batch_len = batch_state.batch_queue.len(); let last_uploaded_batch_block_lock = self.last_uploaded_batch_block.lock().await; @@ -671,6 +713,7 @@ impl Batcher { // once the bug in Lambdaworks merkle tree is fixed. if current_batch_len < 2 { info!("Current batch is empty or length 1. Waiting for more proofs..."); + println!("BATCH STATE UNLOCKED - IS BATCH READY FUNCTION"); return None; } @@ -681,6 +724,7 @@ impl Batcher { "Current batch not ready to be posted. Current block: {} - Last uploaded block: {}. Current batch length: {} - Minimum batch length: {}", block_number, *last_uploaded_batch_block_lock, current_batch_len, self.min_batch_len ); + println!("BATCH STATE UNLOCKED - IS BATCH READY FUNCTION"); return None; } @@ -690,30 +734,54 @@ impl Batcher { info!( "Batch is currently being posted. Waiting for the current batch to be finalized..." ); + println!("BATCH STATE UNLOCKED - IS BATCH READY FUNCTION"); return None; } // Set the batch posting flag to true *batch_posting = true; + + println!("LOCKING USER STATES - IS BATCH READY FUNCTION"); + let user_states = self.user_states.write().await; + println!("USER STATES LOCKED - IS BATCH READY FUNCTION"); + let batch_queue_copy = batch_state.batch_queue.clone(); match batch_queue::try_build_batch(batch_queue_copy, gas_price, self.max_batch_size) { Ok((resulting_batch_queue, finalized_batch)) => { + if finalized_batch.len() == 1 { + println!("BATCH STATE UNLOCKED - IS BATCH READY FUNCTION"); + panic!("FINALIZED BATCH IS LEN 1!!!!!!!!!!!!!!!"); + } batch_state.batch_queue = resulting_batch_queue; + let updated_user_proof_count_and_min_fee = batch_state.get_user_proofs_in_batch_and_min_fee(); - for (addr, (proof_count, min_fee)) in updated_user_proof_count_and_min_fee.iter() { + for addr in user_states.keys() { + let (proof_count, min_fee) = updated_user_proof_count_and_min_fee + .get(addr) + .unwrap_or(&(0, U256::MAX)); + let mut user_state = user_states.get(addr).unwrap().lock().await; user_state.proofs_in_batch = *proof_count; user_state.min_fee = *min_fee; } + // for (addr, (proof_count, min_fee)) in updated_user_proof_count_and_min_fee.iter() { + // let mut user_state = user_states.get(addr).unwrap().lock().await; + // println!("UPDATING OF PROOFS IN BATCH: {}", proof_count); + // user_state.proofs_in_batch = *proof_count; + // user_state.min_fee = *min_fee; + // } + + println!("BATCH STATE UNLOCKED - IS BATCH READY FUNCTION"); Some(finalized_batch) } Err(BatcherError::BatchCostTooHigh) => { // We can't post a batch since users are not willing to pay the needed fee, wait for more proofs info!("No working batch found. Waiting for more proofs..."); *batch_posting = false; + println!("BATCH STATE UNLOCKED - IS BATCH READY FUNCTION"); None } // FIXME: We should refactor this code and instead of returning None, return an error. @@ -721,6 +789,7 @@ impl Batcher { Err(e) => { error!("Unexpected error: {:?}", e); *batch_posting = false; + println!("BATCH STATE UNLOCKED - IS BATCH READY FUNCTION"); None } } @@ -807,7 +876,10 @@ impl Batcher { async fn flush_queue_and_clear_nonce_cache(&self) { warn!("Resetting state... Flushing queue and nonces"); + + println!("LOCKING BATCH STATE - FLUE QUEUE AND CLEAR NONCE FUNCTION"); let mut batch_state = self.batch_state.lock().await; + println!("BATCH STATE LOCKED - FLUE QUEUE AND CLEAR NONCE FUNCTION"); for (entry, _) in batch_state.batch_queue.iter() { if let Some(ws_sink) = entry.messaging_sink.as_ref() { @@ -819,6 +891,7 @@ impl Batcher { batch_state.batch_queue.clear(); self.user_states.write().await.clear(); + println!("BATCH STATE UNLOCKED - FLUE QUEUE AND CLEAR NONCE FUNCTION"); } /// Receives new block numbers, checks if conditions are met for submission and @@ -832,7 +905,7 @@ impl Batcher { } }; - while let Some(finalized_batch) = self.is_batch_ready(block_number, gas_price).await { + if let Some(finalized_batch) = self.is_batch_ready(block_number, gas_price).await { let batch_finalization_result = self .finalize_batch(block_number, finalized_batch, gas_price) .await; @@ -1046,6 +1119,7 @@ impl Batcher { } let user_states = self.user_states.read().await; + let batch_state_lock = self.batch_state.lock().await; // Safe to call `unwrap()` here since at this point we have a non-paying configuration loaded // for sure and the non-paying address nonce cached in the user states. @@ -1069,13 +1143,16 @@ impl Batcher { ) .await; + println!("NONPAYING - ADDING ENTRY TO BATCH"); self.add_to_batch( + batch_state_lock, nonced_verification_data, ws_sink.clone(), client_msg.signature, non_paying_config.address, ) .await; + println!("NONPAYING - ADDED ENTRY TO BATCH SUCCESSFULLY"); // Non-paying user nonce is updated (*non_paying_user_state).nonce = Some(non_paying_nonce + U256::one()); From 85de7d5a89cd382a6686c2ac53decedf7f6bef44 Mon Sep 17 00:00:00 2001 From: Mariano Nicolini Date: Fri, 27 Sep 2024 15:27:52 -0300 Subject: [PATCH 13/47] Save work in progress - Fixed concurrency bugs and refactored batch state --- Makefile | 6 +- batcher/aligned-batcher/src/lib.rs | 153 +++++++++++------- .../aligned-batcher/src/types/batch_state.rs | 49 +++++- 3 files changed, 150 insertions(+), 58 deletions(-) diff --git a/Makefile b/Makefile index 14ee2ff33..04028906c 100644 --- a/Makefile +++ b/Makefile @@ -271,10 +271,11 @@ batcher_send_risc0_burst: --proof ../../scripts/test_files/risc_zero/fibonacci_proof_generator/risc_zero_fibonacci.proof \ --vm_program ../../scripts/test_files/risc_zero/fibonacci_proof_generator/fibonacci_id.bin \ --public_input ../../scripts/test_files/risc_zero/fibonacci_proof_generator/risc_zero_fibonacci.pub \ - --repetitions $(BURST_SIZE) \ + --repetitions 300 \ --proof_generator_addr 0x66f9664f97F2b50F62D13eA064982f936dE76657 \ --rpc_url $(RPC_URL) \ --payment_service_addr $(BATCHER_PAYMENTS_CONTRACT_ADDRESS) + # --private_key 0x4bbbf85ce3377467afe5d46f804f221813b2bb87f24d81f60f1fcdbf7cbf4356 \ batcher_send_plonk_bn254_task: batcher/target/release/aligned @echo "Sending Groth16Bn254 1!=0 task to Batcher..." @@ -361,7 +362,8 @@ batcher_send_halo2_ipa_task_burst_5: batcher/target/release/aligned --proof ../../scripts/test_files/halo2_ipa/proof.bin \ --public_input ../../scripts/test_files/halo2_ipa/pub_input.bin \ --vk ../../scripts/test_files/halo2_ipa/params.bin \ - --repetitions 5 \ + --repetitions 300 \ + --private_key 0xdbda1821b80551c9d65939329250298aa3472ba22feea921c0cf5d620ea67b97 \ --rpc_url $(RPC_URL) \ --payment_service_addr $(BATCHER_PAYMENTS_CONTRACT_ADDRESS) diff --git a/batcher/aligned-batcher/src/lib.rs b/batcher/aligned-batcher/src/lib.rs index e98a28c65..3acd158b1 100644 --- a/batcher/aligned-batcher/src/lib.rs +++ b/batcher/aligned-batcher/src/lib.rs @@ -7,7 +7,7 @@ use connection::{send_message, WsMessageSink}; use dotenvy::dotenv; use ethers::contract::ContractError; use ethers::signers::Signer; -use types::batch_state::BatchState; +use types::batch_state::{self, BatchState}; use types::user_state::UserState; use std::collections::HashMap; @@ -80,7 +80,7 @@ pub struct Batcher { pre_verification_is_enabled: bool, non_paying_config: Option, posting_batch: Mutex, - user_states: RwLock>>, + // user_states: RwLock>>, } impl Batcher { @@ -165,7 +165,7 @@ impl Batcher { .await .expect("Failed to get fallback Batcher Payment Service contract"); - let user_states = RwLock::new(HashMap::new()); + let mut user_states = HashMap::new(); let non_paying_config = if let Some(non_paying_config) = config.batcher.non_paying { warn!("Non-paying address configuration detected. Will replace non-paying address {} with configured address.", non_paying_config.address); @@ -178,7 +178,7 @@ impl Batcher { .expect("Could not get non-paying nonce from Ethereum"); let non_paying_user_state = Mutex::new(UserState::new_non_paying(nonpaying_nonce)); - user_states.write().await.insert( + user_states.insert( non_paying_config.replacement.address(), non_paying_user_state, ); @@ -188,6 +188,8 @@ impl Batcher { None }; + let batch_state = BatchState::new_with_user_states(user_states); + Self { s3_client, s3_bucket_name, @@ -197,7 +199,6 @@ impl Batcher { chain_id, payment_service, payment_service_fallback, - batch_state: Mutex::new(BatchState::new()), max_block_interval: config.batcher.block_interval, min_batch_len: config.batcher.batch_size_interval, max_proof_size: config.batcher.max_proof_size, @@ -206,7 +207,8 @@ impl Batcher { pre_verification_is_enabled: config.batcher.pre_verification_is_enabled, non_paying_config, posting_batch: Mutex::new(false), - user_states, + batch_state: Mutex::new(batch_state), + // user_states: RwLock::new(user_states), } } @@ -394,29 +396,34 @@ impl Batcher { } // Check that we had a user state entry for this user and insert it if not. - if !self.user_states.read().await.contains_key(&addr) { - let new_user_state = Mutex::new(UserState::new()); - self.user_states - .write() - .await - .insert(addr.clone(), new_user_state); + { + let mut batch_state_lock = self.batch_state.lock().await; + if !batch_state_lock.user_states.contains_key(&addr) { + let new_user_state = Mutex::new(UserState::new()); + // let mut batch_state_write_lock = self.batch_state.write().await; + batch_state_lock + .user_states + .insert(addr.clone(), new_user_state); + } } // At this point, we will have a user state for sure, since we have inserted it // if not already present. println!("LOCKING USER STATE FOR ADDR: {addr}"); - let user_state_read_lock = self.user_states.read().await; - let mut user_state = user_state_read_lock.get(&addr).unwrap().lock().await; - println!("USER STATE FOR {addr} LOCKED BABY!"); - - println!("LOCKING BATCH STATE IN HANDLE MESSAGE"); let batch_state_lock = self.batch_state.lock().await; - println!("BATCH STATE LOCKEDIN HANDLE MESSAGE"); + // let mut user_state = batch_state_lock.get_user_state(&addr).unwrap().lock().await; + println!("USER STATE FOR {addr} LOCKED BABY!"); - if !self - .check_user_balance(&addr, user_state.proofs_in_batch + 1) + // println!("LOCKING BATCH STATE IN HANDLE MESSAGE"); + // let batch_state_lock = self.batch_state.lock().await; + // println!("BATCH STATE LOCKEDIN HANDLE MESSAGE"); + let proofs_in_batch = batch_state_lock + .get_user_state(&addr) + .unwrap() + .lock() .await - { + .proofs_in_batch; + if !self.check_user_balance(&addr, proofs_in_batch + 1).await { send_message( ws_conn_sink.clone(), ValidityResponseMessage::InsufficientBalance(addr), @@ -426,7 +433,8 @@ impl Batcher { return Ok(()); } - let expected_nonce = if let Some(cached_nonce) = user_state.nonce { + let user_nonce = batch_state_lock.get_user_nonce(&addr).await; + let expected_nonce = if let Some(cached_nonce) = user_nonce { cached_nonce } else { let ethereum_user_nonce = match self.get_user_nonce_from_ethereum(addr).await { @@ -442,6 +450,7 @@ impl Batcher { } }; + let mut user_state = batch_state_lock.get_user_state(&addr).unwrap().lock().await; user_state.nonce = Some(ethereum_user_nonce); ethereum_user_nonce }; @@ -455,24 +464,25 @@ impl Batcher { send_message(ws_conn_sink.clone(), ValidityResponseMessage::InvalidNonce).await; return Ok(()); } else if expected_nonce == msg_nonce { + // let user_state = batch_state_lock.get_user_state(&addr) let max_fee = nonced_verification_data.max_fee; - if max_fee > user_state.min_fee { + + let (user_min_fee, user_proof_count) = + batch_state_lock.get_user_batch_data(&addr).await.unwrap(); + if max_fee > user_min_fee { warn!( "Invalid max fee for address {addr}, had fee {:?} < {:?}", - user_state.min_fee, max_fee + user_min_fee, max_fee ); send_message(ws_conn_sink.clone(), ValidityResponseMessage::InvalidMaxFee).await; return Ok(()); } - user_state.nonce = Some(msg_nonce + U256::one()); - user_state.min_fee = max_fee; - user_state.proofs_in_batch += 1; + // user_state.nonce = Some(msg_nonce + U256::one()); + // user_state.min_fee = max_fee; + // user_state.proofs_in_batch += 1; - println!( - "USER {addr} PROOFS IN BATCH: {}", - user_state.proofs_in_batch - ); + println!("USER {addr} PROOFS IN BATCH: {}", user_proof_count); println!("ADDING NORMAL ENTRY TO BATCH"); self.add_to_batch( @@ -490,7 +500,6 @@ impl Batcher { if !self .handle_replacement_message( batch_state_lock, - user_state, nonced_verification_data, ws_conn_sink.clone(), client_msg.signature, @@ -539,7 +548,6 @@ impl Batcher { async fn handle_replacement_message( &self, mut batch_state_lock: MutexGuard<'_, BatchState>, - mut user_state_lock: MutexGuard<'_, UserState>, nonced_verification_data: NoncedVerificationData, ws_conn_sink: WsMessageSink, signature: Signature, @@ -630,6 +638,7 @@ impl Batcher { BatchQueueEntryPriority::new(replacement_max_fee, nonce), ); let updated_min_fee_in_batch = batch_state_lock.get_user_min_fee_in_batch(&addr); + let mut user_state_lock = batch_state_lock.get_user_state(&addr).unwrap().lock().await; user_state_lock.min_fee = updated_min_fee_in_batch; assert!(result.is_none()); @@ -655,7 +664,7 @@ impl Batcher { verification_data: NoncedVerificationData, ws_conn_sink: WsMessageSink, proof_submitter_sig: Signature, - proof_submiter_addr: Address, + proof_submitter_addr: Address, ) { info!("Calculating verification data commitments..."); let verification_data_comm = verification_data.clone().into(); @@ -673,7 +682,7 @@ impl Batcher { verification_data_comm, ws_conn_sink, proof_submitter_sig, - proof_submiter_addr, + proof_submitter_addr, ), BatchQueueEntryPriority::new(max_fee, nonce), ); @@ -682,6 +691,28 @@ impl Batcher { "Current batch queue length: {}", batch_state_lock.batch_queue.len() ); + + let mut proof_submitter_addr = proof_submitter_addr; + + proof_submitter_addr = if self.is_nonpaying(&proof_submitter_addr) { + self.non_paying_config + .as_ref() + .unwrap() + .replacement + .address() + } else { + proof_submitter_addr + }; + + let mut user_state = batch_state_lock + .get_user_state(&proof_submitter_addr) + .unwrap() + .lock() + .await; + + user_state.nonce = Some(nonce + U256::one()); + user_state.min_fee = max_fee; + user_state.proofs_in_batch += 1; println!("UNLOCKING BATCH STATE - ADD TO BATCH FUNCTION"); } @@ -703,10 +734,10 @@ impl Batcher { gas_price: U256, ) -> Option> { println!("LOCKING BATCH STATE - IS BATCH READY FUNCTION"); - let mut batch_state = self.batch_state.lock().await; + let mut batch_state_lock = self.batch_state.lock().await; println!("BATCH STATE LOCKED - IS BATCH READY FUNCTION"); - let current_batch_len = batch_state.batch_queue.len(); + let current_batch_len = batch_state_lock.batch_queue.len(); let last_uploaded_batch_block_lock = self.last_uploaded_batch_block.lock().await; // FIXME(marian): This condition should be changed to current_batch_size == 0 @@ -741,28 +772,30 @@ impl Batcher { // Set the batch posting flag to true *batch_posting = true; - println!("LOCKING USER STATES - IS BATCH READY FUNCTION"); - let user_states = self.user_states.write().await; - println!("USER STATES LOCKED - IS BATCH READY FUNCTION"); + // println!("LOCKING USER STATES - IS BATCH READY FUNCTION"); + // let user_states = batch_state_lock.user_states; + // println!("USER STATES LOCKED - IS BATCH READY FUNCTION"); - let batch_queue_copy = batch_state.batch_queue.clone(); + let batch_queue_copy = batch_state_lock.batch_queue.clone(); match batch_queue::try_build_batch(batch_queue_copy, gas_price, self.max_batch_size) { Ok((resulting_batch_queue, finalized_batch)) => { if finalized_batch.len() == 1 { println!("BATCH STATE UNLOCKED - IS BATCH READY FUNCTION"); panic!("FINALIZED BATCH IS LEN 1!!!!!!!!!!!!!!!"); } - batch_state.batch_queue = resulting_batch_queue; + batch_state_lock.batch_queue = resulting_batch_queue; let updated_user_proof_count_and_min_fee = - batch_state.get_user_proofs_in_batch_and_min_fee(); + batch_state_lock.get_user_proofs_in_batch_and_min_fee(); - for addr in user_states.keys() { + for addr in batch_state_lock.user_states.keys() { let (proof_count, min_fee) = updated_user_proof_count_and_min_fee .get(addr) .unwrap_or(&(0, U256::MAX)); - let mut user_state = user_states.get(addr).unwrap().lock().await; + let mut user_state = + batch_state_lock.get_user_state(addr).unwrap().lock().await; + user_state.proofs_in_batch = *proof_count; user_state.min_fee = *min_fee; } @@ -878,10 +911,10 @@ impl Batcher { warn!("Resetting state... Flushing queue and nonces"); println!("LOCKING BATCH STATE - FLUE QUEUE AND CLEAR NONCE FUNCTION"); - let mut batch_state = self.batch_state.lock().await; + let mut batch_state_lock = self.batch_state.lock().await; println!("BATCH STATE LOCKED - FLUE QUEUE AND CLEAR NONCE FUNCTION"); - for (entry, _) in batch_state.batch_queue.iter() { + for (entry, _) in batch_state_lock.batch_queue.iter() { if let Some(ws_sink) = entry.messaging_sink.as_ref() { send_message(ws_sink.clone(), ResponseMessage::BatchReset).await; } else { @@ -889,8 +922,8 @@ impl Batcher { } } - batch_state.batch_queue.clear(); - self.user_states.write().await.clear(); + batch_state_lock.batch_queue.clear(); + batch_state_lock.user_states.clear(); println!("BATCH STATE UNLOCKED - FLUE QUEUE AND CLEAR NONCE FUNCTION"); } @@ -1118,14 +1151,18 @@ impl Batcher { return Ok(()); // Send error message to the client and return } - let user_states = self.user_states.read().await; + // let user_states = self.user_states.read().await; let batch_state_lock = self.batch_state.lock().await; // Safe to call `unwrap()` here since at this point we have a non-paying configuration loaded // for sure and the non-paying address nonce cached in the user states. - let non_paying_user_state_lock = user_states.get(&replacement_addr).unwrap(); - let mut non_paying_user_state = non_paying_user_state_lock.lock().await; - let non_paying_nonce = non_paying_user_state.nonce.unwrap(); + let non_paying_nonce = batch_state_lock + .get_user_nonce(&replacement_addr) + .await + .unwrap(); + + // let mut non_paying_user_state = non_paying_user_state_lock.lock().await; + // let non_paying_nonce = non_paying_user_state.nonce.unwrap(); info!("Non-paying nonce: {:?}", non_paying_nonce); @@ -1154,8 +1191,14 @@ impl Batcher { .await; println!("NONPAYING - ADDED ENTRY TO BATCH SUCCESSFULLY"); - // Non-paying user nonce is updated - (*non_paying_user_state).nonce = Some(non_paying_nonce + U256::one()); + // let mut non_paying_user_state = batch_state_lock + // .get_user_state(&replacement_addr) + // .unwrap() + // .lock() + // .await; + + // // Non-paying user nonce is updated + // (*non_paying_user_state).nonce = Some(non_paying_nonce + U256::one()); info!("Non-paying verification data message handled"); send_message(ws_sink, ValidityResponseMessage::Valid).await; diff --git a/batcher/aligned-batcher/src/types/batch_state.rs b/batcher/aligned-batcher/src/types/batch_state.rs index 497aaebe3..2780b250f 100644 --- a/batcher/aligned-batcher/src/types/batch_state.rs +++ b/batcher/aligned-batcher/src/types/batch_state.rs @@ -1,17 +1,30 @@ use std::collections::HashMap; -use super::batch_queue::{BatchQueue, BatchQueueEntry}; +use super::{ + batch_queue::{BatchQueue, BatchQueueEntry}, + user_state::UserState, +}; use ethers::types::{Address, U256}; use log::debug; +use tokio::sync::Mutex; pub(crate) struct BatchState { pub(crate) batch_queue: BatchQueue, + pub(crate) user_states: HashMap>, } impl BatchState { pub(crate) fn new() -> Self { Self { batch_queue: BatchQueue::new(), + user_states: HashMap::new(), + } + } + + pub(crate) fn new_with_user_states(user_states: HashMap>) -> Self { + Self { + batch_queue: BatchQueue::new(), + user_states, } } @@ -22,6 +35,40 @@ impl BatchState { .find(|entry| entry.sender == sender && entry.nonced_verification_data.nonce == nonce) } + pub(crate) fn get_user_state(&self, addr: &Address) -> Option<&Mutex> { + self.user_states.get(addr) + } + + pub(crate) async fn get_user_nonce(&self, addr: &Address) -> Option { + let Some(user_state) = self.get_user_state(addr) else { + return None; + }; + user_state.lock().await.nonce + } + + pub(crate) async fn get_user_min_fee(&self, addr: &Address) -> Option { + let Some(user_state) = self.get_user_state(addr) else { + return None; + }; + Some(user_state.lock().await.min_fee) + } + + pub(crate) async fn get_user_proof_count(&self, addr: &Address) -> Option { + let Some(user_state) = self.get_user_state(addr) else { + return None; + }; + Some(user_state.lock().await.proofs_in_batch) + } + + pub(crate) async fn get_user_batch_data(&self, addr: &Address) -> Option<(U256, usize)> { + let Some(user_state) = self.get_user_state(addr) else { + return None; + }; + + let user_state_lock = user_state.lock().await; + Some((user_state_lock.min_fee, user_state_lock.proofs_in_batch)) + } + /// Checks if the entry is valid /// An entry is valid if there is no entry with the same sender, lower nonce and a lower fee pub(crate) fn replacement_entry_is_valid( From 5a30e47a20a33007d3a72128b1ee2572a42dfd44 Mon Sep 17 00:00:00 2001 From: Mariano Nicolini Date: Fri, 27 Sep 2024 18:34:17 -0300 Subject: [PATCH 14/47] Remove legacy commented code --- batcher/aligned-batcher/src/lib.rs | 118 +++------------------------- batcher/aligned/send_burst_tasks.sh | 2 +- 2 files changed, 14 insertions(+), 106 deletions(-) diff --git a/batcher/aligned-batcher/src/lib.rs b/batcher/aligned-batcher/src/lib.rs index 3acd158b1..ec24f365c 100644 --- a/batcher/aligned-batcher/src/lib.rs +++ b/batcher/aligned-batcher/src/lib.rs @@ -353,10 +353,16 @@ impl Batcher { }; info!("Message signature verified"); + let nonced_verification_data = client_msg.verification_data.clone(); + if nonced_verification_data.verification_data.proof.len() > self.max_proof_size { + error!("Proof size exceeds the maximum allowed size."); + send_message(ws_conn_sink.clone(), ValidityResponseMessage::ProofTooLarge).await; + return Ok(()); + } + if self.is_nonpaying(&addr) { - println!("Handling nonpaying MESSAGEEE!"); return self - .handle_nonpaying_msg(ws_conn_sink.clone(), client_msg) + .handle_nonpaying_msg(ws_conn_sink.clone(), &client_msg) .await; } @@ -366,14 +372,6 @@ impl Batcher { ValidityResponseMessage::InsufficientBalance(addr), ) .await; - println!("UNLOCKING THE USER {addr} STATEEE!"); - return Ok(()); - } - - let nonced_verification_data = client_msg.verification_data; - if nonced_verification_data.verification_data.proof.len() > self.max_proof_size { - error!("Proof size exceeds the maximum allowed size."); - send_message(ws_conn_sink.clone(), ValidityResponseMessage::ProofTooLarge).await; return Ok(()); } @@ -400,7 +398,6 @@ impl Batcher { let mut batch_state_lock = self.batch_state.lock().await; if !batch_state_lock.user_states.contains_key(&addr) { let new_user_state = Mutex::new(UserState::new()); - // let mut batch_state_write_lock = self.batch_state.write().await; batch_state_lock .user_states .insert(addr.clone(), new_user_state); @@ -409,14 +406,7 @@ impl Batcher { // At this point, we will have a user state for sure, since we have inserted it // if not already present. - println!("LOCKING USER STATE FOR ADDR: {addr}"); let batch_state_lock = self.batch_state.lock().await; - // let mut user_state = batch_state_lock.get_user_state(&addr).unwrap().lock().await; - println!("USER STATE FOR {addr} LOCKED BABY!"); - - // println!("LOCKING BATCH STATE IN HANDLE MESSAGE"); - // let batch_state_lock = self.batch_state.lock().await; - // println!("BATCH STATE LOCKEDIN HANDLE MESSAGE"); let proofs_in_batch = batch_state_lock .get_user_state(&addr) .unwrap() @@ -429,7 +419,6 @@ impl Batcher { ValidityResponseMessage::InsufficientBalance(addr), ) .await; - println!("UNLOCKING THE USER {addr} STATEEE!"); return Ok(()); } @@ -460,15 +449,12 @@ impl Batcher { "Invalid nonce for address {addr}, had nonce {:?} < {:?}", expected_nonce, msg_nonce ); - println!("UNLOCKING THE USER {addr} STATEEE!"); send_message(ws_conn_sink.clone(), ValidityResponseMessage::InvalidNonce).await; return Ok(()); } else if expected_nonce == msg_nonce { - // let user_state = batch_state_lock.get_user_state(&addr) let max_fee = nonced_verification_data.max_fee; - let (user_min_fee, user_proof_count) = - batch_state_lock.get_user_batch_data(&addr).await.unwrap(); + let (user_min_fee, _) = batch_state_lock.get_user_batch_data(&addr).await.unwrap(); if max_fee > user_min_fee { warn!( "Invalid max fee for address {addr}, had fee {:?} < {:?}", @@ -478,13 +464,6 @@ impl Batcher { return Ok(()); } - // user_state.nonce = Some(msg_nonce + U256::one()); - // user_state.min_fee = max_fee; - // user_state.proofs_in_batch += 1; - - println!("USER {addr} PROOFS IN BATCH: {}", user_proof_count); - - println!("ADDING NORMAL ENTRY TO BATCH"); self.add_to_batch( batch_state_lock, nonced_verification_data, @@ -493,7 +472,6 @@ impl Batcher { addr, ) .await; - println!("NORMAL ENTRY ADDED SUCCESSFULLY TO BATCH"); } else { // In this case, the message might be a replacement one. If it is valid, // we replace the old entry with the new from the replacement message. @@ -508,20 +486,10 @@ impl Batcher { ) .await { - println!("UNLOCKING THE USER {addr} STATEEE!"); - // message should not be added to batch return Ok(()); } - - // println!("LOCKING BATCH STATE - HANDLE MESSAGE - REPLACEMENT MESSAGE"); - // let batch_state = self.batch_state.lock().await; - // println!("BATCH STATE LOCKED - HANDLE MESSAGE - REPLACEMENT MESSAGE"); - // let updated_min_fee_in_batch = batch_state.get_user_min_fee_in_batch(&addr); - // user_state.min_fee = updated_min_fee_in_batch; } - println!("BATCH STATE UNLOCKED LOCKED - HANDLE MESSAGE"); - println!("UNLOCKING THE USER {addr} STATEEE!"); info!("Verification data message handled"); send_message(ws_conn_sink, ValidityResponseMessage::Valid).await; Ok(()) @@ -554,10 +522,6 @@ impl Batcher { addr: Address, expected_user_nonce: U256, ) -> bool { - // println!("LOCKING BATCH STATE - HANDLE MESSAGE - REPLACEMENT MESSAGE FUNCTION"); - // let mut batch_state = self.batch_state.lock().await; - // println!("BATCH STATE LOCKED - HANDLE MESSAGE - REPLACEMENT MESSAGE FUNCTION"); - let replacement_max_fee = nonced_verification_data.max_fee; let nonce = nonced_verification_data.nonce; @@ -619,7 +583,6 @@ impl Batcher { ValidityResponseMessage::InvalidReplacementMessage, ) .await; - println!("BATCH STATE UNLOCKED - HANDLE MESSAGE - REPLACEMENT MESSAGE FUNCTION"); return false; } @@ -633,17 +596,15 @@ impl Batcher { // if they have the same nonce and sender, so we can remove the old entry // by calling remove with the new entry batch_state_lock.batch_queue.remove(&replacement_entry); - let result = batch_state_lock.batch_queue.push( + batch_state_lock.batch_queue.push( replacement_entry.clone(), BatchQueueEntryPriority::new(replacement_max_fee, nonce), ); + let updated_min_fee_in_batch = batch_state_lock.get_user_min_fee_in_batch(&addr); let mut user_state_lock = batch_state_lock.get_user_state(&addr).unwrap().lock().await; user_state_lock.min_fee = updated_min_fee_in_batch; - assert!(result.is_none()); - - println!("BATCH STATE UNLOCKED - HANDLE MESSAGE - REPLACEMENT MESSAGE FUNCTION"); true } @@ -672,10 +633,6 @@ impl Batcher { let max_fee = verification_data.max_fee; let nonce = verification_data.nonce; - - // println!("LOCKING BATCH STATE - ADD TO BATCH FUNCTION"); - // let mut batch_state = self.batch_state.lock().await; - // println!("BATCH STATE LOCKED - ADD TO BATCH FUNCTION"); batch_state_lock.batch_queue.push( BatchQueueEntry::new( verification_data, @@ -713,7 +670,6 @@ impl Batcher { user_state.nonce = Some(nonce + U256::one()); user_state.min_fee = max_fee; user_state.proofs_in_batch += 1; - println!("UNLOCKING BATCH STATE - ADD TO BATCH FUNCTION"); } /// Given a new block number listened from the blockchain, checks if the current batch is ready to be posted. @@ -733,10 +689,7 @@ impl Batcher { block_number: u64, gas_price: U256, ) -> Option> { - println!("LOCKING BATCH STATE - IS BATCH READY FUNCTION"); let mut batch_state_lock = self.batch_state.lock().await; - println!("BATCH STATE LOCKED - IS BATCH READY FUNCTION"); - let current_batch_len = batch_state_lock.batch_queue.len(); let last_uploaded_batch_block_lock = self.last_uploaded_batch_block.lock().await; @@ -744,7 +697,6 @@ impl Batcher { // once the bug in Lambdaworks merkle tree is fixed. if current_batch_len < 2 { info!("Current batch is empty or length 1. Waiting for more proofs..."); - println!("BATCH STATE UNLOCKED - IS BATCH READY FUNCTION"); return None; } @@ -755,7 +707,6 @@ impl Batcher { "Current batch not ready to be posted. Current block: {} - Last uploaded block: {}. Current batch length: {} - Minimum batch length: {}", block_number, *last_uploaded_batch_block_lock, current_batch_len, self.min_batch_len ); - println!("BATCH STATE UNLOCKED - IS BATCH READY FUNCTION"); return None; } @@ -765,24 +716,14 @@ impl Batcher { info!( "Batch is currently being posted. Waiting for the current batch to be finalized..." ); - println!("BATCH STATE UNLOCKED - IS BATCH READY FUNCTION"); return None; } // Set the batch posting flag to true *batch_posting = true; - - // println!("LOCKING USER STATES - IS BATCH READY FUNCTION"); - // let user_states = batch_state_lock.user_states; - // println!("USER STATES LOCKED - IS BATCH READY FUNCTION"); - let batch_queue_copy = batch_state_lock.batch_queue.clone(); match batch_queue::try_build_batch(batch_queue_copy, gas_price, self.max_batch_size) { Ok((resulting_batch_queue, finalized_batch)) => { - if finalized_batch.len() == 1 { - println!("BATCH STATE UNLOCKED - IS BATCH READY FUNCTION"); - panic!("FINALIZED BATCH IS LEN 1!!!!!!!!!!!!!!!"); - } batch_state_lock.batch_queue = resulting_batch_queue; let updated_user_proof_count_and_min_fee = @@ -800,21 +741,12 @@ impl Batcher { user_state.min_fee = *min_fee; } - // for (addr, (proof_count, min_fee)) in updated_user_proof_count_and_min_fee.iter() { - // let mut user_state = user_states.get(addr).unwrap().lock().await; - // println!("UPDATING OF PROOFS IN BATCH: {}", proof_count); - // user_state.proofs_in_batch = *proof_count; - // user_state.min_fee = *min_fee; - // } - - println!("BATCH STATE UNLOCKED - IS BATCH READY FUNCTION"); Some(finalized_batch) } Err(BatcherError::BatchCostTooHigh) => { // We can't post a batch since users are not willing to pay the needed fee, wait for more proofs info!("No working batch found. Waiting for more proofs..."); *batch_posting = false; - println!("BATCH STATE UNLOCKED - IS BATCH READY FUNCTION"); None } // FIXME: We should refactor this code and instead of returning None, return an error. @@ -822,7 +754,6 @@ impl Batcher { Err(e) => { error!("Unexpected error: {:?}", e); *batch_posting = false; - println!("BATCH STATE UNLOCKED - IS BATCH READY FUNCTION"); None } } @@ -909,11 +840,7 @@ impl Batcher { async fn flush_queue_and_clear_nonce_cache(&self) { warn!("Resetting state... Flushing queue and nonces"); - - println!("LOCKING BATCH STATE - FLUE QUEUE AND CLEAR NONCE FUNCTION"); let mut batch_state_lock = self.batch_state.lock().await; - println!("BATCH STATE LOCKED - FLUE QUEUE AND CLEAR NONCE FUNCTION"); - for (entry, _) in batch_state_lock.batch_queue.iter() { if let Some(ws_sink) = entry.messaging_sink.as_ref() { send_message(ws_sink.clone(), ResponseMessage::BatchReset).await; @@ -924,7 +851,6 @@ impl Batcher { batch_state_lock.batch_queue.clear(); batch_state_lock.user_states.clear(); - println!("BATCH STATE UNLOCKED - FLUE QUEUE AND CLEAR NONCE FUNCTION"); } /// Receives new block numbers, checks if conditions are met for submission and @@ -1116,7 +1042,7 @@ impl Batcher { async fn handle_nonpaying_msg( &self, ws_sink: WsMessageSink, - client_msg: ClientMessage, + client_msg: &ClientMessage, ) -> Result<(), Error> { let Some(non_paying_config) = self.non_paying_config.as_ref() else { warn!("There isn't a non-paying configuration loaded. This message will be ignored"); @@ -1124,12 +1050,6 @@ impl Batcher { return Ok(()); }; - if client_msg.verification_data.verification_data.proof.len() > self.max_proof_size { - error!("Proof is too large"); - send_message(ws_sink.clone(), ValidityResponseMessage::ProofTooLarge).await; - return Ok(()); - } - let replacement_addr = non_paying_config.replacement.address(); let replacement_user_balance = self.get_user_balance(&replacement_addr).await; if replacement_user_balance == U256::from(0) { @@ -1161,10 +1081,7 @@ impl Batcher { .await .unwrap(); - // let mut non_paying_user_state = non_paying_user_state_lock.lock().await; - // let non_paying_nonce = non_paying_user_state.nonce.unwrap(); - - info!("Non-paying nonce: {:?}", non_paying_nonce); + debug!("Non-paying nonce: {:?}", non_paying_nonce); let nonced_verification_data = NoncedVerificationData::new( client_msg.verification_data.verification_data.clone(), @@ -1180,7 +1097,6 @@ impl Batcher { ) .await; - println!("NONPAYING - ADDING ENTRY TO BATCH"); self.add_to_batch( batch_state_lock, nonced_verification_data, @@ -1189,16 +1105,8 @@ impl Batcher { non_paying_config.address, ) .await; - println!("NONPAYING - ADDED ENTRY TO BATCH SUCCESSFULLY"); - - // let mut non_paying_user_state = batch_state_lock - // .get_user_state(&replacement_addr) - // .unwrap() - // .lock() - // .await; // // Non-paying user nonce is updated - // (*non_paying_user_state).nonce = Some(non_paying_nonce + U256::one()); info!("Non-paying verification data message handled"); send_message(ws_sink, ValidityResponseMessage::Valid).await; diff --git a/batcher/aligned/send_burst_tasks.sh b/batcher/aligned/send_burst_tasks.sh index 25160bf13..f46464aa6 100755 --- a/batcher/aligned/send_burst_tasks.sh +++ b/batcher/aligned/send_burst_tasks.sh @@ -28,6 +28,6 @@ while true do # Run in backaground to be able to run onece per second, and not wait for the previous one to finish . ./batcher/aligned/generate_proof_and_send.sh $counter $burst & - sleep 1 + sleep 3 counter=$((counter + 1)) done From c230382a7482c4d1c139b8016c1c9224b12c8dde Mon Sep 17 00:00:00 2001 From: Mariano Nicolini Date: Fri, 27 Sep 2024 19:52:31 -0300 Subject: [PATCH 15/47] Refactor with Pablo --- batcher/aligned-batcher/src/lib.rs | 200 ++++++++++++++--------------- 1 file changed, 98 insertions(+), 102 deletions(-) diff --git a/batcher/aligned-batcher/src/lib.rs b/batcher/aligned-batcher/src/lib.rs index ec24f365c..7bd64fdbf 100644 --- a/batcher/aligned-batcher/src/lib.rs +++ b/batcher/aligned-batcher/src/lib.rs @@ -407,12 +407,7 @@ impl Batcher { // At this point, we will have a user state for sure, since we have inserted it // if not already present. let batch_state_lock = self.batch_state.lock().await; - let proofs_in_batch = batch_state_lock - .get_user_state(&addr) - .unwrap() - .lock() - .await - .proofs_in_batch; + let proofs_in_batch = batch_state_lock.get_user_proof_count(&addr).await.unwrap(); if !self.check_user_balance(&addr, proofs_in_batch + 1).await { send_message( ws_conn_sink.clone(), @@ -445,21 +440,18 @@ impl Batcher { }; if expected_nonce < msg_nonce { - warn!( - "Invalid nonce for address {addr}, had nonce {:?} < {:?}", - expected_nonce, msg_nonce - ); + warn!("Invalid nonce for address {addr}, had nonce {expected_nonce:?} < {msg_nonce:?}"); send_message(ws_conn_sink.clone(), ValidityResponseMessage::InvalidNonce).await; return Ok(()); } else if expected_nonce == msg_nonce { let max_fee = nonced_verification_data.max_fee; - let (user_min_fee, _) = batch_state_lock.get_user_batch_data(&addr).await.unwrap(); + let Some((user_min_fee, _)) = batch_state_lock.get_user_batch_data(&addr).await else { + send_message(ws_conn_sink.clone(), ValidityResponseMessage::InvalidNonce).await; + return Ok(()); + }; if max_fee > user_min_fee { - warn!( - "Invalid max fee for address {addr}, had fee {:?} < {:?}", - user_min_fee, max_fee - ); + warn!("Invalid max fee for address {addr}, had fee {user_min_fee:?} < {max_fee:?}"); send_message(ws_conn_sink.clone(), ValidityResponseMessage::InvalidMaxFee).await; return Ok(()); } @@ -527,8 +519,7 @@ impl Batcher { let Some(entry) = batch_state_lock.get_entry(addr, nonce) else { warn!( - "Invalid nonce for address {addr} Expected: {:?}, got: {:?}", - expected_user_nonce, nonce + "Invalid nonce for address {addr} Expected: {expected_user_nonce:?}, got: {nonce:?}" ); send_message(ws_conn_sink.clone(), ValidityResponseMessage::InvalidNonce).await; return false; @@ -549,8 +540,7 @@ impl Batcher { } info!( - "Replacing message for address {} with nonce {} and max fee {}", - addr, nonce, replacement_max_fee + "Replacing message for address {addr} with nonce {nonce} and max fee {replacement_max_fee}" ); // The replacement entry is built from the old entry and validated for then to be replaced @@ -566,7 +556,7 @@ impl Batcher { let mut old_sink = messaging_sink.write().await; if let Err(e) = old_sink.close().await { // we dont want to exit here, just log the error - warn!("Error closing sink: {:?}", e); + warn!("Error closing sink: {e:?}"); } } else { warn!( @@ -601,10 +591,12 @@ impl Batcher { BatchQueueEntryPriority::new(replacement_max_fee, nonce), ); + let Some(user_state) = batch_state_lock.get_user_state(&addr) else { + warn!("User state for address {addr:?} was not present in batcher user states, but it should be"); + return false; + }; let updated_min_fee_in_batch = batch_state_lock.get_user_min_fee_in_batch(&addr); - let mut user_state_lock = batch_state_lock.get_user_state(&addr).unwrap().lock().await; - user_state_lock.min_fee = updated_min_fee_in_batch; - + user_state.lock().await.min_fee = updated_min_fee_in_batch; true } @@ -674,12 +666,12 @@ impl Batcher { /// Given a new block number listened from the blockchain, checks if the current batch is ready to be posted. /// There are essentially two conditions to be checked: - /// * Has the current batch reached the minimum size to be posted? - /// * Has the received block number surpassed the maximum interval with respect to the last posted batch block? + /// * Has the current batch reached the minimum size to be posted? + /// * Has the received block number surpassed the maximum interval with respect to the last posted batch block? /// Then the batch will be made as big as possible given this two conditions: - /// * The serialized batch size needs to be smaller than the maximum batch size - /// * The batch submission fee is less than the lowest `max fee` included the batch, - /// * And the batch submission fee is more than the highest `max fee` not included the batch. + /// * The serialized batch size needs to be smaller than the maximum batch size + /// * The batch submission fee is less than the lowest `max fee` included the batch, + /// * And the batch submission fee is more than the highest `max fee` not included the batch. /// An extra sanity check is made to check if the batch size is 0, since it does not make sense to post /// an empty batch, even if the block interval has been reached. /// Once the batch meets the conditions for submission, the finalized batch is then passed to the @@ -722,41 +714,47 @@ impl Batcher { // Set the batch posting flag to true *batch_posting = true; let batch_queue_copy = batch_state_lock.batch_queue.clone(); - match batch_queue::try_build_batch(batch_queue_copy, gas_price, self.max_batch_size) { - Ok((resulting_batch_queue, finalized_batch)) => { - batch_state_lock.batch_queue = resulting_batch_queue; - - let updated_user_proof_count_and_min_fee = - batch_state_lock.get_user_proofs_in_batch_and_min_fee(); - - for addr in batch_state_lock.user_states.keys() { - let (proof_count, min_fee) = updated_user_proof_count_and_min_fee - .get(addr) - .unwrap_or(&(0, U256::MAX)); - - let mut user_state = - batch_state_lock.get_user_state(addr).unwrap().lock().await; + let (resulting_batch_queue, finalized_batch) = + batch_queue::try_build_batch(batch_queue_copy, gas_price, self.max_batch_size) + .inspect_err(|e| { + *batch_posting = false; + match e { + // We can't post a batch since users are not willing to pay the needed fee, wait for more proofs + BatcherError::BatchCostTooHigh => { + info!("No working batch found. Waiting for more proofs") + } + // FIXME: We should refactor this code and instead of returning None, return an error. + // See issue https://github.com/yetanotherco/aligned_layer/issues/1046. + e => error!("Unexpected error: {:?}", e), + } + }) + .ok()?; + + batch_state_lock.batch_queue = resulting_batch_queue; + let updated_user_proof_count_and_min_fee = + batch_state_lock.get_user_proofs_in_batch_and_min_fee(); + + for addr in batch_state_lock.user_states.keys() { + let (proof_count, min_fee) = updated_user_proof_count_and_min_fee + .get(addr) + .unwrap_or(&(0, U256::MAX)); + + // FIXME: The case where a `user_state` is not found in the `user_states` map should not really + // happen here, but doing this check so that we don't unwrap. + // Once https://github.com/yetanotherco/aligned_layer/issues/1046 is done we could return a more + // informative error. + let Some(user_state) = batch_state_lock.get_user_state(addr) else { + return None; + }; - user_state.proofs_in_batch = *proof_count; - user_state.min_fee = *min_fee; - } + // Now we update the user states - Some(finalized_batch) - } - Err(BatcherError::BatchCostTooHigh) => { - // We can't post a batch since users are not willing to pay the needed fee, wait for more proofs - info!("No working batch found. Waiting for more proofs..."); - *batch_posting = false; - None - } - // FIXME: We should refactor this code and instead of returning None, return an error. - // See issue https://github.com/yetanotherco/aligned_layer/issues/1046. - Err(e) => { - error!("Unexpected error: {:?}", e); - *batch_posting = false; - None - } + let mut user_state_lock = user_state.lock().await; + user_state_lock.proofs_in_batch = *proof_count; + user_state_lock.min_fee = *min_fee; } + + Some(finalized_batch) } /// Takes the finalized batch as input and builds the merkle tree, posts verification data batch @@ -1113,54 +1111,52 @@ impl Batcher { Ok(()) } - async fn get_user_balance(&self, addr: &Address) -> U256 { - match self.payment_service.user_balances(*addr).call().await { - Ok(val) => val, - Err(_) => match self - .payment_service_fallback - .user_balances(*addr) - .call() - .await - { - Ok(balance) => balance, - Err(_) => { - warn!("Failed to get balance for address {:?}", addr); - U256::zero() - } - }, - } + /// TODO: ADD DOCUMENTATION ON WHY IT IS AN OPTION + async fn get_user_balance(&self, addr: &Address) -> Option { + if let Ok(balance) = self.payment_service.user_balances(*addr).call().await { + return Some(balance); + }; + + self.payment_service_fallback + .user_balances(*addr) + .call() + .await + .inspect_err(|e| warn!("Failed to get balance for address {:?}", addr)) + .ok() } async fn user_balance_is_unlocked(&self, addr: &Address) -> bool { - let unlock_block = match self.payment_service.user_unlock_block(*addr).call().await { - Ok(val) => val, - Err(_) => match self - .payment_service_fallback - .user_unlock_block(*addr) - .call() - .await - { - Ok(unlock_block) => unlock_block, - Err(_) => { - warn!("Failed to get unlock block for address {:?}", addr); - U256::zero() - } - }, - }; - - unlock_block != U256::zero() + // try with the main + if let Ok(unlock_block) = self.payment_service.user_unlock_block(*addr).call().await { + return unlock_block != U256::zero(); + } + // falback + if let Ok(unlock_block) = self + .payment_service_fallback + .user_unlock_block(*addr) + .call() + .await + { + return unlock_block != U256::zero(); + } + warn!("Could not get user locking state"); + false } async fn get_gas_price(&self) -> Option { - match self.eth_ws_provider.get_gas_price().await { - Ok(gas_price) => Some(gas_price), // this is the block's max priority gas price, not the base fee - Err(_) => match self.eth_ws_provider_fallback.get_gas_price().await { - Ok(gas_price) => Some(gas_price), - Err(_) => { - warn!("Failed to get gas price"); - None - } - }, + if let Ok(gas_price) = self + .eth_ws_provider + .get_gas_price() + .await + .inspect_err(|e| warn!("Failed to get gas price. Trying with fallback: {e:?}")) + { + return Some(gas_price); } + + self.eth_ws_provider_fallback + .get_gas_price() + .await + .inspect_err(|e| warn!("Failed to get gas price: {e:?}")) + .ok() } } From 0c8d6ff52556a9b19a4454b0434dfdc1297cefa8 Mon Sep 17 00:00:00 2001 From: Mariano Nicolini Date: Sat, 28 Sep 2024 13:09:12 -0300 Subject: [PATCH 16/47] Some more refactors and code polishing --- Makefile | 4 +- batcher/aligned-batcher/src/lib.rs | 159 +++++++++--------- .../aligned-batcher/src/types/batch_state.rs | 9 - 3 files changed, 86 insertions(+), 86 deletions(-) diff --git a/Makefile b/Makefile index 04028906c..052857845 100644 --- a/Makefile +++ b/Makefile @@ -271,7 +271,7 @@ batcher_send_risc0_burst: --proof ../../scripts/test_files/risc_zero/fibonacci_proof_generator/risc_zero_fibonacci.proof \ --vm_program ../../scripts/test_files/risc_zero/fibonacci_proof_generator/fibonacci_id.bin \ --public_input ../../scripts/test_files/risc_zero/fibonacci_proof_generator/risc_zero_fibonacci.pub \ - --repetitions 300 \ + --repetitions 500 \ --proof_generator_addr 0x66f9664f97F2b50F62D13eA064982f936dE76657 \ --rpc_url $(RPC_URL) \ --payment_service_addr $(BATCHER_PAYMENTS_CONTRACT_ADDRESS) @@ -362,7 +362,7 @@ batcher_send_halo2_ipa_task_burst_5: batcher/target/release/aligned --proof ../../scripts/test_files/halo2_ipa/proof.bin \ --public_input ../../scripts/test_files/halo2_ipa/pub_input.bin \ --vk ../../scripts/test_files/halo2_ipa/params.bin \ - --repetitions 300 \ + --repetitions 500 \ --private_key 0xdbda1821b80551c9d65939329250298aa3472ba22feea921c0cf5d620ea67b97 \ --rpc_url $(RPC_URL) \ --payment_service_addr $(BATCHER_PAYMENTS_CONTRACT_ADDRESS) diff --git a/batcher/aligned-batcher/src/lib.rs b/batcher/aligned-batcher/src/lib.rs index 7bd64fdbf..64ae2b495 100644 --- a/batcher/aligned-batcher/src/lib.rs +++ b/batcher/aligned-batcher/src/lib.rs @@ -7,7 +7,7 @@ use connection::{send_message, WsMessageSink}; use dotenvy::dotenv; use ethers::contract::ContractError; use ethers::signers::Signer; -use types::batch_state::{self, BatchState}; +use types::batch_state::BatchState; use types::user_state::UserState; use std::collections::HashMap; @@ -166,6 +166,7 @@ impl Batcher { .expect("Failed to get fallback Batcher Payment Service contract"); let mut user_states = HashMap::new(); + let mut batch_state = BatchState::new(); let non_paying_config = if let Some(non_paying_config) = config.batcher.non_paying { warn!("Non-paying address configuration detected. Will replace non-paying address {} with configured address.", non_paying_config.address); @@ -183,13 +184,12 @@ impl Batcher { non_paying_user_state, ); + batch_state = BatchState::new_with_user_states(user_states); Some(non_paying_config) } else { None }; - let batch_state = BatchState::new_with_user_states(user_states); - Self { s3_client, s3_bucket_name, @@ -408,7 +408,7 @@ impl Batcher { // if not already present. let batch_state_lock = self.batch_state.lock().await; let proofs_in_batch = batch_state_lock.get_user_proof_count(&addr).await.unwrap(); - if !self.check_user_balance(&addr, proofs_in_batch + 1).await { + if !self.check_min_balance(&addr, proofs_in_batch + 1).await { send_message( ws_conn_sink.clone(), ValidityResponseMessage::InsufficientBalance(addr), @@ -425,11 +425,9 @@ impl Batcher { Ok(ethereum_user_nonce) => ethereum_user_nonce, Err(e) => { error!( - "Failed to get user nonce from Ethereum for address {:?}. Error: {:?}", - addr, e + "Failed to get user nonce from Ethereum for address {addr:?}. Error: {e:?}" ); send_message(ws_conn_sink.clone(), ValidityResponseMessage::InvalidNonce).await; - return Ok(()); } }; @@ -443,20 +441,13 @@ impl Batcher { warn!("Invalid nonce for address {addr}, had nonce {expected_nonce:?} < {msg_nonce:?}"); send_message(ws_conn_sink.clone(), ValidityResponseMessage::InvalidNonce).await; return Ok(()); - } else if expected_nonce == msg_nonce { - let max_fee = nonced_verification_data.max_fee; - - let Some((user_min_fee, _)) = batch_state_lock.get_user_batch_data(&addr).await else { - send_message(ws_conn_sink.clone(), ValidityResponseMessage::InvalidNonce).await; - return Ok(()); - }; - if max_fee > user_min_fee { - warn!("Invalid max fee for address {addr}, had fee {user_min_fee:?} < {max_fee:?}"); - send_message(ws_conn_sink.clone(), ValidityResponseMessage::InvalidMaxFee).await; - return Ok(()); - } + } - self.add_to_batch( + // In this case, the message might be a replacement one. If it is valid, + // we replace the old entry with the new from the replacement message. + if expected_nonce > msg_nonce { + info!("Possible replacement message received: Expected nonce {expected_nonce:?} - message nonce: {msg_nonce:?}"); + self.handle_replacement_message( batch_state_lock, nonced_verification_data, ws_conn_sink.clone(), @@ -464,24 +455,31 @@ impl Batcher { addr, ) .await; - } else { - // In this case, the message might be a replacement one. If it is valid, - // we replace the old entry with the new from the replacement message. - if !self - .handle_replacement_message( - batch_state_lock, - nonced_verification_data, - ws_conn_sink.clone(), - client_msg.signature, - addr, - expected_nonce, - ) - .await - { - // message should not be added to batch - return Ok(()); - } + + return Ok(()); + } + + let msg_max_fee = nonced_verification_data.max_fee; + let Some(user_min_fee) = batch_state_lock.get_user_min_fee(&addr).await else { + send_message(ws_conn_sink.clone(), ValidityResponseMessage::InvalidNonce).await; + return Ok(()); + }; + + if msg_max_fee > user_min_fee { + warn!("Invalid max fee for address {addr}, had fee {user_min_fee:?} < {msg_max_fee:?}"); + send_message(ws_conn_sink.clone(), ValidityResponseMessage::InvalidMaxFee).await; + return Ok(()); } + + self.add_to_batch( + batch_state_lock, + nonced_verification_data, + ws_conn_sink.clone(), + client_msg.signature, + addr, + ) + .await; + info!("Verification data message handled"); send_message(ws_conn_sink, ValidityResponseMessage::Valid).await; Ok(()) @@ -489,8 +487,10 @@ impl Batcher { // Checks user has sufficient balance // If user has sufficient balance, increments the user's proof count in the batch - async fn check_user_balance(&self, addr: &Address, user_proofs_in_batch: usize) -> bool { - let user_balance = self.get_user_balance(addr).await; + async fn check_min_balance(&self, addr: &Address, user_proofs_in_batch: usize) -> bool { + let Some(user_balance) = self.get_user_balance(addr).await else { + return false; + }; let min_balance = U256::from(user_proofs_in_batch) * U256::from(MIN_FEE_PER_PROOF); if user_balance < min_balance { return false; @@ -512,36 +512,28 @@ impl Batcher { ws_conn_sink: WsMessageSink, signature: Signature, addr: Address, - expected_user_nonce: U256, - ) -> bool { + ) { let replacement_max_fee = nonced_verification_data.max_fee; let nonce = nonced_verification_data.nonce; - let Some(entry) = batch_state_lock.get_entry(addr, nonce) else { - warn!( - "Invalid nonce for address {addr} Expected: {expected_user_nonce:?}, got: {nonce:?}" - ); + warn!("Invalid nonce for address {addr}. Queue entry with nonce {nonce} not found"); send_message(ws_conn_sink.clone(), ValidityResponseMessage::InvalidNonce).await; - return false; + return; }; - if entry.nonced_verification_data.max_fee > replacement_max_fee { - warn!( - "Invalid replacement message for address {addr}, had fee {:?} < {:?}", - entry.nonced_verification_data.max_fee, replacement_max_fee - ); + let original_max_fee = entry.nonced_verification_data.max_fee; + if original_max_fee > replacement_max_fee { + warn!("Invalid replacement message for address {addr}, had fee {original_max_fee:?} < {replacement_max_fee:?}"); send_message( ws_conn_sink.clone(), ValidityResponseMessage::InvalidReplacementMessage, ) .await; - return false; + return; } - info!( - "Replacing message for address {addr} with nonce {nonce} and max fee {replacement_max_fee}" - ); + info!("Replacing message for address {addr} with nonce {nonce} and max fee {replacement_max_fee}"); // The replacement entry is built from the old entry and validated for then to be replaced let mut replacement_entry = entry.clone(); @@ -573,7 +565,7 @@ impl Batcher { ValidityResponseMessage::InvalidReplacementMessage, ) .await; - return false; + return; } info!( @@ -593,11 +585,10 @@ impl Batcher { let Some(user_state) = batch_state_lock.get_user_state(&addr) else { warn!("User state for address {addr:?} was not present in batcher user states, but it should be"); - return false; + return; }; let updated_min_fee_in_batch = batch_state_lock.get_user_min_fee_in_batch(&addr); user_state.lock().await.min_fee = updated_min_fee_in_batch; - true } async fn get_user_nonce_from_ethereum( @@ -847,8 +838,17 @@ impl Batcher { } } + let Some(nonpaying_replacement_addr) = self.get_nonpaying_replacement_addr() else { + batch_state_lock.batch_queue.clear(); + batch_state_lock.user_states.clear(); + return; + }; + batch_state_lock.batch_queue.clear(); - batch_state_lock.user_states.clear(); + warn!("Clearing user states, preserving nonpaying user {nonpaying_replacement_addr}"); + batch_state_lock + .user_states + .retain(|&addr, _| addr == nonpaying_replacement_addr); } /// Receives new block numbers, checks if conditions are met for submission and @@ -1036,6 +1036,13 @@ impl Batcher { .is_some_and(|non_paying_config| non_paying_config.address == *addr) } + fn get_nonpaying_replacement_addr(&self) -> Option
{ + let Some(non_paying_conf) = self.non_paying_config.as_ref() else { + return None; + }; + Some(non_paying_conf.replacement.address()) + } + /// Only relevant for testing and for users to easily use Aligned in testnet. async fn handle_nonpaying_msg( &self, @@ -1049,9 +1056,17 @@ impl Batcher { }; let replacement_addr = non_paying_config.replacement.address(); - let replacement_user_balance = self.get_user_balance(&replacement_addr).await; + let Some(replacement_user_balance) = self.get_user_balance(&replacement_addr).await else { + error!("Could not get balance for non-paying address {replacement_addr:?}"); + send_message( + ws_sink.clone(), + ValidityResponseMessage::InsufficientBalance(replacement_addr), + ) + .await; + return Ok(()); + }; if replacement_user_balance == U256::from(0) { - error!("Insufficient funds for address {:?}", replacement_addr); + error!("Insufficient funds for non-paying address {replacement_addr:?}"); send_message( ws_sink.clone(), ValidityResponseMessage::InsufficientBalance(replacement_addr), @@ -1066,18 +1081,16 @@ impl Batcher { { error!("Invalid proof detected. Verification failed."); send_message(ws_sink.clone(), ValidityResponseMessage::InvalidProof).await; - return Ok(()); // Send error message to the client and return + return Ok(()); } - // let user_states = self.user_states.read().await; let batch_state_lock = self.batch_state.lock().await; - - // Safe to call `unwrap()` here since at this point we have a non-paying configuration loaded - // for sure and the non-paying address nonce cached in the user states. - let non_paying_nonce = batch_state_lock - .get_user_nonce(&replacement_addr) - .await - .unwrap(); + let Some(non_paying_nonce) = batch_state_lock.get_user_nonce(&replacement_addr).await + else { + error!("Nonce for non-paying address {replacement_addr:?} not found in cache."); + send_message(ws_sink.clone(), ValidityResponseMessage::InvalidProof).await; + return Ok(()); + }; debug!("Non-paying nonce: {:?}", non_paying_nonce); @@ -1104,10 +1117,8 @@ impl Batcher { ) .await; - // // Non-paying user nonce is updated info!("Non-paying verification data message handled"); send_message(ws_sink, ValidityResponseMessage::Valid).await; - Ok(()) } @@ -1121,16 +1132,14 @@ impl Batcher { .user_balances(*addr) .call() .await - .inspect_err(|e| warn!("Failed to get balance for address {:?}", addr)) + .inspect_err(|_| warn!("Failed to get balance for address {:?}", addr)) .ok() } async fn user_balance_is_unlocked(&self, addr: &Address) -> bool { - // try with the main if let Ok(unlock_block) = self.payment_service.user_unlock_block(*addr).call().await { return unlock_block != U256::zero(); } - // falback if let Ok(unlock_block) = self .payment_service_fallback .user_unlock_block(*addr) diff --git a/batcher/aligned-batcher/src/types/batch_state.rs b/batcher/aligned-batcher/src/types/batch_state.rs index 2780b250f..ffe9e8054 100644 --- a/batcher/aligned-batcher/src/types/batch_state.rs +++ b/batcher/aligned-batcher/src/types/batch_state.rs @@ -60,15 +60,6 @@ impl BatchState { Some(user_state.lock().await.proofs_in_batch) } - pub(crate) async fn get_user_batch_data(&self, addr: &Address) -> Option<(U256, usize)> { - let Some(user_state) = self.get_user_state(addr) else { - return None; - }; - - let user_state_lock = user_state.lock().await; - Some((user_state_lock.min_fee, user_state_lock.proofs_in_batch)) - } - /// Checks if the entry is valid /// An entry is valid if there is no entry with the same sender, lower nonce and a lower fee pub(crate) fn replacement_entry_is_valid( From 43790343fe6303d3eb878df9d32017beb1a779ab Mon Sep 17 00:00:00 2001 From: Mariano Nicolini Date: Sat, 28 Sep 2024 14:28:49 -0300 Subject: [PATCH 17/47] Polish code, handle errors better --- batcher/aligned-batcher/src/lib.rs | 144 ++++++++++-------- .../aligned-batcher/src/types/batch_state.rs | 12 +- batcher/aligned-batcher/src/types/errors.rs | 9 +- .../src/communication/messaging.rs | 4 + batcher/aligned-sdk/src/core/errors.rs | 2 + batcher/aligned-sdk/src/core/types.rs | 1 + 6 files changed, 95 insertions(+), 77 deletions(-) diff --git a/batcher/aligned-batcher/src/lib.rs b/batcher/aligned-batcher/src/lib.rs index 64ae2b495..7bbd8f8f7 100644 --- a/batcher/aligned-batcher/src/lib.rs +++ b/batcher/aligned-batcher/src/lib.rs @@ -320,11 +320,12 @@ impl Batcher { return Ok(()); } }; + let msg_nonce = client_msg.verification_data.nonce; + info!("Received message with nonce: {msg_nonce:?}",); - info!( - "Received message with nonce: {}", - client_msg.verification_data.nonce - ); + // * ---------------------------------------------------* + // * Perform validations over the message * + // * ---------------------------------------------------* if client_msg.verification_data.chain_id != self.chain_id { warn!( @@ -360,6 +361,15 @@ impl Batcher { return Ok(()); } + // When pre-verification is enabled, batcher will verify proofs for faster feedback with clients + if self.pre_verification_is_enabled + && !zk_utils::verify(&nonced_verification_data.verification_data).await + { + error!("Invalid proof detected. Verification failed."); + send_message(ws_conn_sink.clone(), ValidityResponseMessage::InvalidProof).await; + return Ok(()); + } + if self.is_nonpaying(&addr) { return self .handle_nonpaying_msg(ws_conn_sink.clone(), &client_msg) @@ -375,17 +385,7 @@ impl Batcher { return Ok(()); } - // When pre-verification is enabled, batcher will verify proofs for faster feedback with clients - if self.pre_verification_is_enabled - && !zk_utils::verify(&nonced_verification_data.verification_data).await - { - error!("Invalid proof detected. Verification failed."); - send_message(ws_conn_sink.clone(), ValidityResponseMessage::InvalidProof).await; - return Ok(()); // Send error message to the client and return - } - // Nonce and max fee verification - let msg_nonce = nonced_verification_data.nonce; let max_fee = nonced_verification_data.max_fee; if max_fee < U256::from(MIN_FEE_PER_PROOF) { error!("The max fee signed in the message is less than the accepted minimum fee to be included in the batch."); @@ -396,16 +396,18 @@ impl Batcher { // Check that we had a user state entry for this user and insert it if not. { let mut batch_state_lock = self.batch_state.lock().await; - if !batch_state_lock.user_states.contains_key(&addr) { - let new_user_state = Mutex::new(UserState::new()); - batch_state_lock - .user_states - .insert(addr.clone(), new_user_state); - } + batch_state_lock + .user_states + .entry(addr) + .or_insert_with(|| Mutex::new(UserState::new())); } + // * ---------------------------------------------------* + // * Perform validations over user state * + // * ---------------------------------------------------* + // At this point, we will have a user state for sure, since we have inserted it - // if not already present. + // if not already present. It should be safe to call `unwrap()` here. let batch_state_lock = self.batch_state.lock().await; let proofs_in_batch = batch_state_lock.get_user_proof_count(&addr).await.unwrap(); if !self.check_min_balance(&addr, proofs_in_batch + 1).await { @@ -459,6 +461,10 @@ impl Batcher { return Ok(()); } + // * ---------------------------------------------------------------------* + // * Add message data into the queue and update user state * + // * ---------------------------------------------------------------------* + let msg_max_fee = nonced_verification_data.max_fee; let Some(user_min_fee) = batch_state_lock.get_user_min_fee(&addr).await else { send_message(ws_conn_sink.clone(), ValidityResponseMessage::InvalidNonce).await; @@ -471,14 +477,20 @@ impl Batcher { return Ok(()); } - self.add_to_batch( - batch_state_lock, - nonced_verification_data, - ws_conn_sink.clone(), - client_msg.signature, - addr, - ) - .await; + if let Err(e) = self + .add_to_batch( + batch_state_lock, + nonced_verification_data, + ws_conn_sink.clone(), + client_msg.signature, + addr, + ) + .await + { + error!("Error while adding entry to batch: {e:?}"); + send_message(ws_conn_sink, ValidityResponseMessage::AddToBatchError).await; + return Ok(()); + }; info!("Verification data message handled"); send_message(ws_conn_sink, ValidityResponseMessage::Valid).await; @@ -609,7 +621,7 @@ impl Batcher { ws_conn_sink: WsMessageSink, proof_submitter_sig: Signature, proof_submitter_addr: Address, - ) { + ) -> Result<(), BatcherError> { info!("Calculating verification data commitments..."); let verification_data_comm = verification_data.clone().into(); info!("Adding verification data to batch..."); @@ -634,25 +646,28 @@ impl Batcher { let mut proof_submitter_addr = proof_submitter_addr; + // If the proof submitter is the nonpaying one, we should update the state + // of the replacement address. proof_submitter_addr = if self.is_nonpaying(&proof_submitter_addr) { - self.non_paying_config - .as_ref() - .unwrap() - .replacement - .address() + self.get_nonpaying_replacement_addr() + .unwrap_or(proof_submitter_addr) } else { proof_submitter_addr }; - let mut user_state = batch_state_lock - .get_user_state(&proof_submitter_addr) - .unwrap() - .lock() - .await; + let Some(user_state) = batch_state_lock.get_user_state(&proof_submitter_addr) else { + error!("User state of address {proof_submitter_addr} was not found when trying to update user state. This user state should have been present"); + return Err(BatcherError::AddressNotFoundInUserStates( + proof_submitter_addr, + )); + }; - user_state.nonce = Some(nonce + U256::one()); - user_state.min_fee = max_fee; - user_state.proofs_in_batch += 1; + let mut user_state_lock = user_state.lock().await; + // User state is updated + user_state_lock.nonce = Some(nonce + U256::one()); + user_state_lock.min_fee = max_fee; + user_state_lock.proofs_in_batch += 1; + Ok(()) } /// Given a new block number listened from the blockchain, checks if the current batch is ready to be posted. @@ -734,12 +749,9 @@ impl Batcher { // happen here, but doing this check so that we don't unwrap. // Once https://github.com/yetanotherco/aligned_layer/issues/1046 is done we could return a more // informative error. - let Some(user_state) = batch_state_lock.get_user_state(addr) else { - return None; - }; + let user_state = batch_state_lock.get_user_state(addr)?; // Now we update the user states - let mut user_state_lock = user_state.lock().await; user_state_lock.proofs_in_batch = *proof_count; user_state_lock.min_fee = *min_fee; @@ -1037,9 +1049,7 @@ impl Batcher { } fn get_nonpaying_replacement_addr(&self) -> Option
{ - let Some(non_paying_conf) = self.non_paying_config.as_ref() else { - return None; - }; + let non_paying_conf = self.non_paying_config.as_ref()?; Some(non_paying_conf.replacement.address()) } @@ -1065,6 +1075,7 @@ impl Batcher { .await; return Ok(()); }; + if replacement_user_balance == U256::from(0) { error!("Insufficient funds for non-paying address {replacement_addr:?}"); send_message( @@ -1075,15 +1086,6 @@ impl Batcher { return Ok(()); } - // When pre-verification is enabled, batcher will verify proofs for faster feedback with clients - if self.pre_verification_is_enabled - && !zk_utils::verify(&client_msg.verification_data.verification_data).await - { - error!("Invalid proof detected. Verification failed."); - send_message(ws_sink.clone(), ValidityResponseMessage::InvalidProof).await; - return Ok(()); - } - let batch_state_lock = self.batch_state.lock().await; let Some(non_paying_nonce) = batch_state_lock.get_user_nonce(&replacement_addr).await else { @@ -1108,14 +1110,22 @@ impl Batcher { ) .await; - self.add_to_batch( - batch_state_lock, - nonced_verification_data, - ws_sink.clone(), - client_msg.signature, - non_paying_config.address, - ) - .await; + let signature = client_msg.signature; + let nonpaying_addr = non_paying_config.address; + if let Err(e) = self + .add_to_batch( + batch_state_lock, + nonced_verification_data, + ws_sink.clone(), + signature, + nonpaying_addr, + ) + .await + { + info!("Error while adding nonpaying address entry to batch: {e:?}"); + send_message(ws_sink, ValidityResponseMessage::AddToBatchError).await; + return Ok(()); + }; info!("Non-paying verification data message handled"); send_message(ws_sink, ValidityResponseMessage::Valid).await; diff --git a/batcher/aligned-batcher/src/types/batch_state.rs b/batcher/aligned-batcher/src/types/batch_state.rs index ffe9e8054..528711283 100644 --- a/batcher/aligned-batcher/src/types/batch_state.rs +++ b/batcher/aligned-batcher/src/types/batch_state.rs @@ -40,23 +40,17 @@ impl BatchState { } pub(crate) async fn get_user_nonce(&self, addr: &Address) -> Option { - let Some(user_state) = self.get_user_state(addr) else { - return None; - }; + let user_state = self.get_user_state(addr)?; user_state.lock().await.nonce } pub(crate) async fn get_user_min_fee(&self, addr: &Address) -> Option { - let Some(user_state) = self.get_user_state(addr) else { - return None; - }; + let user_state = self.get_user_state(addr)?; Some(user_state.lock().await.min_fee) } pub(crate) async fn get_user_proof_count(&self, addr: &Address) -> Option { - let Some(user_state) = self.get_user_state(addr) else { - return None; - }; + let user_state = self.get_user_state(addr)?; Some(user_state.lock().await.proofs_in_batch) } diff --git a/batcher/aligned-batcher/src/types/errors.rs b/batcher/aligned-batcher/src/types/errors.rs index 4cf9ecd2a..c62de3576 100644 --- a/batcher/aligned-batcher/src/types/errors.rs +++ b/batcher/aligned-batcher/src/types/errors.rs @@ -1,6 +1,6 @@ use std::fmt; -use ethers::types::SignatureError; +use ethers::types::{Address, SignatureError}; use tokio_tungstenite::tungstenite; pub enum BatcherError { @@ -18,6 +18,7 @@ pub enum BatcherError { GasPriceError, BatchCostTooHigh, WsSinkEmpty, + AddressNotFoundInUserStates(Address), } impl From for BatcherError { @@ -83,6 +84,12 @@ impl fmt::Debug for BatcherError { "Websocket sink was found empty. This should only happen in tests" ) } + BatcherError::AddressNotFoundInUserStates(addr) => { + write!( + f, + "User with address {addr:?} was not found in Batcher user states cache" + ) + } } } } diff --git a/batcher/aligned-sdk/src/communication/messaging.rs b/batcher/aligned-sdk/src/communication/messaging.rs index 336cabe2f..0ae0f6be7 100644 --- a/batcher/aligned-sdk/src/communication/messaging.rs +++ b/batcher/aligned-sdk/src/communication/messaging.rs @@ -114,6 +114,10 @@ pub async fn send_messages( error!("Invalid replacement message!"); return Err(SubmitError::InvalidReplacementMessage); } + ValidityResponseMessage::AddToBatchError => { + error!("Error while pushing the entry to queue"); + return Err(SubmitError::AddToBatchError); + } }; sent_verification_data.push(verification_data.clone()); diff --git a/batcher/aligned-sdk/src/core/errors.rs b/batcher/aligned-sdk/src/core/errors.rs index db2b89969..01a817d3a 100644 --- a/batcher/aligned-sdk/src/core/errors.rs +++ b/batcher/aligned-sdk/src/core/errors.rs @@ -81,6 +81,7 @@ pub enum SubmitError { InvalidReplacementMessage, InsufficientBalance, BatchSubmissionFailed(String), + AddToBatchError, GenericError(String), } @@ -179,6 +180,7 @@ impl fmt::Display for SubmitError { SubmitError::InvalidReplacementMessage => write!(f, "Invalid replacement message"), SubmitError::InsufficientBalance => write!(f, "Insufficient balance"), SubmitError::ProofQueueFlushed => write!(f, "Batch reset"), + SubmitError::AddToBatchError => write!(f, "Error while adding entry to batch"), } } } diff --git a/batcher/aligned-sdk/src/core/types.rs b/batcher/aligned-sdk/src/core/types.rs index f2b961a17..bcfb551e9 100644 --- a/batcher/aligned-sdk/src/core/types.rs +++ b/batcher/aligned-sdk/src/core/types.rs @@ -315,6 +315,7 @@ pub enum ValidityResponseMessage { InvalidProof, InvalidMaxFee, InvalidReplacementMessage, + AddToBatchError, ProofTooLarge, InsufficientBalance(Address), } From f2f363d5f72aa51071bb773d1f48b4990272d0a6 Mon Sep 17 00:00:00 2001 From: Mariano Nicolini Date: Sat, 28 Sep 2024 15:37:48 -0300 Subject: [PATCH 18/47] Fix nonpaying address bug when user states are flushed --- batcher/aligned-batcher/src/lib.rs | 18 +++++++++++++++--- 1 file changed, 15 insertions(+), 3 deletions(-) diff --git a/batcher/aligned-batcher/src/lib.rs b/batcher/aligned-batcher/src/lib.rs index 7bbd8f8f7..721780316 100644 --- a/batcher/aligned-batcher/src/lib.rs +++ b/batcher/aligned-batcher/src/lib.rs @@ -662,8 +662,8 @@ impl Batcher { )); }; - let mut user_state_lock = user_state.lock().await; // User state is updated + let mut user_state_lock = user_state.lock().await; user_state_lock.nonce = Some(nonce + U256::one()); user_state_lock.min_fee = max_fee; user_state_lock.proofs_in_batch += 1; @@ -856,11 +856,23 @@ impl Batcher { return; }; + // If there is a nonpaying address configured, then fetch the correct nonce from Ethereum + // so that it is already loaded + + let Ok(nonpaying_replacement_addr_nonce) = self + .get_user_nonce_from_ethereum(nonpaying_replacement_addr) + .await + else { + batch_state_lock.batch_queue.clear(); + batch_state_lock.user_states.clear(); + return; + }; batch_state_lock.batch_queue.clear(); - warn!("Clearing user states, preserving nonpaying user {nonpaying_replacement_addr}"); + batch_state_lock.user_states.clear(); + let nonpaying_user_state = UserState::new_non_paying(nonpaying_replacement_addr_nonce); batch_state_lock .user_states - .retain(|&addr, _| addr == nonpaying_replacement_addr); + .insert(nonpaying_replacement_addr, Mutex::new(nonpaying_user_state)); } /// Receives new block numbers, checks if conditions are met for submission and From bc12587879cf7e2fddff7d054b6b1a90c05e0c1f Mon Sep 17 00:00:00 2001 From: Mariano Nicolini Date: Mon, 30 Sep 2024 10:42:09 -0300 Subject: [PATCH 19/47] Remove commented code --- batcher/aligned-batcher/src/lib.rs | 2 -- 1 file changed, 2 deletions(-) diff --git a/batcher/aligned-batcher/src/lib.rs b/batcher/aligned-batcher/src/lib.rs index 721780316..63e325755 100644 --- a/batcher/aligned-batcher/src/lib.rs +++ b/batcher/aligned-batcher/src/lib.rs @@ -80,7 +80,6 @@ pub struct Batcher { pre_verification_is_enabled: bool, non_paying_config: Option, posting_batch: Mutex, - // user_states: RwLock>>, } impl Batcher { @@ -208,7 +207,6 @@ impl Batcher { non_paying_config, posting_batch: Mutex::new(false), batch_state: Mutex::new(batch_state), - // user_states: RwLock::new(user_states), } } From 331def9b23a91d50f851883d79c39000e76086a8 Mon Sep 17 00:00:00 2001 From: Mariano Nicolini Date: Mon, 30 Sep 2024 12:53:00 -0300 Subject: [PATCH 20/47] removed mutex from user state --- batcher/aligned-batcher/src/lib.rs | 82 ++++++++++++++----- .../aligned-batcher/src/types/batch_state.rs | 69 ++++++++++++++-- 2 files changed, 124 insertions(+), 27 deletions(-) diff --git a/batcher/aligned-batcher/src/lib.rs b/batcher/aligned-batcher/src/lib.rs index 63e325755..eefea5b5d 100644 --- a/batcher/aligned-batcher/src/lib.rs +++ b/batcher/aligned-batcher/src/lib.rs @@ -177,7 +177,7 @@ impl Batcher { .await .expect("Could not get non-paying nonce from Ethereum"); - let non_paying_user_state = Mutex::new(UserState::new_non_paying(nonpaying_nonce)); + let non_paying_user_state = UserState::new_non_paying(nonpaying_nonce); user_states.insert( non_paying_config.replacement.address(), non_paying_user_state, @@ -397,7 +397,7 @@ impl Batcher { batch_state_lock .user_states .entry(addr) - .or_insert_with(|| Mutex::new(UserState::new())); + .or_insert_with(|| UserState::new()); } // * ---------------------------------------------------* @@ -406,7 +406,7 @@ impl Batcher { // At this point, we will have a user state for sure, since we have inserted it // if not already present. It should be safe to call `unwrap()` here. - let batch_state_lock = self.batch_state.lock().await; + let mut batch_state_lock = self.batch_state.lock().await; let proofs_in_batch = batch_state_lock.get_user_proof_count(&addr).await.unwrap(); if !self.check_min_balance(&addr, proofs_in_batch + 1).await { send_message( @@ -432,8 +432,7 @@ impl Batcher { } }; - let mut user_state = batch_state_lock.get_user_state(&addr).unwrap().lock().await; - user_state.nonce = Some(ethereum_user_nonce); + batch_state_lock.update_user_nonce(&addr, ethereum_user_nonce); ethereum_user_nonce }; @@ -593,12 +592,15 @@ impl Batcher { BatchQueueEntryPriority::new(replacement_max_fee, nonce), ); - let Some(user_state) = batch_state_lock.get_user_state(&addr) else { + // let Some(user_state) = batch_state_lock.get_user_state(&addr) else { + // warn!("User state for address {addr:?} was not present in batcher user states, but it should be"); + // return; + // }; + let updated_min_fee_in_batch = batch_state_lock.get_user_min_fee_in_batch(&addr); + let Some(_) = batch_state_lock.update_user_min_fee(&addr, updated_min_fee_in_batch) else { warn!("User state for address {addr:?} was not present in batcher user states, but it should be"); return; }; - let updated_min_fee_in_batch = batch_state_lock.get_user_min_fee_in_batch(&addr); - user_state.lock().await.min_fee = updated_min_fee_in_batch; } async fn get_user_nonce_from_ethereum( @@ -653,7 +655,17 @@ impl Batcher { proof_submitter_addr }; - let Some(user_state) = batch_state_lock.get_user_state(&proof_submitter_addr) else { + // let Some(user_state) = batch_state_lock.get_user_state(&proof_submitter_addr) else { + // error!("User state of address {proof_submitter_addr} was not found when trying to update user state. This user state should have been present"); + // return Err(BatcherError::AddressNotFoundInUserStates( + // proof_submitter_addr, + // )); + // }; + + let Some(user_proof_count) = batch_state_lock + .get_user_proof_count(&proof_submitter_addr) + .await + else { error!("User state of address {proof_submitter_addr} was not found when trying to update user state. This user state should have been present"); return Err(BatcherError::AddressNotFoundInUserStates( proof_submitter_addr, @@ -661,10 +673,26 @@ impl Batcher { }; // User state is updated - let mut user_state_lock = user_state.lock().await; - user_state_lock.nonce = Some(nonce + U256::one()); - user_state_lock.min_fee = max_fee; - user_state_lock.proofs_in_batch += 1; + // batch_state_lock.update_user_nonce(&proof_submitter_addr, nonce + U256::one()); + // batch_state_lock.update_user_min_fee(&proof_submitter_addr, max_fee); + // batch_state_lock.update_user_proof_count(&proof_submitter_addr, user_proof_count + 1); + + let Some(_) = batch_state_lock.update_user_state( + &proof_submitter_addr, + nonce + U256::one(), + max_fee, + user_proof_count + 1, + ) else { + error!("User state of address {proof_submitter_addr} was not found when trying to update user state. This user state should have been present"); + return Err(BatcherError::AddressNotFoundInUserStates( + proof_submitter_addr, + )); + }; + + // let mut user_state_lock = user_sta + // user_state_lock.nonce = Some(nonce + U256::one()); + // user_state_lock.min_fee = max_fee; + // user_state_lock.proofs_in_batch += 1; Ok(()) } @@ -738,7 +766,24 @@ impl Batcher { let updated_user_proof_count_and_min_fee = batch_state_lock.get_user_proofs_in_batch_and_min_fee(); - for addr in batch_state_lock.user_states.keys() { + // batch_state_lock.user_states.keys().map(|addr| { + // let (proof_count, min_fee) = updated_user_proof_count_and_min_fee + // .get(addr) + // .unwrap_or(&(0, U256::MAX)); + + // // FIXME: The case where a `user_state` is not found in the `user_states` map should not really + // // happen here, but doing this check so that we don't unwrap. + // // Once https://github.com/yetanotherco/aligned_layer/issues/1046 is done we could return a more + // // informative error. + // // let user_state = batch_state_lock.get_user_state(addr)?; + + // // Now we update the user states + // let updated_proof_count = batch_state_lock.update_user_proof_count(addr, *proof_count); + // let update_user_min_fee = batch_state_lock.update_user_min_fee(addr, *min_fee); + // }); + + let user_addresses: Vec
= batch_state_lock.user_states.keys().cloned().collect(); + for addr in user_addresses.iter() { let (proof_count, min_fee) = updated_user_proof_count_and_min_fee .get(addr) .unwrap_or(&(0, U256::MAX)); @@ -747,12 +792,11 @@ impl Batcher { // happen here, but doing this check so that we don't unwrap. // Once https://github.com/yetanotherco/aligned_layer/issues/1046 is done we could return a more // informative error. - let user_state = batch_state_lock.get_user_state(addr)?; + // let user_state = batch_state_lock.get_user_state(addr)?; // Now we update the user states - let mut user_state_lock = user_state.lock().await; - user_state_lock.proofs_in_batch = *proof_count; - user_state_lock.min_fee = *min_fee; + batch_state_lock.update_user_proof_count(addr, *proof_count); + batch_state_lock.update_user_min_fee(addr, *min_fee); } Some(finalized_batch) @@ -870,7 +914,7 @@ impl Batcher { let nonpaying_user_state = UserState::new_non_paying(nonpaying_replacement_addr_nonce); batch_state_lock .user_states - .insert(nonpaying_replacement_addr, Mutex::new(nonpaying_user_state)); + .insert(nonpaying_replacement_addr, nonpaying_user_state); } /// Receives new block numbers, checks if conditions are met for submission and diff --git a/batcher/aligned-batcher/src/types/batch_state.rs b/batcher/aligned-batcher/src/types/batch_state.rs index 528711283..09378ec1f 100644 --- a/batcher/aligned-batcher/src/types/batch_state.rs +++ b/batcher/aligned-batcher/src/types/batch_state.rs @@ -1,4 +1,4 @@ -use std::collections::HashMap; +use std::collections::{hash_map::Entry, HashMap}; use super::{ batch_queue::{BatchQueue, BatchQueueEntry}, @@ -6,11 +6,10 @@ use super::{ }; use ethers::types::{Address, U256}; use log::debug; -use tokio::sync::Mutex; pub(crate) struct BatchState { pub(crate) batch_queue: BatchQueue, - pub(crate) user_states: HashMap>, + pub(crate) user_states: HashMap, } impl BatchState { @@ -21,7 +20,7 @@ impl BatchState { } } - pub(crate) fn new_with_user_states(user_states: HashMap>) -> Self { + pub(crate) fn new_with_user_states(user_states: HashMap) -> Self { Self { batch_queue: BatchQueue::new(), user_states, @@ -35,23 +34,32 @@ impl BatchState { .find(|entry| entry.sender == sender && entry.nonced_verification_data.nonce == nonce) } - pub(crate) fn get_user_state(&self, addr: &Address) -> Option<&Mutex> { + pub(crate) fn get_user_state(&self, addr: &Address) -> Option<&UserState> { self.user_states.get(addr) } pub(crate) async fn get_user_nonce(&self, addr: &Address) -> Option { let user_state = self.get_user_state(addr)?; - user_state.lock().await.nonce + user_state.nonce } pub(crate) async fn get_user_min_fee(&self, addr: &Address) -> Option { let user_state = self.get_user_state(addr)?; - Some(user_state.lock().await.min_fee) + Some(user_state.min_fee) + } + + pub(crate) fn update_user_nonce(&mut self, addr: &Address, new_nonce: U256) -> Option { + if let Entry::Occupied(mut user_state) = self.user_states.entry(*addr) { + let new_nonce = Some(new_nonce); + user_state.get_mut().nonce = new_nonce; + return new_nonce; + } + None } pub(crate) async fn get_user_proof_count(&self, addr: &Address) -> Option { let user_state = self.get_user_state(addr)?; - Some(user_state.lock().await.proofs_in_batch) + Some(user_state.proofs_in_batch) } /// Checks if the entry is valid @@ -86,6 +94,51 @@ impl BatchState { .unwrap_or(U256::max_value()) } + pub(crate) fn update_user_min_fee( + &mut self, + addr: &Address, + new_min_fee: U256, + ) -> Option { + if let Entry::Occupied(mut user_state) = self.user_states.entry(*addr) { + user_state.get_mut().min_fee = new_min_fee; + return Some(new_min_fee); + } + None + } + + pub(crate) fn update_user_proof_count( + &mut self, + addr: &Address, + new_proof_count: usize, + ) -> Option { + if let Entry::Occupied(mut user_state) = self.user_states.entry(*addr) { + user_state.get_mut().proofs_in_batch = new_proof_count; + return Some(new_proof_count); + } + None + } + + /// Updates the user with address `addr` with the provided values of `new_nonce`, `new_min_fee` and + /// `new_proof_count`. + /// If state is updated successfully, returns the updated values inside a `Some()` + /// If the address was not found in the user states, returns `None` + pub(crate) fn update_user_state( + &mut self, + addr: &Address, + new_nonce: U256, + new_min_fee: U256, + new_proof_count: usize, + ) -> Option<(U256, U256, usize)> { + let updated_nonce = self.update_user_nonce(&addr, new_nonce); + let updated_min_fee = self.update_user_min_fee(&addr, new_min_fee); + let updated_proof_count = self.update_user_proof_count(&addr, new_proof_count); + + match (updated_nonce, updated_min_fee, updated_proof_count) { + (Some(_), Some(_), Some(_)) => Some((new_nonce, new_min_fee, new_proof_count)), + _ => None, + } + } + pub(crate) fn get_user_proofs_in_batch_and_min_fee(&self) -> HashMap { let mut updated_user_states = HashMap::new(); for (entry, _) in self.batch_queue.iter() { From d2ce540317742533c6bccc95e0e79a8af87fb0ec Mon Sep 17 00:00:00 2001 From: Mariano Nicolini Date: Mon, 30 Sep 2024 13:11:47 -0300 Subject: [PATCH 21/47] Revert burst size changes in Makefile --- Makefile | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/Makefile b/Makefile index 052857845..92dc6a110 100644 --- a/Makefile +++ b/Makefile @@ -271,7 +271,7 @@ batcher_send_risc0_burst: --proof ../../scripts/test_files/risc_zero/fibonacci_proof_generator/risc_zero_fibonacci.proof \ --vm_program ../../scripts/test_files/risc_zero/fibonacci_proof_generator/fibonacci_id.bin \ --public_input ../../scripts/test_files/risc_zero/fibonacci_proof_generator/risc_zero_fibonacci.pub \ - --repetitions 500 \ + --repetitions $(BURST_SIZE) \ --proof_generator_addr 0x66f9664f97F2b50F62D13eA064982f936dE76657 \ --rpc_url $(RPC_URL) \ --payment_service_addr $(BATCHER_PAYMENTS_CONTRACT_ADDRESS) @@ -297,7 +297,7 @@ batcher_send_plonk_bn254_burst: batcher/target/release/aligned --vk ../../scripts/test_files/gnark_plonk_bn254_script/plonk.vk \ --proof_generator_addr 0x66f9664f97F2b50F62D13eA064982f936dE76657 \ --rpc_url $(RPC_URL) \ - --repetitions 4 \ + --repetitions $(BURST_SIZE) \ --payment_service_addr $(BATCHER_PAYMENTS_CONTRACT_ADDRESS) batcher_send_plonk_bls12_381_task: batcher/target/release/aligned @@ -362,7 +362,7 @@ batcher_send_halo2_ipa_task_burst_5: batcher/target/release/aligned --proof ../../scripts/test_files/halo2_ipa/proof.bin \ --public_input ../../scripts/test_files/halo2_ipa/pub_input.bin \ --vk ../../scripts/test_files/halo2_ipa/params.bin \ - --repetitions 500 \ + --repetitions $(BURST_SIZE) \ --private_key 0xdbda1821b80551c9d65939329250298aa3472ba22feea921c0cf5d620ea67b97 \ --rpc_url $(RPC_URL) \ --payment_service_addr $(BATCHER_PAYMENTS_CONTRACT_ADDRESS) From 5955cb208142681aa07cda70aa0435aea5869201 Mon Sep 17 00:00:00 2001 From: Mariano Nicolini Date: Mon, 30 Sep 2024 13:14:15 -0300 Subject: [PATCH 22/47] Apply clippy suggestions --- batcher/aligned-batcher/src/lib.rs | 2 +- batcher/aligned-batcher/src/types/batch_state.rs | 6 +++--- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/batcher/aligned-batcher/src/lib.rs b/batcher/aligned-batcher/src/lib.rs index eefea5b5d..7cb8cd31a 100644 --- a/batcher/aligned-batcher/src/lib.rs +++ b/batcher/aligned-batcher/src/lib.rs @@ -397,7 +397,7 @@ impl Batcher { batch_state_lock .user_states .entry(addr) - .or_insert_with(|| UserState::new()); + .or_insert_with(UserState::new); } // * ---------------------------------------------------* diff --git a/batcher/aligned-batcher/src/types/batch_state.rs b/batcher/aligned-batcher/src/types/batch_state.rs index 09378ec1f..c1d197f9a 100644 --- a/batcher/aligned-batcher/src/types/batch_state.rs +++ b/batcher/aligned-batcher/src/types/batch_state.rs @@ -129,9 +129,9 @@ impl BatchState { new_min_fee: U256, new_proof_count: usize, ) -> Option<(U256, U256, usize)> { - let updated_nonce = self.update_user_nonce(&addr, new_nonce); - let updated_min_fee = self.update_user_min_fee(&addr, new_min_fee); - let updated_proof_count = self.update_user_proof_count(&addr, new_proof_count); + let updated_nonce = self.update_user_nonce(addr, new_nonce); + let updated_min_fee = self.update_user_min_fee(addr, new_min_fee); + let updated_proof_count = self.update_user_proof_count(addr, new_proof_count); match (updated_nonce, updated_min_fee, updated_proof_count) { (Some(_), Some(_), Some(_)) => Some((new_nonce, new_min_fee, new_proof_count)), From 0a8230642693128724e601afb08c52d38ef54407 Mon Sep 17 00:00:00 2001 From: Mariano Nicolini Date: Mon, 30 Sep 2024 13:17:54 -0300 Subject: [PATCH 23/47] Revert more changes in Makefile --- Makefile | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/Makefile b/Makefile index 92dc6a110..1a34c1488 100644 --- a/Makefile +++ b/Makefile @@ -297,7 +297,7 @@ batcher_send_plonk_bn254_burst: batcher/target/release/aligned --vk ../../scripts/test_files/gnark_plonk_bn254_script/plonk.vk \ --proof_generator_addr 0x66f9664f97F2b50F62D13eA064982f936dE76657 \ --rpc_url $(RPC_URL) \ - --repetitions $(BURST_SIZE) \ + --repetitions 4 \ --payment_service_addr $(BATCHER_PAYMENTS_CONTRACT_ADDRESS) batcher_send_plonk_bls12_381_task: batcher/target/release/aligned @@ -362,8 +362,7 @@ batcher_send_halo2_ipa_task_burst_5: batcher/target/release/aligned --proof ../../scripts/test_files/halo2_ipa/proof.bin \ --public_input ../../scripts/test_files/halo2_ipa/pub_input.bin \ --vk ../../scripts/test_files/halo2_ipa/params.bin \ - --repetitions $(BURST_SIZE) \ - --private_key 0xdbda1821b80551c9d65939329250298aa3472ba22feea921c0cf5d620ea67b97 \ + --repetitions 5 \ --rpc_url $(RPC_URL) \ --payment_service_addr $(BATCHER_PAYMENTS_CONTRACT_ADDRESS) From 770acd1b7d3173d8197381a52033e4db4263df8d Mon Sep 17 00:00:00 2001 From: Mariano Nicolini Date: Mon, 30 Sep 2024 13:18:26 -0300 Subject: [PATCH 24/47] Revert more changes in Makefile --- Makefile | 1 - 1 file changed, 1 deletion(-) diff --git a/Makefile b/Makefile index 1a34c1488..14ee2ff33 100644 --- a/Makefile +++ b/Makefile @@ -275,7 +275,6 @@ batcher_send_risc0_burst: --proof_generator_addr 0x66f9664f97F2b50F62D13eA064982f936dE76657 \ --rpc_url $(RPC_URL) \ --payment_service_addr $(BATCHER_PAYMENTS_CONTRACT_ADDRESS) - # --private_key 0x4bbbf85ce3377467afe5d46f804f221813b2bb87f24d81f60f1fcdbf7cbf4356 \ batcher_send_plonk_bn254_task: batcher/target/release/aligned @echo "Sending Groth16Bn254 1!=0 task to Batcher..." From ca4939ea9b98800975e8f4d922b7f5ec32012b6a Mon Sep 17 00:00:00 2001 From: Mariano Nicolini Date: Mon, 30 Sep 2024 13:35:45 -0300 Subject: [PATCH 25/47] Revert changes in send_burst_tasks.sh script --- batcher/aligned-batcher/src/lib.rs | 60 ++++++++--------------------- batcher/aligned/send_burst_tasks.sh | 2 +- 2 files changed, 17 insertions(+), 45 deletions(-) diff --git a/batcher/aligned-batcher/src/lib.rs b/batcher/aligned-batcher/src/lib.rs index 7cb8cd31a..913af60c0 100644 --- a/batcher/aligned-batcher/src/lib.rs +++ b/batcher/aligned-batcher/src/lib.rs @@ -404,6 +404,10 @@ impl Batcher { // * Perform validations over user state * // * ---------------------------------------------------* + // For now until the message is fully processed, the batch state is locked + // This is needed because we need to query the user state to make validations and + // finally add the proof to the batch queue. + // At this point, we will have a user state for sure, since we have inserted it // if not already present. It should be safe to call `unwrap()` here. let mut batch_state_lock = self.batch_state.lock().await; @@ -494,8 +498,7 @@ impl Batcher { Ok(()) } - // Checks user has sufficient balance - // If user has sufficient balance, increments the user's proof count in the batch + // Checks user has sufficient balance for paying all its the proofs in the current batch. async fn check_min_balance(&self, addr: &Address, user_proofs_in_batch: usize) -> bool { let Some(user_balance) = self.get_user_balance(addr).await else { return false; @@ -592,10 +595,6 @@ impl Batcher { BatchQueueEntryPriority::new(replacement_max_fee, nonce), ); - // let Some(user_state) = batch_state_lock.get_user_state(&addr) else { - // warn!("User state for address {addr:?} was not present in batcher user states, but it should be"); - // return; - // }; let updated_min_fee_in_batch = batch_state_lock.get_user_min_fee_in_batch(&addr); let Some(_) = batch_state_lock.update_user_min_fee(&addr, updated_min_fee_in_batch) else { warn!("User state for address {addr:?} was not present in batcher user states, but it should be"); @@ -655,13 +654,6 @@ impl Batcher { proof_submitter_addr }; - // let Some(user_state) = batch_state_lock.get_user_state(&proof_submitter_addr) else { - // error!("User state of address {proof_submitter_addr} was not found when trying to update user state. This user state should have been present"); - // return Err(BatcherError::AddressNotFoundInUserStates( - // proof_submitter_addr, - // )); - // }; - let Some(user_proof_count) = batch_state_lock .get_user_proof_count(&proof_submitter_addr) .await @@ -673,10 +665,6 @@ impl Batcher { }; // User state is updated - // batch_state_lock.update_user_nonce(&proof_submitter_addr, nonce + U256::one()); - // batch_state_lock.update_user_min_fee(&proof_submitter_addr, max_fee); - // batch_state_lock.update_user_proof_count(&proof_submitter_addr, user_proof_count + 1); - let Some(_) = batch_state_lock.update_user_state( &proof_submitter_addr, nonce + U256::one(), @@ -689,10 +677,6 @@ impl Batcher { )); }; - // let mut user_state_lock = user_sta - // user_state_lock.nonce = Some(nonce + U256::one()); - // user_state_lock.min_fee = max_fee; - // user_state_lock.proofs_in_batch += 1; Ok(()) } @@ -766,37 +750,20 @@ impl Batcher { let updated_user_proof_count_and_min_fee = batch_state_lock.get_user_proofs_in_batch_and_min_fee(); - // batch_state_lock.user_states.keys().map(|addr| { - // let (proof_count, min_fee) = updated_user_proof_count_and_min_fee - // .get(addr) - // .unwrap_or(&(0, U256::MAX)); - - // // FIXME: The case where a `user_state` is not found in the `user_states` map should not really - // // happen here, but doing this check so that we don't unwrap. - // // Once https://github.com/yetanotherco/aligned_layer/issues/1046 is done we could return a more - // // informative error. - // // let user_state = batch_state_lock.get_user_state(addr)?; - - // // Now we update the user states - // let updated_proof_count = batch_state_lock.update_user_proof_count(addr, *proof_count); - // let update_user_min_fee = batch_state_lock.update_user_min_fee(addr, *min_fee); - // }); - let user_addresses: Vec
= batch_state_lock.user_states.keys().cloned().collect(); for addr in user_addresses.iter() { let (proof_count, min_fee) = updated_user_proof_count_and_min_fee .get(addr) .unwrap_or(&(0, U256::MAX)); - // FIXME: The case where a `user_state` is not found in the `user_states` map should not really - // happen here, but doing this check so that we don't unwrap. + // FIXME: The case where a the update functions return `None` can only happen when the user was not found + // in the `user_states` map should not really happen here, but doing this check so that we don't unwrap. // Once https://github.com/yetanotherco/aligned_layer/issues/1046 is done we could return a more // informative error. - // let user_state = batch_state_lock.get_user_state(addr)?; - // Now we update the user states - batch_state_lock.update_user_proof_count(addr, *proof_count); - batch_state_lock.update_user_min_fee(addr, *min_fee); + // Now we update the user states related to the batch (proof count in batch and min fee in batch) + batch_state_lock.update_user_proof_count(addr, *proof_count)?; + batch_state_lock.update_user_min_fee(addr, *min_fee)?; } Some(finalized_batch) @@ -1186,7 +1153,9 @@ impl Batcher { Ok(()) } - /// TODO: ADD DOCUMENTATION ON WHY IT IS AN OPTION + /// Gets the balance of user with address `addr` from Ethereum. + /// Returns `None` if the balance couldn't be returned + /// FIXME: This should return a `Result` instead. async fn get_user_balance(&self, addr: &Address) -> Option { if let Ok(balance) = self.payment_service.user_balances(*addr).call().await { return Some(balance); @@ -1216,6 +1185,9 @@ impl Batcher { false } + /// Gets the current gas price from Ethereum. + /// Returns `None` if the gas price couldn't be returned + /// FIXME: This should return a `Result` instead. async fn get_gas_price(&self) -> Option { if let Ok(gas_price) = self .eth_ws_provider diff --git a/batcher/aligned/send_burst_tasks.sh b/batcher/aligned/send_burst_tasks.sh index f46464aa6..25160bf13 100755 --- a/batcher/aligned/send_burst_tasks.sh +++ b/batcher/aligned/send_burst_tasks.sh @@ -28,6 +28,6 @@ while true do # Run in backaground to be able to run onece per second, and not wait for the previous one to finish . ./batcher/aligned/generate_proof_and_send.sh $counter $burst & - sleep 3 + sleep 1 counter=$((counter + 1)) done From 0a55b8fbdb8d61f3283701f497fc6017a9832cf9 Mon Sep 17 00:00:00 2001 From: Mariano Nicolini Date: Mon, 30 Sep 2024 13:38:51 -0300 Subject: [PATCH 26/47] Apply clippy code docs suggestions --- batcher/aligned-batcher/src/lib.rs | 2 ++ 1 file changed, 2 insertions(+) diff --git a/batcher/aligned-batcher/src/lib.rs b/batcher/aligned-batcher/src/lib.rs index 913af60c0..f6601250e 100644 --- a/batcher/aligned-batcher/src/lib.rs +++ b/batcher/aligned-batcher/src/lib.rs @@ -684,10 +684,12 @@ impl Batcher { /// There are essentially two conditions to be checked: /// * Has the current batch reached the minimum size to be posted? /// * Has the received block number surpassed the maximum interval with respect to the last posted batch block? + /// /// Then the batch will be made as big as possible given this two conditions: /// * The serialized batch size needs to be smaller than the maximum batch size /// * The batch submission fee is less than the lowest `max fee` included the batch, /// * And the batch submission fee is more than the highest `max fee` not included the batch. + /// /// An extra sanity check is made to check if the batch size is 0, since it does not make sense to post /// an empty batch, even if the block interval has been reached. /// Once the batch meets the conditions for submission, the finalized batch is then passed to the From db0a15fb9d683a8ced5d845ae203fe99a140792f Mon Sep 17 00:00:00 2001 From: Mariano Nicolini Date: Mon, 30 Sep 2024 14:31:11 -0300 Subject: [PATCH 27/47] Minor docs fix --- batcher/aligned-batcher/src/lib.rs | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/batcher/aligned-batcher/src/lib.rs b/batcher/aligned-batcher/src/lib.rs index f6601250e..bbf97bbb7 100644 --- a/batcher/aligned-batcher/src/lib.rs +++ b/batcher/aligned-batcher/src/lib.rs @@ -404,7 +404,7 @@ impl Batcher { // * Perform validations over user state * // * ---------------------------------------------------* - // For now until the message is fully processed, the batch state is locked + // For now on until the message is fully processed, the batch state is locked // This is needed because we need to query the user state to make validations and // finally add the proof to the batch queue. @@ -462,10 +462,6 @@ impl Batcher { return Ok(()); } - // * ---------------------------------------------------------------------* - // * Add message data into the queue and update user state * - // * ---------------------------------------------------------------------* - let msg_max_fee = nonced_verification_data.max_fee; let Some(user_min_fee) = batch_state_lock.get_user_min_fee(&addr).await else { send_message(ws_conn_sink.clone(), ValidityResponseMessage::InvalidNonce).await; @@ -478,6 +474,10 @@ impl Batcher { return Ok(()); } + // * ---------------------------------------------------------------------* + // * Add message data into the queue and update user state * + // * ---------------------------------------------------------------------* + if let Err(e) = self .add_to_batch( batch_state_lock, From 03ce15652f5ac78451e55ad134cdd1f2d24a53fe Mon Sep 17 00:00:00 2001 From: Mariano Nicolini Date: Mon, 30 Sep 2024 17:35:52 -0300 Subject: [PATCH 28/47] Add drop on batch state lock when sending error messages --- batcher/aligned-batcher/src/lib.rs | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/batcher/aligned-batcher/src/lib.rs b/batcher/aligned-batcher/src/lib.rs index bbf97bbb7..c802e4121 100644 --- a/batcher/aligned-batcher/src/lib.rs +++ b/batcher/aligned-batcher/src/lib.rs @@ -413,6 +413,7 @@ impl Batcher { let mut batch_state_lock = self.batch_state.lock().await; let proofs_in_batch = batch_state_lock.get_user_proof_count(&addr).await.unwrap(); if !self.check_min_balance(&addr, proofs_in_batch + 1).await { + std::mem::drop(batch_state_lock); send_message( ws_conn_sink.clone(), ValidityResponseMessage::InsufficientBalance(addr), @@ -431,6 +432,7 @@ impl Batcher { error!( "Failed to get user nonce from Ethereum for address {addr:?}. Error: {e:?}" ); + std::mem::drop(batch_state_lock); send_message(ws_conn_sink.clone(), ValidityResponseMessage::InvalidNonce).await; return Ok(()); } @@ -441,6 +443,7 @@ impl Batcher { }; if expected_nonce < msg_nonce { + std::mem::drop(batch_state_lock); warn!("Invalid nonce for address {addr}, had nonce {expected_nonce:?} < {msg_nonce:?}"); send_message(ws_conn_sink.clone(), ValidityResponseMessage::InvalidNonce).await; return Ok(()); @@ -464,11 +467,13 @@ impl Batcher { let msg_max_fee = nonced_verification_data.max_fee; let Some(user_min_fee) = batch_state_lock.get_user_min_fee(&addr).await else { + std::mem::drop(batch_state_lock); send_message(ws_conn_sink.clone(), ValidityResponseMessage::InvalidNonce).await; return Ok(()); }; if msg_max_fee > user_min_fee { + std::mem::drop(batch_state_lock); warn!("Invalid max fee for address {addr}, had fee {user_min_fee:?} < {msg_max_fee:?}"); send_message(ws_conn_sink.clone(), ValidityResponseMessage::InvalidMaxFee).await; return Ok(()); @@ -659,6 +664,7 @@ impl Batcher { .await else { error!("User state of address {proof_submitter_addr} was not found when trying to update user state. This user state should have been present"); + std::mem::drop(batch_state_lock); return Err(BatcherError::AddressNotFoundInUserStates( proof_submitter_addr, )); @@ -672,6 +678,7 @@ impl Batcher { user_proof_count + 1, ) else { error!("User state of address {proof_submitter_addr} was not found when trying to update user state. This user state should have been present"); + std::mem::drop(batch_state_lock); return Err(BatcherError::AddressNotFoundInUserStates( proof_submitter_addr, )); From f4dd5ce5c86dbafb1f80a0d1fe721d6e858cf6a9 Mon Sep 17 00:00:00 2001 From: Mariano Nicolini Date: Mon, 30 Sep 2024 18:05:38 -0300 Subject: [PATCH 29/47] Add some more drops when sending error messages and returnig --- Makefile | 5 +++-- batcher/aligned-batcher/src/lib.rs | 4 ++++ batcher/aligned/send_burst_tasks.sh | 2 +- 3 files changed, 8 insertions(+), 3 deletions(-) diff --git a/Makefile b/Makefile index 14ee2ff33..2b5f6badb 100644 --- a/Makefile +++ b/Makefile @@ -271,7 +271,7 @@ batcher_send_risc0_burst: --proof ../../scripts/test_files/risc_zero/fibonacci_proof_generator/risc_zero_fibonacci.proof \ --vm_program ../../scripts/test_files/risc_zero/fibonacci_proof_generator/fibonacci_id.bin \ --public_input ../../scripts/test_files/risc_zero/fibonacci_proof_generator/risc_zero_fibonacci.pub \ - --repetitions $(BURST_SIZE) \ + --repetitions 500 \ --proof_generator_addr 0x66f9664f97F2b50F62D13eA064982f936dE76657 \ --rpc_url $(RPC_URL) \ --payment_service_addr $(BATCHER_PAYMENTS_CONTRACT_ADDRESS) @@ -361,7 +361,8 @@ batcher_send_halo2_ipa_task_burst_5: batcher/target/release/aligned --proof ../../scripts/test_files/halo2_ipa/proof.bin \ --public_input ../../scripts/test_files/halo2_ipa/pub_input.bin \ --vk ../../scripts/test_files/halo2_ipa/params.bin \ - --repetitions 5 \ + --repetitions 500 \ + --private_key 0xdbda1821b80551c9d65939329250298aa3472ba22feea921c0cf5d620ea67b97 \ --rpc_url $(RPC_URL) \ --payment_service_addr $(BATCHER_PAYMENTS_CONTRACT_ADDRESS) diff --git a/batcher/aligned-batcher/src/lib.rs b/batcher/aligned-batcher/src/lib.rs index c802e4121..c11fef11d 100644 --- a/batcher/aligned-batcher/src/lib.rs +++ b/batcher/aligned-batcher/src/lib.rs @@ -533,6 +533,7 @@ impl Batcher { let replacement_max_fee = nonced_verification_data.max_fee; let nonce = nonced_verification_data.nonce; let Some(entry) = batch_state_lock.get_entry(addr, nonce) else { + std::mem::drop(batch_state_lock); warn!("Invalid nonce for address {addr}. Queue entry with nonce {nonce} not found"); send_message(ws_conn_sink.clone(), ValidityResponseMessage::InvalidNonce).await; return; @@ -540,6 +541,7 @@ impl Batcher { let original_max_fee = entry.nonced_verification_data.max_fee; if original_max_fee > replacement_max_fee { + std::mem::drop(batch_state_lock); warn!("Invalid replacement message for address {addr}, had fee {original_max_fee:?} < {replacement_max_fee:?}"); send_message( ws_conn_sink.clone(), @@ -576,6 +578,7 @@ impl Batcher { replacement_entry.messaging_sink = Some(ws_conn_sink.clone()); if !batch_state_lock.replacement_entry_is_valid(&replacement_entry) { + std::mem::drop(batch_state_lock); warn!("Invalid max fee"); send_message( ws_conn_sink.clone(), @@ -602,6 +605,7 @@ impl Batcher { let updated_min_fee_in_batch = batch_state_lock.get_user_min_fee_in_batch(&addr); let Some(_) = batch_state_lock.update_user_min_fee(&addr, updated_min_fee_in_batch) else { + std::mem::drop(batch_state_lock); warn!("User state for address {addr:?} was not present in batcher user states, but it should be"); return; }; diff --git a/batcher/aligned/send_burst_tasks.sh b/batcher/aligned/send_burst_tasks.sh index 25160bf13..f46464aa6 100755 --- a/batcher/aligned/send_burst_tasks.sh +++ b/batcher/aligned/send_burst_tasks.sh @@ -28,6 +28,6 @@ while true do # Run in backaground to be able to run onece per second, and not wait for the previous one to finish . ./batcher/aligned/generate_proof_and_send.sh $counter $burst & - sleep 1 + sleep 3 counter=$((counter + 1)) done From 84c8e6723269b356da0fc60d3cb39394abcaffe0 Mon Sep 17 00:00:00 2001 From: Mariano Nicolini Date: Mon, 30 Sep 2024 21:37:14 -0300 Subject: [PATCH 30/47] simplify UserState struct --- batcher/aligned-batcher/src/lib.rs | 58 ++++++++++--------- .../aligned-batcher/src/types/batch_state.rs | 5 +- .../aligned-batcher/src/types/user_state.rs | 14 +---- 3 files changed, 37 insertions(+), 40 deletions(-) diff --git a/batcher/aligned-batcher/src/lib.rs b/batcher/aligned-batcher/src/lib.rs index c11fef11d..cb3062ab6 100644 --- a/batcher/aligned-batcher/src/lib.rs +++ b/batcher/aligned-batcher/src/lib.rs @@ -177,7 +177,7 @@ impl Batcher { .await .expect("Could not get non-paying nonce from Ethereum"); - let non_paying_user_state = UserState::new_non_paying(nonpaying_nonce); + let non_paying_user_state = UserState::new(nonpaying_nonce); user_states.insert( non_paying_config.replacement.address(), non_paying_user_state, @@ -394,10 +394,22 @@ impl Batcher { // Check that we had a user state entry for this user and insert it if not. { let mut batch_state_lock = self.batch_state.lock().await; - batch_state_lock - .user_states - .entry(addr) - .or_insert_with(UserState::new); + if !batch_state_lock.user_states.contains_key(&addr) { + let ethereum_user_nonce = match self.get_user_nonce_from_ethereum(addr).await { + Ok(ethereum_user_nonce) => ethereum_user_nonce, + Err(e) => { + error!( + "Failed to get user nonce from Ethereum for address {addr:?}. Error: {e:?}" + ); + std::mem::drop(batch_state_lock); + send_message(ws_conn_sink.clone(), ValidityResponseMessage::InvalidNonce) + .await; + return Ok(()); + } + }; + let user_state = UserState::new(ethereum_user_nonce); + batch_state_lock.user_states.insert(addr, user_state); + } } // * ---------------------------------------------------* @@ -410,8 +422,14 @@ impl Batcher { // At this point, we will have a user state for sure, since we have inserted it // if not already present. It should be safe to call `unwrap()` here. - let mut batch_state_lock = self.batch_state.lock().await; - let proofs_in_batch = batch_state_lock.get_user_proof_count(&addr).await.unwrap(); + let batch_state_lock = self.batch_state.lock().await; + let Some(proofs_in_batch) = batch_state_lock.get_user_proof_count(&addr).await else { + error!("Failed to get user proof count: User not found in user states, but it should have been already inserted"); + std::mem::drop(batch_state_lock); + send_message(ws_conn_sink.clone(), ValidityResponseMessage::InvalidNonce).await; + return Ok(()); + }; + if !self.check_min_balance(&addr, proofs_in_batch + 1).await { std::mem::drop(batch_state_lock); send_message( @@ -422,24 +440,12 @@ impl Batcher { return Ok(()); } - let user_nonce = batch_state_lock.get_user_nonce(&addr).await; - let expected_nonce = if let Some(cached_nonce) = user_nonce { - cached_nonce - } else { - let ethereum_user_nonce = match self.get_user_nonce_from_ethereum(addr).await { - Ok(ethereum_user_nonce) => ethereum_user_nonce, - Err(e) => { - error!( - "Failed to get user nonce from Ethereum for address {addr:?}. Error: {e:?}" - ); - std::mem::drop(batch_state_lock); - send_message(ws_conn_sink.clone(), ValidityResponseMessage::InvalidNonce).await; - return Ok(()); - } - }; - - batch_state_lock.update_user_nonce(&addr, ethereum_user_nonce); - ethereum_user_nonce + let cached_user_nonce = batch_state_lock.get_user_nonce(&addr).await; + let Some(expected_nonce) = cached_user_nonce else { + error!("Failed to get cached user nonce: User not found in user states, but it should have been already inserted"); + std::mem::drop(batch_state_lock); + send_message(ws_conn_sink.clone(), ValidityResponseMessage::InvalidNonce).await; + return Ok(()); }; if expected_nonce < msg_nonce { @@ -891,7 +897,7 @@ impl Batcher { }; batch_state_lock.batch_queue.clear(); batch_state_lock.user_states.clear(); - let nonpaying_user_state = UserState::new_non_paying(nonpaying_replacement_addr_nonce); + let nonpaying_user_state = UserState::new(nonpaying_replacement_addr_nonce); batch_state_lock .user_states .insert(nonpaying_replacement_addr, nonpaying_user_state); diff --git a/batcher/aligned-batcher/src/types/batch_state.rs b/batcher/aligned-batcher/src/types/batch_state.rs index c1d197f9a..72c1ffe54 100644 --- a/batcher/aligned-batcher/src/types/batch_state.rs +++ b/batcher/aligned-batcher/src/types/batch_state.rs @@ -40,7 +40,7 @@ impl BatchState { pub(crate) async fn get_user_nonce(&self, addr: &Address) -> Option { let user_state = self.get_user_state(addr)?; - user_state.nonce + Some(user_state.nonce) } pub(crate) async fn get_user_min_fee(&self, addr: &Address) -> Option { @@ -50,9 +50,8 @@ impl BatchState { pub(crate) fn update_user_nonce(&mut self, addr: &Address, new_nonce: U256) -> Option { if let Entry::Occupied(mut user_state) = self.user_states.entry(*addr) { - let new_nonce = Some(new_nonce); user_state.get_mut().nonce = new_nonce; - return new_nonce; + return Some(new_nonce); } None } diff --git a/batcher/aligned-batcher/src/types/user_state.rs b/batcher/aligned-batcher/src/types/user_state.rs index c5c650ade..a1627b25f 100644 --- a/batcher/aligned-batcher/src/types/user_state.rs +++ b/batcher/aligned-batcher/src/types/user_state.rs @@ -1,7 +1,7 @@ use ethers::types::U256; pub(crate) struct UserState { - pub nonce: Option, + pub nonce: U256, /// The minimum fee of a pending proof for a user. /// This should always be the fee of the biggest pending nonce by the user. /// This is used to check if a user is submitting a proof with a higher nonce and higher fee, @@ -11,17 +11,9 @@ pub(crate) struct UserState { } impl UserState { - pub(crate) fn new() -> Self { + pub(crate) fn new(nonce: U256) -> Self { UserState { - nonce: None, - min_fee: U256::max_value(), - proofs_in_batch: 0, - } - } - - pub(crate) fn new_non_paying(nonce: U256) -> Self { - UserState { - nonce: Some(nonce), + nonce, min_fee: U256::max_value(), proofs_in_batch: 0, } From d77eeac3abfeeb7094de1554c0c81f09fbbc1c56 Mon Sep 17 00:00:00 2001 From: Mariano Nicolini Date: Mon, 30 Sep 2024 23:38:55 -0300 Subject: [PATCH 31/47] Some more refactors on the ordering of validations --- batcher/aligned-batcher/src/lib.rs | 18 +++++++++--------- .../aligned-sdk/src/communication/messaging.rs | 3 +++ batcher/aligned-sdk/src/core/types.rs | 1 + 3 files changed, 13 insertions(+), 9 deletions(-) diff --git a/batcher/aligned-batcher/src/lib.rs b/batcher/aligned-batcher/src/lib.rs index cb3062ab6..379defd60 100644 --- a/batcher/aligned-batcher/src/lib.rs +++ b/batcher/aligned-batcher/src/lib.rs @@ -420,6 +420,12 @@ impl Batcher { // This is needed because we need to query the user state to make validations and // finally add the proof to the batch queue. + let Some(user_balance) = self.get_user_balance(&addr).await else { + error!("Could not get balance for address {addr:?}"); + send_message(ws_conn_sink.clone(), ValidityResponseMessage::EthRpcError).await; + return Ok(()); + }; + // At this point, we will have a user state for sure, since we have inserted it // if not already present. It should be safe to call `unwrap()` here. let batch_state_lock = self.batch_state.lock().await; @@ -430,7 +436,7 @@ impl Batcher { return Ok(()); }; - if !self.check_min_balance(&addr, proofs_in_batch + 1).await { + if !self.check_min_balance(proofs_in_batch + 1, user_balance) { std::mem::drop(batch_state_lock); send_message( ws_conn_sink.clone(), @@ -510,15 +516,9 @@ impl Batcher { } // Checks user has sufficient balance for paying all its the proofs in the current batch. - async fn check_min_balance(&self, addr: &Address, user_proofs_in_batch: usize) -> bool { - let Some(user_balance) = self.get_user_balance(addr).await else { - return false; - }; + fn check_min_balance(&self, user_proofs_in_batch: usize, user_balance: U256) -> bool { let min_balance = U256::from(user_proofs_in_batch) * U256::from(MIN_FEE_PER_PROOF); - if user_balance < min_balance { - return false; - } - true + return user_balance > min_balance; } /// Handles a replacement message diff --git a/batcher/aligned-sdk/src/communication/messaging.rs b/batcher/aligned-sdk/src/communication/messaging.rs index 0ae0f6be7..14d0ad8c5 100644 --- a/batcher/aligned-sdk/src/communication/messaging.rs +++ b/batcher/aligned-sdk/src/communication/messaging.rs @@ -118,6 +118,9 @@ pub async fn send_messages( error!("Error while pushing the entry to queue"); return Err(SubmitError::AddToBatchError); } + ValidityResponseMessage::EthRpcError => { + error!("Batcher suffered connection error with Ethereum RPCs") + } }; sent_verification_data.push(verification_data.clone()); diff --git a/batcher/aligned-sdk/src/core/types.rs b/batcher/aligned-sdk/src/core/types.rs index bcfb551e9..a403dc579 100644 --- a/batcher/aligned-sdk/src/core/types.rs +++ b/batcher/aligned-sdk/src/core/types.rs @@ -318,6 +318,7 @@ pub enum ValidityResponseMessage { AddToBatchError, ProofTooLarge, InsufficientBalance(Address), + EthRpcError, } #[derive(Debug, Clone, Serialize, Deserialize)] From 9f7e6957e93a70c6b7954a1933f5dcd7c3f4f3b6 Mon Sep 17 00:00:00 2001 From: Mariano Nicolini Date: Tue, 1 Oct 2024 10:46:48 -0300 Subject: [PATCH 32/47] Small refactor on check_min_balance --- Makefile | 2 +- batcher/aligned-batcher/src/lib.rs | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/Makefile b/Makefile index 2b5f6badb..9fe58cad3 100644 --- a/Makefile +++ b/Makefile @@ -361,7 +361,7 @@ batcher_send_halo2_ipa_task_burst_5: batcher/target/release/aligned --proof ../../scripts/test_files/halo2_ipa/proof.bin \ --public_input ../../scripts/test_files/halo2_ipa/pub_input.bin \ --vk ../../scripts/test_files/halo2_ipa/params.bin \ - --repetitions 500 \ + --repetitions 1000 \ --private_key 0xdbda1821b80551c9d65939329250298aa3472ba22feea921c0cf5d620ea67b97 \ --rpc_url $(RPC_URL) \ --payment_service_addr $(BATCHER_PAYMENTS_CONTRACT_ADDRESS) diff --git a/batcher/aligned-batcher/src/lib.rs b/batcher/aligned-batcher/src/lib.rs index 379defd60..3ad729036 100644 --- a/batcher/aligned-batcher/src/lib.rs +++ b/batcher/aligned-batcher/src/lib.rs @@ -518,7 +518,7 @@ impl Batcher { // Checks user has sufficient balance for paying all its the proofs in the current batch. fn check_min_balance(&self, user_proofs_in_batch: usize, user_balance: U256) -> bool { let min_balance = U256::from(user_proofs_in_batch) * U256::from(MIN_FEE_PER_PROOF); - return user_balance > min_balance; + return user_balance >= min_balance; } /// Handles a replacement message From 68bb54e1c1e52c0b5ac707f7dcac0d4888fea7f7 Mon Sep 17 00:00:00 2001 From: Mariano Nicolini Date: Tue, 1 Oct 2024 10:48:38 -0300 Subject: [PATCH 33/47] Small refactor on handle_replacement_message --- batcher/aligned-batcher/src/lib.rs | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/batcher/aligned-batcher/src/lib.rs b/batcher/aligned-batcher/src/lib.rs index 3ad729036..449378fcf 100644 --- a/batcher/aligned-batcher/src/lib.rs +++ b/batcher/aligned-batcher/src/lib.rs @@ -610,7 +610,10 @@ impl Batcher { ); let updated_min_fee_in_batch = batch_state_lock.get_user_min_fee_in_batch(&addr); - let Some(_) = batch_state_lock.update_user_min_fee(&addr, updated_min_fee_in_batch) else { + if batch_state_lock + .update_user_min_fee(&addr, updated_min_fee_in_batch) + .is_none() + { std::mem::drop(batch_state_lock); warn!("User state for address {addr:?} was not present in batcher user states, but it should be"); return; From ddb9fd1e5672d344ee842be7ec84419bb54f4543 Mon Sep 17 00:00:00 2001 From: Mariano Nicolini Date: Tue, 1 Oct 2024 10:50:07 -0300 Subject: [PATCH 34/47] Small refactor on add_to_batch --- batcher/aligned-batcher/src/lib.rs | 15 +++++++++------ 1 file changed, 9 insertions(+), 6 deletions(-) diff --git a/batcher/aligned-batcher/src/lib.rs b/batcher/aligned-batcher/src/lib.rs index 449378fcf..9ce116b4c 100644 --- a/batcher/aligned-batcher/src/lib.rs +++ b/batcher/aligned-batcher/src/lib.rs @@ -684,12 +684,15 @@ impl Batcher { }; // User state is updated - let Some(_) = batch_state_lock.update_user_state( - &proof_submitter_addr, - nonce + U256::one(), - max_fee, - user_proof_count + 1, - ) else { + if batch_state_lock + .update_user_state( + &proof_submitter_addr, + nonce + U256::one(), + max_fee, + user_proof_count + 1, + ) + .is_none() + { error!("User state of address {proof_submitter_addr} was not found when trying to update user state. This user state should have been present"); std::mem::drop(batch_state_lock); return Err(BatcherError::AddressNotFoundInUserStates( From a1745f6d9d697e762fedc4897da218643154c71b Mon Sep 17 00:00:00 2001 From: Mariano Nicolini Date: Tue, 1 Oct 2024 10:57:11 -0300 Subject: [PATCH 35/47] Small refactor on update_user_state --- batcher/aligned-batcher/src/types/batch_state.rs | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/batcher/aligned-batcher/src/types/batch_state.rs b/batcher/aligned-batcher/src/types/batch_state.rs index 72c1ffe54..3e09377ac 100644 --- a/batcher/aligned-batcher/src/types/batch_state.rs +++ b/batcher/aligned-batcher/src/types/batch_state.rs @@ -132,10 +132,10 @@ impl BatchState { let updated_min_fee = self.update_user_min_fee(addr, new_min_fee); let updated_proof_count = self.update_user_proof_count(addr, new_proof_count); - match (updated_nonce, updated_min_fee, updated_proof_count) { - (Some(_), Some(_), Some(_)) => Some((new_nonce, new_min_fee, new_proof_count)), - _ => None, + if updated_nonce.is_some() && updated_min_fee.is_some() && updated_proof_count.is_some() { + return Some((new_nonce, new_min_fee, new_proof_count)); } + None } pub(crate) fn get_user_proofs_in_batch_and_min_fee(&self) -> HashMap { From 37a9a0ba03d86a5804a27f272fb84190501d43ba Mon Sep 17 00:00:00 2001 From: Mariano Nicolini Date: Tue, 1 Oct 2024 11:41:29 -0300 Subject: [PATCH 36/47] Revert changes in Makefile --- Makefile | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/Makefile b/Makefile index 9fe58cad3..14ee2ff33 100644 --- a/Makefile +++ b/Makefile @@ -271,7 +271,7 @@ batcher_send_risc0_burst: --proof ../../scripts/test_files/risc_zero/fibonacci_proof_generator/risc_zero_fibonacci.proof \ --vm_program ../../scripts/test_files/risc_zero/fibonacci_proof_generator/fibonacci_id.bin \ --public_input ../../scripts/test_files/risc_zero/fibonacci_proof_generator/risc_zero_fibonacci.pub \ - --repetitions 500 \ + --repetitions $(BURST_SIZE) \ --proof_generator_addr 0x66f9664f97F2b50F62D13eA064982f936dE76657 \ --rpc_url $(RPC_URL) \ --payment_service_addr $(BATCHER_PAYMENTS_CONTRACT_ADDRESS) @@ -361,8 +361,7 @@ batcher_send_halo2_ipa_task_burst_5: batcher/target/release/aligned --proof ../../scripts/test_files/halo2_ipa/proof.bin \ --public_input ../../scripts/test_files/halo2_ipa/pub_input.bin \ --vk ../../scripts/test_files/halo2_ipa/params.bin \ - --repetitions 1000 \ - --private_key 0xdbda1821b80551c9d65939329250298aa3472ba22feea921c0cf5d620ea67b97 \ + --repetitions 5 \ --rpc_url $(RPC_URL) \ --payment_service_addr $(BATCHER_PAYMENTS_CONTRACT_ADDRESS) From c7e0ae1a77e3dcf4d3a5b990c7bcbbdf5643f6fe Mon Sep 17 00:00:00 2001 From: Mariano Nicolini Date: Tue, 1 Oct 2024 11:48:01 -0300 Subject: [PATCH 37/47] Revert changes in Makefile --- Makefile | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Makefile b/Makefile index 14ee2ff33..fae398074 100644 --- a/Makefile +++ b/Makefile @@ -361,7 +361,7 @@ batcher_send_halo2_ipa_task_burst_5: batcher/target/release/aligned --proof ../../scripts/test_files/halo2_ipa/proof.bin \ --public_input ../../scripts/test_files/halo2_ipa/pub_input.bin \ --vk ../../scripts/test_files/halo2_ipa/params.bin \ - --repetitions 5 \ + --repetitions 1000 \ --rpc_url $(RPC_URL) \ --payment_service_addr $(BATCHER_PAYMENTS_CONTRACT_ADDRESS) From 815751b684fd923ea84faf0fea9480006c70b842 Mon Sep 17 00:00:00 2001 From: Mariano Nicolini Date: Tue, 1 Oct 2024 11:48:34 -0300 Subject: [PATCH 38/47] Revert changes in Makefile --- Makefile | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Makefile b/Makefile index fae398074..14ee2ff33 100644 --- a/Makefile +++ b/Makefile @@ -361,7 +361,7 @@ batcher_send_halo2_ipa_task_burst_5: batcher/target/release/aligned --proof ../../scripts/test_files/halo2_ipa/proof.bin \ --public_input ../../scripts/test_files/halo2_ipa/pub_input.bin \ --vk ../../scripts/test_files/halo2_ipa/params.bin \ - --repetitions 1000 \ + --repetitions 5 \ --rpc_url $(RPC_URL) \ --payment_service_addr $(BATCHER_PAYMENTS_CONTRACT_ADDRESS) From 3f0e36277913b9a52859f36b2fe4209879f823e6 Mon Sep 17 00:00:00 2001 From: Mariano Nicolini Date: Tue, 1 Oct 2024 12:01:19 -0300 Subject: [PATCH 39/47] Small changes in logs and removal of verification data cloning before checking its size --- batcher/aligned-batcher/src/lib.rs | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/batcher/aligned-batcher/src/lib.rs b/batcher/aligned-batcher/src/lib.rs index 9ce116b4c..facb63489 100644 --- a/batcher/aligned-batcher/src/lib.rs +++ b/batcher/aligned-batcher/src/lib.rs @@ -319,7 +319,7 @@ impl Batcher { } }; let msg_nonce = client_msg.verification_data.nonce; - info!("Received message with nonce: {msg_nonce:?}",); + debug!("Received message with nonce: {msg_nonce:?}",); // * ---------------------------------------------------* // * Perform validations over the message * @@ -352,13 +352,15 @@ impl Batcher { }; info!("Message signature verified"); - let nonced_verification_data = client_msg.verification_data.clone(); - if nonced_verification_data.verification_data.proof.len() > self.max_proof_size { + let proof_size = client_msg.verification_data.verification_data.proof.len(); + if proof_size > self.max_proof_size { error!("Proof size exceeds the maximum allowed size."); send_message(ws_conn_sink.clone(), ValidityResponseMessage::ProofTooLarge).await; return Ok(()); } + let nonced_verification_data = client_msg.verification_data.clone(); + // When pre-verification is enabled, batcher will verify proofs for faster feedback with clients if self.pre_verification_is_enabled && !zk_utils::verify(&nonced_verification_data.verification_data).await @@ -585,7 +587,7 @@ impl Batcher { replacement_entry.messaging_sink = Some(ws_conn_sink.clone()); if !batch_state_lock.replacement_entry_is_valid(&replacement_entry) { std::mem::drop(batch_state_lock); - warn!("Invalid max fee"); + warn!("Invalid replacement message"); send_message( ws_conn_sink.clone(), ValidityResponseMessage::InvalidReplacementMessage, From e30044f379efd9f19a461f428af6d9e69309b3dd Mon Sep 17 00:00:00 2001 From: Mariano Nicolini Date: Tue, 1 Oct 2024 12:57:46 -0300 Subject: [PATCH 40/47] Apply clippy suggestions --- batcher/aligned-batcher/src/lib.rs | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/batcher/aligned-batcher/src/lib.rs b/batcher/aligned-batcher/src/lib.rs index 97b8e9c96..39a7c6990 100644 --- a/batcher/aligned-batcher/src/lib.rs +++ b/batcher/aligned-batcher/src/lib.rs @@ -10,6 +10,7 @@ use ethers::signers::Signer; use types::batch_state::BatchState; use types::user_state::UserState; +use std::collections::hash_map::Entry; use std::collections::HashMap; use std::env; use std::iter::repeat; @@ -411,7 +412,7 @@ impl Batcher { // Check that we had a user state entry for this user and insert it if not. { let mut batch_state_lock = self.batch_state.lock().await; - if !batch_state_lock.user_states.contains_key(&addr) { + if let Entry::Vacant(user_state_entry) = batch_state_lock.user_states.entry(addr) { let ethereum_user_nonce = match self.get_user_nonce_from_ethereum(addr).await { Ok(ethereum_user_nonce) => ethereum_user_nonce, Err(e) => { @@ -425,7 +426,7 @@ impl Batcher { } }; let user_state = UserState::new(ethereum_user_nonce); - batch_state_lock.user_states.insert(addr, user_state); + user_state_entry.insert(user_state); } } @@ -535,7 +536,7 @@ impl Batcher { // Checks user has sufficient balance for paying all its the proofs in the current batch. fn check_min_balance(&self, user_proofs_in_batch: usize, user_balance: U256) -> bool { let min_balance = U256::from(user_proofs_in_batch) * U256::from(MIN_FEE_PER_PROOF); - return user_balance >= min_balance; + user_balance >= min_balance } /// Handles a replacement message From bfccf27fe7393a6ac283b50af418aef5348ff5f9 Mon Sep 17 00:00:00 2001 From: Mariano Nicolini Date: Tue, 1 Oct 2024 12:58:53 -0300 Subject: [PATCH 41/47] Fix code comment --- batcher/aligned-batcher/src/lib.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/batcher/aligned-batcher/src/lib.rs b/batcher/aligned-batcher/src/lib.rs index 39a7c6990..3b9fd29d9 100644 --- a/batcher/aligned-batcher/src/lib.rs +++ b/batcher/aligned-batcher/src/lib.rs @@ -585,7 +585,7 @@ impl Batcher { nonced_verification_data.verification_data.clone().into(); replacement_entry.nonced_verification_data = nonced_verification_data; - // Close old sink in old entry replace it with new one + // Close old sink in old entry and replace it with the new one { if let Some(messaging_sink) = replacement_entry.messaging_sink { let mut old_sink = messaging_sink.write().await; From 15be4a446a91f0056a1e72e0a1534ecf970f2409 Mon Sep 17 00:00:00 2001 From: Mariano Nicolini Date: Tue, 1 Oct 2024 13:05:34 -0300 Subject: [PATCH 42/47] Remove unneded return statement --- batcher/aligned-batcher/src/lib.rs | 1 - 1 file changed, 1 deletion(-) diff --git a/batcher/aligned-batcher/src/lib.rs b/batcher/aligned-batcher/src/lib.rs index 3b9fd29d9..c6fe73b4f 100644 --- a/batcher/aligned-batcher/src/lib.rs +++ b/batcher/aligned-batcher/src/lib.rs @@ -634,7 +634,6 @@ impl Batcher { { std::mem::drop(batch_state_lock); warn!("User state for address {addr:?} was not present in batcher user states, but it should be"); - return; }; } From e850304f2068900e66633de8b9803504396cca09 Mon Sep 17 00:00:00 2001 From: Urix <43704209+uri-99@users.noreply.github.com> Date: Tue, 1 Oct 2024 18:15:17 -0300 Subject: [PATCH 43/47] feat: observability log of handling paying vs nonpaying --- batcher/aligned-batcher/src/lib.rs | 3 +++ 1 file changed, 3 insertions(+) diff --git a/batcher/aligned-batcher/src/lib.rs b/batcher/aligned-batcher/src/lib.rs index c6fe73b4f..2d4f0966d 100644 --- a/batcher/aligned-batcher/src/lib.rs +++ b/batcher/aligned-batcher/src/lib.rs @@ -392,6 +392,8 @@ impl Batcher { .await; } + info!("Handling paying message"); + if self.user_balance_is_unlocked(&addr).await { send_message( ws_conn_sink.clone(), @@ -1122,6 +1124,7 @@ impl Batcher { ws_sink: WsMessageSink, client_msg: &ClientMessage, ) -> Result<(), Error> { + info!("Handling nonpaying message"); let Some(non_paying_config) = self.non_paying_config.as_ref() else { warn!("There isn't a non-paying configuration loaded. This message will be ignored"); send_message(ws_sink.clone(), ValidityResponseMessage::InvalidNonce).await; From 95897f1863fa9a188f4af93cf37e7489f7fbaa4b Mon Sep 17 00:00:00 2001 From: Mariano Nicolini Date: Wed, 2 Oct 2024 10:42:08 -0300 Subject: [PATCH 44/47] Add drop of batch state lock inside handle_nonpaying_message --- batcher/aligned-batcher/src/lib.rs | 1 + 1 file changed, 1 insertion(+) diff --git a/batcher/aligned-batcher/src/lib.rs b/batcher/aligned-batcher/src/lib.rs index 2d4f0966d..bfd91d1ff 100644 --- a/batcher/aligned-batcher/src/lib.rs +++ b/batcher/aligned-batcher/src/lib.rs @@ -1155,6 +1155,7 @@ impl Batcher { let batch_state_lock = self.batch_state.lock().await; let Some(non_paying_nonce) = batch_state_lock.get_user_nonce(&replacement_addr).await else { + std::mem::drop(batch_state_lock); error!("Nonce for non-paying address {replacement_addr:?} not found in cache."); send_message(ws_sink.clone(), ValidityResponseMessage::InvalidProof).await; return Ok(()); From ccfdacbad78818ac8224854bbb873977797c6616 Mon Sep 17 00:00:00 2001 From: Mariano Nicolini Date: Wed, 2 Oct 2024 11:46:17 -0300 Subject: [PATCH 45/47] Refactor querying of user states to hold the batch state lock the less possible --- batcher/aligned-batcher/src/lib.rs | 42 ++++++++++++++++++------------ 1 file changed, 26 insertions(+), 16 deletions(-) diff --git a/batcher/aligned-batcher/src/lib.rs b/batcher/aligned-batcher/src/lib.rs index bfd91d1ff..bce7f8c23 100644 --- a/batcher/aligned-batcher/src/lib.rs +++ b/batcher/aligned-batcher/src/lib.rs @@ -412,24 +412,34 @@ impl Batcher { } // Check that we had a user state entry for this user and insert it if not. + + // We aquire the lock first only to query if the user is already present and the lock is dropped. + // If it was not present, then the user nonce is queried to the Aligned contract. + // Lastly, we get a lock of the batch state again and insert the user state if it was still missing. + + let is_user_in_state: bool; { + let batch_state_lock = self.batch_state.lock().await; + is_user_in_state = batch_state_lock.user_states.contains_key(&addr); + } + + if !is_user_in_state { + let ethereum_user_nonce = match self.get_user_nonce_from_ethereum(addr).await { + Ok(ethereum_user_nonce) => ethereum_user_nonce, + Err(e) => { + error!( + "Failed to get user nonce from Ethereum for address {addr:?}. Error: {e:?}" + ); + send_message(ws_conn_sink.clone(), ValidityResponseMessage::InvalidNonce).await; + return Ok(()); + } + }; + let user_state = UserState::new(ethereum_user_nonce); let mut batch_state_lock = self.batch_state.lock().await; - if let Entry::Vacant(user_state_entry) = batch_state_lock.user_states.entry(addr) { - let ethereum_user_nonce = match self.get_user_nonce_from_ethereum(addr).await { - Ok(ethereum_user_nonce) => ethereum_user_nonce, - Err(e) => { - error!( - "Failed to get user nonce from Ethereum for address {addr:?}. Error: {e:?}" - ); - std::mem::drop(batch_state_lock); - send_message(ws_conn_sink.clone(), ValidityResponseMessage::InvalidNonce) - .await; - return Ok(()); - } - }; - let user_state = UserState::new(ethereum_user_nonce); - user_state_entry.insert(user_state); - } + batch_state_lock + .user_states + .entry(addr) + .or_insert(user_state); } // * ---------------------------------------------------* From f6febd95d8d0a5e2c8c2a24a004038b9787b9335 Mon Sep 17 00:00:00 2001 From: Mariano Nicolini Date: Wed, 2 Oct 2024 11:51:50 -0300 Subject: [PATCH 46/47] Remove unused import --- batcher/aligned-batcher/src/lib.rs | 1 - 1 file changed, 1 deletion(-) diff --git a/batcher/aligned-batcher/src/lib.rs b/batcher/aligned-batcher/src/lib.rs index bce7f8c23..764e44a5d 100644 --- a/batcher/aligned-batcher/src/lib.rs +++ b/batcher/aligned-batcher/src/lib.rs @@ -10,7 +10,6 @@ use ethers::signers::Signer; use types::batch_state::BatchState; use types::user_state::UserState; -use std::collections::hash_map::Entry; use std::collections::HashMap; use std::env; use std::iter::repeat; From 166509ba09940075c4c6d9a6a8a11990b5513c4b Mon Sep 17 00:00:00 2001 From: Mariano Nicolini Date: Wed, 2 Oct 2024 19:29:22 -0300 Subject: [PATCH 47/47] Add some code comments --- batcher/aligned-batcher/src/lib.rs | 17 +++++++++-------- 1 file changed, 9 insertions(+), 8 deletions(-) diff --git a/batcher/aligned-batcher/src/lib.rs b/batcher/aligned-batcher/src/lib.rs index 764e44a5d..86253dcd1 100644 --- a/batcher/aligned-batcher/src/lib.rs +++ b/batcher/aligned-batcher/src/lib.rs @@ -325,7 +325,7 @@ impl Batcher { // * Perform validations over the message * // * ---------------------------------------------------* - //This check does not save against "Holesky" and "HoleskyStage", since both are chain_id 17000 + // This check does not save against "Holesky" and "HoleskyStage", since both are chain_id 17000 let msg_chain_id = client_msg.verification_data.chain_id; if msg_chain_id != self.chain_id { warn!("Received message with incorrect chain id: {msg_chain_id}"); @@ -338,7 +338,7 @@ impl Batcher { return Ok(()); } - //This checks saves against "Holesky" and "HoleskyStage", since each one has a different payment service address + // This checks saves against "Holesky" and "HoleskyStage", since each one has a different payment service address let msg_payment_service_addr = client_msg.verification_data.payment_service_addr; if msg_payment_service_addr != self.payment_service.address() { warn!("Received message with incorrect payment service address: {msg_payment_service_addr}"); @@ -393,6 +393,9 @@ impl Batcher { info!("Handling paying message"); + // We don't need a batch state lock here, since if the user locks its funds + // after the check, some blocks should pass until he can withdraw. + // It is safe to do just do this here. if self.user_balance_is_unlocked(&addr).await { send_message( ws_conn_sink.clone(), @@ -445,18 +448,16 @@ impl Batcher { // * Perform validations over user state * // * ---------------------------------------------------* - // For now on until the message is fully processed, the batch state is locked - // This is needed because we need to query the user state to make validations and - // finally add the proof to the batch queue. - let Some(user_balance) = self.get_user_balance(&addr).await else { error!("Could not get balance for address {addr:?}"); send_message(ws_conn_sink.clone(), ValidityResponseMessage::EthRpcError).await; return Ok(()); }; - // At this point, we will have a user state for sure, since we have inserted it - // if not already present. It should be safe to call `unwrap()` here. + // For now on until the message is fully processed, the batch state is locked + // This is needed because we need to query the user state to make validations and + // finally add the proof to the batch queue. + let batch_state_lock = self.batch_state.lock().await; let Some(proofs_in_batch) = batch_state_lock.get_user_proof_count(&addr).await else { error!("Failed to get user proof count: User not found in user states, but it should have been already inserted");