diff --git a/crates/node/src/assets.rs b/crates/node/src/assets.rs index 1062cd6c5a..c7d67c389f 100644 --- a/crates/node/src/assets.rs +++ b/crates/node/src/assets.rs @@ -225,18 +225,25 @@ where self.hot_sender.send((id, value)).unwrap() } - pub async fn take_owned(&self) -> (UniqueId, T) { + pub async fn take_owned(&self) -> anyhow::Result<(UniqueId, T)> { // Always query the new condition value before taking an element. // This is to prevent the case where the condition has been updated, // but we're not yet aware of it, and the caller calls this in a loop and // we keep yielding undesired elements, but the caller keeps throwing them // away and we quickly exhaust the available assets. - self.cold_queue.lock().unwrap().update_condition_value(); + self.cold_queue + .lock() + .map_err(|_| anyhow::anyhow!("cold queue lock poisoned"))? + .update_condition_value(); loop { - let taken = self.cold_queue.lock().unwrap().take(); + let taken = self + .cold_queue + .lock() + .map_err(|_| anyhow::anyhow!("cold queue lock poisoned"))? + .take(); match taken { ColdQueueTakeResult::Taken(result) => { - return result; + return Ok(result); } ColdQueueTakeResult::NotTakenButSomeMayBeAvailable => { continue; @@ -254,10 +261,10 @@ where } received = self.hot_receiver.recv_async() => { // can't fail, because self keeps a sender. - let (id, value) = received.unwrap(); - match self.cold_queue.lock().unwrap().add_if_condition_not_satisfied(id, value) { + let (id, value) = received?; + match self.cold_queue.lock().map_err(|_| anyhow::anyhow!("cold queue lock poisoned"))?.add_if_condition_not_satisfied(id, value) { ColdQueueAddIfNotSatisfiedResult::ConditionSatisfied(value) => { - return (id, value); + return Ok((id, value)); } ColdQueueAddIfNotSatisfiedResult::Enqueued => { continue; @@ -272,14 +279,24 @@ where /// Process `num_elements_to_process`, removing any that doesn't satisfy condition. /// Return ids, that were removed from cold storage. - pub async fn maybe_discard_owned(&self, mut num_elements_to_process: usize) -> Vec { - self.cold_queue.lock().unwrap().update_condition_value(); + pub async fn maybe_discard_owned( + &self, + mut num_elements_to_process: usize, + ) -> anyhow::Result> { + self.cold_queue + .lock() + .map_err(|_| anyhow::anyhow!("cold queue lock poisoned"))? + .update_condition_value(); let mut removed_from_cold_queue: Vec = vec![]; // First process elements in the cold queue while num_elements_to_process > 0 { - let discarded = self.cold_queue.lock().unwrap().discard(); + let discarded = self + .cold_queue + .lock() + .map_err(|_| anyhow::anyhow!("cold queue lock poisoned"))? + .discard(); match discarded { ColdQueueDiscardResult::Discarded((id, _)) => { removed_from_cold_queue.push(id); @@ -304,7 +321,7 @@ where let _ = self .cold_queue .lock() - .unwrap() + .map_err(|_| anyhow::anyhow!("cold queue lock poisoned"))? .add_if_condition_satisfied(id, value); } _ => { @@ -314,20 +331,32 @@ where } } - removed_from_cold_queue + Ok(removed_from_cold_queue) } - pub fn available(&self) -> usize { - self.hot_receiver.len() + self.cold_queue.lock().unwrap().cold_available + pub fn available(&self) -> anyhow::Result { + Ok(self.hot_receiver.len() + + self + .cold_queue + .lock() + .map_err(|_| anyhow::anyhow!("cold queue lock poisoned"))? + .cold_available) } - pub fn ready(&self) -> usize { - self.cold_queue.lock().unwrap().cold_ready + pub fn ready(&self) -> anyhow::Result { + Ok(self + .cold_queue + .lock() + .map_err(|_| anyhow::anyhow!("cold queue lock poisoned"))? + .cold_ready) } - pub fn offline(&self) -> usize { - let cold_queue = self.cold_queue.lock().unwrap(); - cold_queue.cold_queue.len() - cold_queue.cold_available + pub fn offline(&self) -> anyhow::Result { + let cold_queue = self + .cold_queue + .lock() + .map_err(|_| anyhow::anyhow!("cold queue lock poisoned"))?; + Ok(cold_queue.cold_queue.len() - cold_queue.cold_available) } } @@ -457,94 +486,91 @@ where /// TODO(#10): This reservation does not persist across restarts, leading to /// the assumption that the clock moves forward at least a second across /// restarts. - pub fn generate_and_reserve_id(&self) -> UniqueId { + pub fn generate_and_reserve_id(&self) -> anyhow::Result { self.generate_and_reserve_id_range(1) } /// Same as `generate_and_reserve_id`, but for a range of IDs. /// The returned ID represents a range that starts from that ID and ending at /// that ID .add_to_counter(count - 1). - pub fn generate_and_reserve_id_range(&self, count: u32) -> UniqueId { + pub fn generate_and_reserve_id_range(&self, count: u32) -> anyhow::Result { assert!(count > 0); - let mut last_id = self.last_id.lock().unwrap(); + let mut last_id = self + .last_id + .lock() + .map_err(|_| anyhow::anyhow!("cold queue lock poisoned"))?; let start = match *last_id { Some(last_id) => last_id.pick_new_after(), None => UniqueId::generate(self.my_participant_id), }; - let end = start.add_to_counter(count - 1).unwrap(); + let end = start.add_to_counter(count - 1)?; *last_id = Some(end); - start + Ok(start) } /// Returns the current number of owned assets in the database. /// Excludes assets which are known to have offline participants. - pub fn num_owned(&self) -> usize { + pub fn num_owned(&self) -> anyhow::Result { self.owned_queue.available() } /// Returns the current number of owned assets in the database which /// are known to have all participants alive. - pub fn num_owned_ready(&self) -> usize { + pub fn num_owned_ready(&self) -> anyhow::Result { self.owned_queue.ready() } /// Returns the current number of owned assets in the database which /// are known to have some participant offline. - pub fn num_owned_offline(&self) -> usize { + pub fn num_owned_offline(&self) -> anyhow::Result { self.owned_queue.offline() } - pub async fn take_owned(&self) -> (UniqueId, T) { - let (id, asset) = self.owned_queue.take_owned().await; + pub async fn take_owned(&self) -> anyhow::Result<(UniqueId, T)> { + let (id, asset) = self.owned_queue.take_owned().await?; let mut update = self.db.update(); update.delete(self.col, &self.make_key(id)); - update - .commit() - .expect("Unrecoverable error writing to database"); - (id, asset) + update.commit()?; + Ok((id, asset)) } /// Adds an owned asset to the storage. - pub fn add_owned(&self, id: UniqueId, value: T) { + pub fn add_owned(&self, id: UniqueId, value: T) -> anyhow::Result<()> { let key = self.make_key(id); - let value_ser = serde_json::to_vec(&value).unwrap(); + let value_ser = serde_json::to_vec(&value)?; let mut update = self.db.update(); update.put(self.col, &key, &value_ser); - update - .commit() - .expect("Unrecoverable error writing to database"); + update.commit()?; // Can't fail, because we keep a receiver alive. - self.owned_queue.add_owned(id, value); + Ok(self.owned_queue.add_owned(id, value)) } /// Examines up to `num_assets_to_process` elements in the storage. /// If any are found not to satisfy the current condition, they are discarded. /// Otherwise, they are kept aside as ready for immediate use. - pub async fn maybe_discard_owned(&self, num_assets_to_process: usize) { + pub async fn maybe_discard_owned(&self, num_assets_to_process: usize) -> anyhow::Result<()> { let removed_cold_ids = self .owned_queue .maybe_discard_owned(num_assets_to_process) - .await; + .await?; if !removed_cold_ids.is_empty() { let mut update = self.db.update(); for id in removed_cold_ids { update.delete(self.col, &self.make_key(id)); } - update - .commit() - .expect("Unrecoverable error writing to database"); + update.commit()?; } + Ok(()) } /// Adds an unowned asset to the storage. - pub fn add_unowned(&self, id: UniqueId, value: T) { + pub fn add_unowned(&self, id: UniqueId, value: T) -> anyhow::Result<()> { let key = self.make_key(id); - let value_ser = serde_json::to_vec(&value).unwrap(); + let value_ser = serde_json::to_vec(&value)?; let mut update = self.db.update(); update.put(self.col, &key, &value_ser); - update - .commit() - .expect("Unrecoverable error writing to database"); + update.commit()?; + Ok(()) } /// Removes an unowned asset from the storage and returns it. Returns @@ -554,7 +580,10 @@ where // Prevent two concurrent callers from both reading the same asset // before either commits the delete (read-then-delete race). { - let mut in_flight = self.unowned_in_flight.lock().unwrap(); + let mut in_flight = self + .unowned_in_flight + .lock() + .map_err(|_| anyhow::anyhow!("cold queue lock poisoned"))?; if !in_flight.insert(id) { anyhow::bail!( "Unowned {} is already being taken by another task: {:?}", @@ -565,7 +594,10 @@ where } let result = self.take_unowned_inner(id); // Always remove from in-flight, whether the take succeeded or not. - self.unowned_in_flight.lock().unwrap().remove(&id); + self.unowned_in_flight + .lock() + .map_err(|_| anyhow::anyhow!("cold queue lock poisoned"))? + .remove(&id); result } @@ -576,9 +608,7 @@ where })?; let mut update = self.db.update(); update.delete(self.col, &key); - update - .commit() - .expect("Unrecoverable error writing to database"); + update.commit()?; Ok(serde_json::from_slice(&value_ser)?) } } @@ -588,7 +618,7 @@ mod tests { use super::{ColdQueue, DistributedAssetStorage, DoubleQueue, UniqueId}; use crate::assets::clean_db; use crate::async_testing::{MaybeReady, run_future_once}; - use crate::db::DBCol; + use crate::db::{DBCol, SecretDB}; use crate::primitives::ParticipantId; use crate::providers::HasParticipants; use borsh::BorshDeserialize; @@ -601,6 +631,22 @@ mod tests { use std::sync::atomic::{AtomicI32, AtomicUsize, Ordering}; use std::sync::{Arc, Mutex}; + fn storage( + db: Arc, + participant: ParticipantId, + domain_id: Vec, + ) -> anyhow::Result> { + DistributedAssetStorage::::new( + FakeClock::default().clock(), + db, + DBCol::Presignature, + domain_id, + participant, + |_, _| true, + Arc::new(Vec::new), + ) + } + /// Adapter used by tests that previously took `Option` to compose /// the equivalent prefix bytes for the generalized [`DistributedAssetStorage`]. fn domain_id_to_prefix(domain_id: Option) -> Vec { @@ -717,7 +763,11 @@ mod tests { }); // Discard should never block, even if the queue is completely empty - queue.maybe_discard_owned(3).now_or_never().unwrap(); + queue + .maybe_discard_owned(3) + .now_or_never() + .unwrap() + .unwrap(); // Add 3 elements, 2 of which don't match the condition let id1 = UniqueId::new(ParticipantId::from_raw(42), 123, 456); @@ -727,25 +777,43 @@ mod tests { queue.add_owned(id1, 1); queue.add_owned(id2, 2); queue.add_owned(id3, 3); - assert_eq!(queue.available(), 3); + assert_eq!(queue.available().unwrap(), 3); - queue.maybe_discard_owned(1).now_or_never().unwrap(); - assert_eq!(queue.available(), 2); + queue + .maybe_discard_owned(1) + .now_or_never() + .unwrap() + .unwrap(); + assert_eq!(queue.available().unwrap(), 2); - assert_eq!(queue.take_owned().now_or_never().unwrap(), (id2, 2)); - assert_eq!(queue.available(), 1); + assert_eq!( + queue.take_owned().now_or_never().unwrap().unwrap(), + (id2, 2) + ); + assert_eq!(queue.available().unwrap(), 1); - queue.maybe_discard_owned(1).now_or_never().unwrap(); - assert_eq!(queue.available(), 0); + queue + .maybe_discard_owned(1) + .now_or_never() + .unwrap() + .unwrap(); + assert_eq!(queue.available().unwrap(), 0); queue.add_owned(id4, 4); - assert_eq!(queue.available(), 1); + assert_eq!(queue.available().unwrap(), 1); - queue.maybe_discard_owned(1).now_or_never().unwrap(); - assert_eq!(queue.available(), 1); + queue + .maybe_discard_owned(1) + .now_or_never() + .unwrap() + .unwrap(); + assert_eq!(queue.available().unwrap(), 1); - assert_eq!(queue.take_owned().now_or_never().unwrap(), (id4, 4)); - assert_eq!(queue.available(), 0); + assert_eq!( + queue.take_owned().now_or_never().unwrap().unwrap(), + (id4, 4) + ); + assert_eq!(queue.available().unwrap(), 0); } // This test covers tricky cases around updates to the condition value @@ -772,7 +840,10 @@ mod tests { // Make condition "% 2 == 1". cond_value.store(1, Ordering::Relaxed); - assert_eq!(queue.take_owned().now_or_never().unwrap(), (id1, 1)); + assert_eq!( + queue.take_owned().now_or_never().unwrap().unwrap(), + (id1, 1) + ); assert_eq!(cond_value_query_count.load(Ordering::Relaxed), 1); // Make condition "% 2 == 0" and start taking an element. @@ -793,7 +864,7 @@ mod tests { // Advance the clock so that the waiting task notices the condition change. clock.advance(near_time::Duration::seconds(1)); - assert_eq!(fut.now_or_never().unwrap(), (id2, 3)); + assert_eq!(fut.now_or_never().unwrap().unwrap(), (id2, 3)); assert_eq!(cond_value_query_count.load(Ordering::Relaxed), 3); // This time change the condition before starting to take an element. @@ -815,11 +886,14 @@ mod tests { // Even though the condition changed, we may get an element returned that satisfied a // stale condition (there's no point to prevent that because there can always be // races). - assert_eq!(fut.now_or_never().unwrap(), (id4, 4)); + assert_eq!(fut.now_or_never().unwrap().unwrap(), (id4, 4)); assert_eq!(cond_value_query_count.load(Ordering::Relaxed), 4); // However, if we take_owned() again, we'll use the correct condition. - assert_eq!(queue.take_owned().now_or_never().unwrap(), (id3, 5)); + assert_eq!( + queue.take_owned().now_or_never().unwrap().unwrap(), + (id3, 5) + ); assert_eq!(cond_value_query_count.load(Ordering::Relaxed), 5); } @@ -827,7 +901,7 @@ mod tests { fn test_distributed_assets_storage() { let clock = FakeClock::default(); let dir = tempfile::tempdir().unwrap(); - let db = crate::db::SecretDB::new(dir.path(), [1; 16]).unwrap(); + let db = SecretDB::new(dir.path(), [1; 16]).unwrap(); let all_participants = vec![ ParticipantId::from_raw(0), ParticipantId::from_raw(1), @@ -848,7 +922,7 @@ mod tests { let store = DistributedAssetStorage::::new( clock.clock(), db, - crate::db::DBCol::TripleV2, + DBCol::TripleV2, Vec::new(), ParticipantId::from_raw(42), |cond, val| val.is_subset_of_active_participants(cond), @@ -858,28 +932,34 @@ mod tests { }, ) .unwrap(); - assert_eq!(store.num_owned(), 0); - - let id1 = store.generate_and_reserve_id(); - let id2 = store.generate_and_reserve_id(); - let id3 = store.generate_and_reserve_id(); - let id4 = store.generate_and_reserve_id(); - let id5 = store.generate_and_reserve_id(); - store.add_owned(id1, ParticipantsWithI32(all_participants.clone(), 123)); - assert_eq!(store.num_owned(), 1); - store.add_owned(id2, ParticipantsWithI32(all_participants.clone(), 456)); - assert_eq!(store.num_owned(), 2); - let asset1 = store.take_owned().now_or_never().unwrap(); + assert_eq!(store.num_owned().unwrap(), 0); + + let id1 = store.generate_and_reserve_id().unwrap(); + let id2 = store.generate_and_reserve_id().unwrap(); + let id3 = store.generate_and_reserve_id().unwrap(); + let id4 = store.generate_and_reserve_id().unwrap(); + let id5 = store.generate_and_reserve_id().unwrap(); + store + .add_owned(id1, ParticipantsWithI32(all_participants.clone(), 123)) + .unwrap(); + assert_eq!(store.num_owned().unwrap(), 1); + store + .add_owned(id2, ParticipantsWithI32(all_participants.clone(), 456)) + .unwrap(); + assert_eq!(store.num_owned().unwrap(), 2); + let asset1 = store.take_owned().now_or_never().unwrap().unwrap(); assert_eq!( asset1, (id1, ParticipantsWithI32(all_participants.clone(), 123)) ); - assert_eq!(store.num_owned(), 1); - store.add_owned( - id3, - ParticipantsWithI32(second_participants_subset.clone(), 789), - ); - assert_eq!(store.num_owned(), 2); + assert_eq!(store.num_owned().unwrap(), 1); + store + .add_owned( + id3, + ParticipantsWithI32(second_participants_subset.clone(), 789), + ) + .unwrap(); + assert_eq!(store.num_owned().unwrap(), 2); *alive_participants.lock().unwrap() = first_participants_subset.clone(); let asset_fut = store.take_owned(); @@ -888,12 +968,14 @@ mod tests { panic!("Cannot take value since set of participants has changed"); }; - store.add_owned( - id4, - ParticipantsWithI32(first_participants_subset.clone(), 101112), - ); + store + .add_owned( + id4, + ParticipantsWithI32(first_participants_subset.clone(), 101112), + ) + .unwrap(); - let asset3 = store.take_owned().now_or_never().unwrap(); + let asset3 = store.take_owned().now_or_never().unwrap().unwrap(); assert_eq!( asset3, ( @@ -906,45 +988,49 @@ mod tests { panic!("Cannot take value since set of participants has changed"); }; - store.add_owned( - id4, - ParticipantsWithI32(first_participants_subset.clone(), 131415), - ); + store + .add_owned( + id4, + ParticipantsWithI32(first_participants_subset.clone(), 131415), + ) + .unwrap(); assert_eq!( - asset_fut.now_or_never().unwrap(), + asset_fut.now_or_never().unwrap().unwrap(), ( id4, ParticipantsWithI32(first_participants_subset.clone(), 131415) ) ); - assert_eq!(store.num_owned(), 0); + assert_eq!(store.num_owned().unwrap(), 0); // Now go back to all participants being available. *alive_participants.lock().unwrap() = all_participants.clone(); - store.add_owned(id5, ParticipantsWithI32(all_participants.clone(), 161718)); - assert_eq!(store.num_owned(), 1); + store + .add_owned(id5, ParticipantsWithI32(all_participants.clone(), 161718)) + .unwrap(); + assert_eq!(store.num_owned().unwrap(), 1); // Previously ineligible assets (456, 789, and 161718) should now be available. assert_eq!( - store.take_owned().now_or_never().unwrap(), + store.take_owned().now_or_never().unwrap().unwrap(), (id2, ParticipantsWithI32(all_participants.clone(), 456)) ); - assert_eq!(store.num_owned(), 2); + assert_eq!(store.num_owned().unwrap(), 2); assert_eq!( - store.take_owned().now_or_never().unwrap(), + store.take_owned().now_or_never().unwrap().unwrap(), ( id3, ParticipantsWithI32(second_participants_subset.clone(), 789) ) ); - assert_eq!(store.num_owned(), 1); + assert_eq!(store.num_owned().unwrap(), 1); assert_eq!( - store.take_owned().now_or_never().unwrap(), + store.take_owned().now_or_never().unwrap().unwrap(), (id5, ParticipantsWithI32(all_participants.clone(), 161718)) ); - assert_eq!(store.num_owned(), 0); + assert_eq!(store.num_owned().unwrap(), 0); } #[test] @@ -979,33 +1065,24 @@ mod tests { #[test] fn test_distributed_store_add_take_owned() { let dir = tempfile::tempdir().unwrap(); - let db = crate::db::SecretDB::new(dir.path(), [1; 16]).unwrap(); - let store = DistributedAssetStorage::::new( - FakeClock::default().clock(), - db, - crate::db::DBCol::TripleV2, - Vec::new(), - ParticipantId::from_raw(42), - |_, _| true, - Arc::new(std::vec::Vec::new), - ) - .unwrap(); - assert_eq!(store.num_owned(), 0); + let db = SecretDB::new(dir.path(), [1; 16]).unwrap(); + let store = storage(db.clone(), ParticipantId::from_raw(42), Vec::new()).unwrap(); + assert_eq!(store.num_owned().unwrap(), 0); // Put in two assets, then dequeue them. - let id1 = store.generate_and_reserve_id(); - let id2 = store.generate_and_reserve_id_range(2); + let id1 = store.generate_and_reserve_id().unwrap(); + let id2 = store.generate_and_reserve_id_range(2).unwrap(); assert!(id2 > id1); - store.add_owned(id1, 123); - assert_eq!(store.num_owned(), 1); - store.add_owned(id2, 456); - assert_eq!(store.num_owned(), 2); - let asset1 = store.take_owned().now_or_never().unwrap(); + store.add_owned(id1, 123).unwrap(); + assert_eq!(store.num_owned().unwrap(), 1); + store.add_owned(id2, 456).unwrap(); + assert_eq!(store.num_owned().unwrap(), 2); + let asset1 = store.take_owned().now_or_never().unwrap().unwrap(); assert_eq!(asset1, (id1, 123)); - assert_eq!(store.num_owned(), 1); - let asset2 = store.take_owned().now_or_never().unwrap(); + assert_eq!(store.num_owned().unwrap(), 1); + let asset2 = store.take_owned().now_or_never().unwrap().unwrap(); assert_eq!(asset2, (id2, 456)); - assert_eq!(store.num_owned(), 0); + assert_eq!(store.num_owned().unwrap(), 0); // Dequeuing an asset before it's available will block. let asset3_fut = store.take_owned(); @@ -1014,33 +1091,24 @@ mod tests { }; let id3 = id2.add_to_counter(1).unwrap(); - store.add_owned(id3, 789); - let asset3 = asset3_fut.now_or_never().unwrap(); + store.add_owned(id3, 789).unwrap(); + let asset3 = asset3_fut.now_or_never().unwrap().unwrap(); assert_eq!(asset3, (id3, 789)); // Sanity check that generated IDs are monotonically increasing. - let id4 = store.generate_and_reserve_id(); + let id4 = store.generate_and_reserve_id().unwrap(); assert!(id4 > id3); } #[test] fn test_distributed_store_add_owned_different_order() { let dir = tempfile::tempdir().unwrap(); - let db = crate::db::SecretDB::new(dir.path(), [1; 16]).unwrap(); - let store = DistributedAssetStorage::::new( - FakeClock::default().clock(), - db.clone(), - crate::db::DBCol::TripleV2, - Vec::new(), - ParticipantId::from_raw(42), - |_, _| true, - Arc::new(std::vec::Vec::new), - ) - .unwrap(); + let db = SecretDB::new(dir.path(), [1; 16]).unwrap(); + let store = storage(db.clone(), ParticipantId::from_raw(42), Vec::new()).unwrap(); // Adding assets in a different order from when the IDs are generated // is fine. They are dequeued in the order that they are queued. - let id1 = store.generate_and_reserve_id_range(3); + let id1 = store.generate_and_reserve_id_range(3).unwrap(); let id2 = id1.add_to_counter(1).unwrap(); let id3 = id1.add_to_counter(2).unwrap(); @@ -1053,56 +1121,56 @@ mod tests { panic!("nothing should not be ready"); }; - store.add_owned(id3, 3); - store.add_owned(id2, 2); + store.add_owned(id3, 3).unwrap(); + store.add_owned(id2, 2).unwrap(); - assert_eq!(asset1_fut.now_or_never().unwrap(), (id3, 3)); - assert_eq!(asset2_fut.now_or_never().unwrap(), (id2, 2)); + assert_eq!(asset1_fut.now_or_never().unwrap().unwrap(), (id3, 3)); + assert_eq!(asset2_fut.now_or_never().unwrap().unwrap(), (id2, 2)); - store.add_owned(id1, 1); - assert_eq!(store.take_owned().now_or_never().unwrap(), (id1, 1)); + store.add_owned(id1, 1).unwrap(); + assert_eq!( + store.take_owned().now_or_never().unwrap().unwrap(), + (id1, 1) + ); // Make sure that ID generation does not depend on the order of adding // them. - let id4 = store.generate_and_reserve_id(); + let id4 = store.generate_and_reserve_id().unwrap(); assert!(id4 > id3); - let id5 = store.generate_and_reserve_id(); - let id6 = store.generate_and_reserve_id(); + let id5 = store.generate_and_reserve_id().unwrap(); + let id6 = store.generate_and_reserve_id().unwrap(); - store.add_owned(id6, 6); - store.add_owned(id5, 5); + store.add_owned(id6, 6).unwrap(); + store.add_owned(id5, 5).unwrap(); // If we reload the store from the db, then the order of the queue would // be based on the key. It doesn't have to be this way, but we test it // here just to clarify the current behavior. drop(store); - let store = DistributedAssetStorage::::new( - FakeClock::default().clock(), - db, - crate::db::DBCol::TripleV2, - Vec::new(), - ParticipantId::from_raw(42), - |_, _| true, - Arc::new(std::vec::Vec::new), - ) - .unwrap(); - assert_eq!(store.take_owned().now_or_never().unwrap(), (id5, 5)); - assert_eq!(store.take_owned().now_or_never().unwrap(), (id6, 6)); + let store = storage(db.clone(), ParticipantId::from_raw(42), Vec::new()).unwrap(); + assert_eq!( + store.take_owned().now_or_never().unwrap().unwrap(), + (id5, 5) + ); + assert_eq!( + store.take_owned().now_or_never().unwrap().unwrap(), + (id6, 6) + ); } #[test] - fn test_distribtued_store_add_take_unowned() { + fn test_distributed_store_add_take_unowned() { let dir = tempfile::tempdir().unwrap(); - let db = crate::db::SecretDB::new(dir.path(), [1; 16]).unwrap(); + let db = SecretDB::new(dir.path(), [1; 16]).unwrap(); let store = DistributedAssetStorage::::new( FakeClock::default().clock(), db, - crate::db::DBCol::TripleV2, + DBCol::TripleV2, Vec::new(), ParticipantId::from_raw(42), |_, _| true, - Arc::new(std::vec::Vec::new), + Arc::new(Vec::new), ) .unwrap(); @@ -1110,9 +1178,9 @@ mod tests { let id1 = UniqueId::new(other, 1, 0); let id2 = UniqueId::new(other, 2, 0); let id3 = UniqueId::new(other, 3, 0); - store.add_unowned(id1, 123); - store.add_unowned(id2, 234); - assert_eq!(store.num_owned(), 0); // does not affect owned + store.add_unowned(id1, 123).unwrap(); + store.add_unowned(id2, 234).unwrap(); + assert_eq!(store.num_owned().unwrap(), 0); // does not affect owned assert_eq!(store.take_unowned(id1).unwrap(), 123); let _ = store @@ -1134,54 +1202,48 @@ mod tests { #[test] fn test_distributed_store_persistence() { let dir = tempfile::tempdir().unwrap(); - let db = crate::db::SecretDB::new(dir.path(), [1; 16]).unwrap(); + let db = SecretDB::new(dir.path(), [1; 16]).unwrap(); let myself = ParticipantId::from_raw(42); let store = DistributedAssetStorage::::new( FakeClock::default().clock(), db.clone(), - crate::db::DBCol::TripleV2, + DBCol::TripleV2, Vec::new(), myself, |_, _| true, - Arc::new(std::vec::Vec::new), + Arc::new(Vec::new), ) .unwrap(); - let id1 = store.generate_and_reserve_id_range(4); - store.add_owned(id1, 1); - store.add_owned(id1.add_to_counter(1).unwrap(), 2); - store.add_owned(id1.add_to_counter(2).unwrap(), 3); - store.add_owned(id1.add_to_counter(3).unwrap(), 4); + let id1 = store.generate_and_reserve_id_range(4).unwrap(); + let _ = store.add_owned(id1, 1).unwrap(); + store.add_owned(id1.add_to_counter(1).unwrap(), 2).unwrap(); + store.add_owned(id1.add_to_counter(2).unwrap(), 3).unwrap(); + store.add_owned(id1.add_to_counter(3).unwrap(), 4).unwrap(); let other = ParticipantId::from_raw(43); - store.add_unowned(UniqueId::new(other, 1, 0), 5); - store.add_unowned(UniqueId::new(other, 2, 0), 6); - store.add_unowned(UniqueId::new(other, 3, 0), 7); - store.add_unowned(UniqueId::new(other, 4, 0), 8); + store.add_unowned(UniqueId::new(other, 1, 0), 5).unwrap(); + store.add_unowned(UniqueId::new(other, 2, 0), 6).unwrap(); + store.add_unowned(UniqueId::new(other, 3, 0), 7).unwrap(); + store.add_unowned(UniqueId::new(other, 4, 0), 8).unwrap(); drop(store); - let store = DistributedAssetStorage::::new( - FakeClock::default().clock(), - db, - crate::db::DBCol::TripleV2, - Vec::new(), - myself, - |_, _| true, - Arc::new(std::vec::Vec::new), - ) - .unwrap(); - assert_eq!(store.num_owned(), 4); - assert_eq!(store.take_owned().now_or_never().unwrap(), (id1, 1)); + let store = storage(db.clone(), myself, Vec::new()).unwrap(); + assert_eq!(store.num_owned().unwrap(), 4); assert_eq!( - store.take_owned().now_or_never().unwrap(), + store.take_owned().now_or_never().unwrap().unwrap(), + (id1, 1) + ); + assert_eq!( + store.take_owned().now_or_never().unwrap().unwrap(), (id1.add_to_counter(1).unwrap(), 2) ); assert_eq!( - store.take_owned().now_or_never().unwrap(), + store.take_owned().now_or_never().unwrap().unwrap(), (id1.add_to_counter(2).unwrap(), 3) ); assert_eq!( - store.take_owned().now_or_never().unwrap(), + store.take_owned().now_or_never().unwrap().unwrap(), (id1.add_to_counter(3).unwrap(), 4) ); @@ -1191,13 +1253,13 @@ mod tests { #[test] fn test_maybe_discard_unowned_persistence() { let dir = tempfile::tempdir().unwrap(); - let db = crate::db::SecretDB::new(dir.path(), [1; 16]).unwrap(); + let db = SecretDB::new(dir.path(), [1; 16]).unwrap(); let myself = ParticipantId::from_raw(42); let store = DistributedAssetStorage::::new( FakeClock::default().clock(), db.clone(), - crate::db::DBCol::TripleV2, + DBCol::TripleV2, Vec::new(), myself, |_, x| *x != 1, @@ -1206,56 +1268,51 @@ mod tests { .unwrap(); // Push asset to the cold queue - let id1 = store.generate_and_reserve_id_range(2); - store.add_owned(id1, 1); - store.add_owned(id1.add_to_counter(1).unwrap(), 2); - assert_eq!(store.take_owned().now_or_never().unwrap().1, 2); - assert_eq!(store.num_owned_offline(), 1); + let id1 = store.generate_and_reserve_id_range(2).unwrap(); + store.add_owned(id1, 1).unwrap(); + store.add_owned(id1.add_to_counter(1).unwrap(), 2).unwrap(); + assert_eq!(store.take_owned().now_or_never().unwrap().unwrap().1, 2); + assert_eq!(store.num_owned_offline().unwrap(), 1); - store.maybe_discard_owned(1).now_or_never().unwrap(); + let _ = store.maybe_discard_owned(1).now_or_never().unwrap(); drop(store); let store = DistributedAssetStorage::::new( FakeClock::default().clock(), db, - crate::db::DBCol::TripleV2, + DBCol::TripleV2, Vec::new(), myself, |_, _| true, - Arc::new(std::vec::Vec::new), + Arc::new(Vec::new), ) .unwrap(); - assert_eq!(store.num_owned(), 0); + assert_eq!(store.num_owned().unwrap(), 0); } #[test] fn test_multiple_domains() { let dir = tempfile::tempdir().unwrap(); - let db = crate::db::SecretDB::new(dir.path(), [1; 16]).unwrap(); + let db = SecretDB::new(dir.path(), [1; 16]).unwrap(); let myself = ParticipantId::from_raw(42); let other = ParticipantId::from_raw(43); for i in 0..4 { let domain_id = Some(DomainId(i)); - let store = DistributedAssetStorage::::new( - FakeClock::default().clock(), - db.clone(), - crate::db::DBCol::Presignature, - domain_id_to_prefix(domain_id), - myself, - |_, _| true, - Arc::new(std::vec::Vec::new), - ) - .unwrap(); + let store = storage(db.clone(), myself, domain_id_to_prefix(domain_id)).unwrap(); for j in 0..10 { - store.add_owned(UniqueId::new(myself, j, 0), 10000 + i * 100 + j); - store.add_unowned(UniqueId::new(other, j, 0), 20000 + i * 100 + j); + store + .add_owned(UniqueId::new(myself, j, 0), 10000 + i * 100 + j) + .unwrap(); + store + .add_unowned(UniqueId::new(other, j, 0), 20000 + i * 100 + j) + .unwrap(); } for j in 0..10 { assert_eq!( - store.take_owned().now_or_never().unwrap().1, + store.take_owned().now_or_never().unwrap().unwrap().1, 10000 + i * 100 + j ); assert_eq!( @@ -1264,8 +1321,12 @@ mod tests { ); } for j in 0..10 { - store.add_owned(UniqueId::new(myself, 100 + j, 0), 30000 + i * 100 + j); - store.add_unowned(UniqueId::new(other, 100 + j, 0), 40000 + i * 100 + j); + store + .add_owned(UniqueId::new(myself, 100 + j, 0), 30000 + i * 100 + j) + .unwrap(); + store + .add_unowned(UniqueId::new(other, 100 + j, 0), 40000 + i * 100 + j) + .unwrap(); } } @@ -1274,17 +1335,17 @@ mod tests { let store = DistributedAssetStorage::::new( FakeClock::default().clock(), db.clone(), - crate::db::DBCol::Presignature, + DBCol::Presignature, domain_id_to_prefix(domain_id), myself, |_, _| true, - Arc::new(std::vec::Vec::new), + Arc::new(Vec::new), ) .unwrap(); for j in 0..10 { assert_eq!( - store.take_owned().now_or_never().unwrap().1, + store.take_owned().now_or_never().unwrap().unwrap().1, 30000 + i * 100 + j ); assert_eq!( @@ -1301,7 +1362,7 @@ mod tests { fn test_distributed_assets_storage_cleanup() { let clock = FakeClock::default(); let dir = tempfile::tempdir().unwrap(); - let db = crate::db::SecretDB::new(dir.path(), [1; 16]).unwrap(); + let db = SecretDB::new(dir.path(), [1; 16]).unwrap(); let all_participants = vec![ ParticipantId::from_raw(0), ParticipantId::from_raw(1), @@ -1340,10 +1401,10 @@ mod tests { }; let assert_db_num_owned = |db_col: DBCol, domain_id: Option, expected: usize| { let store = new_store_from_db(db_col, domain_id); - assert_eq!(store.num_owned(), expected); + assert_eq!(store.num_owned().unwrap(), expected); }; for domain_id in [None, Some(DomainId(0)), Some(DomainId(1))] { - for db_col in [crate::db::DBCol::Presignature, crate::db::DBCol::TripleV2] { + for db_col in [DBCol::Presignature, DBCol::TripleV2] { assert_db_num_owned(db_col, domain_id, 0); { // populate the database @@ -1353,8 +1414,8 @@ mod tests { let subset_c_1 = ParticipantsWithI32(participant_subset_c.clone(), 789); let store = new_store_from_db(db_col, domain_id); for p in [all_1, subset_a_1, subset_b_1, subset_c_1] { - let id = store.generate_and_reserve_id(); - store.add_owned(id, p); + let id = store.generate_and_reserve_id().unwrap(); + store.add_owned(id, p).unwrap(); } } assert_db_num_owned(db_col, domain_id, 4); diff --git a/crates/node/src/assets/cleanup.rs b/crates/node/src/assets/cleanup.rs index 0fb77827c0..d1a12f82b4 100644 --- a/crates/node/src/assets/cleanup.rs +++ b/crates/node/src/assets/cleanup.rs @@ -255,8 +255,10 @@ mod tests { reconstruction_threshold, ) .unwrap(); - let id = triple_store.generate_and_reserve_id(); - triple_store.add_owned(id, make_triple(&all_participants)); + let id = triple_store.generate_and_reserve_id().unwrap(); + triple_store + .add_owned(id, make_triple(&all_participants)) + .unwrap(); let v2_key = triple_v2_key(reconstruction_threshold, id); // Sanity: the row really is in the TripleV2 column. assert!(db.get(DBCol::TripleV2, &v2_key).unwrap().is_some()); @@ -324,8 +326,10 @@ mod tests { reconstruction_threshold, ) .unwrap(); - let id = triple_store.generate_and_reserve_id(); - triple_store.add_owned(id, make_triple(&active_subset)); + let id = triple_store.generate_and_reserve_id().unwrap(); + triple_store + .add_owned(id, make_triple(&active_subset)) + .unwrap(); let v2_key = triple_v2_key(reconstruction_threshold, id); // Seed epoch_data, then run cleanup with one participant's TLS rotated @@ -386,7 +390,9 @@ mod tests { // Even an outright-stale peer-owned triple stays put — the peer cleans // its own. let peer_id = UniqueId::new(peer, 100, 0); - triple_store.add_unowned(peer_id, make_triple(&all_participants)); + triple_store + .add_unowned(peer_id, make_triple(&all_participants)) + .unwrap(); let v2_key = triple_v2_key(reconstruction_threshold, peer_id); delete_stale_triples_and_presignatures( diff --git a/crates/node/src/assets/test_utils.rs b/crates/node/src/assets/test_utils.rs index 59f3749140..020da64bfd 100644 --- a/crates/node/src/assets/test_utils.rs +++ b/crates/node/src/assets/test_utils.rs @@ -128,7 +128,7 @@ impl TestContext { alive_participants: Arc>>, ) -> Self { let dir = tempfile::tempdir().unwrap(); - let db = crate::db::SecretDB::new(dir.path(), [1; 16]).unwrap(); + let db = SecretDB::new(dir.path(), [1; 16]).unwrap(); Self { db, clock: FakeClock::default(), @@ -165,29 +165,29 @@ impl TestContext { pub fn populate(&self, participants: &[ParticipantId]) { // Mirror cleanup's view of triples: per-`t` TripleV2 column. let store = self.new_store::(DBCol::TripleV2, self.triple_prefix()); - let id = store.generate_and_reserve_id(); - store.add_owned(id, make_triple(participants)); + let id = store.generate_and_reserve_id().unwrap(); + store.add_owned(id, make_triple(participants)).unwrap(); for &d in &self.presign_domain_ids { let store = self.new_store::( DBCol::Presignature, d.0.to_be_bytes().to_vec(), ); - let id = store.generate_and_reserve_id(); - store.add_owned(id, make_presign(participants)); + let id = store.generate_and_reserve_id().unwrap(); + store.add_owned(id, make_presign(participants)).unwrap(); } } pub fn assert_owned(&self, expected: usize) { let store = self.new_store::(DBCol::TripleV2, self.triple_prefix()); - assert_eq!(store.num_owned(), expected); + assert_eq!(store.num_owned().unwrap(), expected); for &d in &self.presign_domain_ids { let store = self.new_store::( DBCol::Presignature, d.0.to_be_bytes().to_vec(), ); - assert_eq!(store.num_owned(), expected); + assert_eq!(store.num_owned().unwrap(), expected); } } } diff --git a/crates/node/src/mpc_client.rs b/crates/node/src/mpc_client.rs index 03e4cc3c0f..8fc72fd6dd 100644 --- a/crates/node/src/mpc_client.rs +++ b/crates/node/src/mpc_client.rs @@ -159,7 +159,7 @@ where let metrics_emitter = tracking::spawn("periodically emits metrics", async move { loop { client.emit_metrics(); - tokio::time::sleep(std::time::Duration::from_secs(5)).await; + tokio::time::sleep(Duration::from_secs(5)).await; } }); @@ -169,17 +169,14 @@ where MpcClient::monitor_passive_channels_inner(channel_receiver, self.clone()), ) }; - + let self_clone = self.clone(); let monitor_chain = { let chain_txn_sender = chain_txn_sender.clone(); - tracking::spawn( - "monitor chain", - self.clone().monitor_block_updates( - block_update_receiver, - chain_txn_sender, - debug_receiver, - ), - ) + tracking::spawn("monitor chain", async { + self_clone + .monitor_block_updates(block_update_receiver, chain_txn_sender, debug_receiver) + .await + }) }; let tee_verification_handle = { @@ -237,7 +234,7 @@ where let _ = monitor_passive_channels.await?; metrics_emitter.await?; - monitor_chain.await?; + monitor_chain.await??; let _ = robust_ecdsa_background_tasks.await?; let _ = ecdsa_background_tasks.await?; let _ = eddsa_background_tasks.await?; @@ -254,7 +251,7 @@ where >, chain_txn_sender: impl TransactionSender + 'static, mut debug_receiver: tokio::sync::broadcast::Receiver, - ) { + ) -> anyhow::Result { let mut tasks = AutoAbortTaskCollection::new(); let mut pending_signatures = PendingRequests::::new( @@ -289,7 +286,7 @@ where let Some(block_update) = block_update else { // If this branch hits, it means the channel is closed, meaning the // indexer is being shutdown. So just quit this task. - break; + return Err(anyhow::anyhow!("Block update channel closed")); }; self.client.update_indexer_height(block_update.block.height.into()); @@ -305,7 +302,7 @@ where // TODO(#3031): add batch request and unify stores for request in &signature_requests.requests { - self.sign_request_store.add(request); + self.sign_request_store.add(request)?; } // TODO(#3032): remove completed & finalized requests from store @@ -318,7 +315,7 @@ where block_update.completed_ckds ); for request in &ckd_requests.requests { - self.ckd_request_store.add(request); + self.ckd_request_store.add(request)?; } pending_ckds.notify_new_block(ckd_requests); @@ -331,7 +328,7 @@ where ); for request in &verify_foreign_tx_requests.requests { - self.verify_foreign_tx_request_store.add(request); + self.verify_foreign_tx_request_store.add(request)?; } pending_verify_foreign_txs.notify_new_block(verify_foreign_tx_requests); } @@ -512,7 +509,7 @@ where .unwrap() .last_response_submission = Some(Clock::real().now()); - anyhow::Ok(()) + anyhow::Ok(true) }, ); } diff --git a/crates/node/src/providers/ecdsa/presign.rs b/crates/node/src/providers/ecdsa/presign.rs index e535d04792..51d3901f45 100644 --- a/crates/node/src/providers/ecdsa/presign.rs +++ b/crates/node/src/providers/ecdsa/presign.rs @@ -64,12 +64,12 @@ impl EcdsaSignatureProvider { let parallelism_limiter = Arc::new(tokio::sync::Semaphore::new(config.concurrency)); let mut tasks = AutoAbortTaskCollection::new(); loop { - progress_tracker.update_progress(); + progress_tracker.update_progress().unwrap(); // Since return is ! metrics::MPC_OWNED_NUM_PRESIGNATURES_ONLINE - .set(presignature_store.num_owned_ready() as i64); + .set(presignature_store.num_owned_ready().unwrap() as i64); // Since return is ! metrics::MPC_OWNED_NUM_PRESIGNATURES_WITH_OFFLINE_PARTICIPANT - .set(presignature_store.num_owned_offline() as i64); - let my_presignatures_count: usize = presignature_store.num_owned(); + .set(presignature_store.num_owned_offline().unwrap() as i64); // Since return is ! + let my_presignatures_count: usize = presignature_store.num_owned().unwrap(); // Since return is ! metrics::MPC_OWNED_NUM_PRESIGNATURES_AVAILABLE.set(my_presignatures_count as i64); let should_generate = my_presignatures_count + in_flight_generations.num_in_flight() < config.desired_presignatures_to_buffer; @@ -79,10 +79,11 @@ impl EcdsaSignatureProvider { && in_flight_generations.num_in_flight() < config.concurrency * 2 { - let id = presignature_store.generate_and_reserve_id(); - progress_tracker.set_waiting_for_triples(true); - let (paired_triple_id, (triple0, triple1)) = triple_store.take_owned().await; - progress_tracker.set_waiting_for_triples(false); + let id = presignature_store.generate_and_reserve_id().unwrap(); // Since return is ! + progress_tracker.set_waiting_for_triples(true).unwrap(); // Since return is ! + let (paired_triple_id, (triple0, triple1)) = + triple_store.take_owned().await.unwrap(); // Since return is ! + progress_tracker.set_waiting_for_triples(false).unwrap(); // Since return is ! let participants = participants_from_triples(&triple0, &triple1); let task_id = EcdsaTaskId::Presignature { id, @@ -129,7 +130,7 @@ impl EcdsaSignatureProvider { presignature, participants, }, - ); + )?; anyhow::Ok(()) }), @@ -139,10 +140,10 @@ impl EcdsaSignatureProvider { // If the store is full, try to discard some presignatures which cannot be used right now if my_presignatures_count >= config.desired_presignatures_to_buffer { - presignature_store.maybe_discard_owned(1).await; + presignature_store.maybe_discard_owned(1).await.unwrap(); // Since return is ! } - tokio::time::sleep(std::time::Duration::from_millis(100)).await; + tokio::time::sleep(Duration::from_millis(100)).await; } } @@ -270,7 +271,7 @@ impl MpcLeaderCentricComputation<()> for FollowerPresignComputation { presignature, participants: channel.participants().to_vec(), }, - ); + )?; Ok(()) } @@ -287,16 +288,17 @@ struct PresignatureGenerationProgressTracker { } impl PresignatureGenerationProgressTracker { - pub fn set_waiting_for_triples(&self, waiting: bool) { + pub fn set_waiting_for_triples(&self, waiting: bool) -> anyhow::Result<()> { self.waiting_for_triples .store(waiting, std::sync::atomic::Ordering::Relaxed); - self.update_progress(); + self.update_progress()?; + Ok(()) } - pub fn update_progress(&self) { + pub fn update_progress(&self) -> anyhow::Result<()> { tracking::set_progress(&format!( "Presignatures: available: {}/{}; generating: {}{}", - self.presignature_store.num_owned(), + self.presignature_store.num_owned()?, self.desired_presignatures_to_buffer, self.in_flight_generations .load(std::sync::atomic::Ordering::Relaxed), @@ -308,6 +310,7 @@ impl PresignatureGenerationProgressTracker { } else { "" } - )) + )); + Ok(()) } } diff --git a/crates/node/src/providers/ecdsa/sign.rs b/crates/node/src/providers/ecdsa/sign.rs index 8078e1efa3..4540ddc274 100644 --- a/crates/node/src/providers/ecdsa/sign.rs +++ b/crates/node/src/providers/ecdsa/sign.rs @@ -71,7 +71,7 @@ impl EcdsaSignatureProvider { ) -> anyhow::Result<(Signature, VerifyingKey)> { let sign_request = self.sign_request_store.get(id).await?; let keyshare = self.keyshare(sign_request.domain)?; - let (presignature_id, presignature) = keyshare.presignature_store.take_owned().await; + let (presignature_id, presignature) = keyshare.presignature_store.take_owned().await?; let participants = presignature.participants.clone(); let channel = self.new_channel_for_task( EcdsaTaskId::Signature { diff --git a/crates/node/src/providers/ecdsa/triple.rs b/crates/node/src/providers/ecdsa/triple.rs index 48003d2162..b5a3dc99b9 100644 --- a/crates/node/src/providers/ecdsa/triple.rs +++ b/crates/node/src/providers/ecdsa/triple.rs @@ -94,10 +94,12 @@ impl EcdsaSignatureProvider { let mut offline: i64 = 0; let mut available: i64 = 0; for store in &triple_stores { - online += i64::try_from(store.num_owned_ready()).expect("triple count fits in i64"); - offline += - i64::try_from(store.num_owned_offline()).expect("triple count fits in i64"); - available += i64::try_from(store.num_owned()).expect("triple count fits in i64"); + online += i64::try_from(store.num_owned_ready().unwrap()) + .expect("triple count fits in i64"); + offline += i64::try_from(store.num_owned_offline().unwrap()) + .expect("triple count fits in i64"); + available += + i64::try_from(store.num_owned().unwrap()).expect("triple count fits in i64"); } metrics::MPC_OWNED_NUM_TRIPLES_ONLINE.set(online); metrics::MPC_OWNED_NUM_TRIPLES_WITH_OFFLINE_PARTICIPANT.set(offline); @@ -132,7 +134,7 @@ impl EcdsaSignatureProvider { .collect(); loop { - let my_triples_count = triple_store.num_owned(); + let my_triples_count = triple_store.num_owned().unwrap(); // Since return is ! let should_generate = my_triples_count + in_flight_generations.num_in_flight() < config.desired_triples_to_buffer; @@ -159,7 +161,8 @@ impl EcdsaSignatureProvider { }; let id_start = triple_store - .generate_and_reserve_id_range(SUPPORTED_TRIPLE_GENERATION_BATCH_SIZE as u32); + .generate_and_reserve_id_range(SUPPORTED_TRIPLE_GENERATION_BATCH_SIZE as u32) + .unwrap(); // Since return is ! let task_id = EcdsaTaskId::ManyTriples { start: id_start, count: SUPPORTED_TRIPLE_GENERATION_BATCH_SIZE as u32, @@ -199,10 +202,12 @@ impl EcdsaSignatureProvider { .await?; for (i, paired_triple) in triples.into_iter().enumerate() { - triple_store.add_owned( - id_start.add_to_counter(i.try_into()?)?, - paired_triple, - ); + triple_store + .add_owned( + id_start.add_to_counter(i.try_into()?)?, + paired_triple, + ) + .unwrap(); } anyhow::Ok(()) @@ -212,7 +217,7 @@ impl EcdsaSignatureProvider { // improve throughput by avoiding thundering herd situations. // Further optimization can be done to avoid thundering herd // situations in the first place. - tokio::time::sleep(std::time::Duration::from_secs( + tokio::time::sleep(Duration::from_secs( config.parallel_triple_generation_stagger_time_sec, )) .await; @@ -221,10 +226,10 @@ impl EcdsaSignatureProvider { // If the store is full, try to discard some triples which cannot be used right now if my_triples_count >= config.desired_triples_to_buffer { - triple_store.maybe_discard_owned(32).await; + triple_store.maybe_discard_owned(32).await.unwrap(); } - tokio::time::sleep(std::time::Duration::from_millis(100)).await; + tokio::time::sleep(Duration::from_millis(100)).await; } } @@ -352,7 +357,7 @@ impl MpcLeaderCentricComputation<()> self.out_triple_store.add_unowned( self.out_triple_id_start.add_to_counter(i.try_into()?)?, paired_triple, - ); + )?; } Ok(()) } @@ -559,10 +564,7 @@ mod tests { .flatten() .collect::>(); - Ok(triples - .into_iter() - .chain(passive_triples.await.unwrap()) - .collect()) + Ok(triples.into_iter().chain(passive_triples.await?).collect()) } fn new_triple_store( @@ -623,18 +625,18 @@ mod tests { participants.clone(), ); - let id_a = store_a.generate_and_reserve_id(); - store_a.add_owned(id_a, make_triple(&participants)); - let id_b = store_b.generate_and_reserve_id(); - store_b.add_owned(id_b, make_triple(&participants)); + let id_a = store_a.generate_and_reserve_id().unwrap(); + store_a.add_owned(id_a, make_triple(&participants)).unwrap(); + let id_b = store_b.generate_and_reserve_id().unwrap(); + store_b.add_owned(id_b, make_triple(&participants)).unwrap(); // When taking from `t = 3`. - let (taken_id, _) = store_a.take_owned().now_or_never().unwrap(); + let (taken_id, _) = store_a.take_owned().now_or_never().unwrap().unwrap(); // Then `t = 7`'s store is unaffected. assert_eq!(taken_id, id_a); - assert_eq!(store_a.num_owned(), 0); - assert_eq!(store_b.num_owned(), 1); + assert_eq!(store_a.num_owned().unwrap(), 0); + assert_eq!(store_b.num_owned().unwrap(), 1); } #[test] @@ -647,10 +649,10 @@ mod tests { let participants = vec![me]; let t = ReconstructionThreshold::new(3); let store = new_triple_store(db.clone(), me, t, participants.clone()); - let id = store.generate_and_reserve_id(); + let id = store.generate_and_reserve_id().unwrap(); // When adding an owned triple. - store.add_owned(id, make_triple(&participants)); + store.add_owned(id, make_triple(&participants)).unwrap(); // Then it is persisted under the per-`t` TripleV2 key. assert!( @@ -670,8 +672,8 @@ mod tests { let participants = vec![me]; let t = ReconstructionThreshold::new(3); let store = new_triple_store(db.clone(), me, t, participants.clone()); - let id = store.generate_and_reserve_id(); - store.add_owned(id, make_triple(&participants)); + let id = store.generate_and_reserve_id().unwrap(); + store.add_owned(id, make_triple(&participants)).unwrap(); // When the triple is consumed. let _ = store.take_owned().now_or_never().unwrap(); diff --git a/crates/node/src/providers/robust_ecdsa/presign.rs b/crates/node/src/providers/robust_ecdsa/presign.rs index 17498346e0..8322d6558d 100644 --- a/crates/node/src/providers/robust_ecdsa/presign.rs +++ b/crates/node/src/providers/robust_ecdsa/presign.rs @@ -65,12 +65,12 @@ pub(super) async fn run_background_presignature_generation( .expect("contract validation guarantees a valid threshold"); loop { - progress_tracker.update_progress(); + progress_tracker.update_progress().unwrap(); // Since return is ! metrics::MPC_OWNED_NUM_PRESIGNATURES_ONLINE - .set(presignature_store.num_owned_ready() as i64); + .set(presignature_store.num_owned_ready().unwrap() as i64); // Since return is ! metrics::MPC_OWNED_NUM_PRESIGNATURES_WITH_OFFLINE_PARTICIPANT - .set(presignature_store.num_owned_offline() as i64); - let my_presignatures_count: usize = presignature_store.num_owned(); + .set(presignature_store.num_owned_offline().unwrap() as i64); // Since return is ! + let my_presignatures_count: usize = presignature_store.num_owned().unwrap(); // Since return is ! metrics::MPC_OWNED_NUM_PRESIGNATURES_AVAILABLE.set(my_presignatures_count as i64); let should_generate = my_presignatures_count + in_flight_generations.num_in_flight() < config.desired_presignatures_to_buffer; @@ -80,7 +80,7 @@ pub(super) async fn run_background_presignature_generation( && in_flight_generations.num_in_flight() < config.concurrency * 2 { - let id = presignature_store.generate_and_reserve_id(); + let id = presignature_store.generate_and_reserve_id().unwrap(); // Since return is ! let participants = match client .select_random_active_participants_including_me(num_signers, &running_participants) { @@ -134,7 +134,7 @@ pub(super) async fn run_background_presignature_generation( presignature, participants, }, - ); + )?; anyhow::Ok(()) }), @@ -144,10 +144,10 @@ pub(super) async fn run_background_presignature_generation( // If the store is full, try to discard some presignatures which cannot be used right now if my_presignatures_count >= config.desired_presignatures_to_buffer { - presignature_store.maybe_discard_owned(1).await; + presignature_store.maybe_discard_owned(1).await.unwrap(); // Since return is ! } - tokio::time::sleep(std::time::Duration::from_millis(100)).await; + tokio::time::sleep(Duration::from_millis(100)).await; } } @@ -256,7 +256,7 @@ impl MpcLeaderCentricComputation<()> for FollowerPresignComputation { presignature, participants: channel.participants().to_vec(), }, - ); + )?; Ok(()) } @@ -272,13 +272,14 @@ struct PresignatureGenerationProgressTracker { } impl PresignatureGenerationProgressTracker { - pub fn update_progress(&self) { + pub fn update_progress(&self) -> anyhow::Result<()> { tracking::set_progress(&format!( "Presignatures: available: {}/{}; generating: {}", - self.presignature_store.num_owned(), + self.presignature_store.num_owned()?, self.desired_presignatures_to_buffer, self.in_flight_generations .load(std::sync::atomic::Ordering::Relaxed), - )) + )); + Ok(()) } } diff --git a/crates/node/src/providers/robust_ecdsa/sign.rs b/crates/node/src/providers/robust_ecdsa/sign.rs index 80a88d9eb6..0860533c63 100644 --- a/crates/node/src/providers/robust_ecdsa/sign.rs +++ b/crates/node/src/providers/robust_ecdsa/sign.rs @@ -28,7 +28,7 @@ impl RobustEcdsaSignatureProvider { ) -> anyhow::Result<(Signature, VerifyingKey)> { let sign_request = self.sign_request_store.get(id).await?; let keyshare = self.keyshare(sign_request.domain)?; - let (presignature_id, presignature) = keyshare.presignature_store.take_owned().await; + let (presignature_id, presignature) = keyshare.presignature_store.take_owned().await?; let participants = presignature.participants.clone(); let channel = self.client.new_channel_for_task( RobustEcdsaTaskId::Signature { diff --git a/crates/node/src/providers/verify_foreign_tx/sign.rs b/crates/node/src/providers/verify_foreign_tx/sign.rs index cae149c68d..2173eb0991 100644 --- a/crates/node/src/providers/verify_foreign_tx/sign.rs +++ b/crates/node/src/providers/verify_foreign_tx/sign.rs @@ -79,7 +79,7 @@ where let keyshare = self .ecdsa_signature_provider .keyshare(foreign_tx_request.domain_id)?; - let (presignature_id, presignature) = keyshare.presignature_store.take_owned().await; + let (presignature_id, presignature) = keyshare.presignature_store.take_owned().await?; let participants = presignature.participants.clone(); let channel = self.ecdsa_signature_provider.new_channel_for_task( VerifyForeignTxTaskId::VerifyForeignTx { diff --git a/crates/node/src/storage.rs b/crates/node/src/storage.rs index 378f2526e3..3c3f7af741 100644 --- a/crates/node/src/storage.rs +++ b/crates/node/src/storage.rs @@ -12,30 +12,23 @@ pub struct SignRequestStorage { impl SignRequestStorage { pub fn new(db: Arc) -> anyhow::Result { - let (tx, _) = tokio::sync::broadcast::channel(500); + let (tx, _) = broadcast::channel(500); Ok(Self { db, add_sender: tx }) } /// If given request is already in the database, returns false. /// Otherwise, inserts the request and returns true. - pub fn add(&self, request: &SignatureRequest) -> bool { - let key = borsh::to_vec(&request.id).unwrap(); - if self - .db - .get(DBCol::SignRequest, &key) - .expect("Unrecoverable error reading from database") - .is_some() - { - return false; + pub fn add(&self, request: &SignatureRequest) -> anyhow::Result { + let key = borsh::to_vec(&request.id)?; + if self.db.get(DBCol::SignRequest, &key)?.is_some() { + return Ok(false); } - let value_ser = serde_json::to_vec(&request).unwrap(); + let value_ser = serde_json::to_vec(&request)?; let mut update = self.db.update(); update.put(DBCol::SignRequest, &key, &value_ser); - update - .commit() - .expect("Unrecoverable error writing to database"); + update.commit()?; let _ = self.add_sender.send(request.id); - true + Ok(true) } /// Returns when a signature request with given id is present, then returns it. @@ -78,30 +71,23 @@ pub struct CKDRequestStorage { impl CKDRequestStorage { pub fn new(db: Arc) -> anyhow::Result { - let (tx, _) = tokio::sync::broadcast::channel(500); + let (tx, _) = broadcast::channel(500); Ok(Self { db, add_sender: tx }) } /// If given request is already in the database, returns false. /// Otherwise, inserts the request and returns true. - pub fn add(&self, request: &CKDRequest) -> bool { - let key = borsh::to_vec(&request.id).unwrap(); - if self - .db - .get(DBCol::CKDRequest, &key) - .expect("Unrecoverable error reading from database") - .is_some() - { - return false; + pub fn add(&self, request: &CKDRequest) -> anyhow::Result { + let key = borsh::to_vec(&request.id)?; + if self.db.get(DBCol::CKDRequest, &key)?.is_some() { + return Ok(false); } - let value_ser = serde_json::to_vec(&request).unwrap(); + let value_ser = serde_json::to_vec(&request)?; let mut update = self.db.update(); update.put(DBCol::CKDRequest, &key, &value_ser); - update - .commit() - .expect("Unrecoverable error writing to database"); + update.commit()?; let _ = self.add_sender.send(request.id); - true + Ok(true) } /// Returns when a ckd request with given id is present, then returns it. @@ -144,30 +130,23 @@ pub struct VerifyForeignTransactionRequestStorage { impl VerifyForeignTransactionRequestStorage { pub fn new(db: Arc) -> anyhow::Result { - let (tx, _) = tokio::sync::broadcast::channel(500); + let (tx, _) = broadcast::channel(500); Ok(Self { db, add_sender: tx }) } /// If given request is already in the database, returns false. /// Otherwise, inserts the request and returns true. - pub fn add(&self, request: &VerifyForeignTxRequest) -> bool { - let key = borsh::to_vec(&request.id).unwrap(); - if self - .db - .get(DBCol::VerifyForeignTxRequest, &key) - .expect("Unrecoverable error reading from database") - .is_some() - { - return false; + pub fn add(&self, request: &VerifyForeignTxRequest) -> anyhow::Result { + let key = borsh::to_vec(&request.id)?; + if self.db.get(DBCol::VerifyForeignTxRequest, &key)?.is_some() { + return Ok(false); } - let value_ser = serde_json::to_vec(&request).unwrap(); + let value_ser = serde_json::to_vec(&request)?; let mut update = self.db.update(); update.put(DBCol::VerifyForeignTxRequest, &key, &value_ser); - update - .commit() - .expect("Unrecoverable error writing to database"); + update.commit()?; let _ = self.add_sender.send(request.id); - true + Ok(true) } /// Returns when a verify foreign tx request with given id is present, then returns it. @@ -223,13 +202,9 @@ mod tests { types::SignatureRequest, }; - #[tokio::test] - async fn test_sig_request_storage() { - let dir = tempfile::tempdir().unwrap(); - let db = SecretDB::new(dir.path(), [1; 16]).unwrap(); - let storage = SignRequestStorage::new(db).unwrap(); - - let req1 = SignatureRequest { + fn signature_request() -> SignatureRequest { + // TODO: move this to common test function module since there are more duplicates elsewhere + SignatureRequest { id: CryptoHash(rand::random()), // All other fields are irrelevant for the test. receipt_id: CryptoHash([0; 32]), @@ -238,24 +213,41 @@ mod tests { timestamp_nanosec: 0, tweak: Tweak::new([0; 32]), domain: DomainId::legacy_ecdsa_id(), - }; - assert!(storage.add(&req1)); - assert!(!storage.add(&req1)); - let _ = storage - .get(req1.id) - .await - .expect("Stored signature request should be retrievable"); - let req2 = SignatureRequest { + } + } + + fn cdk_request() -> CKDRequest { + CKDRequest { id: CryptoHash(rand::random()), // All other fields are irrelevant for the test. receipt_id: CryptoHash([0; 32]), + app_public_key: near_mpc_contract_interface::types::CKDAppPublicKey::AppPublicKey( + "bls12381g1:6KtVVcAAGacrjNGePN8bp3KV6fYGrw1rFsyc7cVJCqR16Zc2ZFg3HX3hSZxSfv1oH6" + .parse() + .unwrap(), + ), + app_id: [1u8; 32].into(), entropy: [0; 32], - payload: Payload::from_legacy_ecdsa([0; 32]), timestamp_nanosec: 0, - tweak: Tweak::new([0; 32]), - domain: DomainId::legacy_ecdsa_id(), - }; - storage.add(&req2); + domain_id: DomainId::legacy_ecdsa_id(), + } + } + + #[tokio::test] + async fn test_sig_request_storage() { + let dir = tempfile::tempdir().unwrap(); + let db = SecretDB::new(dir.path(), [1; 16]).unwrap(); + let storage = SignRequestStorage::new(db).unwrap(); + + let req1 = signature_request(); + assert!(storage.add(&req1).unwrap()); + assert!(!storage.add(&req1).unwrap()); + let _ = storage + .get(req1.id) + .await + .expect("Stored signature request should be retrievable"); + let req2 = signature_request(); + storage.add(&req2).unwrap(); let _ = storage .get(req1.id) .await @@ -272,41 +264,15 @@ mod tests { let db = SecretDB::new(dir.path(), [1; 16]).unwrap(); let storage = CKDRequestStorage::new(db).unwrap(); - let req1 = CKDRequest { - id: CryptoHash(rand::random()), - // All other fields are irrelevant for the test. - receipt_id: CryptoHash([0; 32]), - app_public_key: near_mpc_contract_interface::types::CKDAppPublicKey::AppPublicKey( - "bls12381g1:6KtVVcAAGacrjNGePN8bp3KV6fYGrw1rFsyc7cVJCqR16Zc2ZFg3HX3hSZxSfv1oH6" - .parse() - .unwrap(), - ), - app_id: [1u8; 32].into(), - entropy: [0; 32], - timestamp_nanosec: 0, - domain_id: DomainId::legacy_ecdsa_id(), - }; - assert!(storage.add(&req1)); - assert!(!storage.add(&req1)); + let req1 = cdk_request(); + assert!(storage.add(&req1).unwrap()); + assert!(!storage.add(&req1).unwrap()); let _ = storage .get(req1.id) .await .expect("Stored CKD request should be retrievable"); - let req2 = CKDRequest { - id: CryptoHash(rand::random()), - // All other fields are irrelevant for the test. - receipt_id: CryptoHash([0; 32]), - app_public_key: near_mpc_contract_interface::types::CKDAppPublicKey::AppPublicKey( - "bls12381g1:6KtVVcAAGacrjNGePN8bp3KV6fYGrw1rFsyc7cVJCqR16Zc2ZFg3HX3hSZxSfv1oH6" - .parse() - .unwrap(), - ), - app_id: [1u8; 32].into(), - entropy: [0; 32], - timestamp_nanosec: 0, - domain_id: DomainId::legacy_ecdsa_id(), - }; - storage.add(&req2); + let req2 = cdk_request(); + storage.add(&req2).unwrap(); let _ = storage .get(req1.id) .await